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:
parent
acea8cbaac
commit
5f795cb0cb
6 changed files with 71 additions and 24 deletions
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
Loading…
Reference in a new issue