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.
This commit is contained in:
Yiorgis Gozadinos 2026-08-27 15:12:13 +03:00
parent d93cb5c891
commit db7f0e0af6
No known key found for this signature in database
3 changed files with 56 additions and 18 deletions

View file

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

View file

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

View file

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