"""Memory V2 - MemoryManager 与 MemoryProvider 抽象基类测试。""" from __future__ import annotations from types import SimpleNamespace from typing import Any from deerflow.agents.memory.manager import MemoryManager from deerflow.agents.memory.provider import MemoryProvider class _FakeProvider(MemoryProvider): """测试用的可配置 provider。""" def __init__(self, name: str, *, prompt_block: str = "", prefetch_text: str = ""): self._name = name self._prompt_block = prompt_block self._prefetch_text = prefetch_text self.calls: list[str] = [] @property def name(self) -> str: return self._name def is_available(self) -> bool: return True def initialize(self, **kwargs: Any) -> None: self.calls.append("initialize") def shutdown(self) -> None: self.calls.append("shutdown") def system_prompt_block(self) -> str: return self._prompt_block def prefetch(self, query: str, *, thread_id: str = "") -> str: self.calls.append(f"prefetch:{query}") return self._prefetch_text def sync_turn(self, user_content: str, assistant_content: str, *, thread_id: str = "", turn_index: int | None = None) -> None: self.calls.append("sync_turn") self.last_turn_index = turn_index def get_tool_schemas(self) -> list[dict[str, Any]]: return [] def on_session_switch(self, new_thread_id: str, **kwargs: Any) -> None: self.calls.append(f"on_session_switch:{new_thread_id}") def on_memory_write(self, action, target, content, metadata=None) -> None: self.calls.append(f"on_memory_write:{action}:{target}") class _ExplodingProvider(_FakeProvider): """每个 hook 都抛异常,用于验证 manager 的隔离性。""" def prefetch(self, query: str, *, thread_id: str = "") -> str: raise RuntimeError("boom") def system_prompt_block(self) -> str: raise RuntimeError("boom") def sync_turn(self, user_content, assistant_content, *, thread_id: str = "", turn_index: int | None = None) -> None: raise RuntimeError("boom") def initialize(self, **kwargs: Any) -> None: raise RuntimeError("boom") class TestRegistration: def test_builtin_always_accepted(self): manager = MemoryManager() assert manager.add_provider(_FakeProvider("builtin")) is True def test_one_external_accepted(self): manager = MemoryManager() manager.add_provider(_FakeProvider("builtin")) assert manager.add_provider(_FakeProvider("hindsight")) is True def test_second_external_rejected(self): manager = MemoryManager() manager.add_provider(_FakeProvider("builtin")) manager.add_provider(_FakeProvider("hindsight")) assert manager.add_provider(_FakeProvider("mem0")) is False assert len(manager.providers) == 2 def test_knowledge_does_not_consume_external_slot(self): manager = MemoryManager() manager.add_provider(_FakeProvider("builtin")) manager.add_provider(_FakeProvider("knowledge")) assert manager.add_provider(_FakeProvider("hindsight")) is True assert manager.add_provider(_FakeProvider("mem0")) is False assert [p.name for p in manager.providers] == ["builtin", "knowledge", "hindsight"] def test_get_provider(self): manager = MemoryManager() builtin = _FakeProvider("builtin") manager.add_provider(builtin) assert manager.get_provider("builtin") is builtin assert manager.get_provider("missing") is None class TestSystemPromptAndPrefetch: def test_build_system_prompt_concatenates(self): manager = MemoryManager() manager.add_provider(_FakeProvider("builtin", prompt_block="块A")) manager.add_provider(_FakeProvider("hindsight", prompt_block="块B")) result = manager.build_system_prompt() assert "块A" in result and "块B" in result def test_build_system_prompt_skips_empty(self): manager = MemoryManager() manager.add_provider(_FakeProvider("builtin", prompt_block="")) assert manager.build_system_prompt() == "" def test_build_system_prompt_can_exclude_only_legacy_knowledge(self): manager = MemoryManager() manager.add_provider(_FakeProvider("builtin", prompt_block="用户记忆")) manager.add_provider( _FakeProvider("knowledge", prompt_block="/mnt/knowledge/routing.md") ) manager.add_provider(_FakeProvider("hindsight", prompt_block="Hindsight 说明")) result = manager.build_system_prompt( excluded_provider_names={"knowledge"} ) assert "用户记忆" in result assert "Hindsight 说明" in result assert "/mnt/knowledge/routing.md" not in result def test_prefetch_all_concatenates(self): manager = MemoryManager() manager.add_provider(_FakeProvider("builtin", prefetch_text="召回1")) manager.add_provider(_FakeProvider("hindsight", prefetch_text="召回2")) result = manager.prefetch_all("查询") assert "召回1" in result and "召回2" in result def test_build_injection_block_combines_static_and_dynamic(self): manager = MemoryManager() manager.add_provider( _FakeProvider("builtin", prompt_block="静态块", prefetch_text="动态块") ) block = manager.build_injection_block("查询") assert "静态块" in block and "动态块" in block def test_lead_memory_context_keeps_memory_when_legacy_routing_is_disabled( self, monkeypatch, ): from deerflow.agents.lead_agent import prompt as prompt_module from deerflow.agents.memory import manager as manager_module from deerflow.config import memory_config as memory_config_module from deerflow.runtime import user_context as user_context_module manager = MemoryManager() manager.add_provider(_FakeProvider("builtin", prompt_block="用户偏好记忆")) manager.add_provider( _FakeProvider("knowledge", prompt_block="读取 /mnt/knowledge/routing.md") ) manager.add_provider(_FakeProvider("hindsight", prompt_block="历史召回说明")) monkeypatch.setattr( manager_module, "build_memory_manager", lambda **_kwargs: manager, ) monkeypatch.setattr( memory_config_module, "get_effective_memory_config", lambda _user_id: SimpleNamespace(enabled=True), ) monkeypatch.setattr( user_context_module, "get_effective_user_id", lambda: "user-1", ) context = prompt_module._get_memory_context( include_legacy_knowledge_routing=False ) assert "用户偏好记忆" in context assert "历史召回说明" in context assert "/mnt/knowledge/routing.md" not in context class TestHookDispatch: def test_sync_all(self): manager = MemoryManager() p = _FakeProvider("builtin") manager.add_provider(p) manager.sync_all("用户说", "助手答") assert "sync_turn" in p.calls def test_sync_all_passes_turn_index(self): manager = MemoryManager() p = _FakeProvider("builtin") manager.add_provider(p) manager.sync_all("用户说", "助手答", turn_index=5) assert p.last_turn_index == 5 def test_on_session_switch_all(self): manager = MemoryManager() p = _FakeProvider("builtin") manager.add_provider(p) manager.on_session_switch_all("thread-2") assert "on_session_switch:thread-2" in p.calls def test_on_memory_write_skips_builtin(self): manager = MemoryManager() builtin = _FakeProvider("builtin") external = _FakeProvider("hindsight") manager.add_provider(builtin) manager.add_provider(external) manager.on_memory_write("add", "user", "内容") # builtin 是写入源,不应收到镜像通知 assert not any(c.startswith("on_memory_write") for c in builtin.calls) assert "on_memory_write:add:user" in external.calls def test_initialize_and_shutdown_all(self): manager = MemoryManager() p = _FakeProvider("builtin") manager.add_provider(p) manager.initialize_all(user_id="u1") manager.shutdown_all() assert "initialize" in p.calls assert "shutdown" in p.calls class TestFailureIsolation: def test_exploding_provider_does_not_break_others(self): manager = MemoryManager() manager.add_provider(_ExplodingProvider("builtin")) good = _FakeProvider("hindsight", prompt_block="正常块", prefetch_text="正常召回") manager.add_provider(good) # 一个 provider 抛异常,另一个仍正常工作 manager.initialize_all(user_id="u1") assert manager.build_system_prompt() == "正常块" assert manager.prefetch_all("q") == "正常召回" manager.sync_all("u", "a") # 不抛异常 class _MinimalProvider(MemoryProvider): """只实现 abstractmethod 的最小 provider,用于验证基类默认 hook。""" @property def name(self) -> str: return "minimal" def is_available(self) -> bool: return True def initialize(self, **kwargs: Any) -> None: pass def get_tool_schemas(self) -> list[dict[str, Any]]: return [] class TestProviderDefaults: """MemoryProvider 抽象基类的默认 hook 实现。""" def test_default_hooks_are_noop(self): provider = _MinimalProvider() # 未重写的 hook 应有安全的默认实现,不抛异常 assert provider.prefetch("q") == "" assert provider.system_prompt_block() == "" provider.queue_prefetch("q") provider.sync_turn("u", "a") provider.shutdown() provider.on_session_switch("t") provider.on_memory_write("add", "user", "c")