From 5f795cb0cb12d5a1b8f28b3ed7cdfdc17d7c1084 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Thu, 27 Aug 2026 15:31:47 +0300 Subject: [PATCH] Let the MCP server be told which database once `create_mcp_server` promised one database and accepted a scope covering a set, where the write tools exist and fail on use. It refuses that now. Resolving is the public factory's job, as it is `HaikuRAG`'s: `_covering` takes a scope someone already resolved, so the configured name survives without a `DatabaseScope` reaching the public signature. The test that a scope decides the database asserted `all(...)` over a search that could return nothing, which held whatever the server read. It reads the listing instead, so alpha's documents being present and beta's absent both have to be true. --- docs/python.md | 8 +++-- haiku_rag_slim/haiku/rag/app.py | 6 ++-- haiku_rag_slim/haiku/rag/mcp.py | 26 ++++++++++++---- tests/multi_db/test_capabilities.py | 2 +- tests/test_app.py | 6 ++-- tests/test_mcp.py | 47 ++++++++++++++++++++++++----- 6 files changed, 71 insertions(+), 24 deletions(-) diff --git a/docs/python.md b/docs/python.md index 16ce527f..53050c4c 100644 --- a/docs/python.md +++ b/docs/python.md @@ -261,9 +261,11 @@ their documents. none: `search` returns no results, and `ask` and `analyze` run with no evidence from any database. -On the constructor, `sources` is rejected alongside a database path, and `[]` -raises `ValueError`. A selection of nothing to search is a legitimate question; -a client over no database is not. +On the constructor it means something else. Passing `sources` alongside a +database path raises `AmbiguousDatabaseError` immediately, since both say which +database to open. Passing `sources=[]` alone raises `ValueError` on entering the +client: a selection of nothing to search is a legitimate question, a client over +no database is not. #### Inspecting the client scope diff --git a/haiku_rag_slim/haiku/rag/app.py b/haiku_rag_slim/haiku/rag/app.py index 998023bd..54ced988 100644 --- a/haiku_rag_slim/haiku/rag/app.py +++ b/haiku_rag_slim/haiku/rag/app.py @@ -18,7 +18,7 @@ from rich.progress import ( from haiku.rag.client import HaikuRAG, RebuildMode from haiku.rag.config import AppConfig, get_config -from haiku.rag.mcp import create_mcp_server +from haiku.rag.mcp import _covering as _mcp_server_covering from haiku.rag.store.models.chunk import SearchType from haiku.rag.store.models.document import Document @@ -926,9 +926,7 @@ class HaikuRAGApp: # 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( - config=self.config, read_only=self.read_only, scope=self.scope - ) + server = _mcp_server_covering(self.scope, self.config, self.read_only) try: if transport == "stdio": await server.run_stdio_async() diff --git a/haiku_rag_slim/haiku/rag/mcp.py b/haiku_rag_slim/haiku/rag/mcp.py index edcdac02..727bd3f7 100644 --- a/haiku_rag_slim/haiku/rag/mcp.py +++ b/haiku_rag_slim/haiku/rag/mcp.py @@ -28,7 +28,6 @@ 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 over one database. @@ -38,15 +37,30 @@ 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) + return _covering( + DatabaseScope.resolve(config, database_path=db_path), config, read_only + ) + + +def _covering(scope: "DatabaseScope", config: AppConfig, read_only: bool) -> FastMCP: + """An MCP server over databases someone already resolved. + + Internal, as ``HaikuRAG._covering`` is: the public factory takes a path and + resolves it, which is its own job. A caller that resolved already passes the + scope, so the configured name survives, which results and citations carry as + ``source``. + """ + from haiku.rag.store.exceptions import AmbiguousDatabaseError + + if scope.covers_multiple: + raise AmbiguousDatabaseError( + "an MCP server serves one database, and this scope covers " + f"{', '.join(scope.names)}; name the one to serve" + ) client: HaikuRAG | None = None stack = AsyncExitStack() client_lock = asyncio.Lock() diff --git a/tests/multi_db/test_capabilities.py b/tests/multi_db/test_capabilities.py index 2ec676fb..4bb37e2d 100644 --- a/tests/multi_db/test_capabilities.py +++ b/tests/multi_db/test_capabilities.py @@ -226,7 +226,7 @@ class TestLendingANamedClient: run = await capability.for_run(make_context(deps)) assert run.state is not None - # The lent client is what it reads through, whatever it was built with. + # A borrowed client overrides the capability's configured placement. assert await run._ensure_rag() is client run.state.searches["cats"] = await client.search("cats", search_type="fts") diff --git a/tests/test_app.py b/tests/test_app.py index 6714f841..27c278f8 100644 --- a/tests/test_app.py +++ b/tests/test_app.py @@ -620,7 +620,7 @@ def test_search_result_rendering_includes_provenance(app): async def test_run_mcp_stdio(app, client, monkeypatch): server = AsyncMock() - monkeypatch.setattr("haiku.rag.app.create_mcp_server", lambda *a, **kw: server) + monkeypatch.setattr("haiku.rag.app._mcp_server_covering", lambda *a, **kw: server) await app.run_mcp(transport="stdio") @@ -629,7 +629,7 @@ async def test_run_mcp_stdio(app, client, monkeypatch): async def test_run_mcp_http(app, client, monkeypatch): server = AsyncMock() - monkeypatch.setattr("haiku.rag.app.create_mcp_server", lambda *a, **kw: server) + monkeypatch.setattr("haiku.rag.app._mcp_server_covering", lambda *a, **kw: server) await app.run_mcp(host="0.0.0.0", port=9001) @@ -641,7 +641,7 @@ async def test_run_mcp_http(app, client, monkeypatch): async def test_run_mcp_survives_interruption(app, client, monkeypatch): server = AsyncMock() server.run_stdio_async.side_effect = KeyboardInterrupt - monkeypatch.setattr("haiku.rag.app.create_mcp_server", lambda *a, **kw: server) + monkeypatch.setattr("haiku.rag.app._mcp_server_covering", lambda *a, **kw: server) await app.run_mcp(transport="stdio") diff --git a/tests/test_mcp.py b/tests/test_mcp.py index 84167aac..60f105d1 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -1,6 +1,7 @@ import pytest from haiku.rag.client import HaikuRAG +from haiku.rag.mcp import _covering as _mcp_covering from haiku.rag.mcp import create_mcp_server from haiku.rag.store.models import Chunk, Document, SearchResult from haiku.rag.tools.document import DocumentInfo @@ -624,15 +625,46 @@ class TestMCPClientLifetime: ) scope = DatabaseScope.resolve(config, database_name="alpha") - mcp = create_mcp_server(config=config, read_only=True, scope=scope) + mcp = _mcp_covering(scope, config, read_only=True) async with mcp._lifespan_manager(): search = await _get_tool(mcp, "search_documents") results = await search(query="artificial intelligence") - zebras = await search(query="zebras savannah") + listing = await _get_tool(mcp, "list_documents") + documents = await listing() assert results assert {r.source for r in results} == {"alpha"} - assert all(r.source == "alpha" for r in zebras) + titles = {d.title for d in documents} + assert "AI Overview" in titles + assert "Zebras" not in titles + + def test_a_scope_covering_a_set_is_refused(self, tmp_path): + from haiku.rag.client.scope import DatabaseScope + from haiku.rag.config.models import AppConfig, LanceDBConfig + from haiku.rag.store.exceptions import AmbiguousDatabaseError + + config = AppConfig( + lancedb=LanceDBConfig( + databases={"alpha": str(tmp_path / "a"), "beta": str(tmp_path / "b")} + ) + ) + + with pytest.raises(AmbiguousDatabaseError, match="alpha, beta"): + _mcp_covering(DatabaseScope.resolve(config), config, read_only=True) + + def test_the_public_factory_refuses_a_configured_set_too(self, tmp_path): + """It resolves the same scope, so it reaches the same refusal.""" + from haiku.rag.config.models import AppConfig, LanceDBConfig + from haiku.rag.store.exceptions import AmbiguousDatabaseError + + config = AppConfig( + lancedb=LanceDBConfig( + databases={"alpha": str(tmp_path / "a"), "beta": str(tmp_path / "b")} + ) + ) + + with pytest.raises(AmbiguousDatabaseError, match="alpha, beta"): + create_mcp_server(config=config) @pytest.mark.asyncio async def test_the_command_hands_the_server_its_resolved_database( @@ -653,21 +685,22 @@ class TestMCPClientLifetime: async def run_stdio_async(self): return None - def fake_create(db_path=None, config=None, read_only=False, scope=None): - seen.update(db_path=db_path, scope=scope) + def fake_covering(scope, config, read_only): + seen.update(scope=scope, config=config) return _Server() - monkeypatch.setattr("haiku.rag.app.create_mcp_server", fake_create) + monkeypatch.setattr("haiku.rag.app._mcp_server_covering", fake_covering) app = HaikuRAGApp( scope=DatabaseScope.resolve(config, database_name="prod"), config=config ) await app.run_mcp(transport="stdio") - assert seen["db_path"] is None [ref] = seen["scope"].databases assert ref.name == "prod" assert ref.uri == "s3://bucket/prod.lancedb" + # The caller's configuration, not one derived from the ref. + assert seen["config"].lancedb.databases == {"prod": "s3://bucket/prod.lancedb"} @pytest.mark.asyncio async def test_startup_fails_when_the_database_cannot_open(self, tmp_path):