feat(mcp): queue long-running calculations

2026-10-29 18:25:15 +08:00
parent 2865811f1b
commit 63af998045
32 changed files with 1588 additions and 72 deletions
@@ -0,0 +1,103 @@
+"""Durable project calculation queue. Every worker update is generation fenced."""
+from uuid import uuid4
+
+from fastapi import HTTPException
+from psycopg.types.json import Jsonb
+
+
+class CalculationJobRepository:
+    def __init__(self, connection, project_id):
+        self.connection = connection
+        self.project_id = str(project_id)
+
+    async def enqueue(self, *, owner_user_id, request_key, tool, parameters, binding, timeout_seconds):
+        cursor = await self.connection.execute(
+            """INSERT INTO workbench.calculation_jobs
+            (job_id,project_id,owner_user_id,request_key,tool,parameters,binding,timeout_seconds)
+            VALUES (%s,%s,%s,%s,%s,%s,%s,%s)
+            ON CONFLICT(project_id,owner_user_id,request_key) DO NOTHING RETURNING *""",
+            (uuid4(), self.project_id, owner_user_id, request_key, tool,
+             Jsonb(parameters), Jsonb(binding), timeout_seconds),
+        )
+        row = await cursor.fetchone()
+        if row:
+            return dict(row)
+        cursor = await self.connection.execute(
+            "SELECT * FROM workbench.calculation_jobs WHERE project_id=%s AND owner_user_id=%s AND request_key=%s",
+            (self.project_id, owner_user_id, request_key),
+        )
+        row = dict(await cursor.fetchone())
+        if row['tool'] != tool or row['parameters'] != parameters or row['binding'] != binding:
+            raise HTTPException(409, 'Calculation request_key was used with different inputs')
+        return row
+
+    async def get(self, job_id, owner_user_id=None):
+        cursor = await self.connection.execute(
+            """SELECT * FROM workbench.calculation_jobs WHERE project_id=%s AND job_id=%s
+            AND (%s::uuid IS NULL OR owner_user_id=%s)""",
+            (self.project_id, job_id, owner_user_id, owner_user_id),
+        )
+        row = await cursor.fetchone()
+        return dict(row) if row else None
+
+    async def list(self, *, owner_user_id, binding, limit, offset, status=None):
+        cursor = await self.connection.execute(
+            """SELECT * FROM workbench.calculation_jobs WHERE project_id=%s AND owner_user_id=%s
+            AND binding @> %s AND (%s::text IS NULL OR status=%s) ORDER BY created_at DESC,job_id LIMIT %s OFFSET %s""",
+            (self.project_id, owner_user_id, Jsonb(binding), status, status, limit, offset),
+        )
+        return [dict(row) for row in await cursor.fetchall()]
+
+    async def cancel(self, job_id, owner_user_id):
+        await self.connection.execute(
+            """UPDATE workbench.calculation_jobs SET status=CASE WHEN status='queued' THEN 'cancelled'
+            ELSE 'cancel_requested' END,updated_at=clock_timestamp()
+            WHERE project_id=%s AND job_id=%s AND owner_user_id=%s AND status IN ('queued','running')""",
+            (self.project_id, job_id, owner_user_id),
+        )
+        return await self.get(job_id, owner_user_id)
+
+    async def recover(self):
+        # A crashed write is never automatically replayed.
+        await self.connection.execute(
+            """UPDATE workbench.calculation_jobs SET status='interrupted',
+            error='{"code":"WORKER_INTERRUPTED"}'::jsonb,updated_at=clock_timestamp()
+            WHERE project_id=%s AND status IN ('running','cancel_requested')
+            AND lease_expires_at < clock_timestamp()""", (self.project_id,),
+        )
+
+    async def claim(self, owner, lease_seconds):
+        cursor = await self.connection.execute(
+            """UPDATE workbench.calculation_jobs j SET status='running',lease_owner=%s,
+            lease_generation=j.lease_generation+1,
+            lease_expires_at=clock_timestamp()+%s*interval '1 second',updated_at=clock_timestamp()
+            WHERE j.job_id=(SELECT job_id FROM workbench.calculation_jobs
+              WHERE project_id=%s AND status='queued' ORDER BY created_at,job_id
+              FOR UPDATE SKIP LOCKED LIMIT 1) RETURNING j.*""",
+            (owner, lease_seconds, self.project_id),
+        )
+        row = await cursor.fetchone()
+        return dict(row) if row else None
+
+    async def heartbeat(self, row, owner, lease_seconds):
+        cursor = await self.connection.execute(
+            """UPDATE workbench.calculation_jobs SET lease_expires_at=clock_timestamp()+%s*interval '1 second',
+            updated_at=clock_timestamp() WHERE project_id=%s AND job_id=%s AND lease_owner=%s
+            AND lease_generation=%s AND lease_expires_at>clock_timestamp()
+            AND status IN ('running','cancel_requested') RETURNING status""",
+            (lease_seconds, self.project_id, row['job_id'], owner, row['lease_generation']),
+        )
+        result = await cursor.fetchone()
+        return result['status'] if result else None
+
+    async def finish(self, row, owner, status, *, result=None, error=None):
+        cursor = await self.connection.execute(
+            """UPDATE workbench.calculation_jobs SET status=%s,result=%s,error=%s,updated_at=clock_timestamp()
+            WHERE project_id=%s AND job_id=%s AND lease_owner=%s AND lease_generation=%s
+            AND lease_expires_at>clock_timestamp() AND status IN ('running','cancel_requested')
+            AND (status!='cancel_requested' OR %s IN ('cancelled','interrupted')
+                 OR (%s='failed' AND %s='CLEANUP_FAILED')) RETURNING job_id""",
+            (status, Jsonb(result) if result is not None else None, Jsonb(error) if error else None,
+             self.project_id, row['job_id'], owner, row['lease_generation'], status, status, (error or {}).get('code')),
+        )
+        return await cursor.fetchone() is not None
@@ -75,6 +75,20 @@ class OperationalTaskRepository:
                status=409,
            )

+    async def pending_calculation_jobs(
+        self, task_id: str, run_id: str, execution_request_id: str,
+    ) -> list[str]:
+        """Only the exact original run blocks final artifacts and completion."""
+        cursor = await self.connection.execute(
+            """SELECT job_id FROM workbench.calculation_jobs
+               WHERE project_id=%s AND binding->>'task_id'=%s
+                 AND binding->>'run_id'=%s AND binding->>'execution_request_id'=%s
+                 AND status IN ('queued','running','cancel_requested')
+               ORDER BY created_at,job_id LIMIT 100""",
+            (self.project_id, task_id, run_id, execution_request_id),
+        )
+        return [str(row["job_id"]) for row in await cursor.fetchall()]
+
    async def get_task(self, task_id: str) -> dict | None:
        cursor = await self.connection.execute(
            """SELECT document FROM workbench.operational_tasks