diff --git a/CHANGELOG.md b/CHANGELOG.md index 3070696d..555d9ec8 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -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(" failed: ")`; 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 diff --git a/docs/configuration/qa.md b/docs/configuration/qa.md index 32c4c34b..a7efa866 100644 --- a/docs/configuration/qa.md +++ b/docs/configuration/qa.md @@ -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 ``` diff --git a/docs/mcp.md b/docs/mcp.md index d1cb4bac..c2afa2ba 100644 --- a/docs/mcp.md +++ b/docs/mcp.md @@ -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. diff --git a/haiku_rag_slim/haiku/rag/sandbox/sandbox.py b/haiku_rag_slim/haiku/rag/sandbox/sandbox.py index 823f0040..512b5b24 100644 --- a/haiku_rag_slim/haiku/rag/sandbox/sandbox.py +++ b/haiku_rag_slim/haiku/rag/sandbox/sandbox.py @@ -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 diff --git a/tests/sandbox/test_sandbox.py b/tests/sandbox/test_sandbox.py index 98465329..1d3635dc 100644 --- a/tests/sandbox/test_sandbox.py +++ b/tests/sandbox/test_sandbox.py @@ -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): diff --git a/tests/test_mcp.py b/tests/test_mcp.py index 18d7106d..04e7cb0a 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -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(