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
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

View file

@ -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()

View file

@ -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()

View file

@ -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")

View file

@ -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")

View file

@ -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):