From 8a4d488a726dbdbaa837e04292bffebefedd9119 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Tue, 18 Aug 2026 14:20:45 +0300 Subject: [PATCH] Give the MCP server one client for its lifetime All ten tool bodies opened their own `HaikuRAG`, so every tool call paid a connection open and, on object storage, refetched the index the previous call had just cached. The client is now opened once, lazily so that calling a tool function directly still works, and eagerly from the lifespan so an unopenable database fails startup instead of every call. Teardown clears the cached client in a finally, since `_lifespan_manager` can be re-entered and would otherwise hand out a closed connection, including when the close itself fails. `delete_document` no longer opens its own connection with `skip_validation=True`. Keeping it separate broke consistency once connections became long-lived: the delete committed on one connection while reads served from another, which with a 30s consistency interval showed the deleted document as still present. A connection always sees its own writes, so sharing one is what makes delete visible to the next read. So the server no longer opts out of embedding-config validation. Drift that validation rejects now fails MCP startup, where before the server started and only `delete_document` worked while every read returned empty. Same-dimension identity drift still starts a read-only server, matching every other read verb. Delete under drift is now a CLI operation; CLAUDE.md and the CHANGELOG record it. --- CHANGELOG.md | 1 + haiku_rag_slim/haiku/rag/mcp.py | 142 +++++++++++++++++++------------ tests/test_mcp.py | 144 ++++++++++++++++++++++++++++++++ 3 files changed, 232 insertions(+), 55 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index c63d1446..310f5ecd 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -11,6 +11,7 @@ ### Changed - Eval judge pinned to `qwen3.8`: `DEFAULT_JUDGE_MODEL` is `ollama:qwen3.8`, and the reference configs use `Inferact/Qwen3.8-27B-NVFP4` with `extra_body.chat_template_kwargs.reasoning_effort: low`. Results in `docs/benchmarks.md` were judged by `Qwen3.6-35B-A3B-NVFP4` and are not re-judged. +- The MCP server opens one database client for its lifetime instead of one per tool call, and fails at startup if the database cannot be opened. `delete_document` no longer opens its own connection with `skip_validation=True`, so the server no longer opts out of embedding-config validation: drift that validation rejects (any `vector_dim` mismatch, or identity drift on a writable server) now fails MCP startup instead of serving a delete-only server. Use the CLI to delete under drift. - `import_documents` embeds chunks across the whole batch in one pass instead of per document. - `haiku.rag` and `haiku.rag-slim` summaries and keywords; both packages now publish `[project.urls]`. - `server.json` declares `title` and `websiteUrl`, and drops the `keywords` and `license` keys, which are not in the server schema. diff --git a/haiku_rag_slim/haiku/rag/mcp.py b/haiku_rag_slim/haiku/rag/mcp.py index 30c2ed11..a8a44b50 100644 --- a/haiku_rag_slim/haiku/rag/mcp.py +++ b/haiku_rag_slim/haiku/rag/mcp.py @@ -1,3 +1,6 @@ +import asyncio +from collections.abc import AsyncIterator +from contextlib import AsyncExitStack, asynccontextmanager from pathlib import Path from typing import Any @@ -28,7 +31,42 @@ def create_mcp_server( config: Configuration to use. read_only: If True, write tools (add_document_*, delete_document) are not registered. """ - mcp = FastMCP("haiku-rag") + client: HaikuRAG | None = None + stack = AsyncExitStack() + client_lock = asyncio.Lock() + + async def _client() -> HaikuRAG: + """The server's client, opened once. + + Opening cost is per connection, and on object storage the first vector + query loads the index into the session cache, so a client per tool call + pays that repeatedly. + """ + nonlocal client + async with client_lock: + if client is None: + client = await stack.enter_async_context( + HaikuRAG(db_path, config=config, read_only=read_only) + ) + return client + + @asynccontextmanager + async def lifespan(_server: FastMCP) -> AsyncIterator[None]: + # Open eagerly so an unopenable database fails startup rather than + # every tool call. + nonlocal client + await _client() + try: + yield + finally: + # The lifespan can be re-entered; without the reset the next cycle + # hands out the closed client, including when aclose itself fails. + try: + await stack.aclose() + finally: + client = None + + mcp = FastMCP("haiku-rag", lifespan=lifespan) # Write tools - only registered when not in read-only mode if not read_only: @@ -41,14 +79,14 @@ def create_mcp_server( ) -> str | None: """Add a document to the RAG system from a file path.""" try: - async with HaikuRAG(db_path, config=config) as rag: - result = await rag.create_document_from_source( - Path(file_path), title=title, metadata=metadata or {} - ) - # Handle both single document and list of documents (directories) - if isinstance(result, list): - return result[0].id if result else None - return result.id + rag = await _client() + result = await rag.create_document_from_source( + Path(file_path), title=title, metadata=metadata or {} + ) + # Handle both single document and list of documents (directories) + if isinstance(result, list): + return result[0].id if result else None + return result.id except Exception: return None @@ -58,14 +96,14 @@ def create_mcp_server( ) -> str | None: """Add a document to the RAG system from a URL.""" try: - async with HaikuRAG(db_path, config=config) as rag: - result = await rag.create_document_from_source( - url, title=title, metadata=metadata or {} - ) - # Handle both single document and list of documents - if isinstance(result, list): - return result[0].id if result else None - return result.id + rag = await _client() + result = await rag.create_document_from_source( + url, title=title, metadata=metadata or {} + ) + # Handle both single document and list of documents + if isinstance(result, list): + return result[0].id if result else None + return result.id except Exception: return None @@ -78,11 +116,11 @@ def create_mcp_server( ) -> str | None: """Add a document to the RAG system from text content.""" try: - async with HaikuRAG(db_path, config=config) as rag: - document = await rag.create_document( - content, uri, title=title, metadata=metadata or {} - ) - return document.id + rag = await _client() + document = await rag.create_document( + content, uri, title=title, metadata=metadata or {} + ) + return document.id except Exception: return None @@ -90,10 +128,8 @@ def create_mcp_server( async def delete_document(document_id: str) -> bool: """Delete a document by its ID.""" try: - async with HaikuRAG( - db_path, config=config, skip_validation=True - ) as rag: - return await rag.delete_document(document_id) + rag = await _client() + return await rag.delete_document(document_id) except Exception: return False @@ -110,10 +146,8 @@ def create_mcp_server( response (smaller JSON payload for plain-text consumers). """ try: - async with HaikuRAG(db_path, config=config, read_only=read_only) as rag: - return await rag.search( - query, limit=limit, include_images=include_images - ) + rag = await _client() + return await rag.search(query, limit=limit, include_images=include_images) except Exception: return [] @@ -145,10 +179,8 @@ def create_mcp_server( except Exception: return [] try: - async with HaikuRAG(db_path, config=config, read_only=read_only) as rag: - return await rag.search( - raw, limit=limit, include_images=include_images - ) + rag = await _client() + return await rag.search(raw, limit=limit, include_images=include_images) except Exception: return [] @@ -156,8 +188,8 @@ def create_mcp_server( async def get_document(document_id: str) -> Document | None: """Get a document by its ID.""" try: - async with HaikuRAG(db_path, config=config, read_only=read_only) as rag: - return await rag.get_document_by_id(document_id) + rag = await _client() + return await rag.get_document_by_id(document_id) except Exception: return None @@ -175,18 +207,18 @@ def create_mcp_server( filter: Optional SQL WHERE clause to filter documents. """ try: - async with HaikuRAG(db_path, config=config, read_only=read_only) as rag: - documents = await rag.list_documents(limit, offset, filter) + rag = await _client() + documents = await rag.list_documents(limit, offset, filter) - return [ - DocumentInfo( - id=doc.id, - title=doc.title or "Untitled", - uri=doc.uri or "", - created=doc.created_at.strftime("%Y-%m-%d"), - ) - for doc in documents - ] + return [ + DocumentInfo( + id=doc.id, + title=doc.title or "Untitled", + uri=doc.uri or "", + created=doc.created_at.strftime("%Y-%m-%d"), + ) + for doc in documents + ] except Exception: return [] @@ -209,11 +241,11 @@ def create_mcp_server( """ try: images = _decode_images(images_base64) - async with HaikuRAG(db_path, config=config, read_only=read_only) as rag: - answer, citations = await rag.ask(question, images=images) - if cite and citations: - answer += "\n\n" + format_citations(citations) - return answer + rag = await _client() + answer, citations = await rag.ask(question, images=images) + if cite and citations: + answer += "\n\n" + format_citations(citations) + return answer except Exception as e: return f"Error answering question: {e!s}" @@ -240,9 +272,9 @@ def create_mcp_server( """ try: images = _decode_images(images_base64) - async with HaikuRAG(db_path, config=config, read_only=read_only) as rag: - result = await rag.analyze(question, filter=filter, images=images) - return result.answer + rag = await _client() + result = await rag.analyze(question, filter=filter, images=images) + return result.answer except Exception as e: return f"Error running analysis capability: {e!s}" diff --git a/tests/test_mcp.py b/tests/test_mcp.py index f0c36474..9077586e 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -481,3 +481,147 @@ class TestMCPToolsDegradeOnError: assert "AI Overview" in with_cite assert await ask(question="q", cite=False) == "the answer" + + +class TestMCPClientLifetime: + @pytest.mark.asyncio + async def test_tool_calls_share_one_database_open(self, mcp_db, monkeypatch): + from haiku.rag.store.engine import Store + + opens = 0 + initialize = Store._initialize + + async def counted(self): + nonlocal opens + opens += 1 + return await initialize(self) + + monkeypatch.setattr(Store, "_initialize", counted) + + mcp = create_mcp_server(mcp_db, read_only=True) + search = await _get_tool(mcp, "search_documents") + list_docs = await _get_tool(mcp, "list_documents") + await search(query="artificial intelligence") + await list_docs() + await search(query="machine learning") + + assert opens == 1 + + @pytest.mark.asyncio + async def test_concurrent_reads_share_one_open(self, mcp_db, monkeypatch): + import asyncio + + from haiku.rag.store.engine import Store + + opens = 0 + initialize = Store._initialize + + async def counted(self): + nonlocal opens + opens += 1 + return await initialize(self) + + monkeypatch.setattr(Store, "_initialize", counted) + + mcp = create_mcp_server(mcp_db, read_only=True) + list_docs = await _get_tool(mcp, "list_documents") + + results = await asyncio.gather(*(list_docs() for _ in range(5))) + + assert opens == 1 + assert all(len(r) == 2 for r in results) + + @pytest.mark.asyncio + async def test_a_write_is_visible_to_the_next_read(self, mcp_db): + """One connection sees its own writes, whatever the consistency interval.""" + mcp = create_mcp_server(mcp_db, read_only=False) + list_docs = await _get_tool(mcp, "list_documents") + delete_doc = await _get_tool(mcp, "delete_document") + + docs = await list_docs() + assert await delete_doc(document_id=docs[0].id) is True + + assert len(await list_docs()) == len(docs) - 1 + + @pytest.mark.asyncio + async def test_lifespan_opens_and_closes_once(self, mcp_db, monkeypatch): + from haiku.rag.store.engine import Store + + opens = 0 + initialize = Store._initialize + + async def counted(self): + nonlocal opens + opens += 1 + return await initialize(self) + + monkeypatch.setattr(Store, "_initialize", counted) + + mcp = create_mcp_server(mcp_db, read_only=True) + # _lifespan_manager is what every transport enters; the public + # lifespan() combines provider lifespans only. + async with mcp._lifespan_manager(): + assert opens == 1, "startup should open the database, not the first call" + search = await _get_tool(mcp, "search_documents") + await search(query="artificial intelligence") + assert opens == 1 + + assert opens == 1 + + @pytest.mark.asyncio + async def test_startup_fails_when_the_database_cannot_open(self, tmp_path): + mcp = create_mcp_server(tmp_path / "does-not-exist.lancedb", read_only=True) + + with pytest.raises(FileNotFoundError): + async with mcp._lifespan_manager(): + pass + + @pytest.mark.asyncio + async def test_a_second_lifespan_cycle_opens_a_fresh_client( + self, mcp_db, monkeypatch + ): + from haiku.rag.store.engine import Store + + opens = 0 + initialize = Store._initialize + + async def counted(self): + nonlocal opens + opens += 1 + return await initialize(self) + + monkeypatch.setattr(Store, "_initialize", counted) + + mcp = create_mcp_server(mcp_db, read_only=True) + search = await _get_tool(mcp, "search_documents") + + async with mcp._lifespan_manager(): + await search(query="artificial intelligence") + assert opens == 1 + + async with mcp._lifespan_manager(): + results = await search(query="artificial intelligence") + assert opens == 2 + assert len(results) > 0 + + @pytest.mark.asyncio + async def test_same_dim_drift_starts_read_only_but_not_writable(self, mcp_db): + """Validation is unchanged: same-dimension identity drift warns in + read-only mode and raises in writable mode. The MCP server no longer + opts out of it for deletion.""" + from haiku.rag.config import Config + from haiku.rag.store.repositories.settings import ConfigMismatchError + + drifted = Config.model_copy(deep=True) + drifted.embeddings.model.name = "a-different-model" + + async with create_mcp_server( + mcp_db, config=drifted, read_only=True + )._lifespan_manager(): + pass + + with pytest.raises(ConfigMismatchError): + async with create_mcp_server( + mcp_db, config=drifted, read_only=False + )._lifespan_manager(): + pass