diff --git a/backend/packages/harness/deerflow/subagents/batch_service.py b/backend/packages/harness/deerflow/subagents/batch_service.py index 9956e030e2e..d57c9112f55 100644 --- a/backend/packages/harness/deerflow/subagents/batch_service.py +++ b/backend/packages/harness/deerflow/subagents/batch_service.py @@ -42,10 +42,6 @@ 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, *, @@ -67,27 +63,20 @@ def __init__( 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 @@ -103,24 +92,13 @@ async def stop(self) -> None: for task in tasks: task.cancel() - 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 + 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() 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 dcb69fabcf2..1934160fb42 100644 --- a/backend/tests/test_subagent_batch_service.py +++ b/backend/tests/test_subagent_batch_service.py @@ -490,82 +490,3 @@ def execute_async(self, _prompt, task_id=None): 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