diff --git a/tests/test_mcp_server.py b/tests/test_mcp_server.py index 953cd27..27f4251 100644 --- a/tests/test_mcp_server.py +++ b/tests/test_mcp_server.py @@ -9,6 +9,7 @@ via monkeypatch to avoid touching real data. from datetime import datetime import json import os +from types import SimpleNamespace import subprocess import sys from unittest.mock import MagicMock @@ -1068,6 +1069,38 @@ class TestSearchTool: assert "results" in result assert result.get("index_recovered") is True + def test_search_retry_preserves_collection_name(self, monkeypatch, config, kg): + """Retry path must query the same configured collection both times.""" + _patch_mcp_server(monkeypatch, config, kg) + from mempalace import mcp_server + + monkeypatch.setattr( + mcp_server, + "_config", + SimpleNamespace( + palace_path=config.palace_path, + collection_name="custom_drawers", + ), + ) + seen_collection_names = [] + + def fake_search(*args, **kwargs): + seen_collection_names.append(kwargs.get("collection_name")) + if len(seen_collection_names) == 1: + return { + "error": "Search error: Error executing plan: Internal error: Error finding id" + } + return {"results": [{"text": "ok", "wing": "w", "room": "r"}]} + + monkeypatch.setattr(mcp_server, "search_memories", fake_search) + monkeypatch.setattr(mcp_server, "_force_chroma_cache_reset", lambda: None) + monkeypatch.setattr(mcp_server.time, "sleep", lambda _: None) + + result = mcp_server.tool_search(query="anything", wing="wing_api") + + assert "results" in result + assert seen_collection_names == ["custom_drawers", "custom_drawers"] + def test_search_does_not_retry_on_non_transient_error(self, monkeypatch, config, kg): """Validation / unrelated errors must not trigger the retry path.""" _patch_mcp_server(monkeypatch, config, kg)