代码提交

offline-backend-20260512/backend/app/gateway/deep_research_job_dispatcher.py

@@ -0,0 +1,158 @@
+"""Deep-research job dispatcher.
+
+Mirrors :class:`RoundtableJobDispatcher` exactly in shape — same wake-mechanism
+(``asyncio.Event`` + poll), same claim flow (``list_claimable`` → ``claim_job``
+CAS → ``executor.start_job``), same startup-scan, same cancel-reaper — minus the
+roundtable outbox compensator (deep research has no versioned draft to project
+back into; the session row IS the projection target, written directly by the
+executor).
+
+All workers run the same dispatcher; whichever ``claim_job`` returns a row
+(rowcount==1) executes — consistency comes from the DB, not in-process locks.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+from contextlib import suppress
+from datetime import UTC, datetime, timedelta
+
+from app.gateway.deep_research_job_executor import (
+    _LEASE_TTL_SECONDS,
+    WORKER_ID,
+    DeepResearchJobExecutor,
+)
+
+logger = logging.getLogger(__name__)
+
+_DISPATCH_POLL_INTERVAL = 3.0
+_CLAIM_BATCH = 20
+
+
+class DeepResearchJobDispatcher:
+    """Cross-worker scheduler for deep-research jobs (one instance per worker)."""
+
+    def __init__(self, store, executor: DeepResearchJobExecutor, *, worker_id: str | None = None) -> None:
+        self._store = store
+        self._executor = executor
+        self._worker_id = worker_id or WORKER_ID
+        self._wake = asyncio.Event()
+        self._task: asyncio.Task | None = None
+        self._stopping = False
+
+    # ── lifecycle ───────────────────────────────────────────────────────────
+
+    def start(self) -> None:
+        """Start the background loop (called from lifespan startup). Idempotent."""
+        if self._task is not None and not self._task.done():
+            return
+        self._stopping = False
+        self._task = asyncio.create_task(self._loop(), name="deep-research-job-dispatcher")
+        logger.info("[deep-research] dispatcher started (worker=%s)", self._worker_id)
+
+    async def stop(self) -> None:
+        """Stop the background loop (called from lifespan shutdown)."""
+        self._stopping = True
+        self._wake.set()
+        if self._task is not None:
+            self._task.cancel()
+            with suppress(asyncio.CancelledError, Exception):
+                await self._task
+            self._task = None
+
+    def nudge(self) -> None:
+        """Wake the loop immediately (called after start/resume to cut latency)."""
+        self._wake.set()
+
+    # ── main loop ───────────────────────────────────────────────────────────
+
+    async def _loop(self) -> None:
+        # Startup scan: pick up historically queued / expired-lease jobs.
+        await self._safe_dispatch()
+        while not self._stopping:
+            try:
+                await asyncio.wait_for(self._wake.wait(), timeout=_DISPATCH_POLL_INTERVAL)
+            except TimeoutError:
+                pass
+            self._wake.clear()
+            if self._stopping:
+                break
+            await self._safe_dispatch()
+
+    async def _safe_dispatch(self) -> None:
+        try:
+            await self.dispatch_once()
+        except Exception:  # noqa: BLE001
+            logger.exception("[deep-research] dispatch round failed")
+
+    # ── claim ───────────────────────────────────────────────────────────────
+
+    async def dispatch_once(self) -> int:
+        """Scan for claimable jobs and try to claim each.  Returns started count."""
+        now = datetime.now(UTC)
+        await self._reap_cancel_requested(now)
+        try:
+            candidate_ids = await self._store.list_claimable(now=now, limit=_CLAIM_BATCH)
+        except Exception:  # noqa: BLE001
+            logger.warning("[deep-research] list_claimable failed", exc_info=True)
+            return 0
+        if not candidate_ids:
+            return 0
+
+        started = 0
+        lease_until = now + timedelta(seconds=_LEASE_TTL_SECONDS)
+        for job_id in candidate_ids:
+            try:
+                row = await self._store.claim_job(
+                    job_id, lease_owner=self._worker_id, lease_until=lease_until
+                )
+            except Exception:  # noqa: BLE001
+                logger.warning("[deep-research] claim_job failed for %s", job_id, exc_info=True)
+                continue
+            if row is None:
+                continue  # another worker won / status changed
+
+            snapshot = row.get("input_snapshot")
+            if not snapshot or not isinstance(snapshot, dict):
+                logger.error("[deep-research] job %s claimed but snapshot missing; marking error", job_id)
+                with suppress(Exception):
+                    await self._store.update_progress(
+                        job_id,
+                        lease_owner=self._worker_id,
+                        status="failed",
+                        error_code="SNAPSHOT_MISSING",
+                        error_message="input_snapshot 缺失或损坏，无法在后台重建研究",
+                    )
+                continue
+
+            self._executor.start_job(snapshot, job_id=job_id, lease_owner=self._worker_id)
+            started += 1
+            logger.info(
+                "[deep-research] worker=%s claimed job %s (attempt=%s)",
+                self._worker_id,
+                job_id,
+                row.get("attempt"),
+            )
+        return started
+
+    async def _reap_cancel_requested(self, now: datetime) -> int:
+        """Finalize orphaned cancellations whose lease expired (worker crash)."""
+        try:
+            candidate_ids = await self._store.list_cancel_reapable(now=now, limit=_CLAIM_BATCH)
+        except Exception:  # noqa: BLE001
+            return 0
+
+        finalized = 0
+        for job_id in candidate_ids:
+            try:
+                row = await self._store.finalize_cancel(job_id)
+            except Exception:  # noqa: BLE001
+                continue
+            if row is not None:
+                finalized += 1
+                logger.info("[deep-research] reaped orphaned cancellation for job %s", job_id)
+        return finalized
+
+
+__all__ = ["DeepResearchJobDispatcher"]