diff --git a/backend/packages/harness/deerflow/AGENTS.md b/backend/packages/harness/deerflow/AGENTS.md index cdb90112bb..00380b978a 100644 --- a/backend/packages/harness/deerflow/AGENTS.md +++ b/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`. diff --git a/backend/packages/harness/deerflow/projects/summaries.py b/backend/packages/harness/deerflow/projects/summaries.py index 146010ba4b..86c0de3277 100644 --- a/backend/packages/harness/deerflow/projects/summaries.py +++ b/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: diff --git a/backend/tests/test_project_summaries.py b/backend/tests/test_project_summaries.py index e80e3ad9ba..27d203e196 100644 --- a/backend/tests/test_project_summaries.py +++ b/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)