feat: persist authored action idempotency (#201)

src/idempotency.py

@@ -0,0 +1,114 @@
+import json
+import sqlite3
+import time
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Callable
+
+
+@dataclass(frozen=True)
+class Reservation:
+    state: str
+    response: Any = None
+
+
+class IdempotencyLedger:
+    """Small durable ledger for replaying successful Gitea mutations."""
+
+    def __init__(
+        self,
+        path: str | Path,
+        *,
+        ttl_seconds: float,
+        max_entries: int,
+        clock: Callable[[], float] = time.time,
+    ) -> None:
+        self.path = Path(path)
+        self.ttl_seconds = ttl_seconds
+        self.max_entries = max_entries
+        self.clock = clock
+        self.path.parent.mkdir(parents=True, exist_ok=True)
+        self._initialize()
+
+    def _connect(self) -> sqlite3.Connection:
+        connection = sqlite3.connect(self.path, timeout=5)
+        connection.execute("PRAGMA busy_timeout = 5000")
+        return connection
+
+    def _initialize(self) -> None:
+        with self._connect() as connection:
+            connection.execute(
+                """
+                CREATE TABLE IF NOT EXISTS idempotency_operations (
+                    key TEXT PRIMARY KEY,
+                    fingerprint TEXT NOT NULL,
+                    status TEXT NOT NULL CHECK(status IN ('pending', 'completed')),
+                    response_json TEXT,
+                    created_at REAL NOT NULL,
+                    completed_at REAL
+                )
+                """
+            )
+            connection.execute(
+                "CREATE INDEX IF NOT EXISTS idempotency_completed_at_idx "
+                "ON idempotency_operations(completed_at)"
+            )
+
+    @staticmethod
+    def _fingerprint(value: tuple[Any, ...]) -> str:
+        return json.dumps(value, separators=(",", ":"), sort_keys=True)
+
+    def reserve(self, key: str, fingerprint: tuple[Any, ...]) -> Reservation:
+        encoded = self._fingerprint(fingerprint)
+        now = self.clock()
+        with self._connect() as connection:
+            connection.execute("BEGIN IMMEDIATE")
+            connection.execute(
+                "DELETE FROM idempotency_operations "
+                "WHERE status = 'completed' AND completed_at <= ?",
+                (now - self.ttl_seconds,),
+            )
+            row = connection.execute(
+                "SELECT fingerprint, status, response_json FROM idempotency_operations "
+                "WHERE key = ?",
+                (key,),
+            ).fetchone()
+            if row is not None:
+                if row[0] != encoded:
+                    return Reservation("conflict")
+                if row[1] == "completed":
+                    return Reservation("completed", json.loads(row[2]))
+                return Reservation("pending")
+            count = connection.execute(
+                "SELECT COUNT(*) FROM idempotency_operations"
+            ).fetchone()[0]
+            if count >= self.max_entries:
+                completed = connection.execute(
+                    "SELECT key FROM idempotency_operations "
+                    "WHERE status = 'completed' ORDER BY completed_at, rowid LIMIT 1"
+                ).fetchone()
+                if completed is None:
+                    return Reservation("busy")
+                connection.execute(
+                    "DELETE FROM idempotency_operations WHERE key = ?",
+                    (completed[0],),
+                )
+            connection.execute(
+                "INSERT INTO idempotency_operations "
+                "(key, fingerprint, status, created_at) VALUES (?, ?, 'pending', ?)",
+                (key, encoded, now),
+            )
+        return Reservation("reserved")
+
+    def clear(self) -> None:
+        with self._connect() as connection:
+            connection.execute("DELETE FROM idempotency_operations")
+
+    def complete(self, key: str, response: Any) -> None:
+        encoded = json.dumps(response, separators=(",", ":"), sort_keys=True)
+        with self._connect() as connection:
+            connection.execute(
+                "UPDATE idempotency_operations SET status = 'completed', "
+                "response_json = ?, completed_at = ? WHERE key = ?",
+                (encoded, self.clock(), key),
+            )

src/main.py

@@ -1,5 +1,6 @@
 import asyncio
 import math
+import os
 import time
 from collections.abc import Awaitable, Coroutine
 from contextlib import asynccontextmanager
@@ -24,6 +25,7 @@ from src.gitea_proxy import (
     pull_review_detail,
     repos,
 )
+from src.idempotency import IdempotencyLedger
 from src.models import Issue, PullRequest, Repo, User
 from src.suggestion_engine import compute
 from src.views import router as frontend_router
@@ -90,12 +92,18 @@ _live_section_retry_at: dict[str, float | None] = {
 }
 _live_snapshot_refreshing_sections: set[str] = set()
 _read_notification_ids: set[int] = set()
-_issue_creation_operations: dict[
-    str, tuple[tuple[Any, ...], asyncio.Task, float]
-] = {}
 _authored_action_operations: dict[
     str, tuple[tuple[Any, ...], asyncio.Task, float]
 ] = {}
+_state_dir = Path(os.getenv("STACKCHAIN_STATE_DIR", ".stackchain-state"))
+_idempotency_ledger = IdempotencyLedger(
+    os.getenv("STACKCHAIN_IDEMPOTENCY_DB", str(_state_dir / "idempotency.sqlite3")),
+    ttl_seconds=AUTHORED_ACTION_IDEMPOTENCY_TTL_SECONDS,
+    max_entries=(
+        AUTHORED_ACTION_IDEMPOTENCY_MAX_ENTRIES
+        + ISSUE_CREATION_IDEMPOTENCY_MAX_ENTRIES
+    ),
+)
 _available_issue_snapshot_task: asyncio.Task | None = None
 _available_issue_snapshot_value: list[dict] | None = None
 _available_issue_snapshot_created_at: float | None = None
@@ -236,39 +244,52 @@ async def _run_idempotent_authored_action(
     if not idempotency_key:
         return await asyncio.wait_for(operation, timeout=timeout)
 
-    now = time.monotonic()
-    expired = [
-        key for key, (_, task, created_at) in _authored_action_operations.items()
-        if task.done() and now - created_at >= AUTHORED_ACTION_IDEMPOTENCY_TTL_SECONDS
-    ]
-    for key in expired:
-        _authored_action_operations.pop(key, None)
+    reservation = _idempotency_ledger.reserve(idempotency_key, fingerprint)
+    if reservation.state == "conflict":
+        operation.close()
+        raise HTTPException(status_code=409, detail="Idempotency key already used")
+    if reservation.state == "completed":
+        operation.close()
+        return reservation.response
+    if reservation.state == "busy":
+        operation.close()
+        raise HTTPException(
+            status_code=503,
+            detail="Authored action queue is busy; please retry",
+            headers={"Retry-After": "1"},
+        )
 
     existing = _authored_action_operations.get(idempotency_key)
-    if existing is not None:
+    if reservation.state == "pending":
         operation.close()
-        if existing[0] != fingerprint:
-            raise HTTPException(status_code=409, detail="Idempotency key already used")
+        if existing is None or existing[0] != fingerprint:
+            raise HTTPException(
+                status_code=503,
+                detail=(
+                    "This action may still be completing; verify its result before retrying"
+                ),
+                headers={"Retry-After": "5"},
+            )
         task = existing[1]
     else:
-        while len(_authored_action_operations) >= AUTHORED_ACTION_IDEMPOTENCY_MAX_ENTRIES:
-            completed = [
-                key for key, (_, task, _) in _authored_action_operations.items()
-                if task.done()
-            ]
-            if not completed:
-                operation.close()
-                raise HTTPException(
-                    status_code=503,
-                    detail="Authored action queue is busy; please retry",
-                    headers={"Retry-After": "1"},
-                )
-            oldest = min(
-                completed, key=lambda key: _authored_action_operations[key][2]
-            )
-            _authored_action_operations.pop(oldest)
-        task = asyncio.create_task(operation)
-        _authored_action_operations[idempotency_key] = (fingerprint, task, now)
+        async def persist_result():
+            result = await operation
+            _idempotency_ledger.complete(idempotency_key, result)
+            return result
+
+        task = asyncio.create_task(persist_result())
+        _authored_action_operations[idempotency_key] = (
+            fingerprint,
+            task,
+            time.monotonic(),
+        )
+
+        def discard_completed(completed: asyncio.Task) -> None:
+            existing_operation = _authored_action_operations.get(idempotency_key)
+            if existing_operation is not None and existing_operation[1] is completed:
+                _authored_action_operations.pop(idempotency_key, None)
+
+        task.add_done_callback(discard_completed)
     return await asyncio.wait_for(asyncio.shield(task), timeout=timeout)
 
 
@@ -1269,48 +1290,19 @@ async def create_assigned_issue(
             repository, creation.title, creation.body, login, creation.label_ids
         )
 
-    operation = create_issue()
-    if idempotency_key:
-        now = time.monotonic()
-        expired = [
-            key for key, (_, task, created_at) in _issue_creation_operations.items()
-            if task.done()
-            and now - created_at >= ISSUE_CREATION_IDEMPOTENCY_TTL_SECONDS
-        ]
-        for key in expired:
-            _issue_creation_operations.pop(key, None)
-        fingerprint = (
-            repository, creation.title, creation.body, tuple(creation.label_ids)
-        )
-        existing = _issue_creation_operations.get(idempotency_key)
-        if existing is not None:
-            operation.close()
-            if existing[0] != fingerprint:
-                raise HTTPException(status_code=409, detail="Idempotency key already used")
-            task = existing[1]
-        else:
-            while len(_issue_creation_operations) >= ISSUE_CREATION_IDEMPOTENCY_MAX_ENTRIES:
-                completed = [
-                    key for key, (_, task, _) in _issue_creation_operations.items()
-                    if task.done()
-                ]
-                if not completed:
-                    operation.close()
-                    raise HTTPException(
-                        status_code=503,
-                        detail="Issue creation is busy; please retry",
-                        headers={"Retry-After": "1"},
-                    )
-                oldest = min(
-                    completed, key=lambda key: _issue_creation_operations[key][2]
-                )
-                _issue_creation_operations.pop(oldest)
-            task = asyncio.create_task(operation)
-            _issue_creation_operations[idempotency_key] = (fingerprint, task, now)
-        operation = asyncio.shield(task)
-
     try:
-        result = await asyncio.wait_for(operation, timeout=ISSUE_ACTION_TIMEOUT_SECONDS)
+        result = await _run_idempotent_authored_action(
+            create_issue(),
+            idempotency_key=idempotency_key,
+            fingerprint=(
+                "issue-create",
+                repository,
+                creation.title,
+                creation.body,
+                tuple(creation.label_ids),
+            ),
+            timeout=ISSUE_ACTION_TIMEOUT_SECONDS,
+        )
     except HTTPException:
         raise
     except Exception: