feat: extend operational investigation workflows

2026-10-30 17:30:19 +08:00
parent 63af998045
commit 05183aff93
82 changed files with 4644 additions and 330 deletions
@@ -3,6 +3,14 @@ from uuid import uuid4

 from fastapi import HTTPException
 from psycopg.types.json import Jsonb
+from app.mcp.errors import CalculationCapacityExceeded
+
+
+# Explicit metadata projection prevents status polling from decoding large JSON results.
+JOB_COLUMNS = "job_id,project_id,owner_user_id,request_key,tool,parameters,binding,timeout_seconds,status,error,lease_generation,lease_owner,lease_expires_at,created_at,updated_at,started_at,finished_at,queue_expires_at,result_expires_at"
+
+
+PUBLIC_JOB_COLUMNS = "job_id,status,tool,error,created_at,updated_at,started_at,finished_at,queue_expires_at,result_expires_at"


 class CalculationJobRepository:
@@ -10,30 +18,65 @@ class CalculationJobRepository:
        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),
+    async def maintain(self):
+        await self.connection.execute(
+            """UPDATE workbench.calculation_jobs SET status='timed_out',
+            error='{"code":"QUEUE_TIMEOUT"}'::jsonb,finished_at=clock_timestamp(),updated_at=clock_timestamp()
+            WHERE project_id=%s AND status='queued' AND queue_expires_at<=clock_timestamp()""",
+            (self.project_id,),
        )
-        row = await cursor.fetchone()
-        if row:
-            return dict(row)
+        await self.connection.execute(
+            """UPDATE workbench.calculation_jobs SET result=NULL
+            WHERE project_id=%s AND result IS NOT NULL AND result_expires_at<=clock_timestamp()""",
+            (self.project_id,),
+        )
+
+    async def enqueue(self, *, owner_user_id, request_key, tool, parameters, binding, timeout_seconds,
+                      queue_wait_seconds=3600, max_active_jobs=100):
+        async with self.connection.transaction():
+            return await self._enqueue(owner_user_id=owner_user_id, request_key=request_key,
+                tool=tool, parameters=parameters, binding=binding, timeout_seconds=timeout_seconds,
+                queue_wait_seconds=queue_wait_seconds, max_active_jobs=max_active_jobs)
+
+    async def _enqueue(self, *, owner_user_id, request_key, tool, parameters, binding,
+                       timeout_seconds, queue_wait_seconds, max_active_jobs):
+        # Serialize admission within the caller's transaction, across every API process.
+        await self.connection.execute(
+            "SELECT pg_advisory_xact_lock(hashtextextended(%s, 1748531781))", (self.project_id,))
+        await self.maintain()
        cursor = await self.connection.execute(
-            "SELECT * FROM workbench.calculation_jobs WHERE project_id=%s AND owner_user_id=%s AND request_key=%s",
+            f"SELECT {JOB_COLUMNS} 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):
+        existing = await cursor.fetchone()
+        if existing:
+            row = dict(existing)
+            if (row['tool'] != tool
+                    or {key: value for key, value in row['parameters'].items() if key != 'timeout_seconds'}
+                    != {key: value for key, value in parameters.items() if key != 'timeout_seconds'}
+                    or row['binding'] != binding):
+                raise HTTPException(409, 'Calculation request_key was used with different inputs')
+            return row
        cursor = await self.connection.execute(
-            """SELECT * FROM workbench.calculation_jobs WHERE project_id=%s AND job_id=%s
+            """SELECT count(*) AS active FROM workbench.calculation_jobs WHERE project_id=%s
+            AND status IN ('queued','running','cancel_requested')""", (self.project_id,))
+        if (await cursor.fetchone())['active'] >= max_active_jobs:
+            raise CalculationCapacityExceeded()
+        cursor = await self.connection.execute(
+            f"""INSERT INTO workbench.calculation_jobs
+            (job_id,project_id,owner_user_id,request_key,tool,parameters,binding,timeout_seconds,queue_expires_at)
+            VALUES (%s,%s,%s,%s,%s,%s,%s,%s,clock_timestamp()+%s*interval '1 second')
+            RETURNING {JOB_COLUMNS}""",
+            (uuid4(), self.project_id, owner_user_id, request_key, tool,
+             Jsonb(parameters), Jsonb(binding), timeout_seconds, queue_wait_seconds),
+        )
+        return dict(await cursor.fetchone())
+
+    async def get(self, job_id, owner_user_id=None, *, include_result=False):
+        await self.maintain()
+        columns = JOB_COLUMNS + (",result" if include_result else "")
+        cursor = await self.connection.execute(
+            f"""SELECT {columns} 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),
        )
@@ -41,8 +84,9 @@ class CalculationJobRepository:
        return dict(row) if row else None

    async def list(self, *, owner_user_id, binding, limit, offset, status=None):
+        await self.maintain()
        cursor = await self.connection.execute(
-            """SELECT * FROM workbench.calculation_jobs WHERE project_id=%s AND owner_user_id=%s
+            f"""SELECT {PUBLIC_JOB_COLUMNS} 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),
        )
@@ -51,28 +95,30 @@ class CalculationJobRepository:
    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()
+            ELSE 'cancel_requested' END,finished_at=CASE WHEN status='queued' THEN clock_timestamp() ELSE finished_at 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):
+        await self.maintain()
        # 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()
+            error='{"code":"WORKER_INTERRUPTED"}'::jsonb,finished_at=clock_timestamp(),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):
+        await self.maintain()
        cursor = await self.connection.execute(
-            """UPDATE workbench.calculation_jobs j SET status='running',lease_owner=%s,
+            """UPDATE workbench.calculation_jobs j SET status='running',started_at=clock_timestamp(),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
+              WHERE project_id=%s AND status='queued' AND queue_expires_at>clock_timestamp() ORDER BY created_at,job_id
              FOR UPDATE SKIP LOCKED LIMIT 1) RETURNING j.*""",
            (owner, lease_seconds, self.project_id),
        )
@@ -90,14 +136,15 @@ class CalculationJobRepository:
        result = await cursor.fetchone()
        return result['status'] if result else None

-    async def finish(self, row, owner, status, *, result=None, error=None):
+    async def finish(self, row, owner, status, *, result=None, error=None, result_retention_seconds=604800):
        cursor = await self.connection.execute(
-            """UPDATE workbench.calculation_jobs SET status=%s,result=%s,error=%s,updated_at=clock_timestamp()
+            """UPDATE workbench.calculation_jobs SET status=%s,result=%s,error=%s,finished_at=clock_timestamp(),
+            result_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')
            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,
+            (status, Jsonb(result) if result is not None else None, Jsonb(error) if error else None, result_retention_seconds,
             self.project_id, row['job_id'], owner, row['lease_generation'], status, status, (error or {}).get('code')),
        )
        return await cursor.fetchone() is not None