deer-flow/backend/tests/test_memory_retrieval_adapter.py

"""Regression coverage for DeerMem's #4279 RetrievalPort integration."""

from __future__ import annotations

from concurrent.futures import ThreadPoolExecutor
from pathlib import Path
from unittest.mock import MagicMock

import pytest

from deerflow.agents.memory.backends.deermem.deer_mem import DeerMem
from deerflow.agents.memory.backends.deermem.deermem.config import DeerMemConfig
from deerflow.agents.memory.backends.deermem.deermem.core.retrieval import (
    FTS5Retrieval,
    FTS5RetrievalAdapter,
    _is_advanced_query,
    _jieba_available,
    create_fts5_retrieval,
)
from deerflow.agents.memory.backends.deermem.deermem.core.storage import FileMemoryStorage


def _fact(fact_id: str, content: str, *, category: str = "context") -> dict:
    return {
        "id": fact_id,
        "content": content,
        "category": category,
        "confidence": 0.8,
        "createdAt": "2026-07-21T00:00:00Z",
        "source": {"type": "test", "threadId": None},
    }


def test_natural_language_hyphens_are_not_treated_as_fts5_syntax() -> None:
    assert not _is_advanced_query("real-time co-pilot node -js +python")


@pytest.mark.skipif(not _jieba_available, reason="install the optional memory-zh extra")
def test_chinese_subphrase_search_uses_jieba_tokenization(tmp_path: Path) -> None:
    adapter = FTS5RetrievalAdapter(tmp_path / "facts.sqlite3")
    scope = {"userId": "alice", "agentName": "agent-a"}
    try:
        adapter.upsert(_fact("zh", "中文检索支持验证"), scope=scope, path="")
        results = adapter.search("检索", scopes=[scope], top_k=5, mode="fts5", filters=None)
        assert [item["fact"]["id"] for item in results] == ["zh"]
    finally:
        adapter.close()


def test_score_time_decay_and_advanced_phrase_query(tmp_path: Path) -> None:
    engine = FTS5Retrieval(tmp_path / "facts.sqlite3")
    try:
        assert engine._compute_final_score(1.0, 0.5, "2026-07-21T00:00:00Z") > engine._compute_final_score(1.0, 0.5, "2020-01-01T00:00:00Z")
        engine.index_fact("phrase", "alpha beta", scope_user='"alice"', scope_agent='"agent-a"')
        assert engine.search('"alpha beta"', scope_user='"alice"', scope_agent='"agent-a"')
    finally:
        engine.close()


def test_warm_retrieval_rebuilds_the_complete_index(tmp_path: Path) -> None:
    manager = DeerMem(backend_config={"storage_path": str(tmp_path), "token_counting": "char"})
    _, fact_id = manager.create_fact("warm retrieval fact", user_id="alice")
    assert manager.warm_retrieval()
    assert any(fact["id"] == fact_id for fact in manager.search("warm", user_id="alice"))


def test_warm_retrieval_accepts_partial_fact_failures(tmp_path: Path) -> None:
    manager = DeerMem(backend_config={"storage_path": str(tmp_path), "token_counting": "char"})
    manager._storage.rebuild_index = lambda scopes=None: {"supported": True, "indexed": 2, "failed": 1}  # type: ignore[method-assign]

    assert manager.warm_retrieval()
    assert manager._retrieval_fully_warmed


def test_warm_retrieval_retries_after_fatal_rebuild_failure(tmp_path: Path) -> None:
    manager = DeerMem(backend_config={"storage_path": str(tmp_path), "token_counting": "char"})
    manager._storage.rebuild_index = lambda scopes=None: {"supported": True, "indexed": 0, "failed": 1, "fatal": True}  # type: ignore[method-assign]

    assert not manager.warm_retrieval()
    assert not manager._retrieval_fully_warmed


def test_lazy_warm_rebuilds_each_requested_scope(tmp_path: Path) -> None:
    manager = DeerMem(backend_config={"storage_path": str(tmp_path), "token_counting": "char"})
    calls: list[list[dict[str, str | None]]] = []
    manager._storage.rebuild_index = lambda scopes=None: calls.append(scopes or []) or {"supported": True, "failed": 0}  # type: ignore[method-assign]
    scopes = [{"userId": "alice", "agentName": "a"}, {"userId": "bob", "agentName": "b"}]
    manager._ensure_retrieval_scopes(scopes)
    assert calls == [[scopes[0]], [scopes[1]]]


def test_lazy_warm_does_not_retry_partial_fact_failures(tmp_path: Path) -> None:
    manager = DeerMem(backend_config={"storage_path": str(tmp_path), "token_counting": "char"})
    rebuild = MagicMock(return_value={"supported": True, "indexed": 2, "failed": 1})
    manager._storage.rebuild_index = rebuild  # type: ignore[method-assign]
    scopes = [{"userId": "alice", "agentName": "a"}]

    manager._ensure_retrieval_scopes(scopes)
    manager._ensure_retrieval_scopes(scopes)

    rebuild.assert_called_once_with(scopes)


def test_deermem_close_releases_retrieval_connection(tmp_path: Path) -> None:
    manager = DeerMem(backend_config={"storage_path": str(tmp_path), "token_counting": "char"})
    close_storage = MagicMock()
    manager._storage.close = close_storage  # type: ignore[attr-defined]

    manager.close()

    close_storage.assert_called_once_with()


def test_file_storage_close_releases_retrieval_connection(tmp_path: Path) -> None:
    retrieval = MagicMock()
    storage = FileMemoryStorage(DeerMemConfig(storage_path=str(tmp_path)), retrieval=retrieval)

    storage.close()

    retrieval.close.assert_called_once_with()


def test_factory_recreates_corrupt_persistent_index(tmp_path: Path) -> None:
    index_dir = tmp_path / ".retrieval"
    index_dir.mkdir()
    db_path = index_dir / "memory-fts5.sqlite3"
    db_path.write_bytes(b"not a sqlite database")

    adapter = create_fts5_retrieval(DeerMemConfig(storage_path=str(tmp_path)))

    assert adapter is not None
    try:
        adapter.upsert(
            _fact("recovered", "recreated derived index"),
            scope={"userId": "alice", "agentName": "agent-a"},
            path="",
        )
        assert adapter.search(
            "recreated",
            scopes=[{"userId": "alice", "agentName": "agent-a"}],
            top_k=5,
            mode="fts5",
            filters=None,
        )
    finally:
        adapter.close()


def test_adapter_isolates_scopes_even_when_fact_ids_repeat(tmp_path: Path) -> None:
    adapter = FTS5RetrievalAdapter(tmp_path / "facts.sqlite3")
    try:
        adapter.upsert(_fact("same", "Alice private alpha"), scope={"userId": "alice", "agentName": "__default__"}, path="")
        adapter.upsert(_fact("same", "Bob private beta"), scope={"userId": "bob", "agentName": "__default__"}, path="")

        alice = adapter.search("alpha", scopes=[{"userId": "alice", "agentName": "__default__"}], top_k=5, mode="hybrid", filters=None)
        bob = adapter.search("alpha", scopes=[{"userId": "bob", "agentName": "__default__"}], top_k=5, mode="hybrid", filters=None)

        assert [item["fact"]["content"] for item in alice] == ["Alice private alpha"]
        assert bob == []
    finally:
        adapter.close()


def test_adapter_persists_across_restart_and_supports_category_filter(tmp_path: Path) -> None:
    db_path = tmp_path / "facts.sqlite3"
    scope = {"userId": "alice", "agentName": "agent-a"}
    first = FTS5RetrievalAdapter(db_path)
    first.upsert(_fact("one", "Python deployment preference", category="preference"), scope=scope, path="")
    first.upsert(_fact("two", "Python deployment context", category="context"), scope=scope, path="")
    first.close()

    second = FTS5RetrievalAdapter(db_path)
    try:
        results = second.search("Python deployment", scopes=[scope], top_k=5, mode="hybrid", filters={"category": "preference"})
        assert [item["fact"]["id"] for item in results] == ["one"]
        assert results[0]["matchType"] == "fts5"
    finally:
        second.close()


def test_file_storage_incremental_notifications_update_and_remove_index(tmp_path: Path) -> None:
    adapter = FTS5RetrievalAdapter(tmp_path / "facts.sqlite3")
    storage = FileMemoryStorage(DeerMemConfig(storage_path=str(tmp_path)), retrieval=adapter)
    scope = {"userId": "alice", "agentName": "agent-a"}
    try:
        storage.upsert_fact(_fact("one", "initial searchable value"), user_id="alice", agent_name="agent-a", expected_manifest_revision=0)
        assert storage.search_facts("initial", scopes=[scope])[0]["fact"]["id"] == "one"

        updated = storage.get_fact("one", user_id="alice", agent_name="agent-a")
        assert updated is not None
        updated["content"] = "replacement searchable value"
        storage.upsert_fact(
            updated,
            user_id="alice",
            agent_name="agent-a",
            expected_manifest_revision=1,
            expected_fact_revision=updated["revision"],
        )
        assert storage.search_facts("replacement", scopes=[scope])[0]["fact"]["id"] == "one"
        assert storage.search_facts("initial", scopes=[scope]) == []

        storage.delete_fact("one", user_id="alice", agent_name="agent-a")
        assert storage.search_facts("replacement", scopes=[scope]) == []
    finally:
        adapter.close()


def test_full_rebuild_removes_stale_rows_after_markdown_delete(tmp_path: Path) -> None:
    adapter = FTS5RetrievalAdapter(tmp_path / "facts.sqlite3")
    storage = FileMemoryStorage(DeerMemConfig(storage_path=str(tmp_path)), retrieval=adapter)
    scope = {"userId": "alice", "agentName": "agent-a"}
    try:
        storage.upsert_fact(_fact("one", "stale markdown value"), user_id="alice", agent_name="agent-a", expected_manifest_revision=0)
        memory_path = storage._get_memory_file_path("agent-a", user_id="alice")
        fact_path = next(memory_path.parent.glob("agents/agent-a/facts/**/*.md"))
        fact_path.unlink()

        result = storage.rebuild_index()
        assert result == {"supported": True, "indexed": 0, "failed": 0}
        assert storage.search_facts("stale", scopes=[scope]) == []
    finally:
        adapter.close()


def test_reload_rebuilds_index_after_out_of_band_markdown_edit(tmp_path: Path) -> None:
    adapter = FTS5RetrievalAdapter(tmp_path / "facts.sqlite3")
    storage = FileMemoryStorage(DeerMemConfig(storage_path=str(tmp_path)), retrieval=adapter)
    scope = {"userId": "alice", "agentName": "agent-a"}
    try:
        storage.upsert_fact(_fact("one", "original markdown value"), user_id="alice", agent_name="agent-a", expected_manifest_revision=0)
        memory_path = storage._get_memory_file_path("agent-a", user_id="alice")
        fact_path = next(memory_path.parent.glob("agents/agent-a/facts/**/*.md"))
        raw = fact_path.read_text(encoding="utf-8").replace("original markdown value", "edited markdown value")
        fact_path.write_text(raw, encoding="utf-8")

        storage.reload("agent-a", user_id="alice")
        assert storage.search_facts("edited", scopes=[scope])[0]["fact"]["id"] == "one"
        assert storage.search_facts("original", scopes=[scope]) == []
    finally:
        adapter.close()


def test_failed_notification_marks_scope_dirty_and_rebuilds_on_search(tmp_path: Path) -> None:
    adapter = FTS5RetrievalAdapter(tmp_path / "facts.sqlite3")
    storage = FileMemoryStorage(DeerMemConfig(storage_path=str(tmp_path)), retrieval=adapter)
    original_upsert = adapter.upsert
    failed_once = True

    def fail_once(fact, *, scope, path):
        nonlocal failed_once
        if failed_once:
            failed_once = False
            raise RuntimeError("simulated index outage")
        original_upsert(fact, scope=scope, path=path)

    adapter.upsert = fail_once  # type: ignore[method-assign]
    try:
        storage.upsert_fact(_fact("one", "recoverable indexed value"), user_id="alice", agent_name="agent-a", expected_manifest_revision=0)
        results = storage.search_facts("recoverable", scopes=[{"userId": "alice", "agentName": "agent-a"}])
        assert results[0]["fact"]["id"] == "one"
    finally:
        adapter.close()


def test_concurrent_upserts_are_searchable(tmp_path: Path) -> None:
    adapter = FTS5RetrievalAdapter(tmp_path / "facts.sqlite3")
    scope = {"userId": "alice", "agentName": "agent-a"}
    try:

        def write(index: int) -> None:
            adapter.upsert(_fact(f"fact-{index}", f"concurrent memory item {index}"), scope=scope, path="")

        with ThreadPoolExecutor(max_workers=8) as pool:
            list(pool.map(write, range(20)))
        results = adapter.search("concurrent memory", scopes=[scope], top_k=20, mode="hybrid", filters=None)
        assert {item["fact"]["id"] for item in results} == {f"fact-{index}" for index in range(20)}
    finally:
        adapter.close()


def test_deermem_create_and_restart_use_retrieval_adapter(tmp_path: Path) -> None:
    config = {"storage_path": str(tmp_path), "token_counting": "char"}
    manager = DeerMem(backend_config=config)
    _, fact_id = manager.create_fact("BM25 retrieval survives restart", category="knowledge", user_id="alice")
    assert fact_id is not None
    assert any(fact["id"] == fact_id for fact in manager.search("retrieval BM25", user_id="alice"))

    restarted = DeerMem(backend_config=config)
    assert any(fact["id"] == fact_id for fact in restarted.search("restart retrieval", user_id="alice"))


def test_indexed_zero_confidence_keeps_its_retrieval_weight(tmp_path: Path) -> None:
    adapter = FTS5RetrievalAdapter(tmp_path / "facts.sqlite3")
    scope = {"userId": "alice", "agentName": "agent-a"}
    try:
        adapter.upsert({**_fact("zero", "confidence weight probe"), "confidence": 0.0}, scope=scope, path="")
        adapter.upsert({**_fact("half", "confidence weight probe"), "confidence": 0.5}, scope=scope, path="")
        unset_fact = _fact("unset", "confidence weight probe")
        del unset_fact["confidence"]
        adapter.upsert(unset_fact, scope=scope, path="")
        scores = {item["fact"]["id"]: item["score"] for item in adapter.search("confidence", scopes=[scope], top_k=5, mode="fts5", filters=None)}

        assert scores["zero"] < scores["half"]
        # A missing confidence still falls back to the default, so only the explicit 0.0 is discounted.
        assert scores["unset"] == pytest.approx(scores["half"])
    finally:
        adapter.close()


def test_substring_fallback_keeps_zero_confidence_weight(tmp_path: Path) -> None:
    storage = FileMemoryStorage(DeerMemConfig(storage_path=str(tmp_path)))
    scope = {"userId": "alice", "agentName": "agent-a"}
    try:
        storage.upsert_fact({**_fact("zero", "substring confidence probe"), "confidence": 0.0}, user_id="alice", agent_name="agent-a")
        storage.upsert_fact({**_fact("half", "substring confidence probe"), "confidence": 0.5}, user_id="alice", agent_name="agent-a")
        scores = {item["fact"]["id"]: item["score"] for item in storage.search_facts("substring", scopes=[scope], top_k=5)}

        assert scores == {"zero": 0.0, "half": 0.5}
    finally:
        storage.close()


def test_substring_fallback_searches_every_fact_of_a_large_scope(tmp_path: Path) -> None:
    """The fallback must read the whole scope, not the first page of 100 facts.

    With 150 facts (max_facts allows up to 500) and no retrieval adapter, a
    top_k of 200 returned only the first 100 matches: the fallback iterated
    list_facts' default page instead of the complete scope, so facts past
    the page silently vanished from search results.
    """
    storage = FileMemoryStorage(DeerMemConfig(storage_path=str(tmp_path), max_facts=200))
    scope = {"userId": "alice", "agentName": "agent-a"}
    try:
        facts = [_fact(f"fact-{index:03d}", f"zeta common memory number {index:03d}") for index in range(1, 151)]
        storage.apply_changes({"upserts": facts}, user_id="alice", agent_name="agent-a")

        results = storage.search_facts("zeta", scopes=[scope], top_k=200)

        assert len(results) == 150
        returned_ids = {item["fact"]["id"] for item in results}
        # fact-001..fact-100 sat inside the old first page too (``matched[0:100]``
        # in insertion order), so fact-101 and fact-150 are the ids only a
        # complete-scope read can return.
        assert {"fact-101", "fact-150"} <= returned_ids
    finally:
        storage.close()


@pytest.mark.parametrize("top_k", [1, 5])
@pytest.mark.parametrize("reversal", ["confidence", "time_decay"])
def test_engine_selects_composite_top_k_before_truncating(top_k: int, reversal: str, caplog: pytest.LogCaptureFixture) -> None:
    from datetime import UTC, datetime, timedelta

    now = datetime.now(UTC)
    recent = now.isoformat()
    older = (now - timedelta(days=3650)).isoformat()
    engine = FTS5Retrieval()
    try:
        for index in range(30):
            engine.index_fact(
                f"short-{index:02d}",
                f"Python Python tool{index:02d} reminder",
                confidence=0.7 if reversal == "confidence" else 0.8,
                created_at=recent if reversal == "confidence" else older,
            )
        engine.index_fact(
            "best",
            "Python deployment preference uses project virtual environments and pinned dependencies",
            confidence=0.95 if reversal == "confidence" else 0.8,
            created_at=recent,
        )
        complete = engine.search("Python", top_k=31)
        assert len(complete) == 31
        assert complete[0]["id"] == "best"
        # The longer fact has weaker BM25 but the highest composite score.
        assert complete[0]["bm25_score"] < complete[1]["bm25_score"]
        assert complete[0]["score"] > complete[1]["score"]

        with caplog.at_level("DEBUG", logger=FTS5Retrieval.__module__):
            limited = engine.search("Python", top_k=top_k)
        assert [fact["id"] for fact in limited] == [fact["id"] for fact in complete[:top_k]]
        assert any(f"matched=31 retained={top_k}" in message for message in caplog.messages)
    finally:
        engine.close()


@pytest.mark.parametrize("category", [None, "preference"])
@pytest.mark.parametrize("top_k", [1, 5])
def test_adapter_composite_ranking_preserves_category_and_limit(category: str | None, top_k: int) -> None:
    from datetime import UTC, datetime

    adapter = FTS5RetrievalAdapter()
    scope = {"userId": "alice", "agentName": "agent-a"}
    recent = datetime.now(UTC).isoformat()
    try:
        for index in range(45):
            adapter.upsert(
                {**_fact(f"short-{index:02d}", f"Python Python tool{index:02d} reminder", category="preference"), "confidence": 0.7, "createdAt": recent},
                scope=scope,
                path="",
            )
        adapter.upsert(
            {**_fact("best", "Python deployment preference uses project virtual environments and pinned dependencies", category="preference"), "confidence": 0.95, "createdAt": recent},
            scope=scope,
            path="",
        )
        adapter.upsert({**_fact("other-category", "Python", category="context"), "confidence": 1.0, "createdAt": recent}, scope=scope, path="")
        adapter.upsert({**_fact("nonmatch", "Coffee brewing preference", category="preference"), "confidence": 1.0, "createdAt": recent}, scope=scope, path="")
        filters = {"category": category} if category else None
        complete = adapter.search("Python", scopes=[scope], top_k=100, mode="fts5", filters=filters)
        assert complete[0]["fact"]["id"] == ("best" if category else "other-category")
        limited = adapter.search("Python", scopes=[scope], top_k=top_k, mode="fts5", filters=filters)
        assert [row["fact"]["id"] for row in limited] == [row["fact"]["id"] for row in complete[:top_k]]
        assert all(row["matchType"] == "fts5" for row in limited)
        assert all("Python" in row["fact"]["content"] for row in limited)
        if category:
            assert all(row["fact"]["category"] == category for row in limited)
    finally:
        adapter.close()


def test_engine_composite_score_ties_keep_bm25_and_insertion_order() -> None:
    from datetime import UTC, datetime

    engine = FTS5Retrieval()
    recent = datetime.now(UTC).isoformat()
    try:
        for fact_id in ("z-first", "a-second", "middle"):
            engine.index_fact(fact_id, "Python same content", confidence=0.8, created_at=recent)
        for top_k in (1, 2, 3, 10):
            assert [fact["id"] for fact in engine.search("Python", top_k=top_k)] == ["z-first", "a-second", "middle"][:top_k]
    finally:
        engine.close()


def test_deermem_default_tool_limit_keeps_best_composite_fact(tmp_path: Path) -> None:
    manager = DeerMem(backend_config={"storage_path": str(tmp_path), "token_counting": "char"})
    try:
        for index in range(90):
            manager.create_fact(f"Python Python tool{index:02d} reminder", confidence=0.7, user_id="alice")
        _, best_id = manager.create_fact(
            "Python deployment preference uses project virtual environments and pinned dependencies",
            confidence=0.95,
            category="preference",
            user_id="alice",
        )
        complete = manager.search("Python", top_k=100, user_id="alice")
        assert len(complete) == 91
        assert complete[0]["id"] == best_id
        # memory_search's default limit is 10. Changing only the requested
        # count must not decide whether the highest-scoring fact is reachable.
        for top_k in (10, 12):
            assert [fact["id"] for fact in manager.search("Python", top_k=top_k, user_id="alice")] == [fact["id"] for fact in complete[:top_k]]
    finally:
        manager.close()