fix(memory): drain shutdown workers across cancellation (#5531)

* fix(memory): drain shutdown workers across cancellation

* fix(memory): contain shutdown config resolution failures

---------

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

README.md

@@ -1622,6 +1622,14 @@ and the [request contract](backend/docs/API.md#referencing-a-previous-conversati
 
 ### Long-Term Memory
 
+Gateway shutdown drains memory updates before closing the backend, even when
+shutdown is cancelled. Config reload failures are logged without aborting runtime
+teardown. For Kubernetes, budget `terminationGracePeriodSeconds` for all shutdown
+hooks, config/backend resolution, `memory.shutdown_flush_timeout_seconds`, and
+backend close plus a safety margin. The flush timeout does not bound `close()`:
+custom backends must make close quick or internally bounded, or shutdown can wait
+until the process is forcibly terminated.
+
 Most agents forget everything the moment a conversation ends. DeerFlow remembers.
 
 DeerMem can optionally suppress near-duplicate extracted facts with

backend/app/gateway/AGENTS.md

@@ -1,5 +1,13 @@
 ### Gateway API (`app/gateway/`)
 
+Memory shutdown resolves hot-reloaded config and the backend, flushes, then
+closes as one `await_drained` operation. Keep config resolution inside the
+best-effort error handler and off the event loop so malformed config edits do
+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.
+
 `conversation_access.py` binds an opt-in read-only tool to a run request's
 explicit `conversation_references` and effective `runs:read` permission. Never
 derive grants from message contents or checkpoints. The callback travels through

backend/app/gateway/app.py

@@ -53,6 +53,7 @@ from deerflow.config import app_config as deerflow_app_config
 from deerflow.logging_config import DEFAULT_LOG_DATE_FORMAT, DEFAULT_LOG_FORMAT, configure_logging
 from deerflow.tracing.monocle import setup_monocle_tracing_if_enabled
 from deerflow.uploads.manager import cleanup_stale_upload_staging_files
+from deerflow.utils.file_io import await_drained
 
 AppConfig = deerflow_app_config.AppConfig
 get_app_config = deerflow_app_config.get_app_config
@@ -197,6 +198,39 @@ async def _warm_memory_retrieval(manager) -> None:
         logger.warning("Memory retrieval index rebuild skipped", exc_info=True)
 
 
+async def _shutdown_memory_backend(*, retrieval_warm_finished: bool) -> None:
+    """Resolve, drain, and close memory within the caller's cancellation shield.
+
+    Backend ``close()`` overrides must be quick or internally bounded: unlike
+    ``shutdown_flush``, close has no host timeout and is drained even on cancellation.
+    """
+    manager = None
+    try:
+        app_cfg: AppConfig = await asyncio.to_thread(get_app_config)
+        if app_cfg.memory.enabled:
+            from deerflow.agents.memory import get_memory_manager
+
+            manager = await asyncio.to_thread(get_memory_manager)
+            flush_timeout = app_cfg.memory.shutdown_flush_timeout_seconds
+            completed = await asyncio.to_thread(manager.shutdown_flush, flush_timeout)
+            if completed:
+                logger.info("Memory queue flush completed within %.1fs", flush_timeout)
+            else:
+                logger.warning(
+                    "Memory queue flush did not finish within %.1fs; remaining updates may be lost",
+                    flush_timeout,
+                )
+    except Exception:
+        logger.exception("Failed to flush memory queue on shutdown")
+    finally:
+        close = getattr(manager, "close", None)
+        if callable(close) and retrieval_warm_finished:
+            try:
+                await asyncio.to_thread(close)
+            except Exception:
+                logger.exception("Failed to close memory backend on shutdown")
+
+
 async def _run_startup_trash_sweep(app: FastAPI, startup_config) -> None:
     """One trash retention sweep at gateway startup (Phase-2 spec §8.3).
 
@@ -615,9 +649,12 @@ async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]:
         #
         # K8s caveat: ``shutdown_flush_timeout_seconds`` must fit inside the
         # pod's ``terminationGracePeriodSeconds`` (channel stop + browser
-        # session close + the brief retrieval-warm wait + this drain + buffer),
+        # session close + the brief retrieval-warm wait + config/backend
+        # resolution + this drain + backend close + buffer),
         # set on the gateway Helm deployment -- or K8s SIGKILLs the drain
         # mid-flight and the loss this is fixing is silently re-introduced.
+        # Backend close has no host timeout: overrides must be quick or
+        # internally bounded, since cancellation waits for that worker too.
         # The retrieval index is derived from canonical memory files, so its
         # wait is independently capped and never consumes the flush budget.
         retrieval_warm_finished = True
@@ -634,9 +671,8 @@ async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]:
                 retrieval_warm_finished = False
                 logger.warning("Memory retrieval index rebuild is still running; leaving its connection open during shutdown")
 
-        manager = None
         try:
-            # Memory shutdown runs on a worker thread and can trigger detached
+            # Memory shutdown runs on worker threads and can trigger detached
             # system-model callbacks. Stop accepting those callbacks before
             # flushing, while keeping the registered loop alive for awaited
             # task hooks until langgraph_runtime drains runs and subagents.
@@ -646,30 +682,15 @@ async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]:
         except Exception:
             logger.debug("Failed to suspend extension system observations (non-fatal)", exc_info=True)
 
-        try:
-            app_cfg = get_app_config()
-            if app_cfg.memory.enabled:
-                from deerflow.agents.memory import get_memory_manager
-
-                manager = await asyncio.to_thread(get_memory_manager)
-                flush_timeout = app_cfg.memory.shutdown_flush_timeout_seconds
-                completed = await asyncio.to_thread(manager.shutdown_flush, flush_timeout)
-                if completed:
-                    logger.info("Memory queue flush completed within %.1fs", flush_timeout)
-                else:
-                    logger.warning(
-                        "Memory queue flush did not finish within %.1fs; remaining updates may be lost",
-                        flush_timeout,
-                    )
-        except Exception:
-            logger.exception("Failed to flush memory queue on shutdown")
-        finally:
-            close = getattr(manager, "close", None)
-            if callable(close) and retrieval_warm_finished:
-                try:
-                    await asyncio.to_thread(close)
-                except Exception:
-                    logger.exception("Failed to close memory backend on shutdown")
+        # ``asyncio.to_thread`` cancellation only detaches the awaiter; the
+        # worker keeps running. Treat resolve + flush + close as one owned
+        # shutdown operation so lifespan cancellation cannot close the backend
+        # underneath an in-flight flush or return while either worker is live.
+        await await_drained(
+            _shutdown_memory_backend(
+                retrieval_warm_finished=retrieval_warm_finished,
+            )
+        )
 
     logger.info("Shutting down API Gateway")
 

backend/tests/test_gateway_memory_shutdown_cancellation.py

@@ -0,0 +1,143 @@
+"""Cancellation ownership regression for Gateway memory shutdown."""
+
+from __future__ import annotations
+
+import asyncio
+import threading
+from contextlib import asynccontextmanager
+from types import SimpleNamespace
+from unittest.mock import AsyncMock, MagicMock, patch
+
+import pytest
+from fastapi import FastAPI
+
+
+@asynccontextmanager
+async def _noop_langgraph_runtime(_app, _startup_config):
+    yield
+
+
+@pytest.mark.parametrize("config_outcome", ["ready", "blocked", "invalid"])
+def test_lifespan_cancellation_drains_memory_flush_and_close(config_outcome: str, caplog: pytest.LogCaptureFixture) -> None:
+    """Config errors stay best-effort; cancellation cannot detach shutdown workers."""
+
+    async def scenario() -> None:
+        from app.gateway.app import lifespan
+
+        app = FastAPI()
+        startup_config = SimpleNamespace(
+            log_level="INFO",
+            memory=SimpleNamespace(
+                token_counting="char",
+                enabled=True,
+                shutdown_flush_timeout_seconds=5.0,
+            ),
+        )
+        fake_service = MagicMock()
+        fake_service.get_status.return_value = {}
+        flush_started = threading.Event()
+        flush_finished = threading.Event()
+        allow_flush = threading.Event()
+        close_started = threading.Event()
+        close_finished = threading.Event()
+        allow_close = threading.Event()
+        config_started = threading.Event()
+        allow_config = threading.Event()
+        loop_thread = threading.get_ident()
+        manager = MagicMock()
+        manager.warm_retrieval = None
+        manager.warm.return_value = True
+
+        def blocking_flush(_timeout: float) -> bool:
+            flush_started.set()
+            assert allow_flush.wait(5.0)
+            flush_finished.set()
+            return True
+
+        def blocking_close() -> None:
+            close_started.set()
+            assert flush_finished.is_set(), "memory close raced the still-running shutdown flush"
+            assert allow_close.wait(5.0)
+            close_finished.set()
+
+        manager.shutdown_flush.side_effect = blocking_flush
+        manager.close.side_effect = blocking_close
+
+        async def fake_start(_startup_config, **_kwargs):
+            return fake_service
+
+        def shutdown_config():
+            if config_outcome == "invalid":
+                raise ValueError("invalid shutdown config")
+            assert threading.get_ident() != loop_thread, "shutdown config resolution blocked the event loop"
+            config_started.set()
+            if config_outcome == "blocked":
+                assert allow_config.wait(5.0)
+            return startup_config
+
+        with (
+            patch("app.gateway.app.get_app_config", return_value=startup_config) as get_config,
+            patch("app.gateway.app.get_gateway_config", return_value=MagicMock(host="x", port=0)),
+            patch("app.gateway.app.langgraph_runtime", _noop_langgraph_runtime),
+            patch("deerflow.skills.projection.ensure_public_skill_projection"),
+            patch("app.gateway.app.auth.close_oidc_service", AsyncMock()),
+            patch("app.channels.service.start_channel_service", side_effect=fake_start),
+            patch("app.channels.service.stop_channel_service", AsyncMock()),
+            patch("deerflow.agents.memory.get_memory_manager", return_value=manager),
+            patch("deerflow.extensions.notify.suspend_extension_system_observations"),
+        ):
+            context = lifespan(app)
+            await context.__aenter__()
+            get_config.side_effect = shutdown_config
+            shutdown_task = asyncio.create_task(context.__aexit__(None, None, None))
+            try:
+                if config_outcome == "invalid":
+                    await shutdown_task
+                    assert "Failed to flush memory queue on shutdown" in caplog.text
+                    assert "invalid shutdown config" in caplog.text
+                    manager.shutdown_flush.assert_not_called()
+                    manager.close.assert_not_called()
+                    return
+
+                assert await asyncio.to_thread(config_started.wait, 1.0)
+                if config_outcome == "blocked":
+                    shutdown_task.cancel("first shutdown cancellation")
+                    await asyncio.sleep(0)
+                    shutdown_task.cancel("cancellation during config resolution")
+                    for _ in range(10):
+                        await asyncio.sleep(0)
+                    assert not shutdown_task.done(), "Gateway shutdown abandoned config resolution"
+                    assert not flush_started.is_set()
+                    assert not close_started.is_set()
+                    allow_config.set()
+
+                assert await asyncio.to_thread(flush_started.wait, 1.0)
+                shutdown_task.cancel("first shutdown cancellation")
+                await asyncio.sleep(0)
+                shutdown_task.cancel("second shutdown cancellation")
+                for _ in range(10):
+                    await asyncio.sleep(0)
+
+                assert not close_started.is_set()
+                assert not shutdown_task.done(), "Gateway shutdown abandoned the in-flight memory flush"
+
+                allow_flush.set()
+                assert await asyncio.to_thread(close_started.wait, 1.0)
+                assert flush_finished.is_set()
+                shutdown_task.cancel("third shutdown cancellation")
+                for _ in range(10):
+                    await asyncio.sleep(0)
+                assert not shutdown_task.done(), "Gateway shutdown abandoned the in-flight memory close"
+
+                allow_close.set()
+                with pytest.raises(asyncio.CancelledError) as exc_info:
+                    await shutdown_task
+                assert exc_info.value.args == ("first shutdown cancellation",)
+                assert close_finished.is_set()
+            finally:
+                allow_config.set()
+                allow_flush.set()
+                allow_close.set()
+                await asyncio.gather(shutdown_task, return_exceptions=True)
+
+    asyncio.run(scenario())