fix(gateway): bound subagent batch shutdown (#6434)

* fix(gateway): bound subagent batch shutdown

* test(gateway): cover bounded subagent batch shutdown

* fix(subagents): cancel batch work before stop drain

* test(subagents): pin cancellation before stop drain

* test(gateway): account for subagent shutdown budget

* docs(helm): refresh gateway shutdown budget

* docs(compose): refresh gateway shutdown budget

* docs(compose): correct gateway shutdown total

---------

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

backend/app/gateway/app.py

@@ -416,6 +416,27 @@ async def _shutdown_scheduled_task_service(app: FastAPI) -> None:
         logger.exception("Failed to stop scheduled task service")
 
 
+async def _shutdown_subagent_batch_service(app: FastAPI) -> None:
+    """Bound durable subagent batch stop so Gateway exit cannot hang forever."""
+    service = getattr(app.state, "subagent_batch_service", None)
+    if service is None:
+        return
+    app.state.subagent_batches_available = False
+    try:
+        await asyncio.wait_for(service.stop(), timeout=_SHUTDOWN_HOOK_TIMEOUT_SECONDS)
+    except TimeoutError:
+        logger.warning(
+            "Subagent batch service shutdown exceeded %.1fs; proceeding with worker exit.",
+            _SHUTDOWN_HOOK_TIMEOUT_SECONDS,
+        )
+    except Exception:
+        logger.exception("Failed to stop subagent batch service")
+    finally:
+        from deerflow.subagents.batch_runtime import set_subagent_batch_submitter
+
+        set_subagent_batch_submitter(None)
+
+
 def _scheduled_task_notification_repos(startup_config: AppConfig):
     """Return ``(connection_repo, notification_repo)`` for the scheduled-run outbox (issue #4254).
 
@@ -848,16 +869,7 @@ async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]:
 
         set_mcp_task_config_snapshot(None)
 
-        if getattr(app.state, "subagent_batch_service", None) is not None:
-            app.state.subagent_batches_available = False
-            try:
-                await app.state.subagent_batch_service.stop()
-            except Exception:
-                logger.exception("Failed to stop subagent batch service")
-            finally:
-                from deerflow.subagents.batch_runtime import set_subagent_batch_submitter
-
-                set_subagent_batch_submitter(None)
+        await _shutdown_subagent_batch_service(app)
 
         # Browser sessions have their own bounded teardown. MCP sessions close
         # after the runtime drains runs, since those runs may still call tools.

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

@@ -77,15 +77,23 @@ class SubagentBatchService:
         self._stop.set()
         poller = self._poller
         self._poller = None
-        if poller is not None:
-            poller.cancel()
-            await asyncio.gather(poller, return_exceptions=True)
+
+        # Issue every owned-work cancellation before the first await. The
+        # Gateway wraps this stop hook in a deadline; if poller teardown is
+        # slow, cancellation of this coroutine must not prevent native/item
+        # cancellation from being requested.
         execution_ids = list(self._execution_ids.values())
         for execution_id in execution_ids:
             request_cancel_background_task(execution_id)
         tasks = list(self._executions.values())
+
+        if poller is not None:
+            poller.cancel()
         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()

backend/tests/_gateway_shutdown_budget.py

@@ -24,6 +24,7 @@ HOOKS_BOUNDED_BY_SHUTDOWN_HOOK_TIMEOUT: tuple[str, ...] = (
     "notification delivery worker stop",
     "channel service stop",
     "scheduled task service stop",
+    "subagent batch service stop",
     "browser session close",
     "MCP session pool close",
 )
@@ -62,8 +63,8 @@ def memory_flush_seconds() -> float:
 def lifespan_shutdown_seconds() -> float:
     """Worst-case lifespan shutdown: every bounded hook + retrieval-warm wait + run drain + memory flush.
 
-    Unbounded steps (OIDC close, MCP/subagent task-service stops, backend close)
-    are expected to be quick and are not modeled; deployments keep slack on top
-    of this figure for them.
+    Unbounded steps (OIDC close, MCP task-service stop, backend close) are
+    expected to be quick and are not modeled; deployments keep slack on top of
+    this figure for them.
     """
     return bounded_hooks_seconds() + retrieval_warm_wait_seconds() + run_drain_seconds() + memory_flush_seconds()

backend/tests/test_subagent_batch_shutdown_timeout.py

@@ -0,0 +1,86 @@
+from __future__ import annotations
+
+import asyncio
+import inspect
+import logging
+from unittest.mock import AsyncMock, MagicMock
+
+import pytest
+from fastapi import FastAPI
+
+from app.gateway import app as gateway_app
+from deerflow.subagents import batch_service as batch_service_module
+from deerflow.subagents.batch_service import SubagentBatchService
+
+
+@pytest.mark.asyncio
+async def test_subagent_batch_service_shutdown_is_bounded(
+    monkeypatch: pytest.MonkeyPatch,
+    caplog: pytest.LogCaptureFixture,
+) -> None:
+    async def hang() -> None:
+        await asyncio.sleep(3600)
+
+    app = FastAPI()
+    service = MagicMock()
+    service.stop = AsyncMock(side_effect=hang)
+    app.state.subagent_batch_service = service
+    app.state.subagent_batches_available = True
+    monkeypatch.setattr(gateway_app, "_SHUTDOWN_HOOK_TIMEOUT_SECONDS", 0.01)
+
+    started = asyncio.get_running_loop().time()
+    with caplog.at_level(logging.WARNING, logger="app.gateway.app"):
+        await gateway_app._shutdown_subagent_batch_service(app)
+    elapsed = asyncio.get_running_loop().time() - started
+
+    assert elapsed < 0.5
+    assert app.state.subagent_batches_available is False
+    service.stop.assert_awaited_once()
+    assert "Subagent batch service shutdown exceeded" in caplog.text
+
+
+def test_lifespan_uses_bounded_subagent_batch_shutdown() -> None:
+    source = inspect.getsource(gateway_app.lifespan)
+    assert "await _shutdown_subagent_batch_service(app)" in source
+
+
+@pytest.mark.asyncio
+async def test_batch_stop_requests_owned_cancellation_before_poller_drain(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    release_poller = asyncio.Event()
+    poller_cancelled = asyncio.Event()
+    requested: list[str] = []
+
+    async def slow_poller() -> None:
+        try:
+            await asyncio.sleep(3600)
+        except asyncio.CancelledError:
+            poller_cancelled.set()
+            await release_poller.wait()
+
+    async def item_work() -> None:
+        await asyncio.sleep(3600)
+
+    service = SubagentBatchService.__new__(SubagentBatchService)
+    service._stop = asyncio.Event()
+    service._poller = asyncio.create_task(slow_poller())
+    item_task = asyncio.create_task(item_work())
+    service._executions = {"item-1": item_task}
+    service._execution_ids = {"item-1": "execution-1"}
+    service._item_batches = {"item-1": "batch-1"}
+
+    monkeypatch.setattr(
+        batch_service_module,
+        "request_cancel_background_task",
+        requested.append,
+    )
+
+    stop_task = asyncio.create_task(service.stop())
+    await poller_cancelled.wait()
+
+    assert requested == ["execution-1"]
+    assert item_task.cancelled() or item_task.cancelling()
+
+    release_poller.set()
+    await stop_task

deploy/helm/deer-flow/README.md

@@ -277,7 +277,7 @@ kubectl -n deer-flow exec deploy/deer-flow-provisioner -- curl -s localhost:8002
   `%40`). The chart uses an external `databaseUrl` verbatim and does not
   rewrite the DSN in a user-managed Secret.
 
-- **Graceful shutdown & memory drain.** The gateway pod sets `terminationGracePeriodSeconds` (default 90s, overridable via `gateway.terminationGracePeriodSeconds`) plus an optional `preStop` sleep (`gateway.preStopSleepSeconds`, default 5s), and bounds uvicorn's `--timeout-graceful-shutdown` (`gateway.uvicornGracefulShutdownSeconds`, default 10s) so an idle SSE connection cannot hold up lifespan shutdown indefinitely. The grace period MUST exceed the Gateway's graceful-shutdown work — the preStop sleep, the uvicorn timeout, and the lifespan's worst case: six hooks bounded at 5s each (startup trash sweep, notification delivery worker, channel service, scheduled task service, browser sessions, MCP session pool), the 1s retrieval-warm wait, the in-flight run drain (5s) and the memory-queue drain (`memory.shutdown_flush_timeout_seconds`, default 30s), about 66s in total, plus a buffer — because the drain runs on a daemon thread and K8s SIGKILLs anything still running at the end of the grace window. K8s defaults to 30s, which SIGKILLs the drain mid-flight and silently re-introduces the memory loss the drain is fixing. `backend/tests/_gateway_shutdown_budget.py` reads these bounds from the Gateway and pins the chart and compose budgets against them. **When you raise `memory.shutdown_flush_timeout_seconds`, raise `gateway.terminationGracePeriodSeconds` to match** (preStop + uvicorn timeout + ~36s of bounded hooks and drains + memory drain + buffer).
+- **Graceful shutdown & memory drain.** The gateway pod sets `terminationGracePeriodSeconds` (default 90s, overridable via `gateway.terminationGracePeriodSeconds`) plus an optional `preStop` sleep (`gateway.preStopSleepSeconds`, default 5s), and bounds uvicorn's `--timeout-graceful-shutdown` (`gateway.uvicornGracefulShutdownSeconds`, default 10s) so an idle SSE connection cannot hold up lifespan shutdown indefinitely. The grace period MUST exceed the Gateway's graceful-shutdown work — the preStop sleep, the uvicorn timeout, and the lifespan's worst case: seven hooks bounded at 5s each (startup trash sweep, notification delivery worker, channel service, scheduled task service, subagent batch service, browser sessions, MCP session pool), the 1s retrieval-warm wait, the in-flight run drain (5s) and the memory-queue drain (`memory.shutdown_flush_timeout_seconds`, default 30s), about 71s in total, plus a buffer — because the drain runs on a daemon thread and K8s SIGKILLs anything still running at the end of the grace window. K8s defaults to 30s, which SIGKILLs the drain mid-flight and silently re-introduces the memory loss the drain is fixing. `backend/tests/_gateway_shutdown_budget.py` reads these bounds from the Gateway and pins the chart and compose budgets against them. **When you raise `memory.shutdown_flush_timeout_seconds`, raise `gateway.terminationGracePeriodSeconds` to match** (preStop + uvicorn timeout + ~41s of bounded hooks and drains + memory drain + buffer).
 - **Gateway replicas.** Run control is cross-pod-safe since the work tracked
   by [issue #3948](https://github.com/bytedance/deer-flow/issues/3948) landed
   (#4003, #4064, #4500): admission is a durable one-active-run-per-thread

docker/docker-compose.yaml

@@ -124,11 +124,12 @@ services:
     # QR login and browser tools stay worker-local, so keep one worker while
     # they are enabled. --timeout-graceful-shutdown bounds how long open SSE
     # connections may delay lifespan shutdown after SIGTERM; stop_grace_period
-    # below must cover that bound plus the lifespan's worst case: five hooks
+    # below must cover that bound plus the lifespan's worst case: seven hooks
     # bounded at 5s each (trash sweep, notification worker, channel service,
-    # browser sessions, MCP session pool), the 1s retrieval-warm wait, the 5s
+    # scheduled task service, subagent batch service, browser sessions,
+    # MCP session pool), the 1s retrieval-warm wait, the 5s
     # run drain and the 30s memory flush (memory.shutdown_flush_timeout_seconds)
-    # -- about 61s -- plus slack. backend/tests/_gateway_shutdown_budget.py
+    # -- about 71s -- plus slack. backend/tests/_gateway_shutdown_budget.py
     # pins the arithmetic.
     command: sh -c "cd backend && PYTHONPATH=. uv run --no-sync uvicorn app.gateway.app:app --host 0.0.0.0 --port 8001 --workers ${GATEWAY_WORKERS:-1} --timeout-graceful-shutdown 10"
     # Docker's 10s default SIGKILLed the 30s memory queue flush