mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-08-10 14:58:46 +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). * 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.
194 lines
8.6 KiB
Python
194 lines
8.6 KiB
Python
"""Pluggable memory manager: the factory resolves the configured backend.
|
|
|
|
Covers the drop-in contract end-to-end:
|
|
- short name -> registered backend (deermem / noop);
|
|
- dotted path (``module.Attr`` and ``module:Attr``) -> the same class;
|
|
- unknown value -> raise (fail-fast: a wrong store is a silent data-integrity footgun).
|
|
|
|
Also pins the noop empty-memory behaviour and the ``hasattr`` capability
|
|
probing surface (reload_memory + fact CRUD) that the gateway/client rely on.
|
|
|
|
Each test resets the singleton + backend cache and sets the config, so they
|
|
are independent of order.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from pathlib import Path
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
import pytest
|
|
|
|
from deerflow.agents.memory import (
|
|
MemoryManager,
|
|
get_memory_manager,
|
|
reset_memory_manager,
|
|
)
|
|
from deerflow.agents.memory.backends.deermem.deer_mem import DeerMem
|
|
from deerflow.agents.memory.backends.noop.noop_manager import NoopMemoryManager
|
|
from deerflow.config.memory_config import MemoryConfig, get_memory_config, set_memory_config
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def _isolate_memory_manager():
|
|
"""Reset the singleton + restore config around every test."""
|
|
orig = get_memory_config()
|
|
reset_memory_manager()
|
|
yield
|
|
set_memory_config(orig)
|
|
reset_memory_manager()
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"manager_class, expected",
|
|
[
|
|
("deermem", DeerMem),
|
|
("noop", NoopMemoryManager),
|
|
("deerflow.agents.memory.backends.deermem.deer_mem.DeerMem", DeerMem),
|
|
("deerflow.agents.memory.backends.deermem.deer_mem:DeerMem", DeerMem),
|
|
("deerflow.agents.memory.backends.noop.noop_manager.NoopMemoryManager", NoopMemoryManager),
|
|
("deerflow.agents.memory.backends.noop.noop_manager:NoopMemoryManager", NoopMemoryManager),
|
|
],
|
|
)
|
|
def test_resolves_configured_backend(manager_class: str, expected: type[MemoryManager]) -> None:
|
|
set_memory_config(MemoryConfig(manager_class=manager_class))
|
|
manager = get_memory_manager()
|
|
assert isinstance(manager, expected)
|
|
# singleton: a second call returns the same instance
|
|
assert get_memory_manager() is manager
|
|
|
|
|
|
def test_unknown_backend_raises_instead_of_falling_back() -> None:
|
|
"""An unknown manager_class is a config error: raise, don't silently fall
|
|
back to DeerMem (memory is persistent state -- a wrong store is a silent
|
|
data-integrity footgun)."""
|
|
set_memory_config(MemoryConfig(manager_class="bogus-backend"))
|
|
with pytest.raises(ValueError, match="bogus-backend"):
|
|
get_memory_manager()
|
|
|
|
|
|
def test_noop_runs_with_empty_memory() -> None:
|
|
set_memory_config(MemoryConfig(manager_class="noop"))
|
|
manager = get_memory_manager()
|
|
assert manager.get_context(user_id="u") == ""
|
|
assert manager.search("anything") == []
|
|
assert manager.get_memory(user_id="u") == {"facts": []}
|
|
# writes are no-ops; memory stays empty
|
|
manager.add("t", [], agent_name=None, user_id="u")
|
|
manager.add_nowait("t", [], agent_name=None, user_id="u")
|
|
assert manager.get_memory(user_id="u") == {"facts": []}
|
|
|
|
|
|
def test_internal_capabilities_are_hasattr_probeable() -> None:
|
|
"""reload_memory + fact CRUD + warm exist on DeerMem but not on noop (the ABC omits them)."""
|
|
set_memory_config(MemoryConfig(manager_class="deermem"))
|
|
deermem = get_memory_manager()
|
|
for cap in ("warm", "reload_memory", "create_fact", "delete_fact", "update_fact"):
|
|
assert hasattr(deermem, cap), cap
|
|
|
|
reset_memory_manager()
|
|
set_memory_config(MemoryConfig(manager_class="noop"))
|
|
noop = get_memory_manager()
|
|
for cap in ("warm", "reload_memory", "create_fact", "delete_fact", "update_fact"):
|
|
assert not hasattr(noop, cap), cap
|
|
|
|
|
|
def test_deermem_search_works_delete_export_are_stubs() -> None:
|
|
set_memory_config(MemoryConfig(manager_class="deermem"))
|
|
deermem = get_memory_manager()
|
|
# search is implemented (substring match) -- returns a list, does not raise.
|
|
assert isinstance(deermem.search("q", user_id="u"), list)
|
|
# delete_memory / export_memory remain unimplemented stubs this phase.
|
|
with pytest.raises(NotImplementedError):
|
|
deermem.delete_memory(user_id="u")
|
|
with pytest.raises(NotImplementedError):
|
|
deermem.export_memory(user_id="u")
|
|
|
|
|
|
def test_factory_raises_when_storage_path_is_existing_file(tmp_path) -> None:
|
|
"""A storage_path that resolves to an existing FILE is a config error: DeerMem
|
|
treats storage_path as a root directory, so a file would make save's mkdir
|
|
raise NotADirectoryError (silent write failure). Fail loud at startup (#1)."""
|
|
file_path = tmp_path / "mem.json"
|
|
file_path.write_text("{}", encoding="utf-8")
|
|
set_memory_config(MemoryConfig(manager_class="deermem", backend_config={"storage_path": str(file_path)}))
|
|
with pytest.raises(ValueError, match="existing file"):
|
|
get_memory_manager()
|
|
|
|
|
|
def test_migration_drops_file_style_legacy_storage_path(caplog) -> None:
|
|
"""A legacy top-level storage_path that looks like a file (ends in .json) is
|
|
dropped, not carried verbatim -- DeerMem now treats storage_path as a root
|
|
directory, so carrying 'memory.json' would orphan per-user memory / hit
|
|
NotADirectoryError. Dropping lets the factory inject runtime_home (per-user
|
|
location unchanged). Non-file legacy fields still migrate; empty values are
|
|
skipped silently (#1, #6)."""
|
|
from deerflow.config.memory_config import load_memory_config_from_dict
|
|
|
|
with caplog.at_level("WARNING", logger="deerflow.config.memory_config"):
|
|
load_memory_config_from_dict({"storage_path": "memory.json", "max_facts": 50})
|
|
cfg = get_memory_config()
|
|
assert "storage_path" not in cfg.backend_config # file-style dropped
|
|
assert cfg.backend_config.get("max_facts") == 50 # non-file legacy still migrates
|
|
assert any("looks like a file path" in r.message for r in caplog.records)
|
|
|
|
|
|
def test_empty_storage_path_factory_injects_runtime_home(tmp_path, monkeypatch) -> None:
|
|
"""Empty/absent storage_path -> factory injects runtime_home() as the root, so
|
|
per-user memory lands at {runtime_home}/users/{uid}/memory.json (matches
|
|
pre-abstraction per-user location). Pins the zero-config default (reviewer #1)."""
|
|
import deerflow.config.runtime_paths as rp
|
|
|
|
monkeypatch.setattr(rp, "runtime_home", lambda: tmp_path)
|
|
set_memory_config(MemoryConfig(manager_class="deermem")) # no storage_path
|
|
manager = get_memory_manager()
|
|
assert Path(manager._config.storage_path) == tmp_path
|
|
manager.create_fact("hello", user_id="u1")
|
|
# 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()
|