fix(subagents): drain owned batch stop across cancellation (#5525)

* fix(subagents): drain owned batch stop across cancellation

* fix(subagents): preserve cancellation across stop failures

---------

Co-authored-by: Willem Jiang <willem.jiang@gmail.com>

README.md

@@ -1408,6 +1408,8 @@ async with runtime:
     # Serve or invoke graph while the durable worker is running.
 ```
 
+`SubagentRuntime.stop()` waits for its owned batch service to finish before propagating caller cancellation, including repeated cancellation. This drain has no timeout; repository operations and child cleanup must terminate. If service shutdown also fails or is cancelled, the first caller cancellation is preserved with the service failure as its cause.
+
 The factory still does not load YAML or create SQL infrastructure: the caller supplies the config snapshot, repository, and lifecycle. Because it accepts a caller-owned `system_prompt`, direct integrations also own any model-visible wording about those limits; the default middleware enforces the runtime limits regardless. The factory does not mount the Gateway owner-scoped HTTP routes or Web UI, so direct applications must expose their own result API/UI if they need those surfaces. For ordinary delegation only, `SubagentRuntime(...)` needs no asynchronous startup.
 
 Administrators can add, edit, disable, and delete reusable worker definitions from **Settings → Subagents**. Built-in and `config.yaml` definitions remain visible there as read-only entries. The default Lead Agent can use every enabled runtime sub-agent; each page-created Custom Agent can instead allow all, none, or a selected set. That selection is enforced both in the model-visible directory and by the server-side `task` tool. Managed definitions are deployment-wide in this version and follow `agent_storage.backend`: atomic files for a local deployment or the shared application database for multiple instances.

backend/packages/harness/deerflow/subagents/AGENTS.md

@@ -2,6 +2,8 @@
 
 **Context**: Capture after validation, before setup. Keep genuine replies, even hidden clarifications; exclude framework state and unpaired calls. Mark unserializable media as omitted.
 
+**Direct runtime shutdown**: `SubagentRuntime.stop()` holds its lifecycle lock until the owned service stop task terminates, then propagates the first caller cancellation with any service failure/cancellation as its cause. The drain is intentionally unbounded: repository awaits and child cleanup must terminate; a timeout must not detach still-owned work. Keep terminal-outcome and repeated-cancellation coverage in `tests/test_subagent_runtime.py`.
+
 **Durable batch acceptance**: `batch_task` normalizes optional per-item criteria
 before persistence (empty becomes null; 20 items × 500 neutralized characters),
 sharing `normalize_acceptance_criteria` with the executor and checker.

backend/packages/harness/deerflow/subagents/runtime.py

@@ -129,15 +129,42 @@ class SubagentRuntime:
             self._batch_started = True
 
     async def stop(self) -> None:
-        """Stop the owned worker and hide its bound tools from new graphs."""
+        """Stop the owned worker before propagating caller cancellation.
 
-        if self._owned_batch_service is None:
+        The drain intentionally has no timeout: releasing lifecycle ownership
+        while the service is still stopping would allow work to outlive this
+        runtime. The owned service's stop() must terminate, so its repository
+        awaits and child cleanup must not suppress cancellation indefinitely.
+        """
+
+        service = self._owned_batch_service
+        if service is None:
             return
         async with self._lifecycle_lock:
             if not self._batch_started:
                 return
+            # Hide the submitter immediately, but keep this lifecycle operation
+            # alive until the owned worker has actually finished stopping.
             self._batch_started = False
-            await self._owned_batch_service.stop()
+            stop_task = asyncio.create_task(service.stop(), name="subagent-runtime-batch-stop")
+            cancellation: asyncio.CancelledError | None = None
+            while not stop_task.done():
+                try:
+                    # wait() neither forwards caller cancellation to the owned
+                    # task nor raises that task's exception. Inspect its outcome
+                    # below so a service failure cannot replace cancellation.
+                    await asyncio.wait({stop_task})
+                except asyncio.CancelledError as exc:
+                    if cancellation is None:
+                        cancellation = exc
+
+            if cancellation is not None:
+                try:
+                    stop_task.result()
+                except (asyncio.CancelledError, Exception) as exc:
+                    raise cancellation from exc
+                raise cancellation
+            stop_task.result()
 
     async def __aenter__(self) -> SubagentRuntime:
         await self.start()

backend/tests/test_subagent_runtime.py

@@ -1,3 +1,4 @@
+import asyncio
 from types import SimpleNamespace
 from unittest.mock import AsyncMock, MagicMock, patch
 
@@ -75,3 +76,102 @@ async def test_runtime_owns_batch_worker_lifecycle_and_shared_capacity() -> None
     )
     service.start.assert_awaited_once_with()
     service.stop.assert_awaited_once_with()
+
+
+@pytest.mark.asyncio
+async def test_runtime_stop_drains_owned_batch_worker_across_repeated_cancellation() -> None:
+    stop_started = asyncio.Event()
+    allow_stop = asyncio.Event()
+
+    async def blocking_stop() -> None:
+        stop_started.set()
+        await allow_stop.wait()
+
+    service = MagicMock()
+    service.start = AsyncMock()
+    service.stop = AsyncMock(side_effect=blocking_stop)
+    repository = MagicMock()
+    app_config = MagicMock()
+
+    with patch(
+        "deerflow.subagents.batch_service.SubagentBatchService",
+        return_value=service,
+    ):
+        runtime = SubagentRuntime(
+            SubagentRuntimeConfig(max_running=1),
+            batch_repository=repository,
+            batch_config=SubagentBatchesConfig(enabled=True),
+            app_config=app_config,
+        )
+        await runtime.start()
+
+        stop_task = asyncio.create_task(runtime.stop())
+        await asyncio.wait_for(stop_started.wait(), timeout=1)
+
+        stop_task.cancel()
+        for _ in range(10):
+            await asyncio.sleep(0)
+
+        assert runtime.batch_submitter is None
+        assert not stop_task.done(), "runtime stop released ownership after the first cancellation"
+
+        stop_task.cancel()
+        for _ in range(10):
+            await asyncio.sleep(0)
+        assert not stop_task.done(), "runtime stop released ownership after repeated cancellation"
+
+        allow_stop.set()
+        with pytest.raises(asyncio.CancelledError):
+            await stop_task
+
+        service.stop.assert_awaited_once_with()
+        assert runtime.batch_submitter is None
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("cancel_caller", [False, True])
+@pytest.mark.parametrize("stop_error_type", [RuntimeError, asyncio.CancelledError])
+async def test_runtime_stop_preserves_caller_cancellation_when_service_fails(cancel_caller: bool, stop_error_type: type[BaseException]) -> None:
+    stop_error = stop_error_type("service shutdown failed")
+    stop_started = asyncio.Event()
+    allow_stop = asyncio.Event()
+
+    async def failing_stop() -> None:
+        stop_started.set()
+        await allow_stop.wait()
+        raise stop_error
+
+    service = MagicMock()
+    service.start = AsyncMock()
+    service.stop = AsyncMock(side_effect=failing_stop)
+    with patch("deerflow.subagents.batch_service.SubagentBatchService", return_value=service):
+        runtime = SubagentRuntime(
+            batch_repository=MagicMock(),
+            batch_config=SubagentBatchesConfig(enabled=True),
+            app_config=MagicMock(),
+        )
+        await runtime.start()
+        stop_task = asyncio.create_task(runtime.stop())
+        try:
+            await asyncio.wait_for(stop_started.wait(), timeout=1)
+            if cancel_caller:
+                stop_task.cancel("first caller cancellation")
+                await asyncio.sleep(0)
+                stop_task.cancel("second caller cancellation")
+                await asyncio.sleep(0)
+                assert not stop_task.done()
+            allow_stop.set()
+
+            expected_error = asyncio.CancelledError if cancel_caller else type(stop_error)
+            with pytest.raises(expected_error) as raised:
+                await stop_task
+            if cancel_caller:
+                assert raised.value.args == ("first caller cancellation",)
+                assert raised.value.__cause__ is stop_error
+            else:
+                assert raised.value is stop_error
+            service.stop.assert_awaited_once_with()
+            assert runtime.batch_submitter is None
+        finally:
+            allow_stop.set()
+            await asyncio.gather(stop_task, return_exceptions=True)