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.
This commit is contained in:
Tu Naichao 2026-07-15 20:25:41 +08:00 committed by GitHub
parent fa419b3e4c
commit 959bf13406
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
16 changed files with 538 additions and 14 deletions

View File

@ -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. - `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) - `storage_path` - Path to memory.json (absolute path opts out of per-user isolation)
- `debounce_seconds` - Wait time before processing (default: 30) - `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; 1300). 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) - `model_name` - LLM for updates (null = default model)
- `max_facts` / `fact_confidence_threshold` - Fact storage limits (100 / 0.7) - `max_facts` / `fact_confidence_threshold` - Fact storage limits (100 / 0.7)
- `max_injection_tokens` - Token limit for prompt injection (2000) - `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) - `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) - `summarization` - Context summarization (enabled, trigger conditions, keep policy)
- `subagents.enabled` - Master switch for subagent delegation - `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`**: **`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 `<mcp_routing_hints>` 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. - `mcpServers` - Map of server name → config (enabled, type, command, args, env, url, headers, oauth, description, `routing`, `tools`, `tool_call_timeout`). `routing.mode="prefer"` emits `<mcp_routing_hints>` 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.

View File

@ -301,6 +301,43 @@ async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]:
except Exception: except Exception:
logger.exception("Failed to stop scheduled task service") 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") logger.info("Shutting down API Gateway")

View File

@ -130,6 +130,7 @@ class MemoryConfigResponse(BaseModel):
enabled: bool = Field(..., description="Whether the memory mechanism is enabled (call-site gate).") 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.") 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).") 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).") 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).") 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, "enabled": true,
"injection_enabled": true, "injection_enabled": true,
"shutdown_flush_timeout_seconds": 30.0,
"mode": "middleware", "mode": "middleware",
"manager_class": "deermem", "manager_class": "deermem",
"backend_config": { "backend_config": {
@ -380,6 +382,7 @@ async def get_memory_config_endpoint() -> MemoryConfigResponse:
enabled=config.enabled, enabled=config.enabled,
mode=config.mode, mode=config.mode,
injection_enabled=config.injection_enabled, injection_enabled=config.injection_enabled,
shutdown_flush_timeout_seconds=config.shutdown_flush_timeout_seconds,
manager_class=config.manager_class, manager_class=config.manager_class,
backend_config=config.backend_config, backend_config=config.backend_config,
) )
@ -406,6 +409,7 @@ async def get_memory_status(http_request: Request) -> MemoryStatusResponse:
enabled=config.enabled, enabled=config.enabled,
mode=config.mode, mode=config.mode,
injection_enabled=config.injection_enabled, injection_enabled=config.injection_enabled,
shutdown_flush_timeout_seconds=config.shutdown_flush_timeout_seconds,
manager_class=config.manager_class, manager_class=config.manager_class,
backend_config=config.backend_config, backend_config=config.backend_config,
), ),

View File

@ -237,6 +237,19 @@ class DeerMem(MemoryManager):
"""Not implemented this phase (no distinct export yet; /export routes via get_memory).""" """Not implemented this phase (no distinct export yet; /export routes via get_memory)."""
raise NotImplementedError("DeerMem.export_memory is not implemented yet") 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) ─── # ── DeerMem-internal (NOT on the ABC; reached via hasattr probing) ───
def warm(self) -> bool: def warm(self) -> bool:
"""Pre-warm DeerMem-specific resources (the tiktoken encoding cache). """Pre-warm DeerMem-specific resources (the tiktoken encoding cache).

View File

@ -47,6 +47,11 @@ class MemoryUpdateQueue:
self._lock = threading.Lock() self._lock = threading.Lock()
self._timer: threading.Timer | None = None self._timer: threading.Timer | None = None
self._processing = False 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 self._reprocess_pending = False
@staticmethod @staticmethod
@ -172,8 +177,15 @@ class MemoryUpdateQueue:
self._timer.daemon = True self._timer.daemon = True
self._timer.start() self._timer.start()
def _process_queue(self) -> None: def _process_queue(self, *, skip_inter_item_delay: bool = False) -> None:
"""Process all queued conversation contexts.""" """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: with self._lock:
if self._processing: if self._processing:
# Another worker is already draining the queue. Instead of # Another worker is already draining the queue. Instead of
@ -188,12 +200,15 @@ class MemoryUpdateQueue:
return return
self._processing = True self._processing = True
self._processing_thread = threading.current_thread()
contexts_to_process = self._queue.copy() contexts_to_process = self._queue.copy()
self._queue.clear() self._queue.clear()
self._timer = None self._timer = None
logger.info("Processing %d queued memory updates", len(contexts_to_process)) logger.info("Processing %d queued memory updates", len(contexts_to_process))
succeeded = 0
failed = 0
try: try:
for context in contexts_to_process: for context in contexts_to_process:
try: try:
@ -208,35 +223,122 @@ class MemoryUpdateQueue:
trace_id=context.trace_id, trace_id=context.trace_id,
) )
if success: if success:
succeeded += 1
logger.info("Memory updated successfully for thread %s (trace_id=%s)", context.thread_id, context.trace_id) logger.info("Memory updated successfully for thread %s (trace_id=%s)", context.thread_id, context.trace_id)
else: else:
failed += 1
logger.warning("Memory update skipped/failed for thread %s (trace_id=%s)", context.thread_id, context.trace_id) logger.warning("Memory update skipped/failed for thread %s (trace_id=%s)", context.thread_id, context.trace_id)
except Exception as e: 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) 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 # Small delay between updates to avoid rate limiting.
if len(contexts_to_process) > 1: # 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) time.sleep(0.5)
finally: 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: with self._lock:
self._processing = False self._processing = False
self._processing_thread = None
if self._reprocess_pending: if self._reprocess_pending:
self._reprocess_pending = False self._reprocess_pending = False
if self._queue: if self._queue:
self._schedule_timer(0) self._schedule_timer(0)
def flush(self) -> None: def flush(self, *, skip_inter_item_delay: bool = False) -> None:
"""Force immediate processing of the queue. """Force immediate processing of the queue.
This is useful for testing or graceful shutdown. 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: with self._lock:
if self._timer is not None: if self._timer is not None:
self._timer.cancel() self._timer.cancel()
self._timer = None 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: def flush_nowait(self) -> None:
"""Start queue processing immediately in a background thread.""" """Start queue processing immediately in a background thread."""
@ -256,6 +358,7 @@ class MemoryUpdateQueue:
self._timer = None self._timer = None
self._queue.clear() self._queue.clear()
self._processing = False self._processing = False
self._processing_thread = None
self._reprocess_pending = False self._reprocess_pending = False
@property @property

View File

@ -157,6 +157,11 @@ class NoopMemoryManager(MemoryManager):
) -> dict[str, Any]: ) -> dict[str, Any]:
return _empty_memory() 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) ────────── # ── Optional DeerMem-internal capabilities (NOT on the ABC) ──────────
# The host gateway discovers these via ``hasattr(manager, "<name>")`` and # The host gateway discovers these via ``hasattr(manager, "<name>")`` and
# returns 501 when absent. Implement the ones your backend supports so the # returns 501 when absent. Implement the ones your backend supports so the

View File

@ -188,6 +188,34 @@ class MemoryManager(ABC):
) -> dict[str, Any]: ) -> dict[str, Any]:
"""Export the memory document for the bucket. *stub* this phase (no caller yet).""" """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) ─────────────────────────────────────────── # ── Backend discovery (drop-in) ───────────────────────────────────────────
def _scan_backends() -> dict[str, type[MemoryManager]]: def _scan_backends() -> dict[str, type[MemoryManager]]:

View File

@ -1342,6 +1342,7 @@ class DeerFlowClient:
"enabled": config.enabled, "enabled": config.enabled,
"mode": config.mode, "mode": config.mode,
"injection_enabled": config.injection_enabled, "injection_enabled": config.injection_enabled,
"shutdown_flush_timeout_seconds": config.shutdown_flush_timeout_seconds,
"manager_class": config.manager_class, "manager_class": config.manager_class,
"backend_config": config.backend_config, "backend_config": config.backend_config,
} }

View File

@ -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 reached via ``backend_config`` (a dict the factory passes to the backend's
``__init__``). This module holds ONLY the host-shared fields every backend / ``__init__``). This module holds ONLY the host-shared fields every backend /
call site / factory reads: ``enabled`` / ``injection_enabled`` / 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 makes backends swappable and portable (DeerMem's knobs do not leak onto the
shared contract). shared contract).
""" """
@ -17,7 +18,7 @@ from pydantic import BaseModel, Field
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
# Host-shared MemoryConfig fields (read by every backend / call site / factory). # 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 # 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 # config.yaml (pre-abstraction). On load they are auto-migrated into
@ -66,6 +67,23 @@ class MemoryConfig(BaseModel):
default=True, default=True,
description="Whether to inject memory into the system prompt (call-site gate).", 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( manager_class: str = Field(
default="deermem", default="deermem",
description=( description=(

View File

@ -10,6 +10,7 @@ signal-reentrancy deadlock described in
from __future__ import annotations from __future__ import annotations
import asyncio import asyncio
import logging
from contextlib import asynccontextmanager from contextlib import asynccontextmanager
from types import SimpleNamespace from types import SimpleNamespace
from unittest.mock import AsyncMock, MagicMock, patch from unittest.mock import AsyncMock, MagicMock, patch
@ -35,6 +36,9 @@ async def _run_lifespan_with_hanging_stop() -> float:
app = FastAPI() app = FastAPI()
startup_config = MagicMock() startup_config = MagicMock()
startup_config.log_level = "INFO" 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 = MagicMock()
fake_service.get_status = MagicMock(return_value={}) 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 from app.gateway.app import lifespan
app = FastAPI() 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 = MagicMock()
fake_service.get_status = MagicMock(return_value={}) fake_service.get_status = MagicMock(return_value={})
cleanup_upload_staging_files = MagicMock(return_value=2) 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() cleanup_upload_staging_files.assert_called_once_with()
close_oidc_service.assert_awaited_once() close_oidc_service.assert_awaited_once()
stop_channel_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()

View File

@ -15,6 +15,7 @@ are independent of order.
from __future__ import annotations from __future__ import annotations
from pathlib import Path from pathlib import Path
from unittest.mock import MagicMock, patch
import pytest 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 # 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()] user_dirs = [p.name for p in (tmp_path / "users").iterdir() if p.is_dir()]
assert len(user_dirs) == 1 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()

View File

@ -247,3 +247,159 @@ def test_process_queue_forwards_trace_id_to_updater() -> None:
user_id=None, user_id=None,
trace_id="trace-memory-1", 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

View File

@ -15,7 +15,7 @@
# ============================================================================ # ============================================================================
# Bump this number when the config schema changes. # Bump this number when the config schema changes.
# Run `make config-upgrade` to merge new fields into your local config.yaml. # Run `make config-upgrade` to merge new fields into your local config.yaml.
config_version: 25 config_version: 26
# ============================================================================ # ============================================================================
# Logging # Logging
@ -1453,6 +1453,7 @@ summarization:
# Shared fields (host level, backend-agnostic): # Shared fields (host level, backend-agnostic):
# enabled - Master switch for the memory mechanism (call-site gate) # enabled - Master switch for the memory mechanism (call-site gate)
# injection_enabled - Whether to inject memory into the system prompt (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 # manager_class - Backend selector: registered name (deermem/noop) or dotted path
# backend_config - Backend-private config dict (passthrough; each backend self-interprets) # backend_config - Backend-private config dict (passthrough; each backend self-interprets)
# #
@ -1476,6 +1477,12 @@ summarization:
memory: memory:
enabled: true enabled: true
injection_enabled: true # Whether to inject memory into system prompt 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 # Memory backend selector. Either a registered backend name (matching a
# backends/<name>/ folder that exposes MANAGER_CLASS, e.g. deermem / noop) # backends/<name>/ folder that exposes MANAGER_CLASS, e.g. deermem / noop)
# or a dotted import path to a MemoryManager subclass. # or a dotted import path to a MemoryManager subclass.

View File

@ -124,7 +124,7 @@ they resolve from the `secrets` map):
```yaml ```yaml
config: | config: |
config_version: 25 config_version: 26
models: models:
- name: gpt-4 - name: gpt-4
use: langchain_openai:ChatOpenAI use: langchain_openai:ChatOpenAI
@ -227,6 +227,7 @@ kubectl -n deer-flow exec deploy/deer-flow-provisioner -- curl -s localhost:8002
username: deerflow username: deerflow
password: changeme 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 replicas.** Postgres + the Redis stream bridge together make the
gateway's *persisted* state (checkpointer + run/thread metadata) and *live gateway's *persisted* state (checkpointer + run/thread metadata) and *live
stream* path cross-pod-safe. The default is still 1 replica: **do not raise stream* path cross-pod-safe. The default is still 1 replica: **do not raise

View File

@ -24,6 +24,10 @@ spec:
checksum/extensions: {{ include "deer-flow.extensionsChecksum" . }} checksum/extensions: {{ include "deer-flow.extensionsChecksum" . }}
spec: spec:
{{- include "deer-flow.imagePullSecrets" . | nindent 6 }} {{- 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: securityContext:
# Non-root: the gateway image has no USER directive and would otherwise # 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 # run as root. uid/gid 1000; fsGroup 1000 so the .deer-flow PVC is
@ -145,6 +149,16 @@ spec:
port: http port: http
initialDelaySeconds: 30 initialDelaySeconds: 30
periodSeconds: 20 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 }} {{- with .Values.gateway.resources }}
resources: resources:
{{- toYaml . | nindent 12 }} {{- toYaml . | nindent 12 }}

View File

@ -28,6 +28,18 @@ gateway:
replicas: 1 # Safe default. Postgres + the Redis stream bridge are replicas: 1 # Safe default. Postgres + the Redis stream bridge are
# wired, but multi-replica needs issue #3948's run-control # wired, but multi-replica needs issue #3948's run-control
# work (cancel/dedup/reconcile) — see README "Gateway replicas". # 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: resources:
requests: requests:
cpu: 200m cpu: 200m
@ -221,7 +233,7 @@ ingress:
# -- DeerFlow config.yaml content. Secrets MUST stay as $VAR references — never # -- DeerFlow config.yaml content. Secrets MUST stay as $VAR references — never
# inline literal secret values here. The default enables provisioner sandbox. # inline literal secret values here. The default enables provisioner sandbox.
config: | config: |
config_version: 25 config_version: 26
log_level: info log_level: info
models: [] models: []