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

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

Browse source 
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
Show stats Download patch file Download diff file Expand all files Collapse all files

6
CHANGELOG.md
View file

@ -50,9 +50,6 @@
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.
- 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
`search_documents`, `search_documents_by_image` and `execute_code`; `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
`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

2
docs/configuration/qa.md
View file

@ -50,7 +50,7 @@ analysis:
provider: anthropic
name: claude-sonnet-4-20250514
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_executions: 15 # Max execute_code calls per question
```

13
docs/mcp.md
View file

@ -171,9 +171,10 @@ Useful modules include `json`, `re`, `math`, `pathlib`, `datetime`,
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
there is no network and no filesystem
beyond `/documents`. `analysis.code_timeout` bounds a call, counted separately
for compute and for document reads, and `analysis.max_output_chars` bounds its
output.
beyond `/documents`. `analysis.code_timeout` is the call's budget: compute is
stopped at it, and past it no further host call starts, a file read or an
in-code search alike, though one already running finishes.
`analysis.max_output_chars` bounds the output.
### Filters
@ -192,10 +193,8 @@ title = 'Q3 report'
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`. A failure on the
server inside a program, a database read or an in-code search raising, reaches
the program and the client as its exception type only; the traceback goes to
the server log.
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.

93
haiku_rag_slim/haiku/rag/sandbox/sandbox.py
View file

@ -1,6 +1,5 @@
import asyncio
import json
import logging
import os
from collections.abc import AsyncIterator, Callable, Coroutine
from contextlib import asynccontextmanager, suppress
@ -31,21 +30,9 @@ if TYPE_CHECKING:
from haiku.rag.client.scope import DatabaseScope
logger = logging.getLogger(__name__)
_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
class SandboxResult:
"""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
inside the callback answers the worker's suspension, which keeps the
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, (
"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()
raise TimeoutError(
"time limit exceeded: no further document reads after "
f"{self._config.analysis.code_timeout}s"
)
try:
return asyncio.run_coroutine_threadsafe(coro, self._loop).result()
except Exception as e:
raise _host_failure("document read", e) from None
raise self._time_limit()
return asyncio.run_coroutine_threadsafe(coro, self._loop).result()
def _past_deadline(self) -> bool:
return (
self._deadline is not None
and self._loop is not 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:
"""Drop a session whose worker is gone.
@ -334,6 +347,7 @@ class Sandbox:
context = self._context
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
# results: the Monty interpreter has no PIL/base64/hashlib, so the
# agent's Python can't do anything with them. The driving model
@ -373,6 +387,7 @@ class Sandbox:
return out
async def list_documents() -> list[dict[str, Any]]:
self._check_deadline()
docs, _ = await self._documents()
return [
{
@ -387,22 +402,10 @@ class Sandbox:
]
return {
"search": self._guarded("search()", search),
"list_documents": self._guarded("list_documents()", list_documents),
"search": search,
"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:
"""Build the virtual filesystem with document data.
@ -554,7 +557,7 @@ class Sandbox:
files.append(
CallbackFile(
f"{doc_dir}/metadata.json",
read=lambda _path, text=metadata: text,
read=self._timed(lambda _path, text=metadata: text),
write=_deny_write,
)
)
@ -575,21 +578,21 @@ class Sandbox:
files.append(
CallbackFile(
f"{doc_dir}/content.txt",
read=_make_content_reader(doc_id),
read=self._timed(_make_content_reader(doc_id)),
write=_deny_write,
)
)
files.append(
CallbackFile(
f"{doc_dir}/items.jsonl",
read=_make_items_reader(doc_id),
read=self._timed(_make_items_reader(doc_id)),
write=_deny_write,
)
)
files.append(
CallbackFile(
f"{doc_dir}/chunks.jsonl",
read=_make_chunks_reader(doc_id),
read=self._timed(_make_chunks_reader(doc_id)),
write=_deny_write,
)
)
@ -600,7 +603,7 @@ class Sandbox:
files.append(
CallbackFile(
f"{doc_dir}/toc.json",
read=_make_toc_reader(doc_id),
read=self._timed(_make_toc_reader(doc_id)),
write=_deny_write,
)
)
@ -612,9 +615,9 @@ class Sandbox:
Monty spends ``max_duration_secs`` across the session's whole life, and
the session is reused so variables persist between calls: the budget
covers the whole run. ``code_timeout`` is enforced per call elsewhere: the read
deadline in ``_run_on_loop`` bounds a call that reads, and the pool's
``request_timeout`` bounds one that computes.
covers the whole run. ``code_timeout`` is enforced per call elsewhere: past
its deadline no further host call starts (``_check_deadline``), and the
pool's ``request_timeout`` bounds compute.
``max_suspensions`` counts host callbacks per session, document reads
included, defaults to 1000 and cannot be disabled. The time budgets are

124
tests/sandbox/test_sandbox.py
View file

@ -1,5 +1,4 @@
import asyncio
import logging
import threading
from pathlib import Path
@ -334,69 +333,21 @@ class TestSandboxExternalFunctionEdgeCases:
assert "external error" in result.stderr
@pytest.mark.asyncio
async def test_a_failing_search_reaches_the_program_by_type_only(
self, sandbox, monkeypatch, caplog
async def test_a_failing_search_keeps_its_message_for_the_program(
self, sandbox, monkeypatch
):
"""A host-side failure inside search() names its exception type to
the program; the message and traceback go to the log."""
"""A host-side failure inside search() reaches the program with its
message, which the agent reads to repair its code."""
async def boom(self, *args, **kwargs):
raise ValueError("failed at /secret/path")
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 "search() failed: ValueError" 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
)
assert "ValueError: failed at /secret/path" in result.stderr
class TestSandboxOutputTruncation:
@ -1013,24 +964,59 @@ class TestSandboxReadDeadline:
enforces the budget itself, before each read."""
@pytest.mark.asyncio
async def test_a_failed_read_reaches_the_program_by_type_only(
self, sandbox, caplog
async def test_the_deadline_covers_reads_from_memory_and_in_code_calls(
self, temp_db_path, monkeypatch
):
"""The bridged read hands the program the exception type, not the
message, and logs the traceback."""
sandbox._loop = asyncio.get_running_loop()
"""Once a call's time is up, a file served from memory and an in-code
listing are refused like a database read. A slow first read spends the
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():
raise ValueError("failed at /secret/path")
config = AppConfig()
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"):
with pytest.raises(RuntimeError, match="document read failed: ValueError"):
await asyncio.to_thread(sandbox._run_on_loop, failing_read())
async def slow_content(self, *args, **kwargs):
await asyncio.sleep(1.3)
return "body"
assert any(
r.exc_info and "failed at /secret/path" in str(r.exc_info[1])
for r in caplog.records
)
monkeypatch.setattr(repository, "get_content", slow_content)
sb = Sandbox(db_path=temp_db_path, config=config, context=AnalysisContext())
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
async def test_read_after_deadline_raises_without_scheduling(self, sandbox):

21
tests/test_mcp.py
View file

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