services: update comfy agent, senpai md consumer, and swapper deps

2026-02-19 00:14:18 -08:00
parent c201d105f6
commit c57e6ed96b
15 changed files with 586 additions and 40 deletions
--- a/services/comfy-agent/app/idempotency.py
+++ b/services/comfy-agent/app/idempotency.py
@@ -0,0 +1,111 @@
+import sqlite3
+import threading
+import time
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Literal, Optional
+
+from .config import settings
+
+
+Decision = Literal["created", "exists", "conflict"]
+
+
+@dataclass
+class IdempotencyResult:
+    decision: Decision
+    job_id: str
+
+
+class IdempotencyStore:
+    def __init__(self, db_path: str, ttl_s: int) -> None:
+        self.db_path = db_path
+        self.ttl_s = max(60, int(ttl_s))
+        self._lock = threading.Lock()
+        self._init_db()
+
+    def _connect(self) -> sqlite3.Connection:
+        conn = sqlite3.connect(self.db_path, check_same_thread=False)
+        conn.execute("PRAGMA journal_mode=WAL")
+        conn.execute("PRAGMA synchronous=NORMAL")
+        return conn
+
+    def _init_db(self) -> None:
+        Path(self.db_path).parent.mkdir(parents=True, exist_ok=True)
+        with self._connect() as conn:
+            conn.execute(
+                """
+                CREATE TABLE IF NOT EXISTS idempotency_jobs (
+                    idem_key TEXT NOT NULL,
+                    gen_type TEXT NOT NULL,
+                    req_hash TEXT NOT NULL,
+                    job_id TEXT NOT NULL,
+                    created_at INTEGER NOT NULL,
+                    expires_at INTEGER NOT NULL,
+                    PRIMARY KEY (idem_key, gen_type)
+                )
+                """
+            )
+            conn.execute(
+                "CREATE INDEX IF NOT EXISTS idx_idem_expires_at ON idempotency_jobs(expires_at)"
+            )
+            conn.commit()
+
+    def reserve(
+        self,
+        *,
+        idem_key: str,
+        gen_type: str,
+        req_hash: str,
+        job_id: str,
+    ) -> IdempotencyResult:
+        now = int(time.time())
+        expires = now + self.ttl_s
+
+        with self._lock, self._connect() as conn:
+            conn.execute("DELETE FROM idempotency_jobs WHERE expires_at < ?", (now,))
+            row = conn.execute(
+                """
+                SELECT req_hash, job_id, expires_at
+                FROM idempotency_jobs
+                WHERE idem_key = ? AND gen_type = ?
+                """,
+                (idem_key, gen_type),
+            ).fetchone()
+
+            if row:
+                existing_hash, existing_job_id, existing_exp = row
+                if existing_exp >= now:
+                    if existing_hash == req_hash:
+                        return IdempotencyResult(decision="exists", job_id=existing_job_id)
+                    return IdempotencyResult(decision="conflict", job_id=existing_job_id)
+
+                conn.execute(
+                    "DELETE FROM idempotency_jobs WHERE idem_key = ? AND gen_type = ?",
+                    (idem_key, gen_type),
+                )
+
+            conn.execute(
+                """
+                INSERT INTO idempotency_jobs
+                    (idem_key, gen_type, req_hash, job_id, created_at, expires_at)
+                VALUES (?, ?, ?, ?, ?, ?)
+                """,
+                (idem_key, gen_type, req_hash, job_id, now, expires),
+            )
+            conn.commit()
+            return IdempotencyResult(decision="created", job_id=job_id)
+
+
+IDEMPOTENCY_STORE: Optional[IdempotencyStore] = None
+
+
+def init_idempotency_store() -> None:
+    global IDEMPOTENCY_STORE
+    if IDEMPOTENCY_STORE is not None:
+        return
+
+    db_path = settings.IDEMPOTENCY_DB_PATH
+    if not db_path:
+        db_path = str(Path(settings.STORAGE_PATH) / "idempotency.sqlite3")
+    IDEMPOTENCY_STORE = IdempotencyStore(db_path=db_path, ttl_s=settings.IDEMPOTENCY_TTL_S)