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.
This commit is contained in:
Yiorgis Gozadinos 2026-08-27 15:31:47 +03:00
parent acea8cbaac
commit 5f795cb0cb
No known key found for this signature in database
6 changed files with 71 additions and 24 deletions

View file

@ -261,9 +261,11 @@ their documents.
none: `search` returns no results, and `ask` and `analyze` run with no evidence none: `search` returns no results, and `ask` and `analyze` run with no evidence
from any database. from any database.
On the constructor, `sources` is rejected alongside a database path, and `[]` On the constructor it means something else. Passing `sources` alongside a
raises `ValueError`. A selection of nothing to search is a legitimate question; database path raises `AmbiguousDatabaseError` immediately, since both say which
a client over no database is not. 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 #### Inspecting the client scope

View file

@ -18,7 +18,7 @@ from rich.progress import (
from haiku.rag.client import HaikuRAG, RebuildMode from haiku.rag.client import HaikuRAG, RebuildMode
from haiku.rag.config import AppConfig, get_config 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.chunk import SearchType
from haiku.rag.store.models.document import Document 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 # The resolved scope, not a derived path and configuration: a path
# would override a configured URI, and deriving drops the name results # would override a configured URI, and deriving drops the name results
# and citations carry. # and citations carry.
server = create_mcp_server( server = _mcp_server_covering(self.scope, self.config, self.read_only)
config=self.config, read_only=self.read_only, scope=self.scope
)
try: try:
if transport == "stdio": if transport == "stdio":
await server.run_stdio_async() await server.run_stdio_async()

View file

@ -28,7 +28,6 @@ def create_mcp_server(
db_path: Path | None = None, db_path: Path | None = None,
config: AppConfig | None = None, config: AppConfig | None = None,
read_only: bool = False, read_only: bool = False,
scope: "DatabaseScope | None" = None,
) -> FastMCP: ) -> FastMCP:
"""Create an MCP server over one database. """Create an MCP server over one database.
@ -38,15 +37,30 @@ def create_mcp_server(
must pass None rather than a local stand-in. must pass None rather than a local stand-in.
config: Configuration to use. config: Configuration to use.
read_only: If True, write tools (add_document_*, delete_document) are not registered. 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 from haiku.rag.client.scope import DatabaseScope
config = config if config is not None else get_config() config = config if config is not None else get_config()
if scope is None: return _covering(
scope = DatabaseScope.resolve(config, database_path=db_path) 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 client: HaikuRAG | None = None
stack = AsyncExitStack() stack = AsyncExitStack()
client_lock = asyncio.Lock() client_lock = asyncio.Lock()

View file

@ -226,7 +226,7 @@ class TestLendingANamedClient:
run = await capability.for_run(make_context(deps)) run = await capability.for_run(make_context(deps))
assert run.state is not None 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 assert await run._ensure_rag() is client
run.state.searches["cats"] = await client.search("cats", search_type="fts") run.state.searches["cats"] = await client.search("cats", search_type="fts")

View file

@ -620,7 +620,7 @@ def test_search_result_rendering_includes_provenance(app):
async def test_run_mcp_stdio(app, client, monkeypatch): async def test_run_mcp_stdio(app, client, monkeypatch):
server = AsyncMock() 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") 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): async def test_run_mcp_http(app, client, monkeypatch):
server = AsyncMock() 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) 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): async def test_run_mcp_survives_interruption(app, client, monkeypatch):
server = AsyncMock() server = AsyncMock()
server.run_stdio_async.side_effect = KeyboardInterrupt 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") await app.run_mcp(transport="stdio")

View file

@ -1,6 +1,7 @@
import pytest import pytest
from haiku.rag.client import HaikuRAG 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.mcp import create_mcp_server
from haiku.rag.store.models import Chunk, Document, SearchResult from haiku.rag.store.models import Chunk, Document, SearchResult
from haiku.rag.tools.document import DocumentInfo from haiku.rag.tools.document import DocumentInfo
@ -624,15 +625,46 @@ class TestMCPClientLifetime:
) )
scope = DatabaseScope.resolve(config, database_name="alpha") 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(): async with mcp._lifespan_manager():
search = await _get_tool(mcp, "search_documents") search = await _get_tool(mcp, "search_documents")
results = await search(query="artificial intelligence") 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 results
assert {r.source for r in results} == {"alpha"} 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 @pytest.mark.asyncio
async def test_the_command_hands_the_server_its_resolved_database( async def test_the_command_hands_the_server_its_resolved_database(
@ -653,21 +685,22 @@ class TestMCPClientLifetime:
async def run_stdio_async(self): async def run_stdio_async(self):
return None return None
def fake_create(db_path=None, config=None, read_only=False, scope=None): def fake_covering(scope, config, read_only):
seen.update(db_path=db_path, scope=scope) seen.update(scope=scope, config=config)
return _Server() return _Server()
monkeypatch.setattr("haiku.rag.app.create_mcp_server", fake_create) monkeypatch.setattr("haiku.rag.app._mcp_server_covering", fake_covering)
app = HaikuRAGApp( app = HaikuRAGApp(
scope=DatabaseScope.resolve(config, database_name="prod"), config=config scope=DatabaseScope.resolve(config, database_name="prod"), config=config
) )
await app.run_mcp(transport="stdio") await app.run_mcp(transport="stdio")
assert seen["db_path"] is None
[ref] = seen["scope"].databases [ref] = seen["scope"].databases
assert ref.name == "prod" assert ref.name == "prod"
assert ref.uri == "s3://bucket/prod.lancedb" 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 @pytest.mark.asyncio
async def test_startup_fails_when_the_database_cannot_open(self, tmp_path): async def test_startup_fails_when_the_database_cannot_open(self, tmp_path):