fix(gateway): bound scheduled task shutdown (#6409)

* fix(gateway): rebase scheduler shutdown bound onto latest main

* style(gateway): format scheduled task shutdown

---------

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

backend/app/gateway/app.py

@@ -394,6 +394,28 @@ async def _shutdown_startup_trash_sweep(app: FastAPI) -> None:
         logger.exception("Startup trash sweep failed during shutdown")
 
 
+async def _shutdown_scheduled_task_service(app: FastAPI) -> None:
+    """Bound scheduler stop so a stuck poll cannot block Gateway exit.
+
+    ``asyncio.wait_for`` cancels the ``stop()`` coroutine at the deadline.
+    If the scheduler is already stuck inside an iteration, its owned poll task
+    may remain pending until event-loop teardown; this helper bounds Gateway
+    shutdown latency rather than promising completion of stuck scheduler work.
+    """
+    service = getattr(app.state, "scheduled_task_service", None)
+    if service is None:
+        return
+    try:
+        await asyncio.wait_for(service.stop(), timeout=_SHUTDOWN_HOOK_TIMEOUT_SECONDS)
+    except TimeoutError:
+        logger.warning(
+            "Scheduled task service shutdown exceeded %.1fs; proceeding with worker exit.",
+            _SHUTDOWN_HOOK_TIMEOUT_SECONDS,
+        )
+    except Exception:
+        logger.exception("Failed to stop scheduled task service")
+
+
 def _scheduled_task_notification_repos(startup_config: AppConfig):
     """Return ``(connection_repo, notification_repo)`` for the scheduled-run outbox (issue #4254).
 
@@ -810,11 +832,7 @@ async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]:
         except Exception:
             logger.exception("Failed to stop channel service")
 
-        if getattr(app.state, "scheduled_task_service", None) is not None:
-            try:
-                await app.state.scheduled_task_service.stop()
-            except Exception:
-                logger.exception("Failed to stop scheduled task service")
+        await _shutdown_scheduled_task_service(app)
 
         if getattr(app.state, "mcp_task_service", None) is not None:
             app.state.mcp_tasks_available = False

backend/tests/_gateway_shutdown_budget.py

@@ -23,6 +23,7 @@ HOOKS_BOUNDED_BY_SHUTDOWN_HOOK_TIMEOUT: tuple[str, ...] = (
     "startup trash sweep",
     "notification delivery worker stop",
     "channel service stop",
+    "scheduled task service stop",
     "browser session close",
     "MCP session pool close",
 )
@@ -61,7 +62,7 @@ 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, scheduler / task-service stops, backend close)
+    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.
     """

backend/tests/test_scheduler_shutdown_timeout.py

@@ -0,0 +1,37 @@
+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
+
+
+@pytest.mark.asyncio
+async def test_scheduled_task_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.scheduled_task_service = service
+    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_scheduled_task_service(app)
+    elapsed = asyncio.get_running_loop().time() - started
+
+    assert elapsed < 0.5
+    service.stop.assert_awaited_once()
+    assert "Scheduled task service shutdown exceeded" in caplog.text
+
+
+def test_lifespan_uses_bounded_scheduled_task_shutdown() -> None:
+    source = inspect.getsource(gateway_app.lifespan)
+    assert "await _shutdown_scheduled_task_service(app)" in source

deploy/helm/deer-flow/README.md

@@ -270,7 +270,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: five hooks bounded at 5s each (startup trash sweep, notification delivery worker, channel 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 61s 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 + ~31s 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: 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).
 - **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