From ada1b9ba6b2d0c47f8517d9e7a635f7fe9c5cd38 Mon Sep 17 00:00:00 2001 From: NanPan <111261006+poijygfdyy@users.noreply.github.com> Date: Thu, 8 Oct 2026 23:21:19 +0800 Subject: [PATCH] fix(subagents): fence batch poller start while stop drains (#6517) * fix(subagents): fence batch poller restart until stop drains * test(subagents): pin concurrent start during batch stop drain * fix(subagents): keep restart fenced across overlapping stop drains * fix(subagents): reconcile aborted stop drains and support direct fixtures * test(subagents): retry cancellation-interrupted poller shutdown --- .../deerflow/subagents/batch_service.py | 38 +++++++-- backend/tests/test_subagent_batch_service.py | 79 +++++++++++++++++++ 2 files changed, 109 insertions(+), 8 deletions(-) diff --git a/backend/packages/harness/deerflow/subagents/batch_service.py b/backend/packages/harness/deerflow/subagents/batch_service.py index d57c9112f5..9956e030e2 100644 --- a/backend/packages/harness/deerflow/subagents/batch_service.py +++ b/backend/packages/harness/deerflow/subagents/batch_service.py @@ -42,6 +42,10 @@ def _usage(records: list[dict[str, Any]] | None) -> dict[str, int] | None: class SubagentBatchService: """Lease, execute, and recover durable native-subagent batch items.""" + # Also used by focused shutdown tests constructing an instance via __new__. + _stopping: bool = False + _stop_drains: int = 0 + def __init__( self, *, @@ -63,20 +67,27 @@ class SubagentBatchService: self._lease_owner = f"{socket.gethostname()}:{uuid.uuid4().hex}" self._stop = asyncio.Event() self._poller: asyncio.Task[None] | None = None + self._stopping = False + self._stop_drains = 0 self._executions: dict[str, asyncio.Task[None]] = {} self._execution_ids: dict[str, str] = {} self._item_batches: dict[str, str] = {} async def start(self) -> None: + if self._stopping: + raise RuntimeError("cannot start subagent batch poller before stop completes") if self._poller is not None: return self._stop.clear() self._poller = asyncio.create_task(self._run(), name="subagent-batch-poller") async def stop(self) -> None: + # Keep poller ownership visible until the entire drain finishes. + # Otherwise start() could create a fresh poller during this await. + self._stop_drains += 1 + self._stopping = True self._stop.set() poller = self._poller - self._poller = None # Issue every owned-work cancellation before the first await. The # Gateway wraps this stop hook in a deadline; if poller teardown is @@ -92,13 +103,24 @@ class SubagentBatchService: for task in tasks: task.cancel() - if poller is not None: - await asyncio.gather(poller, return_exceptions=True) - if tasks: - await asyncio.gather(*tasks, return_exceptions=True) - self._executions.clear() - self._execution_ids.clear() - self._item_batches.clear() + completed = False + try: + if poller is not None: + await asyncio.gather(poller, return_exceptions=True) + if tasks: + await asyncio.gather(*tasks, return_exceptions=True) + self._executions.clear() + self._execution_ids.clear() + self._item_batches.clear() + if self._poller is poller: + self._poller = None + completed = True + finally: + # A cancelled caller cannot leave a phantom drain count. Keep + # start fenced until a subsequent successful stop retries cleanup. + self._stop_drains -= 1 + if completed and self._stop_drains == 0: + self._stopping = False async def _run(self) -> None: while not self._stop.is_set(): diff --git a/backend/tests/test_subagent_batch_service.py b/backend/tests/test_subagent_batch_service.py index 1934160fb4..dcb69fabcf 100644 --- a/backend/tests/test_subagent_batch_service.py +++ b/backend/tests/test_subagent_batch_service.py @@ -490,3 +490,82 @@ async def test_cancel_during_tool_assembly_skips_launch(monkeypatch, tmp_path) - assert batch["counts"]["cancelled"] == 1 finally: await close_engine() + + +@pytest.mark.asyncio +async def test_start_cannot_replace_poller_while_stop_is_draining(monkeypatch: pytest.MonkeyPatch) -> None: + service = SubagentBatchService( + repository=SimpleNamespace(), + config=SubagentBatchesConfig(), + runtime_config=SubagentRuntimeConfig(), + ) + poller_entered = asyncio.Event() + allow_poller_exit = asyncio.Event() + + async def blocking_poller() -> None: + poller_entered.set() + try: + await asyncio.Future() + except asyncio.CancelledError: + await allow_poller_exit.wait() + + monkeypatch.setattr(service, "_run", blocking_poller) + await service.start() + original_poller = service._poller + assert original_poller is not None + await asyncio.wait_for(poller_entered.wait(), timeout=1) + + stop_task = asyncio.create_task(service.stop()) + await asyncio.sleep(0) + assert not stop_task.done() + assert service._poller is original_poller + with pytest.raises(RuntimeError, match="before stop completes"): + await service.start() + assert service._poller is original_poller + + allow_poller_exit.set() + await asyncio.wait_for(stop_task, timeout=1) + assert original_poller.done() + assert service._poller is None + assert not service._stopping + + await service.start() + assert service._poller is not original_poller + await service.stop() + + +@pytest.mark.asyncio +async def test_cancelled_batch_stop_allows_successful_retry(monkeypatch: pytest.MonkeyPatch) -> None: + service = SubagentBatchService( + repository=SimpleNamespace(), + config=SubagentBatchesConfig(), + runtime_config=SubagentRuntimeConfig(), + ) + poller_entered = asyncio.Event() + release_poller = asyncio.Event() + + async def reluctant_poller() -> None: + poller_entered.set() + try: + await asyncio.Future() + except asyncio.CancelledError: + await release_poller.wait() + + monkeypatch.setattr(service, "_run", reluctant_poller) + await service.start() + await asyncio.wait_for(poller_entered.wait(), timeout=1) + first_stop = asyncio.create_task(service.stop()) + await asyncio.sleep(0) + first_stop.cancel() + with pytest.raises(asyncio.CancelledError): + await first_stop + assert service._stop_drains == 0 + assert service._stopping + with pytest.raises(RuntimeError, match="before stop completes"): + await service.start() + + release_poller.set() + await asyncio.wait_for(service.stop(), timeout=1) + assert service._poller is None + assert service._stop_drains == 0 + assert not service._stopping