Check the sandbox deadline on every host call, one error contract

analysis.code_timeout was enforced only in _run_on_loop, the bridge for
database-bound reads; metadata.json, the cached JSONL files and in-code
search() and list_documents() never looked at the clock, and Monty's
watchdog counts compute only. Every host call now checks the deadline
before it starts.

Host errors keep their message for every caller: the masking added for the
MCP server goes, and the sandbox is one path for the capability and the
server alike.
This commit is contained in:
Yiorgis Gozadinos 2026-09-07 12:47:28 +03:00
parent bae359e3bf
commit 95fdb46c3a
No known key found for this signature in database
6 changed files with 122 additions and 137 deletions

View file

@ -50,9 +50,6 @@
Unknown document, unknown collection, invalid filter, invalid base64 and a Unknown document, unknown collection, invalid filter, invalid base64 and a
failing program carry a message. Anything else is masked failing program carry a message. Anything else is masked
(`mask_error_details=True`) and logged server-side. (`mask_error_details=True`) and logged server-side.
- A host-side failure inside the analysis sandbox (a document read or an
in-code `search()` raising) reaches the program as
`RuntimeError("<call> failed: <ExceptionType>")`; the traceback is logged.
- `haiku-rag mcp` covers the configured `lancedb.databases` set. `sources` on - `haiku-rag mcp` covers the configured `lancedb.databases` set. `sources` on
`search_documents`, `search_documents_by_image` and `execute_code`; `source` `search_documents`, `search_documents_by_image` and `execute_code`; `source`
on `get_document`; an unknown name is a tool error. `DocumentInfo.source`. on `get_document`; an unknown name is a tool error. `DocumentInfo.source`.
@ -61,6 +58,9 @@
- `toc.json` `item_range` in the analysis sandbox is a line slice into - `toc.json` `item_range` in the analysis sandbox is a line slice into
`items.jsonl`, as documented; it held item positions. `items.jsonl`, as documented; it held item positions.
- Past `analysis.code_timeout` a sandbox program starts no further host call.
Files served from memory and in-code `search()` / `list_documents()` were
not checked against the deadline.
### Removed ### Removed

View file

@ -50,7 +50,7 @@ analysis:
provider: anthropic provider: anthropic
name: claude-sonnet-4-20250514 name: claude-sonnet-4-20250514
temperature: 0.0 # Default: 0.0 (deterministic for code generation) temperature: 0.0 # Default: 0.0 (deterministic for code generation)
code_timeout: 60.0 # Max seconds a call may spend reading documents code_timeout: 60.0 # Per call: compute stops, no read or search starts past it
max_output_chars: 50000 # Truncate output after this many chars max_output_chars: 50000 # Truncate output after this many chars
max_executions: 15 # Max execute_code calls per question max_executions: 15 # Max execute_code calls per question
``` ```

View file

@ -171,9 +171,10 @@ Useful modules include `json`, `re`, `math`, `pathlib`, `datetime`,
reached for: `decimal` and `statistics`. No generator functions, class reached for: `decimal` and `statistics`. No generator functions, class
inheritance or `match` statements, and a file object cannot be iterated. Files are read-only, and inheritance or `match` statements, and a file object cannot be iterated. Files are read-only, and
there is no network and no filesystem there is no network and no filesystem
beyond `/documents`. `analysis.code_timeout` bounds a call, counted separately beyond `/documents`. `analysis.code_timeout` is the call's budget: compute is
for compute and for document reads, and `analysis.max_output_chars` bounds its stopped at it, and past it no further host call starts, a file read or an
output. in-code search alike, though one already running finishes.
`analysis.max_output_chars` bounds the output.
### Filters ### Filters
@ -192,10 +193,8 @@ title = 'Q3 report'
A failure is an MCP error, never an empty result. Expected failures carry a 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 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), server does not cover, a filter the query engine rejects (with its message),
invalid base64, and a program that fails in `execute_code`. A failure on the invalid base64, and a program that fails in `execute_code`, with the error the
server inside a program, a database read or an in-code search raising, reaches program hit.
the program and the client as its exception type only; the traceback goes to
the server log.
Anything else reaches the client as `Error calling tool 'name'` and its Anything else reaches the client as `Error calling tool 'name'` and its
traceback goes to the server log. traceback goes to the server log.

View file

@ -1,6 +1,5 @@
import asyncio import asyncio
import json import json
import logging
import os import os
from collections.abc import AsyncIterator, Callable, Coroutine from collections.abc import AsyncIterator, Callable, Coroutine
from contextlib import asynccontextmanager, suppress from contextlib import asynccontextmanager, suppress
@ -31,21 +30,9 @@ if TYPE_CHECKING:
from haiku.rag.client.scope import DatabaseScope from haiku.rag.client.scope import DatabaseScope
logger = logging.getLogger(__name__)
_MAX_HOST_CALLS = 10_000_000 _MAX_HOST_CALLS = 10_000_000
def _host_failure(where: str, e: Exception) -> RuntimeError:
"""The error a program gets for a failure on the host side of a call.
The message and traceback go to the log. The program, and through the MCP
server its client, learn the exception type only.
"""
logger.exception("%s failed inside the sandbox", where)
return RuntimeError(f"{where} failed: {type(e).__name__}")
@dataclass @dataclass
class SandboxResult: class SandboxResult:
"""Result of executing code in the sandbox.""" """Result of executing code in the sandbox."""
@ -283,21 +270,47 @@ class Sandbox:
loop overruns it by however long the outstanding reads take. Raising from loop overruns it by however long the outstanding reads take. Raising from
inside the callback answers the worker's suspension, which keeps the inside the callback answers the worker's suspension, which keeps the
session usable cancelling ``feed_run`` from outside does not, and wedges session usable cancelling ``feed_run`` from outside does not, and wedges
the protocol. A failed read reaches the program by type only. the protocol.
""" """
assert self._loop is not None, ( assert self._loop is not None, (
"VFS reads happen during execute(); the loop must be captured first." "VFS reads happen during execute(); the loop must be captured first."
) )
if self._deadline is not None and self._loop.time() > self._deadline: if self._past_deadline():
coro.close() coro.close()
raise TimeoutError( raise self._time_limit()
"time limit exceeded: no further document reads after " return asyncio.run_coroutine_threadsafe(coro, self._loop).result()
f"{self._config.analysis.code_timeout}s"
) def _past_deadline(self) -> bool:
try: return (
return asyncio.run_coroutine_threadsafe(coro, self._loop).result() self._deadline is not None
except Exception as e: and self._loop is not None
raise _host_failure("document read", e) from None and self._loop.time() > self._deadline
)
def _time_limit(self) -> TimeoutError:
return TimeoutError(
"time limit exceeded: no further document reads or calls after "
f"{self._config.analysis.code_timeout}s"
)
def _check_deadline(self) -> None:
"""Refuse a host call once the call's time is up.
Monty's watchdog counts only time the worker spends computing, so every
host call, a file served from memory and an in-code search included,
checks the deadline before it runs.
"""
if self._past_deadline():
raise self._time_limit()
def _timed(
self, read: Callable[["PurePosixPath"], str]
) -> Callable[["PurePosixPath"], str]:
def call(path: "PurePosixPath") -> str:
self._check_deadline()
return read(path)
return call
async def _discard_session(self) -> None: async def _discard_session(self) -> None:
"""Drop a session whose worker is gone. """Drop a session whose worker is gone.
@ -334,6 +347,7 @@ class Sandbox:
context = self._context context = self._context
async def search(query: str, limit: int = 10) -> list[dict[str, Any]]: async def search(query: str, limit: int = 10) -> list[dict[str, Any]]:
self._check_deadline()
# Picture bytes are deliberately not attached to in-code search # Picture bytes are deliberately not attached to in-code search
# results: the Monty interpreter has no PIL/base64/hashlib, so the # results: the Monty interpreter has no PIL/base64/hashlib, so the
# agent's Python can't do anything with them. The driving model # agent's Python can't do anything with them. The driving model
@ -373,6 +387,7 @@ class Sandbox:
return out return out
async def list_documents() -> list[dict[str, Any]]: async def list_documents() -> list[dict[str, Any]]:
self._check_deadline()
docs, _ = await self._documents() docs, _ = await self._documents()
return [ return [
{ {
@ -387,22 +402,10 @@ class Sandbox:
] ]
return { return {
"search": self._guarded("search()", search), "search": search,
"list_documents": self._guarded("list_documents()", list_documents), "list_documents": list_documents,
} }
@staticmethod
def _guarded(
where: str, fn: Callable[..., Coroutine[Any, Any, Any]]
) -> Callable[..., Coroutine[Any, Any, Any]]:
async def call(*args: Any, **kwargs: Any) -> Any:
try:
return await fn(*args, **kwargs)
except Exception as e:
raise _host_failure(where, e) from None
return call
async def _build_vfs(self) -> OSAccess: async def _build_vfs(self) -> OSAccess:
"""Build the virtual filesystem with document data. """Build the virtual filesystem with document data.
@ -554,7 +557,7 @@ class Sandbox:
files.append( files.append(
CallbackFile( CallbackFile(
f"{doc_dir}/metadata.json", f"{doc_dir}/metadata.json",
read=lambda _path, text=metadata: text, read=self._timed(lambda _path, text=metadata: text),
write=_deny_write, write=_deny_write,
) )
) )
@ -575,21 +578,21 @@ class Sandbox:
files.append( files.append(
CallbackFile( CallbackFile(
f"{doc_dir}/content.txt", f"{doc_dir}/content.txt",
read=_make_content_reader(doc_id), read=self._timed(_make_content_reader(doc_id)),
write=_deny_write, write=_deny_write,
) )
) )
files.append( files.append(
CallbackFile( CallbackFile(
f"{doc_dir}/items.jsonl", f"{doc_dir}/items.jsonl",
read=_make_items_reader(doc_id), read=self._timed(_make_items_reader(doc_id)),
write=_deny_write, write=_deny_write,
) )
) )
files.append( files.append(
CallbackFile( CallbackFile(
f"{doc_dir}/chunks.jsonl", f"{doc_dir}/chunks.jsonl",
read=_make_chunks_reader(doc_id), read=self._timed(_make_chunks_reader(doc_id)),
write=_deny_write, write=_deny_write,
) )
) )
@ -600,7 +603,7 @@ class Sandbox:
files.append( files.append(
CallbackFile( CallbackFile(
f"{doc_dir}/toc.json", f"{doc_dir}/toc.json",
read=_make_toc_reader(doc_id), read=self._timed(_make_toc_reader(doc_id)),
write=_deny_write, write=_deny_write,
) )
) )
@ -612,9 +615,9 @@ class Sandbox:
Monty spends ``max_duration_secs`` across the session's whole life, and Monty spends ``max_duration_secs`` across the session's whole life, and
the session is reused so variables persist between calls: the budget the session is reused so variables persist between calls: the budget
covers the whole run. ``code_timeout`` is enforced per call elsewhere: the read covers the whole run. ``code_timeout`` is enforced per call elsewhere: past
deadline in ``_run_on_loop`` bounds a call that reads, and the pool's its deadline no further host call starts (``_check_deadline``), and the
``request_timeout`` bounds one that computes. pool's ``request_timeout`` bounds compute.
``max_suspensions`` counts host callbacks per session, document reads ``max_suspensions`` counts host callbacks per session, document reads
included, defaults to 1000 and cannot be disabled. The time budgets are included, defaults to 1000 and cannot be disabled. The time budgets are

View file

@ -1,5 +1,4 @@
import asyncio import asyncio
import logging
import threading import threading
from pathlib import Path from pathlib import Path
@ -334,69 +333,21 @@ class TestSandboxExternalFunctionEdgeCases:
assert "external error" in result.stderr assert "external error" in result.stderr
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_a_failing_search_reaches_the_program_by_type_only( async def test_a_failing_search_keeps_its_message_for_the_program(
self, sandbox, monkeypatch, caplog self, sandbox, monkeypatch
): ):
"""A host-side failure inside search() names its exception type to """A host-side failure inside search() reaches the program with its
the program; the message and traceback go to the log.""" message, which the agent reads to repair its code."""
async def boom(self, *args, **kwargs): async def boom(self, *args, **kwargs):
raise ValueError("failed at /secret/path") raise ValueError("failed at /secret/path")
monkeypatch.setattr(HaikuRAG, "search", boom) monkeypatch.setattr(HaikuRAG, "search", boom)
with caplog.at_level(logging.ERROR, logger="haiku.rag.sandbox.sandbox"): result = await sandbox.execute("await search('hello')")
result = await sandbox.execute("await search('hello')")
assert not result.success assert not result.success
assert "search() failed: ValueError" in result.stderr assert "ValueError: failed at /secret/path" in result.stderr
assert "/secret/path" not in result.stderr
assert any(
r.exc_info and "failed at /secret/path" in str(r.exc_info[1])
for r in caplog.records
)
@pytest.mark.asyncio
async def test_a_failing_document_read_reaches_the_program_by_type_only(
self, temp_db_path, monkeypatch, caplog
):
"""A program can catch a failed file read, and what it catches names
the exception type only."""
from haiku.rag.store.models.document import Document
async with HaikuRAG(temp_db_path, create=True) as client:
doc = await client.document_repository.create(
Document(content="x", uri="test://read", title="Read")
)
repository = type(client.document_repository)
async def boom(self, *args, **kwargs):
raise ValueError("failed at /secret/path")
monkeypatch.setattr(repository, "get_content", boom)
sb = Sandbox(
db_path=temp_db_path, config=AppConfig(), context=AnalysisContext()
)
try:
with caplog.at_level(logging.ERROR, logger="haiku.rag.sandbox.sandbox"):
result = await sb.execute(
"from pathlib import Path\n"
"try:\n"
f" Path('/documents/{doc.id}/content.txt').read_text()\n"
"except Exception as e:\n"
" print('caught:', e)"
)
finally:
await sb.close()
assert result.success, result.stderr
assert "caught:" in result.stdout
assert "ValueError" in result.stdout
assert "/secret/path" not in result.stdout
assert any(
r.exc_info and "failed at /secret/path" in str(r.exc_info[1])
for r in caplog.records
)
class TestSandboxOutputTruncation: class TestSandboxOutputTruncation:
@ -1013,24 +964,59 @@ class TestSandboxReadDeadline:
enforces the budget itself, before each read.""" enforces the budget itself, before each read."""
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_a_failed_read_reaches_the_program_by_type_only( async def test_the_deadline_covers_reads_from_memory_and_in_code_calls(
self, sandbox, caplog self, temp_db_path, monkeypatch
): ):
"""The bridged read hands the program the exception type, not the """Once a call's time is up, a file served from memory and an in-code
message, and logs the traceback.""" listing are refused like a database read. A slow first read spends the
sandbox._loop = asyncio.get_running_loop() budget; the watchdog does not count time spent waiting on the host."""
from docling_core.types.doc.document import DoclingDocument
from docling_core.types.doc.labels import DocItemLabel
async def failing_read(): config = AppConfig()
raise ValueError("failed at /secret/path") config.analysis.code_timeout = 1.0
docling = DoclingDocument(name="d")
docling.add_text(label=DocItemLabel.TEXT, text="Foxes and dogs.")
async with HaikuRAG(temp_db_path, create=True) as client:
doc = await client.import_document(
docling,
[
Chunk(
content="Foxes and dogs.",
embedding=[0.1] * config.embeddings.model.vector_dim,
order=0,
)
],
uri="test://deadline-paths",
)
repository = type(client.document_repository)
with caplog.at_level(logging.ERROR, logger="haiku.rag.sandbox.sandbox"): async def slow_content(self, *args, **kwargs):
with pytest.raises(RuntimeError, match="document read failed: ValueError"): await asyncio.sleep(1.3)
await asyncio.to_thread(sandbox._run_on_loop, failing_read()) return "body"
assert any( monkeypatch.setattr(repository, "get_content", slow_content)
r.exc_info and "failed at /secret/path" in str(r.exc_info[1]) sb = Sandbox(db_path=temp_db_path, config=config, context=AnalysisContext())
for r in caplog.records try:
) result = await sb.execute(
"from pathlib import Path\n"
f"root = Path('/documents/{doc.id}')\n"
"print(len((root / 'content.txt').read_text()))\n"
"try:\n"
" (root / 'metadata.json').read_text()\n"
" print('static: read')\n"
"except Exception as e:\n"
" print('static:', type(e).__name__)\n"
"await list_documents()\n"
"print('listed')"
)
finally:
await sb.close()
assert "static: TimeoutError" in result.stdout
assert "listed" not in result.stdout
assert not result.success
assert "time limit exceeded" in result.stderr
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_read_after_deadline_raises_without_scheduling(self, sandbox): async def test_read_after_deadline_raises_without_scheduling(self, sandbox):

View file

@ -1098,25 +1098,22 @@ class TestMCPErrorContract:
assert "base64" in result.content[0].text assert "base64" in result.content[0].text
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_a_host_failure_inside_a_program_names_only_its_type( async def test_a_host_failure_inside_a_program_carries_its_message(
self, mcp_db, monkeypatch, caplog self, mcp_db, monkeypatch
): ):
"""One contract for the sandbox: the client reads the same error the
program did, message included."""
async def boom(self, *args, **kwargs): async def boom(self, *args, **kwargs):
raise RuntimeError("boom at /secret/path") raise RuntimeError("boom at /secret/path")
monkeypatch.setattr(HaikuRAG, "search", boom) monkeypatch.setattr(HaikuRAG, "search", boom)
with caplog.at_level(logging.ERROR, logger="haiku.rag.sandbox.sandbox"): result = await _call(
result = await _call( create_mcp_server(mcp_db), "execute_code", code="await search('x')"
create_mcp_server(mcp_db), "execute_code", code="await search('x')" )
)
assert result.is_error assert result.is_error
assert "RuntimeError" in result.content[0].text assert "RuntimeError: boom at /secret/path" in result.content[0].text
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
)
@pytest.mark.asyncio @pytest.mark.asyncio
@pytest.mark.parametrize( @pytest.mark.parametrize(