From 959bf13406eadfc22d10086bcc18fbcb4751aeb1 Mon Sep 17 00:00:00 2001 From: Tu Naichao <102199610+TuNaiChao@users.noreply.github.com> Date: Wed, 15 Jul 2026 20:25:41 +0800 Subject: [PATCH] fix(memory): flush memory queue on graceful shutdown to prevent loss (#4181) * fix(memory): bounded shutdown flush via MemoryManager.shutdown_flush Re-applies the memory-queue shutdown drain on top of the pluggable MemoryManager abstraction (#4122): the old top-level MemoryUpdateQueue singleton is gone, so the drain is now a backend contract instead of host code reaching into the queue. - MemoryManager ABC: shutdown_flush(timeout) -> bool. Every backend implements a bounded graceful-shutdown drain. - DeerMem: queue.flush_sync (daemon-thread + Event.wait hard timeout for the uninterruptible sync LLM call; joins an in-flight worker first so contexts a debounce Timer already pulled out are not lost on exit; skips inter-item sleep on the drain path; per-item succeeded/failed count), exposed via shutdown_flush. - noop: shutdown_flush is a clean no-op success. - Gateway lifespan: call get_memory_manager().shutdown_flush(timeout) after channels/scheduler stop, via asyncio.to_thread, try/except bounded. No host-level pending/processing guard -- the backend short-circuits on an idle buffer, so the host cannot "forget" the in-flight case (structurally eliminates the guard race flagged on the prior revision). - shutdown_flush_timeout_seconds added to the shared MemoryConfig (host-owned lifecycle budget, default 30, 1-300) + exposed on MemoryConfigResponse and the embedded client; config_version 25 -> 26. Tests: queue flush_sync (7), lifespan drain incl. False-branch caplog assertion + disabled gate (3), ABC contract noop/deermem (3). * fix(chart): gateway grace period so memory drain is not SIGKILLed K8s defaults terminationGracePeriodSeconds to 30s, shorter than the Gateway's graceful-shutdown work (channel stop ~5s + memory queue drain default 30s). Without an explicit grace period, K8s SIGKILLs the memory drain mid-flight and silently re-introduces the loss shutdown_flush is fixing (flagged on the prior revision). - gateway pod: terminationGracePeriodSeconds (default 45, configurable). - gateway container: preStop sleep (default 5, 0 disables) so the Service/ingress deregisters the pod before SIGTERM begins the drain. - values.yaml + README: both configurable; README documents that the grace period must track memory.shutdown_flush_timeout_seconds. * docs(memory): document shutdown_flush_timeout_seconds + lifespan drain Add the host-shared field to the memory config list and Config Schema summary in backend/AGENTS.md, noting the lifespan drain and the K8s grace-period relationship. * fix(chart): bump embedded config_version to 26 The chart's embedded `config:` block (values.yaml + README example) still had config_version: 25 after commit f3ca8e9f raised config.example.yaml to 26, failing the validate-chart config_version drift check. Bump both to 26. --- backend/AGENTS.md | 3 +- backend/app/gateway/app.py | 37 +++++ backend/app/gateway/routers/memory.py | 4 + .../memory/backends/deermem/deer_mem.py | 13 ++ .../backends/deermem/deermem/core/queue.py | 117 ++++++++++++- .../memory/backends/noop/noop_manager.py | 5 + .../harness/deerflow/agents/memory/manager.py | 28 ++++ backend/packages/harness/deerflow/client.py | 1 + .../harness/deerflow/config/memory_config.py | 22 ++- .../tests/test_gateway_lifespan_shutdown.py | 81 ++++++++- .../tests/test_memory_manager_pluggable.py | 45 +++++ backend/tests/test_memory_queue.py | 156 ++++++++++++++++++ config.example.yaml | 9 +- deploy/helm/deer-flow/README.md | 3 +- .../templates/gateway-deployment.yaml | 14 ++ deploy/helm/deer-flow/values.yaml | 14 +- 16 files changed, 538 insertions(+), 14 deletions(-) diff --git a/backend/AGENTS.md b/backend/AGENTS.md index 7250dd567..f4a13158b 100644 --- a/backend/AGENTS.md +++ b/backend/AGENTS.md @@ -596,6 +596,7 @@ Focused regression coverage for the updater lives in `backend/tests/test_memory_ - `mode` - Operation mode: `middleware` (default passive background extraction) or `tool` (experimental model-driven memory tools). Modes are mutually exclusive. - `storage_path` - Path to memory.json (absolute path opts out of per-user isolation) - `debounce_seconds` - Wait time before processing (default: 30) +- `shutdown_flush_timeout_seconds` - Host-shared hard budget (seconds) to drain the memory backend's pending-update buffer on Gateway graceful shutdown (default: 30; 1–300). Each pending item does one LLM call, so large IM batches may need more. The Gateway lifespan calls `MemoryManager.shutdown_flush(timeout)` after channels/scheduler stop; the backend short-circuits on an idle buffer, so the host calls it unconditionally (no pending/processing gate). Must fit inside the pod's K8s `terminationGracePeriodSeconds` (gateway Helm chart sets this; default 45s) or K8s SIGKILLs the drain mid-flight. - `model_name` - LLM for updates (null = default model) - `max_facts` / `fact_confidence_threshold` - Fact storage limits (100 / 0.7) - `max_injection_tokens` - Token limit for prompt injection (2000) @@ -732,7 +733,7 @@ Config is env-driven like the others — `MonocleTracingConfig`, built in `get_t - `title` - Auto-title generation (enabled, max_words, max_chars, model_name; null model_name uses fast local fallback, explicit model_name uses the prompt_template LLM path) - `summarization` - Context summarization (enabled, trigger conditions, keep policy) - `subagents.enabled` - Master switch for subagent delegation -- `memory` - Memory system (enabled, storage_path, debounce_seconds, model_name, max_facts, fact_confidence_threshold, injection_enabled, max_injection_tokens, staleness_review_enabled, staleness_age_days, staleness_min_candidates, staleness_max_removals_per_cycle, staleness_protected_categories) +- `memory` - Memory system (enabled, storage_path, debounce_seconds, shutdown_flush_timeout_seconds, model_name, max_facts, fact_confidence_threshold, injection_enabled, max_injection_tokens, staleness_review_enabled, staleness_age_days, staleness_min_candidates, staleness_max_removals_per_cycle, staleness_protected_categories) **`extensions_config.json`**: - `mcpServers` - Map of server name → config (enabled, type, command, args, env, url, headers, oauth, description, `routing`, `tools`, `tool_call_timeout`). `routing.mode="prefer"` emits `` prompt guidance; if `tool_search` defers the hinted tool, `McpRoutingMiddleware` can also auto-promote matching deferred schemas before the model call. It does not hard-disable other tools. diff --git a/backend/app/gateway/app.py b/backend/app/gateway/app.py index ff93bf8e7..65b7a35e3 100644 --- a/backend/app/gateway/app.py +++ b/backend/app/gateway/app.py @@ -301,6 +301,43 @@ async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]: except Exception: logger.exception("Failed to stop scheduled task service") + # Drain the memory backend's pending-update buffer before the worker + # exits (best-effort, bounded). IM channels and the scheduler are + # already stopped above, so no new IM/scheduler updates arrive during + # the drain; the LangGraph runtime / in-flight HTTP requests can still + # complete memory enqueues in a narrow window, but anything added after + # the drain copies the buffer only resets the debounce Timer + # (best-effort, same as today). + # + # No host-level pending/processing guard: ``shutdown_flush`` + # short-circuits on a truly idle buffer (returns True immediately), so + # calling it unconditionally is cheap and keeps the in-flight-worker + # race entirely inside the backend (where the buffer lives) -- the host + # cannot "forget" that case the way a ``pending_count > 0``-only guard + # would (review #6 on the original PR). + # + # K8s caveat: ``shutdown_flush_timeout_seconds`` must fit inside the + # pod's ``terminationGracePeriodSeconds`` (channel stop + this drain + + # buffer), set on the gateway Helm deployment -- or K8s SIGKILLs the + # drain mid-flight and the loss this is fixing is silently re-introduced. + try: + app_cfg = get_app_config() + if app_cfg.memory.enabled: + from deerflow.agents.memory import get_memory_manager + + manager = 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") + logger.info("Shutting down API Gateway") diff --git a/backend/app/gateway/routers/memory.py b/backend/app/gateway/routers/memory.py index c5d32640a..8c701005e 100644 --- a/backend/app/gateway/routers/memory.py +++ b/backend/app/gateway/routers/memory.py @@ -130,6 +130,7 @@ class MemoryConfigResponse(BaseModel): enabled: bool = Field(..., description="Whether the memory mechanism is enabled (call-site gate).") mode: Literal["middleware", "tool"] = Field(..., description="Memory operation mode: 'middleware' (passive per-turn LLM summarization) or 'tool' (model calls memory tools directly). Mechanism-level, applies to any backend.") injection_enabled: bool = Field(..., description="Whether memory is injected into the system prompt (call-site gate).") + shutdown_flush_timeout_seconds: float = Field(..., description="Hard budget (s) to drain pending memory updates on Gateway graceful shutdown; must fit inside the pod's K8s terminationGracePeriodSeconds.") manager_class: str = Field(..., description="Active memory backend selector (backend name or dotted path).") backend_config: dict = Field(..., description="Backend-private config (self-interpreted by the backend).") @@ -362,6 +363,7 @@ async def get_memory_config_endpoint() -> MemoryConfigResponse: { "enabled": true, "injection_enabled": true, + "shutdown_flush_timeout_seconds": 30.0, "mode": "middleware", "manager_class": "deermem", "backend_config": { @@ -380,6 +382,7 @@ async def get_memory_config_endpoint() -> MemoryConfigResponse: enabled=config.enabled, mode=config.mode, injection_enabled=config.injection_enabled, + shutdown_flush_timeout_seconds=config.shutdown_flush_timeout_seconds, manager_class=config.manager_class, backend_config=config.backend_config, ) @@ -406,6 +409,7 @@ async def get_memory_status(http_request: Request) -> MemoryStatusResponse: enabled=config.enabled, mode=config.mode, injection_enabled=config.injection_enabled, + shutdown_flush_timeout_seconds=config.shutdown_flush_timeout_seconds, manager_class=config.manager_class, backend_config=config.backend_config, ), diff --git a/backend/packages/harness/deerflow/agents/memory/backends/deermem/deer_mem.py b/backend/packages/harness/deerflow/agents/memory/backends/deermem/deer_mem.py index 758ece764..8ecacbc0f 100644 --- a/backend/packages/harness/deerflow/agents/memory/backends/deermem/deer_mem.py +++ b/backend/packages/harness/deerflow/agents/memory/backends/deermem/deer_mem.py @@ -237,6 +237,19 @@ class DeerMem(MemoryManager): """Not implemented this phase (no distinct export yet; /export routes via get_memory).""" raise NotImplementedError("DeerMem.export_memory is not implemented yet") + # ── Lifecycle ─────────────────────────────────────────────────────── + def shutdown_flush(self, timeout: float) -> bool: + """Drain the debounce queue within ``timeout`` on graceful shutdown. + + Delegates to the queue's bounded synchronous flush, which joins an + in-flight worker first (so contexts a debounce Timer already pulled out + of the queue are not lost on exit) and otherwise drains the queue on a + daemon thread with a real hard timeout (the memory-update LLM call is + synchronous and cannot be interrupted). Returns ``True`` only when the + drain genuinely finished within ``timeout``. + """ + return self._queue.flush_sync(timeout) + # ── DeerMem-internal (NOT on the ABC; reached via hasattr probing) ─── def warm(self) -> bool: """Pre-warm DeerMem-specific resources (the tiktoken encoding cache). diff --git a/backend/packages/harness/deerflow/agents/memory/backends/deermem/deermem/core/queue.py b/backend/packages/harness/deerflow/agents/memory/backends/deermem/deermem/core/queue.py index b35cb59a8..14a8f0f1c 100644 --- a/backend/packages/harness/deerflow/agents/memory/backends/deermem/deermem/core/queue.py +++ b/backend/packages/harness/deerflow/agents/memory/backends/deermem/deermem/core/queue.py @@ -47,6 +47,11 @@ class MemoryUpdateQueue: self._lock = threading.Lock() self._timer: threading.Timer | None = None self._processing = False + # Thread currently running ``_process_queue`` (None when idle). ``flush_sync`` + # joins an in-flight worker instead of reporting a false-positive "completed" + # while contexts it already pulled out of the queue are still being processed + # (and would be lost on exit). See ``flush_sync`` step (1). + self._processing_thread: threading.Thread | None = None self._reprocess_pending = False @staticmethod @@ -172,8 +177,15 @@ class MemoryUpdateQueue: self._timer.daemon = True self._timer.start() - def _process_queue(self) -> None: - """Process all queued conversation contexts.""" + def _process_queue(self, *, skip_inter_item_delay: bool = False) -> None: + """Process all queued conversation contexts. + + Args: + skip_inter_item_delay: When set, skip the inter-item rate-limit + ``time.sleep``. Intended for the shutdown-drain path + (:meth:`flush_sync`), which races a bounded timeout and should + not waste budget sleeping between items. + """ with self._lock: if self._processing: # Another worker is already draining the queue. Instead of @@ -188,12 +200,15 @@ class MemoryUpdateQueue: return self._processing = True + self._processing_thread = threading.current_thread() contexts_to_process = self._queue.copy() self._queue.clear() self._timer = None logger.info("Processing %d queued memory updates", len(contexts_to_process)) + succeeded = 0 + failed = 0 try: for context in contexts_to_process: try: @@ -208,35 +223,122 @@ class MemoryUpdateQueue: trace_id=context.trace_id, ) if success: + succeeded += 1 logger.info("Memory updated successfully for thread %s (trace_id=%s)", context.thread_id, context.trace_id) else: + failed += 1 logger.warning("Memory update skipped/failed for thread %s (trace_id=%s)", context.thread_id, context.trace_id) except Exception as e: + failed += 1 logger.error("Error updating memory for thread %s (trace_id=%s): %s", context.thread_id, context.trace_id, e) - # Small delay between updates to avoid rate limiting - if len(contexts_to_process) > 1: + # Small delay between updates to avoid rate limiting. + # Skipped on the shutdown-drain path, which races a bounded + # timeout and should spend that budget on LLM calls, not on + # sleeping between items. + if not skip_inter_item_delay and len(contexts_to_process) > 1: time.sleep(0.5) - finally: + # Summary count disambiguates "drained" (queue emptied) from "saved" + # (every extraction persisted): per-item ``update_memory`` failures are + # swallowed above, so without this an operator debugging missing + # memories would see only the happy-path "Processing N" line. + if succeeded or failed: + logger.info("Memory update batch done: %d succeeded, %d failed", succeeded, failed) with self._lock: self._processing = False + self._processing_thread = None if self._reprocess_pending: self._reprocess_pending = False if self._queue: self._schedule_timer(0) - def flush(self) -> None: + def flush(self, *, skip_inter_item_delay: bool = False) -> None: """Force immediate processing of the queue. This is useful for testing or graceful shutdown. + + Args: + skip_inter_item_delay: Forwarded to :meth:`_process_queue`; skip the + inter-item rate-limit sleep. Intended for the shutdown-drain + path (:meth:`flush_sync`). """ with self._lock: if self._timer is not None: self._timer.cancel() self._timer = None - self._process_queue() + self._process_queue(skip_inter_item_delay=skip_inter_item_delay) + + def flush_sync(self, timeout: float) -> bool: + """Best-effort synchronous flush bounded by ``timeout`` seconds. + + Unlike :meth:`flush_nowait` (which only schedules a daemon timer that + is killed on process exit), this runs :meth:`flush` on a daemon thread + and waits up to ``timeout`` seconds for it to finish. Intended for + graceful shutdown: without it, any updates enqueued since the last + timer fire are lost on restart / rolling deploy / SIGTERM, because the + queue is pure in-memory and the debounce Timer is a daemon thread. + + The drain accounts for two races a naive ``flush()`` would miss: + + - **In-flight worker.** If the debounce Timer already fired, an + ``_process_queue`` worker is mid-LLM-call holding contexts it already + pulled out of the queue (``_processing=True``, queue empty). ``flush`` + alone would see ``_processing=True``, no-op, and report success while + that worker is still running and likely killed on exit. So we join + the in-flight worker first (bounded by the remaining budget). + - **Failed flush.** ``flush`` makes a synchronous LLM call that can + raise; success is tracked on the happy path only, so the return value + matches the docstring's "completed". + + Note: steps (1) and (3) share the same ``deadline`` budget. A slow + in-flight worker can consume most/all of it, leaving step (3) to no-op; + ``timeout`` must therefore cover both a slow in-flight worker *and* the + remaining queue (best-effort: any tail not drained in budget is dropped, + same failure direction as no flush, scoped to the tail). + + Returns ``True`` only if the drain genuinely finished (queue empty, no + worker still running, flush did not raise) within ``timeout``. + """ + deadline = time.monotonic() + timeout + + # (1) Wait for an in-flight _process_queue first (bounded). Otherwise + # flush() would see _processing=True, no-op, and we would report + # success while that worker is still mid-LLM-call on a daemon thread + # that exit will kill — losing the contexts it already pulled out. + with self._lock: + in_flight = self._processing_thread + if in_flight is not None: + in_flight.join(timeout=max(0.0, deadline - time.monotonic())) + + # (2) Genuine idle: nothing pending and no worker still running. + if self.pending_count == 0 and not self.is_processing: + return True + + # (3) Drain the queue on a daemon thread so the timeout is a real hard + # stop: flush() makes a synchronous LLM call that cannot be + # interrupted, so we wait on Event.wait, not on Thread.join. + success = False + done = threading.Event() + + def _run() -> None: + nonlocal success + try: + self.flush(skip_inter_item_delay=True) + success = True + except Exception: + logger.exception("Memory queue flush failed during shutdown drain") + finally: + done.set() + + worker = threading.Thread(target=_run, name="memory-shutdown-flush", daemon=True) + worker.start() + finished = done.wait(timeout=max(0.0, deadline - time.monotonic())) + if not finished: + return False + # flush() returned; only report success if no worker raced back in. + return bool(success) and not self.is_processing def flush_nowait(self) -> None: """Start queue processing immediately in a background thread.""" @@ -256,6 +358,7 @@ class MemoryUpdateQueue: self._timer = None self._queue.clear() self._processing = False + self._processing_thread = None self._reprocess_pending = False @property diff --git a/backend/packages/harness/deerflow/agents/memory/backends/noop/noop_manager.py b/backend/packages/harness/deerflow/agents/memory/backends/noop/noop_manager.py index 8e06c856c..028dd97d9 100644 --- a/backend/packages/harness/deerflow/agents/memory/backends/noop/noop_manager.py +++ b/backend/packages/harness/deerflow/agents/memory/backends/noop/noop_manager.py @@ -157,6 +157,11 @@ class NoopMemoryManager(MemoryManager): ) -> dict[str, Any]: return _empty_memory() + # ── Lifecycle ─────────────────────────────────────────────────────── + def shutdown_flush(self, timeout: float) -> bool: + """Nothing is ever queued, so shutdown drain is a clean no-op success.""" + return True + # ── Optional DeerMem-internal capabilities (NOT on the ABC) ────────── # The host gateway discovers these via ``hasattr(manager, "")`` and # returns 501 when absent. Implement the ones your backend supports so the diff --git a/backend/packages/harness/deerflow/agents/memory/manager.py b/backend/packages/harness/deerflow/agents/memory/manager.py index 95352fa0b..85e8036a9 100644 --- a/backend/packages/harness/deerflow/agents/memory/manager.py +++ b/backend/packages/harness/deerflow/agents/memory/manager.py @@ -188,6 +188,34 @@ class MemoryManager(ABC): ) -> dict[str, Any]: """Export the memory document for the bucket. *stub* this phase (no caller yet).""" + # ── Lifecycle ─────────────────────────────────────────────────────── + @abstractmethod + def shutdown_flush(self, timeout: float) -> bool: + """Best-effort bounded drain of pending updates on graceful shutdown. + + Runs on the Gateway shutdown path (after IM channels and the scheduler + stop, so no new IM/scheduler updates arrive during the drain) to flush + updates still sitting in the backend's debounce buffer. Without it, any + update enqueued since the last timer fire is lost on restart / rolling + deploy / SIGTERM, because the buffer is pure in-memory and the debounce + worker is a daemon thread killed on process exit. + + Implementations must honour a *hard* ``timeout``: the drain makes a + synchronous LLM call that cannot be interrupted, so the caller (the + Gateway lifespan) needs a real upper bound that lines up with the K8s + ``terminationGracePeriodSeconds`` (the drain must finish inside the pod + grace window, or K8s SIGKILLs it mid-drain and the loss the drain is + fixing is silently re-introduced). + + Returns ``True`` if the drain genuinely finished within ``timeout`` + (buffer empty, no worker still running, no exception); ``False`` on + timeout or failure (the caller logs a warning and proceeds to exit -- + any unfinished tail is dropped, strictly better than no flush). A + backend with no pending work (or no buffer at all) returns ``True`` + immediately, so the host may call this unconditionally when memory is + enabled without gating on backend-private queue state. + """ + # ── Backend discovery (drop-in) ─────────────────────────────────────────── def _scan_backends() -> dict[str, type[MemoryManager]]: diff --git a/backend/packages/harness/deerflow/client.py b/backend/packages/harness/deerflow/client.py index 25e91daae..c8046c0c7 100644 --- a/backend/packages/harness/deerflow/client.py +++ b/backend/packages/harness/deerflow/client.py @@ -1342,6 +1342,7 @@ class DeerFlowClient: "enabled": config.enabled, "mode": config.mode, "injection_enabled": config.injection_enabled, + "shutdown_flush_timeout_seconds": config.shutdown_flush_timeout_seconds, "manager_class": config.manager_class, "backend_config": config.backend_config, } diff --git a/backend/packages/harness/deerflow/config/memory_config.py b/backend/packages/harness/deerflow/config/memory_config.py index af1b5bb07..80fb4e4bf 100644 --- a/backend/packages/harness/deerflow/config/memory_config.py +++ b/backend/packages/harness/deerflow/config/memory_config.py @@ -4,7 +4,8 @@ DeerMem-private fields live in ``backends/deermem/config.py`` (``DeerMemConfig`` reached via ``backend_config`` (a dict the factory passes to the backend's ``__init__``). This module holds ONLY the host-shared fields every backend / call site / factory reads: ``enabled`` / ``injection_enabled`` / -``manager_class`` / ``backend_config``. Keeping the shared schema slim is what +``shutdown_flush_timeout_seconds`` / ``manager_class`` / ``backend_config``. +Keeping the shared schema slim is what makes backends swappable and portable (DeerMem's knobs do not leak onto the shared contract). """ @@ -17,7 +18,7 @@ from pydantic import BaseModel, Field logger = logging.getLogger(__name__) # Host-shared MemoryConfig fields (read by every backend / call site / factory). -_SHARED_FIELDS = frozenset({"enabled", "mode", "injection_enabled", "manager_class", "backend_config"}) +_SHARED_FIELDS = frozenset({"enabled", "mode", "injection_enabled", "shutdown_flush_timeout_seconds", "manager_class", "backend_config"}) # DeerMem-private fields that used to live at the top level of `memory:` in # config.yaml (pre-abstraction). On load they are auto-migrated into @@ -66,6 +67,23 @@ class MemoryConfig(BaseModel): default=True, description="Whether to inject memory into the system prompt (call-site gate).", ) + shutdown_flush_timeout_seconds: float = Field( + default=30.0, + ge=1.0, + le=300.0, + description=( + "Hard time budget (seconds) for draining the memory backend's " + "pending-update buffer during Gateway graceful shutdown. The drain " + "makes one LLM call per pending item, so large IM batches may need " + "a higher value. Must fit inside the pod's K8s " + "terminationGracePeriodSeconds (together with channel/scheduler " + "stop) or K8s SIGKILLs the drain mid-flight. The drain runs on a " + "daemon thread, so on timeout the process proceeds to exit and any " + "unfinished tail is dropped (same failure direction as no flush, " + "scoped to the tail). Host-shared (not backend-private): the host " + "owns the lifespan budget and the K8s grace relationship." + ), + ) manager_class: str = Field( default="deermem", description=( diff --git a/backend/tests/test_gateway_lifespan_shutdown.py b/backend/tests/test_gateway_lifespan_shutdown.py index 79d328155..070542cd1 100644 --- a/backend/tests/test_gateway_lifespan_shutdown.py +++ b/backend/tests/test_gateway_lifespan_shutdown.py @@ -10,6 +10,7 @@ signal-reentrancy deadlock described in from __future__ import annotations import asyncio +import logging from contextlib import asynccontextmanager from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock, patch @@ -35,6 +36,9 @@ async def _run_lifespan_with_hanging_stop() -> float: app = FastAPI() startup_config = MagicMock() startup_config.log_level = "INFO" + # Keep this test focused on the channel-hang timing: skip the memory drain. + startup_config.memory.enabled = False + startup_config.memory.shutdown_flush_timeout_seconds = 5.0 fake_service = MagicMock() fake_service.get_status = MagicMock(return_value={}) @@ -79,7 +83,7 @@ async def _run_lifespan_with_upload_staging_cleanup(): from app.gateway.app import lifespan app = FastAPI() - startup_config = SimpleNamespace(log_level="INFO", memory=SimpleNamespace(token_counting="char")) + startup_config = SimpleNamespace(log_level="INFO", memory=SimpleNamespace(token_counting="char", enabled=False, shutdown_flush_timeout_seconds=30.0)) fake_service = MagicMock() fake_service.get_status = MagicMock(return_value={}) cleanup_upload_staging_files = MagicMock(return_value=2) @@ -110,3 +114,78 @@ def test_lifespan_sweeps_upload_staging_files_on_startup(): cleanup_upload_staging_files.assert_called_once_with() close_oidc_service.assert_awaited_once() stop_channel_service.assert_awaited_once() + + +async def _run_lifespan_with_memory_flush(*, enabled: bool, flush_return: bool) -> MagicMock: + """Drive lifespan with a spied memory manager.shutdown_flush. + + Returns the manager mock so the caller can assert the shutdown flush was + reached (and with what timeout). The host calls ``shutdown_flush`` + unconditionally when memory is enabled -- there is no host-level + ``pending_count/is_processing`` gate, because the backend short-circuits on + an idle buffer and keeping the in-flight race inside the backend means the + host cannot "forget" it (review #6 on the original PR). + """ + from app.gateway.app import lifespan + + app = FastAPI() + startup_config = SimpleNamespace( + log_level="INFO", + memory=SimpleNamespace( + token_counting="char", + enabled=enabled, + shutdown_flush_timeout_seconds=5.0, + ), + ) + fake_service = MagicMock() + fake_service.get_status = MagicMock(return_value={}) + close_oidc_service = AsyncMock() + stop_channel_service = AsyncMock() + + async def fake_start(_startup_config): + return fake_service + + manager = MagicMock() + manager.shutdown_flush.return_value = flush_return + + with ( + patch("app.gateway.app.get_app_config", return_value=startup_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("app.gateway.app.auth.close_oidc_service", close_oidc_service), + patch("app.channels.service.start_channel_service", side_effect=fake_start), + patch("app.channels.service.stop_channel_service", stop_channel_service), + patch("deerflow.agents.memory.get_memory_manager", return_value=manager), + ): + async with lifespan(app): + pass + + return manager + + +def test_lifespan_drains_memory_on_shutdown_with_configured_timeout(caplog) -> None: + """When memory is enabled, shutdown calls manager.shutdown_flush with the + configured timeout (asserts the timeout is forwarded, review #3) and logs + 'completed' at INFO when the drain finishes.""" + caplog.set_level(logging.INFO, logger="app.gateway.app") + manager = asyncio.run(_run_lifespan_with_memory_flush(enabled=True, flush_return=True)) + manager.shutdown_flush.assert_called_once_with(5.0) + assert any(r.levelno == logging.INFO and "flush completed" in r.message for r in caplog.records) + + +def test_lifespan_warns_when_memory_flush_does_not_finish(caplog) -> None: + """A False return (timeout/failure) is the path operators actually see when + K8s SIGKILLs the drain; the host must log a WARNING (not 'completed'), so + the loss risk is visible (review #3 False-branch coverage; review #2/#4 + failed-flush semantics).""" + caplog.set_level(logging.WARNING, logger="app.gateway.app") + manager = asyncio.run(_run_lifespan_with_memory_flush(enabled=True, flush_return=False)) + manager.shutdown_flush.assert_called_once_with(5.0) + assert any(r.levelno == logging.WARNING and "did not finish" in r.message for r in caplog.records) + assert not any("flush completed" in r.message for r in caplog.records) + + +def test_lifespan_skips_memory_flush_when_disabled() -> None: + """memory.enabled=False skips the drain entirely.""" + manager = asyncio.run(_run_lifespan_with_memory_flush(enabled=False, flush_return=True)) + manager.shutdown_flush.assert_not_called() diff --git a/backend/tests/test_memory_manager_pluggable.py b/backend/tests/test_memory_manager_pluggable.py index 501f50b87..8431536e1 100644 --- a/backend/tests/test_memory_manager_pluggable.py +++ b/backend/tests/test_memory_manager_pluggable.py @@ -15,6 +15,7 @@ are independent of order. from __future__ import annotations from pathlib import Path +from unittest.mock import MagicMock, patch import pytest @@ -146,3 +147,47 @@ def test_empty_storage_path_factory_injects_runtime_home(tmp_path, monkeypatch) # per-user dir created under the injected runtime_home root user_dirs = [p.name for p in (tmp_path / "users").iterdir() if p.is_dir()] assert len(user_dirs) == 1 + + +def test_shutdown_flush_is_on_abc_and_noop_is_noop_success() -> None: + """``shutdown_flush`` is part of the MemoryManager ABC (every backend must + implement a bounded graceful-shutdown drain); noop has no buffer, so it + returns True immediately.""" + assert "shutdown_flush" in MemoryManager.__abstractmethods__ + reset_memory_manager() + set_memory_config(MemoryConfig(manager_class="noop")) + noop = get_memory_manager() + assert noop.shutdown_flush(1.0) is True + reset_memory_manager() + + +def test_deermem_shutdown_flush_delegates_to_queue_flush_sync() -> None: + """DeerMem.shutdown_flush delegates to its queue's bounded flush_sync, + forwarding the host-owned timeout budget unchanged.""" + reset_memory_manager() + set_memory_config(MemoryConfig(manager_class="deermem")) + deermem = get_memory_manager() + with patch.object(deermem._queue, "flush_sync", return_value=True) as spy: + assert deermem.shutdown_flush(7.0) is True + spy.assert_called_once_with(7.0) + reset_memory_manager() + + +def test_deermem_shutdown_flush_drains_a_pending_update() -> None: + """End-to-end: a pending update sitting in DeerMem's queue is drained + within the timeout on shutdown_flush (the loss-on-exit bug the ABC method + fixes). The drain skips inter-item sleep so the budget goes to LLM calls.""" + from deerflow.agents.memory.backends.deermem.deermem.core.queue import ConversationContext + + reset_memory_manager() + set_memory_config(MemoryConfig(manager_class="deermem")) + deermem = get_memory_manager() + # Inject a mock updater so no real LLM call is made; both items "succeed". + mock_updater = MagicMock() + mock_updater.update_memory.return_value = True + deermem._queue._updater = mock_updater + deermem._queue._queue = [ConversationContext(thread_id=f"t{i}", messages=["m"], agent_name="lead_agent") for i in range(3)] + assert deermem.shutdown_flush(5.0) is True + assert deermem._queue.pending_count == 0 + assert mock_updater.update_memory.call_count == 3 + reset_memory_manager() diff --git a/backend/tests/test_memory_queue.py b/backend/tests/test_memory_queue.py index d86103c5d..13d3c5e34 100644 --- a/backend/tests/test_memory_queue.py +++ b/backend/tests/test_memory_queue.py @@ -247,3 +247,159 @@ def test_process_queue_forwards_trace_id_to_updater() -> None: user_id=None, trace_id="trace-memory-1", ) + + +# --------------------------------------------------------------------------- +# shutdown_flush / flush_sync (graceful-shutdown drain) — review carry-overs. +# The queue is a daemon-timer + in-memory buffer, so anything pending at +# process exit is lost. flush_sync drains it within a hard timeout, joining an +# in-flight worker first so contexts a debounce Timer already pulled out of the +# queue are not lost either. +# --------------------------------------------------------------------------- + +_QUEUE_MODULE = "deerflow.agents.memory.backends.deermem.deermem.core.queue" + + +def test_flush_sync_noop_on_empty_queue() -> None: + """flush_sync short-circuits and returns True when there is nothing to drain.""" + queue = _queue() + assert queue.pending_count == 0 + assert queue.flush_sync(timeout=5.0) is True + + +def test_flush_sync_drains_pending_queue_and_returns_true() -> None: + """flush_sync runs the synchronous flush() and waits for it to finish.""" + mock_updater = MagicMock() + mock_updater.update_memory.return_value = True + queue = _queue(mock_updater) + queue._queue = [ConversationContext(thread_id="thread-1", messages=["conversation"], agent_name="lead_agent")] + + with ( + patch(_QUEUE_MODULE + ".MemoryUpdater", create=True), + patch(_QUEUE_MODULE + ".time.sleep"), + ): + completed = queue.flush_sync(timeout=5.0) + + assert completed is True + assert queue.pending_count == 0 + mock_updater.update_memory.assert_called_once_with( + messages=["conversation"], + thread_id="thread-1", + agent_name="lead_agent", + correction_detected=False, + reinforcement_detected=False, + user_id=None, + trace_id=None, + ) + + +def test_flush_sync_returns_false_when_flush_exceeds_timeout() -> None: + """flush_sync does not block past ``timeout``; a slow flush returns False.""" + queue = _queue() + queue._queue = [ConversationContext(thread_id="thread-1", messages=["conversation"], agent_name="lead_agent")] + release = threading.Event() + + def _slow_flush() -> None: + # Block until the test releases us (well past the flush_sync timeout). + release.wait(timeout=5.0) + + with patch.object(queue, "flush", side_effect=_slow_flush): + completed = queue.flush_sync(timeout=0.1) + + assert completed is False + # The queue was not drained because flush() never returned. + assert queue.pending_count == 1 + # Release the daemon thread so it does not linger past the test. + release.set() + + +def _run_inflight_worker(queue: MemoryUpdateQueue, release: threading.Event) -> threading.Thread: + """Start a thread that mimics _process_queue's "pulled contexts, mid-LLM" state. + + It claims ``_processing`` / ``_processing_thread`` (so the queue looks idle + by ``pending_count`` but a worker is in flight), blocks on ``release``, + then clears the flags on the way out. + """ + + def _inflight() -> None: + with queue._lock: + queue._processing = True + queue._processing_thread = threading.current_thread() + release.wait(timeout=5.0) + with queue._lock: + queue._processing = False + queue._processing_thread = None + + thread = threading.Thread(target=_inflight, name="fake-inflight-worker", daemon=True) + thread.start() + # Wait until the fake worker has claimed _processing. + while not queue.is_processing: + time.sleep(0.005) + return thread + + +def test_flush_sync_waits_for_inflight_worker_and_returns_false_if_unfinished() -> None: + """flush_sync must not report success while an in-flight _process_queue is + still mid-LLM-call — the contexts it already pulled out would be lost on + exit. It joins the in-flight worker (bounded) and returns False when the + worker does not finish within the budget (review comment #1).""" + queue = _queue() + release = threading.Event() + inflight = _run_inflight_worker(queue, release) + + try: + completed = queue.flush_sync(timeout=0.2) + finally: + release.set() + inflight.join(timeout=5.0) + + assert completed is False + + +def test_flush_sync_returns_true_when_inflight_worker_finishes_in_budget() -> None: + """When the in-flight worker finishes within the budget, flush_sync joins it + and reports success (review comment #1, positive case).""" + queue = _queue() + release = threading.Event() + inflight = _run_inflight_worker(queue, release) + + # Let the in-flight worker finish well within the budget. + release.set() + + completed = queue.flush_sync(timeout=5.0) + inflight.join(timeout=5.0) + + assert completed is True + assert queue.is_processing is False + assert queue._processing_thread is None + + +def test_flush_sync_returns_false_when_flush_raises() -> None: + """flush_sync reports failure (not success) when flush() raises, so the + caller never logs a contradictory 'completed' next to the exception + (review comment #2).""" + queue = _queue() + queue._queue = [ConversationContext(thread_id="thread-1", messages=["conversation"], agent_name="lead_agent")] + + with patch.object(queue, "flush", side_effect=RuntimeError("boom")): + completed = queue.flush_sync(timeout=5.0) + + assert completed is False + + +def test_flush_sync_skips_inter_item_delay_on_drain_path() -> None: + """On the shutdown-drain path the per-item rate-limit sleep is skipped so + the bounded timeout covers as many items as possible (review comment #5).""" + mock_updater = MagicMock() + mock_updater.update_memory.return_value = True + queue = _queue(mock_updater) + queue._queue = [ConversationContext(thread_id=f"thread-{i}", messages=["conversation"], agent_name="lead_agent") for i in range(3)] + + with patch(_QUEUE_MODULE + ".time.sleep") as mock_sleep: + completed = queue.flush_sync(timeout=5.0) + + assert completed is True + assert queue.pending_count == 0 + # No inter-item rate-limit sleep on the drain path. + mock_sleep.assert_not_called() + assert mock_updater.update_memory.call_count == 3 diff --git a/config.example.yaml b/config.example.yaml index 3c1c458f9..e4e28733d 100644 --- a/config.example.yaml +++ b/config.example.yaml @@ -15,7 +15,7 @@ # ============================================================================ # Bump this number when the config schema changes. # Run `make config-upgrade` to merge new fields into your local config.yaml. -config_version: 25 +config_version: 26 # ============================================================================ # Logging @@ -1453,6 +1453,7 @@ summarization: # Shared fields (host level, backend-agnostic): # enabled - Master switch for the memory mechanism (call-site gate) # injection_enabled - Whether to inject memory into the system prompt (call-site gate) +# shutdown_flush_timeout_seconds - Hard budget (s) to drain pending updates on Gateway graceful shutdown (default: 30) # manager_class - Backend selector: registered name (deermem/noop) or dotted path # backend_config - Backend-private config dict (passthrough; each backend self-interprets) # @@ -1476,6 +1477,12 @@ summarization: memory: enabled: true injection_enabled: true # Whether to inject memory into system prompt + # Hard budget (seconds) to drain pending memory updates on Gateway graceful + # shutdown. Each pending item does one LLM call, so large IM batches may need + # more. Must fit inside the pod's K8s terminationGracePeriodSeconds (channel + # stop + this drain + buffer) or K8s SIGKILLs the drain -- set that on the + # gateway Helm deployment (see deploy/helm/deer-flow). Default 30s. + shutdown_flush_timeout_seconds: 30.0 # Memory backend selector. Either a registered backend name (matching a # backends// folder that exposes MANAGER_CLASS, e.g. deermem / noop) # or a dotted import path to a MemoryManager subclass. diff --git a/deploy/helm/deer-flow/README.md b/deploy/helm/deer-flow/README.md index cb72565e5..2c6aa9bd5 100644 --- a/deploy/helm/deer-flow/README.md +++ b/deploy/helm/deer-flow/README.md @@ -124,7 +124,7 @@ they resolve from the `secrets` map): ```yaml config: | - config_version: 25 + config_version: 26 models: - name: gpt-4 use: langchain_openai:ChatOpenAI @@ -227,6 +227,7 @@ kubectl -n deer-flow exec deploy/deer-flow-provisioner -- curl -s localhost:8002 username: deerflow password: changeme ``` +- **Graceful shutdown & memory drain.** The gateway pod sets `terminationGracePeriodSeconds` (default 45s, overridable via `gateway.terminationGracePeriodSeconds`) plus an optional `preStop` sleep (`gateway.preStopSleepSeconds`, default 5s). The grace period MUST exceed the Gateway's graceful-shutdown work — channel stop (~5s) plus the memory-queue drain (`memory.shutdown_flush_timeout_seconds`, default 30s) 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. **When you raise `memory.shutdown_flush_timeout_seconds`, raise `gateway.terminationGracePeriodSeconds` to match** (channel stop + drain + buffer). - **Gateway replicas.** Postgres + the Redis stream bridge together make the gateway's *persisted* state (checkpointer + run/thread metadata) and *live stream* path cross-pod-safe. The default is still 1 replica: **do not raise diff --git a/deploy/helm/deer-flow/templates/gateway-deployment.yaml b/deploy/helm/deer-flow/templates/gateway-deployment.yaml index 6227ebc93..4c907a7d2 100644 --- a/deploy/helm/deer-flow/templates/gateway-deployment.yaml +++ b/deploy/helm/deer-flow/templates/gateway-deployment.yaml @@ -24,6 +24,10 @@ spec: checksum/extensions: {{ include "deer-flow.extensionsChecksum" . }} spec: {{- include "deer-flow.imagePullSecrets" . | nindent 6 }} + # MUST exceed channel stop (~5s) + memory shutdown_flush_timeout_seconds + # (default 30s) + buffer, or K8s SIGKILLs the memory drain mid-flight + # (values.yaml gateway.terminationGracePeriodSeconds). K8s defaults to 30s. + terminationGracePeriodSeconds: {{ .Values.gateway.terminationGracePeriodSeconds | default 45 }} securityContext: # Non-root: the gateway image has no USER directive and would otherwise # run as root. uid/gid 1000; fsGroup 1000 so the .deer-flow PVC is @@ -145,6 +149,16 @@ spec: port: http initialDelaySeconds: 30 periodSeconds: 20 + {{- /* preStop sleep so the Service/ingress deregisters the pod and + in-flight requests drain before SIGTERM begins graceful shutdown + (which runs the memory queue drain). 0 / unset disables it. + Counts against terminationGracePeriodSeconds. */ -}} + {{- if and .Values.gateway.preStopSleepSeconds (gt (.Values.gateway.preStopSleepSeconds | int) 0) }} + lifecycle: + preStop: + exec: + command: ["sleep", "{{ .Values.gateway.preStopSleepSeconds | int }}"] + {{- end }} {{- with .Values.gateway.resources }} resources: {{- toYaml . | nindent 12 }} diff --git a/deploy/helm/deer-flow/values.yaml b/deploy/helm/deer-flow/values.yaml index c7c815159..a9b9ef5e0 100644 --- a/deploy/helm/deer-flow/values.yaml +++ b/deploy/helm/deer-flow/values.yaml @@ -28,6 +28,18 @@ gateway: replicas: 1 # Safe default. Postgres + the Redis stream bridge are # wired, but multi-replica needs issue #3948's run-control # work (cancel/dedup/reconcile) — see README "Gateway replicas". + # Pod termination grace period (seconds). MUST exceed the Gateway's graceful + # shutdown work: channel stop (~5s) + memory queue drain + # (memory.shutdown_flush_timeout_seconds, default 30s) + buffer. K8s defaults + # to 30s, which SIGKILLs the memory drain mid-flight and silently re-introduces + # the memory loss that shutdown_flush is fixing. Raise this when you raise the + # memory timeout. See gateway-deployment.yaml + README "Graceful shutdown". + terminationGracePeriodSeconds: 45 + # preStop hook sleep (seconds) before SIGTERM, so the Service/ingress + # deregisters the pod and in-flight HTTP requests drain before graceful + # shutdown begins (which runs the memory queue drain). 0 disables the hook. + # Counts against terminationGracePeriodSeconds. + preStopSleepSeconds: 5 resources: requests: cpu: 200m @@ -221,7 +233,7 @@ ingress: # -- DeerFlow config.yaml content. Secrets MUST stay as $VAR references — never # inline literal secret values here. The default enables provisioner sandbox. config: | - config_version: 25 + config_version: 26 log_level: info models: []