fix(extensions): drain service shutdown across cancellation (#5549)

* fix(extensions): drain service shutdown across cancellation

* docs(gateway): document extension shutdown drain

backend/app/gateway/AGENTS.md

@@ -7,6 +7,10 @@ not abort runtime teardown. Cancellation waits for the owned workers before
 propagating; backend `close()` overrides must be quick or internally bounded
 because close has no host timeout. Budget resolution and close in the pod grace
 period in addition to the configured flush timeout and other shutdown hooks.
+Extension-service teardown follows the same ownership rule: `stop_services()`
+is drained across host cancellation before later runtime resources unwind; each
+service `stop()` remains bounded by its 30-second `asyncio.timeout`, so include
+that bound in pod grace-period budgeting.
 
 `conversation_access.py` binds an opt-in read-only tool to a run request's
 explicit `conversation_references` and effective `runs:read` permission. Never

backend/app/gateway/deps.py

@@ -33,6 +33,7 @@ from deerflow.persistence.feedback import FeedbackRepository
 from deerflow.runtime import ORPHAN_RECOVERY_STOP_REASON, STARTUP_ORPHAN_RECOVERY_ERROR, RunContext, RunManager, StreamBridge
 from deerflow.runtime.events.store.base import RunEventStore
 from deerflow.runtime.runs.store.base import RunStore
+from deerflow.utils.file_io import await_drained
 
 logger = logging.getLogger(__name__)
 
@@ -520,9 +521,11 @@ async def langgraph_runtime(app: FastAPI, startup_config: AppConfig) -> AsyncGen
 
         async def stop_extension_services() -> None:
             record_runtime_diagnostics(
-                await stop_services(
-                    extensions,
-                    service_entries=attempted_services,
+                await await_drained(
+                    stop_services(
+                        extensions,
+                        service_entries=attempted_services,
+                    )
                 )
             )
 

backend/tests/test_gateway_extension_service_lifecycle.py

@@ -201,3 +201,71 @@ async def test_cancellation_during_service_start_propagates_after_cleanup(monkey
         "stop:first",
         "engine_close",
     ]
+
+
+@pytest.mark.asyncio
+async def test_host_cancellation_does_not_abandon_extension_service_shutdown(monkeypatch):
+    from app.gateway.deps import langgraph_runtime
+
+    events: list[str] = []
+    blocking_stop_entered = asyncio.Event()
+    allow_blocking_stop = asyncio.Event()
+
+    class _Service:
+        def __init__(self, name: str, *, block_stop: bool = False) -> None:
+            self.name = name
+            self.block_stop = block_stop
+
+        async def start(self, _deps) -> None:
+            events.append(f"start:{self.name}")
+
+        async def stop(self) -> None:
+            events.append(f"stop:{self.name}")
+            if self.block_stop:
+                blocking_stop_entered.set()
+                await allow_blocking_stop.wait()
+
+    registry = ExtensionRegistry()
+    with registry.attributed_to("first:install"):
+        registry.service(_Service("first"))
+    with registry.attributed_to("blocking:install"):
+        registry.service(_Service("blocking", block_stop=True))
+
+    _patch_runtime_resources(monkeypatch, events)
+    monkeypatch.setattr(
+        "deerflow.persistence.thread_meta.make_thread_store",
+        lambda _sf, _store: (_ for _ in ()).throw(RuntimeError("later startup failure")),
+    )
+    app = FastAPI()
+    app.state.extensions = registry.build()
+
+    async def run_runtime() -> None:
+        async with langgraph_runtime(
+            app,
+            SimpleNamespace(database=_database_config()),
+        ):
+            pytest.fail("runtime must not yield")
+
+    task = asyncio.create_task(run_runtime())
+    await asyncio.wait_for(blocking_stop_entered.wait(), timeout=1.0)
+
+    task.cancel()
+    await asyncio.sleep(0)
+    task.cancel()
+    for _ in range(5):
+        await asyncio.sleep(0)
+
+    assert not task.done(), "host cancellation abandoned extension shutdown"
+    assert "stop:first" not in events
+
+    allow_blocking_stop.set()
+    with pytest.raises(asyncio.CancelledError):
+        await task
+
+    assert events == [
+        "start:first",
+        "start:blocking",
+        "stop:blocking",
+        "stop:first",
+        "engine_close",
+    ]