deerflow-code/offline-backend-20260512/backend/tests/test_memory_manager.py

"""Memory V2 - MemoryManager 与 MemoryProvider 抽象基类测试。"""

from __future__ import annotations

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_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


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")