danvics/haiku.rag
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

One error contract for the MCP server

Browse source 
Every failure reaches the client as an MCP error carrying its message;
mask_error_details is set off explicitly, since FastMCP also reads it from
the environment. The masking goes, and with it the filter pre-check that
ran a count before every filtered call and the UnknownDatabaseError
translations that existed only to survive it. Explicit domain errors stay.
This commit is contained in:
Yiorgis Gozadinos 2026-09-07 13:14:16 +03:00
parent 95fdb46c3a
commit 4b1ec096c6
No known key found for this signature in database
5 changed files with 50 additions and 151 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

6
CHANGELOG.md
View file

@ -46,10 +46,8 @@
`ImageContent` blocks, with no structured content.
`SearchResult.format_for_agent(include_document_id=, include_chunk_meta=)`;
`collect_pictures` in `haiku.rag.tools.search`.
- MCP tools raise on failure; an empty result no longer doubles as an error.
Unknown document, unknown collection, invalid filter, invalid base64 and a
failing program carry a message. Anything else is masked
(`mask_error_details=True`) and logged server-side.
- MCP tools raise on failure, with the error's message; an empty result no
longer doubles as an error.
- `haiku-rag mcp` covers the configured `lancedb.databases` set. `sources` on
`search_documents`, `search_documents_by_image` and `execute_code`; `source`
on `get_document`; an unknown name is a tool error. `DocumentInfo.source`.

2
docs/configuration/qa.md
View file

@ -56,7 +56,7 @@ analysis:
```
- **model**: LLM configuration (see [Providers](providers.md#model-settings)). When unset, falls back to `qa.model`.
- **code_timeout**: Seconds a single `execute_code` call may spend reading documents (default: 60). The sandbox refuses further reads past this point. Code that computes without reading is killed by the worker watchdog at the same limit. `code_timeout * max_executions` is the cumulative ceiling across all calls in one question.
- **code_timeout**: Seconds a single `execute_code` call has (default: 60). Past it the sandbox starts no further host call, a document read or an in-code `search()` / `list_documents()`; one already running finishes. Code that computes without host calls is killed by the worker watchdog at the same limit. `code_timeout * max_executions` is the cumulative ceiling across all calls in one question.
- **max_output_chars**: Truncate code output after this many characters (default: 50000)
- **max_executions**: Maximum `execute_code` calls per question before the capability is told to answer from what it has (default: 15)

16
docs/mcp.md
View file

@ -31,8 +31,8 @@ The server opens the database read-only. Ingestion goes through the CLI
## Collections
With several databases in `lancedb.databases`, the server covers all of
them, as `haiku-rag search` does. Results, documents and citations name
theirs in `source`. `sources` on `search_documents`, `search_documents_by_image`
them, as `haiku-rag search` does. Results and documents name theirs in
`source`. `sources` on `search_documents`, `search_documents_by_image`
and `execute_code` restricts a call to a subset; `source` on `get_document` names the database holding the
document. A name the server does not cover is an error.
`haiku-rag --db-name NAME mcp` serves one. See
@ -190,13 +190,11 @@ title = 'Q3 report'
### Errors
A failure is an MCP error, never an empty result. Expected failures carry a
message: a document or section id that matches nothing, a collection the
server does not cover, a filter the query engine rejects (with its message),
invalid base64, and a program that fails in `execute_code`, with the error the
program hit.
Anything else reaches the client as `Error calling tool 'name'` and its
traceback goes to the server log.
A failure is an MCP error carrying its message, never an empty result: a
document or section id that matches nothing, a collection the server does not
cover, a filter the query engine rejects, invalid base64, a program that fails
in `execute_code` with the error it hit, and anything unexpected with its own
message.
### Instructions

108
haiku_rag_slim/haiku/rag/mcp.py
View file

@ -1,6 +1,5 @@
import asyncio
import logging
import re
import base64
from collections.abc import AsyncIterator
from contextlib import AsyncExitStack, asynccontextmanager
from importlib import metadata
@ -17,9 +16,7 @@ from haiku.rag.client import HaikuRAG
from haiku.rag.config import AppConfig, get_config
from haiku.rag.context import build_toc
from haiku.rag.sandbox import AnalysisContext, Sandbox, recovery_hint
from haiku.rag.store.exceptions import UnknownDatabaseError
from haiku.rag.store.models import Document, SearchResult
from haiku.rag.store.models.document_item import DocumentItem
from haiku.rag.store.schema import DocumentMetaRecord
from haiku.rag.tools.document import DocumentInfo, DocumentSection, OutlineNode
from haiku.rag.tools.search import collect_pictures
@ -28,8 +25,7 @@ if TYPE_CHECKING:
from typing import Any
from haiku.rag.client.scope import DatabaseScope
logger = logging.getLogger(__name__)
from haiku.rag.store.models.document_item import DocumentItem
_FILTER_COLUMNS = ", ".join(DocumentMetaRecord.model_fields)
@ -55,8 +51,6 @@ def _read_only(title: str) -> ToolAnnotations:
def _decode_image(image_base64: str) -> bytes:
import base64
try:
return base64.b64decode(image_base64, validate=True)
except ValueError as e:
@ -65,33 +59,6 @@ def _decode_image(image_base64: str) -> bytes:
raise ToolError("Invalid base64 image") from e
async def _check_filter(
rag: HaikuRAG, filter: str | None, sources: list[str] | None = None
) -> None:
"""Evaluate a filter on its own before the read that would use it.
A filtered count on one selected database runs the same predicate on the
same table and nothing else, so a ValueError here is the query engine
rejecting the filter; its message names columns and the statement, never
a location. A ValueError raised later in the read stays masked. Only the
selection is touched: every database shares the schema, so one suffices.
"""
if filter is None:
return
selected = await rag.clients_covering(sources)
if not selected:
return
try:
await selected[0].count_documents(filter=filter)
except ValueError as e:
# The engine lists its own columns too, lance internals among them.
reason = re.sub(r"\s*Valid fields are .*", "", str(e), flags=re.DOTALL)
raise ToolError(
f"Invalid filter {filter!r}: {reason.rstrip('. ')}. "
f"Columns: {_FILTER_COLUMNS}."
) from e
def _instructions(scope: "DatabaseScope", config: AppConfig) -> str:
"""What the server is for, naming no tools: the client has every tool's
description from the listing."""
@ -107,7 +74,7 @@ def _instructions(scope: "DatabaseScope", config: AppConfig) -> str:
if scope.covers_multiple:
lines.append(
f"It holds several collections: {', '.join(scope.names)}. Results "
"and citations name theirs in `source`; pass `sources` to use a subset."
"name theirs in `source`; pass `sources` to use a subset."
)
if config.prompts.domain_preamble:
lines.append(config.prompts.domain_preamble)
@ -119,8 +86,6 @@ def _search_result(results: list[SearchResult], covers_multiple: bool) -> ToolRe
metadata, then each distinct picture as an image block labelled with its
result. No structured content: a client given both shows the model the
JSON and drops the text, or shows both."""
import base64
total = len(results)
text = "\n\n".join(
result.format_for_agent(
@ -196,8 +161,7 @@ def _covering(scope: "DatabaseScope", config: AppConfig) -> FastMCP:
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``.
scope, so the configured name survives, which results carry as ``source``.
"""
client: HaikuRAG | None = None
stack = AsyncExitStack()
@ -233,14 +197,14 @@ def _covering(scope: "DatabaseScope", config: AppConfig) -> FastMCP:
finally:
client = None
# Masking keeps paths and provider URLs out of an unexpected error's text;
# the traceback goes to the server log. A ToolError reaches the client as is.
# Explicit: the setting is also read from the environment, and the contract
# is that every failure reaches the client with its message.
mcp = FastMCP(
"haiku-rag",
instructions=_instructions(scope, config),
version=metadata.version("haiku.rag-slim"),
lifespan=lifespan,
mask_error_details=True,
mask_error_details=False,
)
@mcp.tool(annotations=_read_only("Search documents"))
@ -272,17 +236,13 @@ def _covering(scope: "DatabaseScope", config: AppConfig) -> FastMCP:
False for a smaller response.
"""
rag = await _client()
try:
await _check_filter(rag, filter, sources)
results = await rag.search(
query,
limit=limit,
filter=filter,
include_images=include_images,
sources=sources,
)
except UnknownDatabaseError as e:
raise ToolError(str(e)) from e
results = await rag.search(
query,
limit=limit,
filter=filter,
include_images=include_images,
sources=sources,
)
return _search_result(await rag.expand_context(results), rag.covers_multiple)
# Image-as-query tool, only registered when the configured embedder
@ -317,17 +277,13 @@ def _covering(scope: "DatabaseScope", config: AppConfig) -> FastMCP:
"""
raw = _decode_image(image_base64)
rag = await _client()
try:
await _check_filter(rag, filter, sources)
results = await rag.search(
raw,
limit=limit,
filter=filter,
include_images=include_images,
sources=sources,
)
except UnknownDatabaseError as e:
raise ToolError(str(e)) from e
results = await rag.search(
raw,
limit=limit,
filter=filter,
include_images=include_images,
sources=sources,
)
return _search_result(
await rag.expand_context(results), rag.covers_multiple
)
@ -346,24 +302,18 @@ def _covering(scope: "DatabaseScope", config: AppConfig) -> FastMCP:
is asked.
"""
rag = await _client()
try:
document = await rag.get_document_by_id(document_id, source)
except UnknownDatabaseError as e:
raise ToolError(str(e)) from e
document = await rag.get_document_by_id(document_id, source)
if document is None:
raise ToolError(f"No document with id {document_id!r}")
return document
async def _items_of(document_id: str, source: str | None) -> list[DocumentItem]:
async def _items_of(document_id: str, source: str | None) -> list["DocumentItem"]:
"""A document's items in reading order, from the database holding it."""
rag = await _client()
try:
document = await rag.get_document_by_id(document_id, source)
if document is None:
raise ToolError(f"No document with id {document_id!r}")
owner = await rag.reader_for(source or document.source)
except UnknownDatabaseError as e:
raise ToolError(str(e)) from e
document = await rag.get_document_by_id(document_id, source)
if document is None:
raise ToolError(f"No document with id {document_id!r}")
owner = await rag.reader_for(source or document.source)
assert owner is not None, "a stored document names its database"
return await owner.document_item_repository.get_all_items(document_id)
@ -434,7 +384,6 @@ def _covering(scope: "DatabaseScope", config: AppConfig) -> FastMCP:
offset: How many documents to skip, for paging.
"""
rag = await _client()
await _check_filter(rag, filter)
documents = await rag.list_documents(limit, offset, filter)
return [
DocumentInfo(
@ -498,10 +447,7 @@ def _covering(scope: "DatabaseScope", config: AppConfig) -> FastMCP:
scope, config, AnalysisContext(filter=filter, sources=sources), rag=rag
)
try:
await _check_filter(rag, filter, sources)
result = await sandbox.execute(code)
except UnknownDatabaseError as e:
raise ToolError(str(e)) from e
finally:
await sandbox.close()
if not result.success:

69
tests/test_mcp.py
View file

@ -1,4 +1,3 @@
import logging
import re
from pathlib import Path
@ -8,6 +7,7 @@ from fastmcp.exceptions import ToolError
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.exceptions import UnknownDatabaseError
from haiku.rag.store.models import Chunk, Document, SearchResult
from haiku.rag.tools.document import DocumentInfo
from tests.multi_db.helpers import _config, _seed, _seed_expandable
@ -485,7 +485,7 @@ class TestMCPDocumentNavigation:
assert (
await section(document_id=doc.id, section_id="#/texts/0", source="beta")
).title == "Only in beta"
with pytest.raises(ToolError, match="nope"):
with pytest.raises(UnknownDatabaseError, match="nope"):
await outline(document_id=doc.id, source="nope")
with pytest.raises(ToolError, match=doc.id):
await outline(document_id=doc.id, source="alpha")
@ -908,11 +908,10 @@ class TestMCPCoversTheConfiguredSet:
async def test_an_unknown_database_is_an_error_not_an_empty_result(
self, two_dbs, multimodal_embedder, tool_name, kwargs
):
mcp = _covering_all(two_dbs)
tool = await _get_tool(mcp, tool_name)
result = await _call(_covering_all(two_dbs), tool_name, **kwargs)
with pytest.raises(ToolError, match="nope"):
await tool(**kwargs)
assert result.is_error
assert "nope" in result.content[0].text
@pytest.mark.asyncio
async def test_a_filtered_search_touches_only_the_selected_databases(self, two_dbs):
@ -1029,9 +1028,8 @@ class TestMCPImageQuery:
@pytest.mark.filterwarnings("ignore:Found propagated trace context:RuntimeWarning")
class TestMCPErrorContract:
"""A failure is an error on the wire, never an empty result. Expected
failures say what went wrong; anything else is masked and logged on the
server."""
"""A failure is an error on the wire carrying its message, never an empty
result."""
@pytest.mark.asyncio
async def test_an_unknown_document_is_an_error(self, mcp_db):
@ -1051,37 +1049,13 @@ class TestMCPErrorContract:
("execute_code", {"code": "print(1)"}),
],
)
async def test_an_invalid_filter_is_an_error_naming_the_filter(
self, mcp_db, tool_name, kwargs
):
async def test_an_invalid_filter_is_an_error(self, mcp_db, tool_name, kwargs):
result = await _call(
create_mcp_server(mcp_db), tool_name, filter="no_such_column = 1", **kwargs
)
assert result.is_error
assert "no_such_column = 1" in result.content[0].text
assert "created_at" in result.content[0].text
assert "_rowid" not in result.content[0].text
@pytest.mark.asyncio
@pytest.mark.parametrize("filter", [None, "title = 'AI Overview'"])
async def test_a_value_error_from_the_read_is_not_an_invalid_filter(
self, mcp_db, monkeypatch, filter
):
"""Only the filter check translates ValueError; one raised by the read
itself, with or without a valid filter, stays masked."""
async def boom(self, *args, **kw):
raise ValueError("boom at /secret/path")
monkeypatch.setattr(HaikuRAG, "search", boom)
result = await _call(
create_mcp_server(mcp_db), "search_documents", query="x", filter=filter
)
assert result.is_error
assert "filter" not in result.content[0].text
assert "/secret/path" not in result.content[0].text
assert "no_such_column" in result.content[0].text
@pytest.mark.asyncio
@pytest.mark.parametrize(
@ -1125,34 +1099,17 @@ class TestMCPErrorContract:
("list_documents", "list_documents", {}),
],
)
async def test_an_unexpected_failure_is_masked_and_logged(
self,
mcp_db,
multimodal_embedder,
monkeypatch,
caplog,
client_method,
tool_name,
kwargs,
async def test_an_unexpected_failure_carries_its_message(
self, mcp_db, multimodal_embedder, monkeypatch, client_method, tool_name, kwargs
):
async def boom(self, *args, **kw):
raise RuntimeError("boom at /secret/path")
monkeypatch.setattr(HaikuRAG, client_method, boom)
# fastmcp's logger does not propagate, so listen to it directly.
fastmcp_logger = logging.getLogger("fastmcp")
fastmcp_logger.addHandler(caplog.handler)
try:
result = await _call(create_mcp_server(mcp_db), tool_name, **kwargs)
finally:
fastmcp_logger.removeHandler(caplog.handler)
result = await _call(create_mcp_server(mcp_db), tool_name, **kwargs)
assert result.is_error
assert "/secret/path" not in result.content[0].text
assert any(
r.exc_info and "boom at /secret/path" in str(r.exc_info[1])
for r in caplog.records
)
assert "boom at /secret/path" in result.content[0].text
class TestAgentPlugins: