mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-07-21 10:15:47 +00:00
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).
This commit is contained in:
@@ -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")
|
||||
|
||||
|
||||
|
||||
@@ -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,
|
||||
),
|
||||
|
||||
@@ -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).
|
||||
|
||||
+110
-7
@@ -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
|
||||
|
||||
@@ -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, "<name>")`` and
|
||||
# returns 501 when absent. Implement the ones your backend supports so the
|
||||
|
||||
@@ -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]]:
|
||||
|
||||
@@ -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,
|
||||
}
|
||||
|
||||
@@ -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=(
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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
|
||||
|
||||
+8
-1
@@ -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/<name>/ folder that exposes MANAGER_CLASS, e.g. deermem / noop)
|
||||
# or a dotted import path to a MemoryManager subclass.
|
||||
|
||||
Reference in New Issue
Block a user