diff --git a/mempalace/mcp_server.py b/mempalace/mcp_server.py index 2169255..81c9b50 100644 --- a/mempalace/mcp_server.py +++ b/mempalace/mcp_server.py @@ -39,13 +39,21 @@ logger = logging.getLogger("mempalace_mcp") _config = MempalaceConfig() +_client_cache = None +_collection_cache = None + + def _get_collection(create=False): - """Return the ChromaDB collection, or None on failure.""" + """Return the ChromaDB collection, caching the client between calls.""" + global _client_cache, _collection_cache try: - client = chromadb.PersistentClient(path=_config.palace_path) + if _client_cache is None: + _client_cache = chromadb.PersistentClient(path=_config.palace_path) if create: - return client.get_or_create_collection(_config.collection_name) - return client.get_collection(_config.collection_name) + _collection_cache = _client_cache.get_or_create_collection(_config.collection_name) + elif _collection_cache is None: + _collection_cache = _client_cache.get_collection(_config.collection_name) + return _collection_cache except Exception: return None diff --git a/tests/conftest.py b/tests/conftest.py index 22b5e42..fa08802 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -34,6 +34,19 @@ from mempalace.config import MempalaceConfig # noqa: E402 from mempalace.knowledge_graph import KnowledgeGraph # noqa: E402 +@pytest.fixture(autouse=True) +def _reset_mcp_cache(): + """Reset the MCP server's cached ChromaDB client/collection between tests.""" + yield + try: + from mempalace import mcp_server + + mcp_server._client_cache = None + mcp_server._collection_cache = None + except (ImportError, AttributeError): + pass + + @pytest.fixture(scope="session", autouse=True) def _isolate_home(): """Ensure HOME points to a temp dir for the entire test session.