fix(projects): drain summary shutdown across caller cancellation (#7684)

backend/packages/harness/deerflow/AGENTS.md

@@ -135,3 +135,7 @@ continue, and consume/log its outcome in a completion callback. Prepare-stage
 cancellation retains its handoff/reclaim path. Regressions in
 `tests/blocking_io/test_workspace_changes_cancellation.py` must cover prompt
 metadata cancellation and text-cache drain/cleanup.
+
+### Project Summary Worker Shutdown (`projects/summaries.py`)
+
+`stop()` closes admission synchronously and shares one shielded `_stop_task`; `start()` reopens admission only after it completes and is cleared. Preserve queue-drop/sentinel/grace semantics. Regressions: `test_stop_drains_workers_before_propagating_repeated_cancellation` and `test_start_after_stop_restarts_worker_pool` in `backend/tests/test_project_summaries.py`.

backend/packages/harness/deerflow/projects/summaries.py

@@ -163,6 +163,7 @@ class SummaryGenerator:
         self._queue: asyncio.Queue[tuple[str, str, str] | None] = asyncio.Queue(maxsize=pcfg.summary_queue_size)
         self._workers: list[asyncio.Task] = []
         self._stopping = False
+        self._stop_task: asyncio.Task[None] | None = None
         self._metrics: dict[str, int] = {key: 0 for key in _METRIC_KEYS}
 
     @property
@@ -173,9 +174,15 @@ class SummaryGenerator:
         return dict(self._metrics)
 
     async def start(self) -> None:
-        """Spawn the worker pool; a no-op when disabled or already running."""
+        """Spawn or restart the worker pool when no shutdown is in progress."""
         if not self._enabled or self._workers:
             return
+        if self._stop_task is not None:
+            if not self._stop_task.done():
+                return
+            self._stop_task.result()
+            self._stop_task = None
+            self._stopping = False
         self._workers = [asyncio.create_task(self._worker(), name=f"summary-generator-{i}") for i in range(self._worker_count)]
 
     def enqueue(self, *, document_id: str, project_id: str, user_id: str) -> bool:
@@ -195,14 +202,57 @@ class SummaryGenerator:
             return False
 
     async def stop(self) -> None:
+        """Await owned shutdown before propagating caller cancellation.
+
+        Close admission synchronously so enqueue() cannot slip work in before the
+        shutdown task starts. Repeated callers share the same task; cancellation
+        is propagated only after owned cleanup has finished.
+        """
+        self._stopping = True
+        if self._stop_task is None:
+            self._stop_task = asyncio.create_task(self._stop_impl(), name="summary-generator-stop")
+        stop_task = self._stop_task
+        current = asyncio.current_task()
+        initial_cancellations = current.cancelling() if current is not None else 0
+        cancellation: asyncio.CancelledError | None = None
+
+        while True:
+            try:
+                await asyncio.shield(stop_task)
+                break
+            except asyncio.CancelledError as exc:
+                current = asyncio.current_task()
+                if current is None or current.cancelling() <= initial_cancellations:
+                    raise
+                if cancellation is None:
+                    cancellation = exc
+                if stop_task.done():
+                    break
+            except Exception:
+                break
+
+        # A cancellation can be requested just as the owned task completes, in
+        # which case awaiting an already-done task may not suspend to deliver it.
+        current = asyncio.current_task()
+        if cancellation is None and current is not None and current.cancelling() > initial_cancellations:
+            cancellation = asyncio.CancelledError()
+        if cancellation is not None:
+            try:
+                stop_task.result()
+            except BaseException as error:
+                raise cancellation from error
+            raise cancellation
+        stop_task.result()
+
+    async def _stop_impl(self) -> None:
         """Stop admitting, discard queued jobs, let idle workers exit, drain the rest.
 
-        Queued-but-unstarted jobs are dropped (counted ``queue_dropped``).
-        One sentinel per worker then wakes every idle loop so an empty queue
-        shuts down immediately instead of waiting out the grace period; the
-        grace bounds only genuinely in-flight jobs. A cancelled job drains
-        its conversion worker per ``await_drained`` semantics,
-        so shutdown may exceed the grace period while a conversion finishes.
+        Queued-but-unstarted jobs are dropped (counted as `queue_dropped`).
+        One sentinel per worker wakes every idle loop so an empty queue shuts
+        down immediately instead of waiting out the grace period; the grace
+        bounds only genuinely in-flight jobs. A cancelled job drains its
+        conversion worker per `await_drained` semantics, so shutdown may
+        exceed the grace period while a conversion finishes.
         """
         self._stopping = True
         while True:

backend/tests/test_project_summaries.py

@@ -412,3 +412,95 @@ class TestGenerator:
 
 async def _fake_model(text: str, metadata: dict | None = None) -> OneshotLLMResult:
     return _result(text, metadata)
+
+
+async def test_stop_drains_workers_before_propagating_repeated_cancellation(repos, text_file, monkeypatch):
+    sf, projects, docs = repos
+    await _shelf_row(projects, docs)
+    _patch_source(monkeypatch, (text_file, None))
+    blocker = asyncio.Event()
+    model_started = asyncio.Event()
+
+    async def blocking_model(**kwargs):
+        model_started.set()
+        await blocker.wait()
+        return _result("finished before shutdown")
+
+    monkeypatch.setattr(summaries, "run_oneshot_llm_detailed", blocking_model)
+    generator = SummaryGenerator(
+        session_factory=sf,
+        paths=None,
+        app_config=_app_config(summary_concurrency=1),
+    )
+    await generator.start()
+    row = await docs.get("doc-1", user_id="u1")
+    assert generator.enqueue(document_id="doc-1", project_id=row["project_id"], user_id="u1") is True
+    await model_started.wait()
+    workers = list(generator._workers)
+
+    stop_task = asyncio.create_task(generator.stop())
+    waited_after_first_cancel = False
+    waited_after_second_cancel = False
+    try:
+
+        async def wait_until_stopping():
+            while not generator._stopping:
+                await asyncio.sleep(0)
+
+        await asyncio.wait_for(wait_until_stopping(), timeout=1)
+        stop_task.cancel()
+        await asyncio.sleep(0)
+        waited_after_first_cancel = not stop_task.done()
+        stop_task.cancel()
+        await asyncio.sleep(0)
+        waited_after_second_cancel = not stop_task.done()
+    finally:
+        # Always let the fake provider finish so a failing regression cannot
+        # leave a background worker behind in the test process.
+        blocker.set()
+        if not stop_task.done():
+            stop_task.cancel()
+
+    with pytest.raises(asyncio.CancelledError):
+        await asyncio.wait_for(stop_task, timeout=2)
+    await asyncio.wait_for(asyncio.gather(*workers, return_exceptions=True), timeout=2)
+
+    assert waited_after_first_cancel
+    assert waited_after_second_cancel
+    assert generator._workers == []
+    assert (await docs.get("doc-1", user_id="u1"))["summary"] == "finished before shutdown"
+
+
+async def test_start_after_stop_restarts_worker_pool(repos, text_file, monkeypatch):
+    sf, projects, docs = repos
+    await _shelf_row(projects, docs)
+    _patch_source(monkeypatch, (text_file, None))
+    monkeypatch.setattr(summaries, "run_oneshot_llm_detailed", lambda **kwargs: _fake_model("restarted summary"))
+    generator = SummaryGenerator(
+        session_factory=sf,
+        paths=None,
+        app_config=_app_config(summary_concurrency=1),
+    )
+
+    await generator.start()
+    first_workers = list(generator._workers)
+    await generator.stop()
+    assert generator._workers == []
+    assert all(task.done() for task in first_workers)
+
+    # A completed shutdown belongs to the previous lifecycle only. Starting
+    # again must clear it and reopen admission for the new worker pool.
+    await generator.start()
+    second_workers = list(generator._workers)
+    assert len(second_workers) == 1
+    assert not generator._stopping
+    assert all(not task.done() for task in second_workers)
+
+    row = await docs.get("doc-1", user_id="u1")
+    assert generator.enqueue(document_id="doc-1", project_id=row["project_id"], user_id="u1") is True
+    await _drain(generator)
+    await generator.stop()
+
+    assert (await docs.get("doc-1", user_id="u1"))["summary"] == "restarted summary"
+    assert generator._workers == []
+    assert all(task.done() for task in second_workers)