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,74 @@
+"""Record only job-owned native databases so abrupt cancellation is recoverable."""
+import asyncio
+import json
+import os
+from pathlib import Path
+import re
+from uuid import UUID
+
+
+def _names(directory: Path) -> list[str]:
+    path = directory / "databases.json"
+    if not path.exists():
+        return []
+    if path.is_symlink() or path.stat().st_size > 65536:
+        raise ValueError("Invalid calculation database ledger")
+    names = json.loads(path.read_text())
+    if not isinstance(names, list) or len(names) > 100 or any(
+        not isinstance(name, str) or not re.fullmatch(r"tjw_tmp_[a-z0-9_]+_[a-f0-9]{32}", name)
+        or len(name) > 63 for name in names
+    ):
+        raise ValueError("Invalid calculation database ledger")
+    return names
+
+
+def _write(directory: Path, names: list[str]) -> None:
+    path = directory / "databases.json"
+    staging = directory / "databases.json.new"
+    with staging.open("w") as output:
+        json.dump(names, output)
+        output.flush()
+        os.fsync(output.fileno())
+    staging.replace(path)
+
+
+def record_calculation_database(name: str, *, removed: bool = False) -> None:
+    """Called before clone creation and after successful normal cleanup."""
+    root = os.environ.get("TJWATER_CALCULATION_JOB_DIR")
+    if not root:
+        return
+    directory = Path(root)
+    names = _names(directory)
+    if removed:
+        names = [value for value in names if value != name]
+    elif name not in names:
+        names.append(name)
+    _write(directory, names)
+
+
+async def cleanup_calculation_databases(project_id: str, job_directory: str) -> None:
+    directory = Path(job_directory)
+    names = _names(directory)
+    if not names:
+        return
+    from app.infra.db.metadb.database import SessionLocal
+    from app.infra.db.metadb.repositories.metadata_repository import MetadataRepository
+    from app.infra.db.project_routing import ActiveProjectRouting, activate_project_routing
+    from app.native.wndb.core.projects import have_project, delete_project
+    async with SessionLocal() as session:
+        metadata = MetadataRepository(session)
+        project = await metadata.get_project_by_id(UUID(str(project_id)))
+        route = await metadata.get_project_db_routing(UUID(str(project_id)), "biz_data")
+        if not project or not route or route.db_type != "postgresql":
+            raise ValueError("Calculation cleanup routing unavailable")
+        routing = ActiveProjectRouting(project.code, route.dsn, project_id=project.id)
+        await session.commit()
+
+    def cleanup():
+        with activate_project_routing(routing):
+            for name in names[:]:
+                if have_project(name):
+                    delete_project(name)
+                names.remove(name)
+                _write(directory, names)
+    await asyncio.to_thread(cleanup)
@@ -0,0 +1,318 @@
+"""Independent, singleton calculation supervisor and isolated job subprocesses.
+
+Never imported by API lifespan. Child output contains the original QueryResult;
+algorithm stdout/stderr is discarded rather than interpreted as a protocol.
+"""
+from __future__ import annotations
+
+import asyncio
+from contextlib import asynccontextmanager, contextmanager
+from dataclasses import dataclass
+from datetime import datetime, timezone
+import json
+import os
+from pathlib import Path
+import signal
+import shutil
+import sys
+import tempfile
+import time
+from uuid import UUID, uuid4
+
+
+@dataclass(frozen=True)
+class WorkerSettings:
+    project_ids: tuple[str, ...]
+    max_concurrency: int = 2
+    poll_seconds: float = 2
+    lease_seconds: int = 30
+    heartbeat_seconds: float = 5
+
+    @classmethod
+    def from_environment(cls):
+        projects = tuple(dict.fromkeys(str(UUID(p.strip())) for p in
+            os.environ['CALCULATION_WORKER_PROJECT_IDS'].split(',')))
+        result = cls(projects, int(os.getenv('CALCULATION_WORKER_MAX_CONCURRENCY', '2')),
+                     float(os.getenv('CALCULATION_WORKER_POLL_SECONDS', '2')),
+                     int(os.getenv('CALCULATION_WORKER_LEASE_SECONDS', '30')),
+                     float(os.getenv('CALCULATION_WORKER_HEARTBEAT_SECONDS', '5')))
+        if not projects or not 1 <= result.max_concurrency <= 32:
+            raise ValueError('Configure explicit project IDs and concurrency between 1 and 32')
+        if not 0 < result.heartbeat_seconds < result.lease_seconds / 3 or result.poll_seconds <= 0:
+            raise ValueError('Heartbeat must be positive and below one third of the lease')
+        return result
+
+
+@asynccontextmanager
+async def queue_repository(project_id):
+    """Resolve current routing from metadata, never from a job payload."""
+    from psycopg import AsyncConnection
+    from psycopg.rows import dict_row
+    from app.infra.db.metadb.database import SessionLocal
+    from app.infra.db.metadb.repositories.metadata_repository import MetadataRepository
+    from app.infra.db.postgresql.calculation_jobs import CalculationJobRepository
+    async with SessionLocal() as session:
+        metadata = MetadataRepository(session)
+        project = await metadata.get_project_by_id(UUID(str(project_id)))
+        if not project or project.status != 'active':
+            raise ValueError('Calculation project is not active')
+        routing = await metadata.get_project_db_routing(UUID(str(project_id)), 'biz_data')
+        if not routing or routing.db_type != 'postgresql':
+            raise ValueError('Project business PostgreSQL routing is unavailable')
+        await session.commit()
+        async with await AsyncConnection.connect(routing.dsn, row_factory=dict_row) as conn:
+            yield CalculationJobRepository(conn, project_id)
+
+
+@asynccontextmanager
+async def execution_scope(row):
+    from fastapi import HTTPException
+    from app.mcp.scope import QueryMetadataRepository, QueryScope
+    from app.auth.project_dependencies import (resolve_project_context,
+        resolve_project_business_routing, get_project_pg_connection)
+    from app.infra.db.metadb.database import SessionLocal
+    from app.infra.db.project_routing import activate_project_routing
+    async with SessionLocal() as session:
+        metadata = QueryMetadataRepository(session)
+        user = await metadata.get_user_by_id(UUID(str(row['owner_user_id'])))
+        if not user or not user.is_active:
+            raise HTTPException(403, 'Calculation owner is inactive')
+        project = await resolve_project_context(str(row['project_id']), user, metadata)
+        routing = await resolve_project_business_routing(project, metadata)
+        await session.commit()
+        async with asynccontextmanager(get_project_pg_connection)(project, metadata) as conn:
+            await session.commit()
+            with activate_project_routing(routing):
+                yield QueryScope(project, user, conn, routing, metadata)
+
+
+async def execute_child(input_path: str, output_path: str):
+    from app.mcp.calculation_jobs import execute_calculation_job
+    row = json.loads(Path(input_path).read_text())
+    try:
+        async with execution_scope(row) as scope:
+            result = await execute_calculation_job(scope, row)
+        payload = {'status': 'succeeded', 'result': result.model_dump(mode='json')}
+    except Exception as exc:
+        # Do not persist exception text: DSNs, SQL, credentials or model data can occur there.
+        category = 'CALCULATION_FAILED'
+        if getattr(exc, 'status_code', None) in (401, 403):
+            category = 'AUTHORIZATION_DENIED'
+        elif getattr(exc, 'status_code', None) == 409:
+            category = 'EXECUTION_CONFLICT'
+        elif getattr(exc, 'status_code', None) in (404, 410):
+            category = 'EXECUTION_CONTEXT_UNAVAILABLE'
+        payload = {'status': 'failed', 'error': {'code': category}}
+    encoded = json.dumps(payload, allow_nan=False)
+    if len(encoded.encode()) > 8 * 1024 * 1024:
+        encoded = json.dumps({'status': 'failed', 'error': {'code': 'RESULT_TOO_LARGE'}})
+    Path(output_path).write_text(encoded)
+
+
+async def stop_process_group(process):
+    """Always kill the group, even if its original leader already exited."""
+    try:
+        os.killpg(process.pid, signal.SIGTERM)
+    except ProcessLookupError:
+        pass
+    try:
+        await asyncio.wait_for(process.wait(), 2)
+    except TimeoutError:
+        pass
+    # Native pools can ignore SIGTERM or survive the Python leader.
+    try:
+        os.killpg(process.pid, signal.SIGKILL)
+    except ProcessLookupError:
+        pass
+    await process.wait()
+    if process.stdin is not None:
+        process.stdin.close()
+
+
+@contextmanager
+def job_directory(row):
+    root = Path(os.getenv('CALCULATION_WORKER_STATE_DIR', 'temp/calculation-jobs'))
+    root.mkdir(parents=True, exist_ok=True)
+    directory = Path(tempfile.mkdtemp(prefix=f"{row.get('job_id', 'job')}-", dir=root))
+    directory.chmod(0o700)
+    (directory / 'scope.json').write_text(json.dumps({'project_id': str(row['project_id']), 'job_id': str(row.get('job_id', 'job')),
+        'lease_generation': row.get('lease_generation')}))
+    try:
+        yield str(directory.resolve())
+    finally:
+        # Failed external resource cleanup must remain operator-recoverable.
+        if not (directory / 'cleanup-failed').exists():
+            shutil.rmtree(directory)
+
+
+async def cleanup_job_resources(row, directory):
+    from app.services.calculation_cleanup import cleanup_calculation_databases
+    try:
+        await cleanup_calculation_databases(row['project_id'], directory)
+    except BaseException:
+        (Path(directory) / 'cleanup-failed').touch()
+        raise
+
+
+async def recover_job_resources(project_ids):
+    """Reconcile only this worker's explicit scopes after obtaining the singleton."""
+    root = Path(os.getenv('CALCULATION_WORKER_STATE_DIR', 'temp/calculation-jobs'))
+    if not root.exists():
+        return
+    for directory in root.iterdir():
+        if directory.is_symlink() or not directory.is_dir():
+            continue
+        scope_file = directory / 'scope.json'
+        if not scope_file.exists():
+            continue
+        row = json.loads(scope_file.read_text())
+        if str(row.get('project_id')) not in project_ids:
+            continue
+        job_id = str(UUID(row['job_id']))
+        if not directory.name.startswith(job_id + '-'):
+            raise RuntimeError('Calculation state directory identity is invalid')
+        async with queue_repository(row['project_id']) as repository:
+            stored = await repository.get(job_id)
+        if not stored or stored['lease_generation'] != row['lease_generation']:
+            raise RuntimeError('Calculation state ownership cannot be verified')
+        expiry = stored.get('lease_expires_at')
+        if expiry and expiry > datetime.now(timezone.utc) and stored['status'] in ('running', 'cancel_requested'):
+            raise RuntimeError('Previous calculation lease has not expired')
+        await cleanup_job_resources(row, str(directory.resolve()))
+        async with queue_repository(row['project_id']) as repository:
+            await repository.recover()
+        shutil.rmtree(directory)
+
+
+class CalculationWorker:
+    def __init__(self, settings, *, repository=queue_repository):
+        self.settings = settings
+        self.repository = repository
+        self.owner = str(uuid4())
+
+    async def run_job(self, row):
+        from fastapi.encoders import jsonable_encoder
+        with job_directory(row) as directory:
+            input_path = Path(directory) / 'input.json'
+            output_path = Path(directory) / 'result.json'
+            input_path.write_text(json.dumps(jsonable_encoder(row)))
+            process = await asyncio.create_subprocess_exec(
+                sys.executable, '-m', 'app.services.calculation_worker',
+                '--guard', str(input_path), str(output_path), start_new_session=True,
+                stdin=asyncio.subprocess.PIPE, stdout=asyncio.subprocess.DEVNULL,
+                stderr=asyncio.subprocess.DEVNULL,
+                env={**os.environ, 'TMPDIR': directory, 'TJWATER_CALCULATION_JOB_DIR': directory},
+            )
+            started = time.monotonic()
+            cleaned = False
+            try:
+                while True:
+                    async with asyncio.timeout(self.settings.heartbeat_seconds):
+                        async with self.repository(row['project_id']) as repository:
+                            status = await repository.heartbeat(row, self.owner, self.settings.lease_seconds)
+                    if status is None:
+                        return  # Finally stops all owned computation on a lost lease.
+                    terminal = None
+                    if status == 'cancel_requested':
+                        terminal = {'status': 'cancelled'}
+                    elif time.monotonic() - started >= row['timeout_seconds']:
+                        terminal = {'status': 'timed_out', 'error': {'code': 'EXECUTION_TIMEOUT'}}
+                    elif process.returncode is not None:
+                        try:
+                            if output_path.stat().st_size > 8 * 1024 * 1024:
+                                raise ValueError('Result exceeds persistence limit')
+                            terminal = json.loads(output_path.read_text())
+                        except (OSError, ValueError):
+                            terminal = {'status': 'failed', 'error': {'code': 'CHILD_PROCESS_FAILED'}}
+                    if terminal:
+                        await stop_process_group(process)
+                        try:
+                            await cleanup_job_resources(row, directory)
+                            cleaned = True
+                        except Exception:
+                            terminal = {'status': 'failed', 'error': {'code': 'CLEANUP_FAILED'}}
+                            cleaned = True
+                        async with self.repository(row['project_id']) as repository:
+                            finished = await repository.finish(row, self.owner, **terminal)
+                            if not finished:
+                                status = await repository.heartbeat(row, self.owner, self.settings.lease_seconds)
+                                if status == 'cancel_requested':
+                                    await repository.finish(row, self.owner, 'cancelled')
+                        return
+                    await asyncio.sleep(self.settings.heartbeat_seconds)
+            finally:
+                await stop_process_group(process)
+                if not cleaned:
+                    await cleanup_job_resources(row, directory)
+
+    async def run(self):
+        from sqlalchemy import text
+        from app.infra.db.metadb.database import engine
+        # Dedicated session-level metadata lock: independent API workers never
+        # participate; duplicate supervisors exit before claiming any jobs.
+        async with engine.connect() as lock:
+            acquired = await lock.scalar(text("SELECT pg_try_advisory_lock(1748531781, 2)"))
+            await lock.commit()
+            if not acquired:
+                raise RuntimeError('Another calculation supervisor is already active')
+            tasks = set()
+            project_cursor = 0
+            try:
+                # Allow any old supervisor to detect a broken lock connection and
+                # stop its children before this replacement claims capacity.
+                await asyncio.sleep(self.settings.lease_seconds)
+                await recover_job_resources(self.settings.project_ids)
+                while True:
+                    async with asyncio.timeout(self.settings.heartbeat_seconds):
+                        await lock.execute(text('SELECT 1'))
+                        await lock.commit()
+                    for task in tuple(tasks):
+                        if task.done():
+                            tasks.remove(task)
+                            task.result()  # Fail closed if heartbeat/persistence failed.
+                    projects = self.settings.project_ids
+                    ordered = projects[project_cursor:] + projects[:project_cursor]
+                    project_cursor = (project_cursor + 1) % len(projects)
+                    for project_id in ordered:
+                        if len(tasks) >= self.settings.max_concurrency:
+                            break
+                        async with self.repository(project_id) as repository:
+                            await repository.recover()
+                            row = await repository.claim(self.owner, self.settings.lease_seconds)
+                        if row:
+                            tasks.add(asyncio.create_task(self.run_job(row)))
+                    await asyncio.sleep(min(self.settings.poll_seconds, self.settings.heartbeat_seconds))
+            finally:
+                for task in tasks:
+                    task.cancel()
+                await asyncio.gather(*tasks, return_exceptions=True)
+                await lock.execute(text('SELECT pg_advisory_unlock(1748531781, 2)'))
+                await lock.commit()
+
+
+async def guard_child(input_path, output_path):
+    """A parent-death pipe kills the whole group even when the supervisor is SIGKILLed."""
+    reader = asyncio.StreamReader()
+    protocol = asyncio.StreamReaderProtocol(reader)
+    transport, _ = await asyncio.get_running_loop().connect_read_pipe(lambda: protocol, sys.stdin)
+    process = await asyncio.create_subprocess_exec(
+        sys.executable, '-m', 'app.services.calculation_worker', input_path, output_path,
+        stdin=asyncio.subprocess.DEVNULL,
+    )
+    completed = asyncio.create_task(process.wait())
+    parent_gone = asyncio.create_task(reader.read(1))
+    try:
+        done, _ = await asyncio.wait((completed, parent_gone), return_when=asyncio.FIRST_COMPLETED)
+        if parent_gone in done:
+            os.killpg(os.getpgrp(), signal.SIGKILL)
+        await completed
+    finally:
+        parent_gone.cancel()
+        transport.close()
+
+
+if __name__ == '__main__':
+    if sys.argv[1] == '--guard':
+        asyncio.run(guard_child(sys.argv[2], sys.argv[3]))
+    else:
+        asyncio.run(execute_child(sys.argv[1], sys.argv[2]))
@@ -778,6 +778,10 @@ class OperationalTaskService:
                    "Only a running run accepts Agent artifacts", status=409)
            now = self._now()
            self._require_active_run_lease(active_run, now)
+            if command.report and (
+                command.report.stage == "final" or command.report.outcome == "insufficient_data"
+            ):
+                await self._require_no_pending_calculations(task_id, active_run)
            previous_report = await self.repo.get_latest_report(str(task_id))
            require(not previous_report or previous_report.get("execution_generation") != task["execution_generation"]
                    or previous_report["report"].get("outcome") != "insufficient_data",
@@ -1866,6 +1870,15 @@ class OperationalTaskService:
                task_ids.append(task_id)
        return DispatchDueSchedulesResult(task_ids=task_ids, dispatched_at=now)

+    async def _require_no_pending_calculations(self, task_id: UUID, run: dict) -> None:
+        if hasattr(self.repo, "pending_calculation_jobs"):
+            pending = await self.repo.pending_calculation_jobs(
+                str(task_id), str(run["run_id"]), str(run["execution_request_id"]),
+            )
+            require(not pending, "PENDING_CALCULATIONS",
+                    "Wait for submitted calculations before publishing a final report or completing the run",
+                    refs=pending, status=409, recovery="retry_same_request")
+
    async def _finish_run(
        self, task_id: UUID, run_id: UUID, command: AgentRunCommand,
        *, outcome: str, failure_code: str | None = None,
@@ -1890,6 +1903,7 @@ class OperationalTaskService:
            task = await self._task(task_id)
            final_outcome = outcome
            if outcome == "succeeded":
+                await self._require_no_pending_calculations(task_id, run)
                if hasattr(self.repo, "pending_run_messages"):
                    pending_messages = await self.repo.pending_run_messages(str(run_id))
                    require(not pending_messages, "PENDING_AGENT_MESSAGES",
@@ -2,6 +2,7 @@ from collections.abc import Iterator
 from contextlib import contextmanager
 from pathlib import Path
 from tempfile import NamedTemporaryFile
+import os

 from app.native.wndb.inp.exporter import dump_inp

@@ -17,9 +18,11 @@ def temporary_project_inp(
    version: str = "2",
 ) -> Iterator[Path]:
    """Export one project model to an isolated INP and remove it afterwards."""
-    PROJECT_INP_DIRECTORY.mkdir(parents=True, exist_ok=True)
+    job_directory = os.environ.get("TJWATER_CALCULATION_JOB_DIR")
+    directory = Path(job_directory) / "inp" if job_directory else PROJECT_INP_DIRECTORY
+    directory.mkdir(parents=True, exist_ok=True)
    with NamedTemporaryFile(
-        dir=PROJECT_INP_DIRECTORY,
+        dir=directory,
        prefix=f"{purpose}_",
        suffix=".inp",
        delete=False,