From 2bdae7518d85081ae24cc1364704d4c9f61a7a9d Mon Sep 17 00:00:00 2001 From: NanPan <111261006+poijygfdyy@users.noreply.github.com> Date: Fri, 18 Sep 2026 17:19:10 +0800 Subject: [PATCH] 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 --- README.md | 8 + backend/app/gateway/AGENTS.md | 8 + backend/app/gateway/app.py | 75 +++++---- ...st_gateway_memory_shutdown_cancellation.py | 143 ++++++++++++++++++ 4 files changed, 207 insertions(+), 27 deletions(-) create mode 100644 backend/tests/test_gateway_memory_shutdown_cancellation.py diff --git a/README.md b/README.md index 91119d6b5..6a8db16d7 100644 --- a/README.md +++ b/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 diff --git a/backend/app/gateway/AGENTS.md b/backend/app/gateway/AGENTS.md index bd620c336..4ed8c8127 100644 --- a/backend/app/gateway/AGENTS.md +++ b/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 diff --git a/backend/app/gateway/app.py b/backend/app/gateway/app.py index 07c3669f8..74094e022 100644 --- a/backend/app/gateway/app.py +++ b/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") diff --git a/backend/tests/test_gateway_memory_shutdown_cancellation.py b/backend/tests/test_gateway_memory_shutdown_cancellation.py new file mode 100644 index 000000000..c9d9059ec --- /dev/null +++ b/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())