From db7f0e0af6855bfc68d3fc3ba4c6ddfdf21b423f Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Thu, 27 Aug 2026 15:12:13 +0300 Subject: [PATCH] Give the MCP server the database, not a description of it `run_mcp` derived a path and a configuration for the server, and deriving drops the configured name: MCP results and citations carried `source=None` where the same database named through any other path carried "alpha". `create_mcp_server` takes the resolved scope and opens through it, so the name survives. Passing a path and configuration still works and resolves to the same place. --- haiku_rag_slim/haiku/rag/app.py | 9 +++---- haiku_rag_slim/haiku/rag/mcp.py | 17 +++++++++--- tests/test_mcp.py | 48 ++++++++++++++++++++++++++------- 3 files changed, 56 insertions(+), 18 deletions(-) diff --git a/haiku_rag_slim/haiku/rag/app.py b/haiku_rag_slim/haiku/rag/app.py index 95e99fc5..998023bd 100644 --- a/haiku_rag_slim/haiku/rag/app.py +++ b/haiku_rag_slim/haiku/rag/app.py @@ -923,12 +923,11 @@ class HaikuRAGApp: The server opens its own client and validates it on startup, so nothing is opened here first. """ - # The ref's own path, not `_path`: a URI-backed database has none, and - # the local stand-in would override the URI in `_store_config`. + # The resolved scope, not a derived path and configuration: a path + # would override a configured URI, and deriving drops the name results + # and citations carry. server = create_mcp_server( - self._one.db_path, - config=self._store_config, - read_only=self.read_only, + config=self.config, read_only=self.read_only, scope=self.scope ) try: if transport == "stdio": diff --git a/haiku_rag_slim/haiku/rag/mcp.py b/haiku_rag_slim/haiku/rag/mcp.py index 929b2d62..edcdac02 100644 --- a/haiku_rag_slim/haiku/rag/mcp.py +++ b/haiku_rag_slim/haiku/rag/mcp.py @@ -2,7 +2,7 @@ import asyncio from collections.abc import AsyncIterator from contextlib import AsyncExitStack, asynccontextmanager from pathlib import Path -from typing import Any +from typing import TYPE_CHECKING, Any from fastmcp import FastMCP @@ -12,6 +12,9 @@ from haiku.rag.store.models import Document, SearchResult from haiku.rag.tools.document import DocumentInfo from haiku.rag.utils import format_citations +if TYPE_CHECKING: + from haiku.rag.client.scope import DatabaseScope + def _decode_images(images_base64: list[str] | None) -> list[bytes] | None: if not images_base64: @@ -25,8 +28,9 @@ def create_mcp_server( db_path: Path | None = None, config: AppConfig | None = None, read_only: bool = False, + scope: "DatabaseScope | None" = None, ) -> FastMCP: - """Create an MCP server with the specified database path. + """Create an MCP server over one database. Args: db_path: Path to the database file, or None to let `config` place it. A @@ -34,8 +38,15 @@ def create_mcp_server( must pass None rather than a local stand-in. config: Configuration to use. read_only: If True, write tools (add_document_*, delete_document) are not registered. + scope: The database, already resolved. Pass this rather than a derived + path and configuration to keep its configured name, which results + and citations carry as `source`. """ + from haiku.rag.client.scope import DatabaseScope + config = config if config is not None else get_config() + if scope is None: + scope = DatabaseScope.resolve(config, database_path=db_path) client: HaikuRAG | None = None stack = AsyncExitStack() client_lock = asyncio.Lock() @@ -51,7 +62,7 @@ def create_mcp_server( async with client_lock: if client is None: client = await stack.enter_async_context( - HaikuRAG(db_path, config=config, read_only=read_only) + HaikuRAG._covering(scope, config, read_only=read_only) ) return client diff --git a/tests/test_mcp.py b/tests/test_mcp.py index 6835d920..84167aac 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -604,12 +604,42 @@ class TestMCPClientLifetime: assert opens == 1 @pytest.mark.asyncio - async def test_a_uri_backed_database_is_not_replaced_by_a_local_path( + async def test_the_scope_decides_the_database_and_names_its_results( + self, mcp_db, tmp_path + ): + """The scope is the selection, so the server reads the one database it + names and results carry that name. The configuration alone would place + every database it configures.""" + from haiku.rag.client.scope import DatabaseScope + from haiku.rag.config.models import AppConfig, LanceDBConfig + + other = tmp_path / "beta.lancedb" + async with HaikuRAG(other, create=True) as rag: + await rag.create_document( + "Zebras graze on the savannah.", title="Zebras", uri="test://zebras" + ) + + config = AppConfig( + lancedb=LanceDBConfig(databases={"alpha": str(mcp_db), "beta": str(other)}) + ) + scope = DatabaseScope.resolve(config, database_name="alpha") + + mcp = create_mcp_server(config=config, read_only=True, scope=scope) + async with mcp._lifespan_manager(): + search = await _get_tool(mcp, "search_documents") + results = await search(query="artificial intelligence") + zebras = await search(query="zebras savannah") + + assert results + assert {r.source for r in results} == {"alpha"} + assert all(r.source == "alpha" for r in zebras) + + @pytest.mark.asyncio + async def test_the_command_hands_the_server_its_resolved_database( self, monkeypatch ): - """A path overrides `lancedb.uri`, so the server gets the database's own - path — None where a URI placed it — not the local stand-in a URI-backed - ref resolves to for display.""" + """A path would override a configured URI, and a derived configuration + would lose the name, so `run_mcp` passes neither.""" from haiku.rag.app import HaikuRAGApp from haiku.rag.client.scope import DatabaseScope from haiku.rag.config.models import AppConfig, LanceDBConfig @@ -623,9 +653,8 @@ class TestMCPClientLifetime: async def run_stdio_async(self): return None - def fake_create(db_path=None, config=None, read_only=False): - seen["db_path"] = db_path - seen["config"] = config + def fake_create(db_path=None, config=None, read_only=False, scope=None): + seen.update(db_path=db_path, scope=scope) return _Server() monkeypatch.setattr("haiku.rag.app.create_mcp_server", fake_create) @@ -636,9 +665,8 @@ class TestMCPClientLifetime: await app.run_mcp(transport="stdio") assert seen["db_path"] is None - [ref] = DatabaseScope.resolve( - seen["config"], database_path=seen["db_path"] - ).databases + [ref] = seen["scope"].databases + assert ref.name == "prod" assert ref.uri == "s3://bucket/prod.lancedb" @pytest.mark.asyncio