feat(mcp): queue long-running calculations

tests/integration/test_calculation_queue_storage.py

@@ -0,0 +1,70 @@
+"""Opt-in real PostgreSQL queue races, using an explicitly disposable database."""
+import asyncio
+import os
+from pathlib import Path
+from uuid import uuid4
+
+import pytest
+from fastapi import HTTPException
+from psycopg import AsyncConnection
+from psycopg.conninfo import conninfo_to_dict
+from psycopg.rows import dict_row
+
+from app.infra.db.postgresql.calculation_jobs import CalculationJobRepository
+
+
+def test_concurrent_submission_claims_fencing_and_project_isolation():
+    dsn = os.getenv("CALCULATION_QUEUE_TEST_DSN")
+    if not dsn:
+        pytest.skip("Set CALCULATION_QUEUE_TEST_DSN to a disposable PostgreSQL database")
+    assert conninfo_to_dict(dsn)["dbname"].startswith("tjwater_calc_test_")
+
+    async def run():
+        project, other, owner = uuid4(), uuid4(), uuid4()
+        async with await AsyncConnection.connect(dsn, row_factory=dict_row) as conn:
+            await conn.execute((Path(__file__).parents[2] / "resources/sql/026_calculation_jobs.sql").read_text())
+
+        async def submit(key, parameters=None, selected_project=project):
+            async with await AsyncConnection.connect(dsn, row_factory=dict_row) as conn:
+                return await CalculationJobRepository(conn, selected_project).enqueue(
+                    owner_user_id=owner, request_key=key, tool="run_burst_location",
+                    parameters=parameters or {"x": 1}, binding={"native_session_id": "session"}, timeout_seconds=600)
+
+        # Four independent API connections, same logical submission.
+        submissions = await asyncio.gather(*(submit("same") for _ in range(4)))
+        assert len({str(row["job_id"]) for row in submissions}) == 1
+        with pytest.raises(HTTPException) as conflict:
+            await submit("same", {"x": 2})
+        assert conflict.value.status_code == 409
+        second = await submit("second")
+        foreign = await submit("foreign", selected_project=other)
+
+        async def claim(worker):
+            async with await AsyncConnection.connect(dsn, row_factory=dict_row) as conn:
+                return await CalculationJobRepository(conn, project).claim(worker, 60)
+
+        claimed = await asyncio.gather(*(claim(f"worker-{index}") for index in range(4)))
+        rows = [row for row in claimed if row]
+        assert len(rows) == 2
+        assert len({str(row["job_id"]) for row in rows}) == 2
+        assert str(foreign["job_id"]) not in {str(row["job_id"]) for row in rows}
+
+        async with await AsyncConnection.connect(dsn, row_factory=dict_row) as conn:
+            repo = CalculationJobRepository(conn, project)
+            assert await repo.get(foreign["job_id"], owner) is None
+            assert await repo.get(second["job_id"], uuid4()) is None
+            row = rows[0]
+            stale = {**row, "lease_generation": row["lease_generation"] - 1}
+            assert await repo.heartbeat(stale, row["lease_owner"], 60) is None
+            assert not await repo.finish(stale, row["lease_owner"], "succeeded", result={})
+            cancelled = await repo.cancel(row["job_id"], owner)
+            assert cancelled["status"] == "cancel_requested"
+            assert not await repo.finish(row, row["lease_owner"], "succeeded", result={})
+            assert await repo.finish(row, row["lease_owner"], "cancelled")
+            expired = rows[1]
+            await conn.execute("UPDATE workbench.calculation_jobs SET lease_expires_at=now()-interval '1 second' WHERE job_id=%s", (expired["job_id"],))
+            await repo.recover()
+            assert (await repo.get(expired["job_id"], owner))["status"] == "interrupted"
+            assert not await repo.finish(expired, expired["lease_owner"], "succeeded", result={})
+            assert await repo.claim("replacement", 60) is None
+    asyncio.run(run())