From d9e4d1e57c4e012d8d693a0d04012af95dc0441a Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Tue, 28 Jul 2026 17:33:08 +0300 Subject: [PATCH] Bound a call that never reads and restore the iteration test The read deadline only gets control at a read, so code that computes without reading escaped it. Give the pool a request_timeout above code_timeout. The watchdog kills the worker, and execute() already replaces a dead session, so a runaway call fails and the next call recovers. A read refusal keeps the variables, so it has to win the race whenever code does read. Restore test_open_file_objects_are_not_iterable. Replacing the neighbouring write test by text range deleted it, which left the instructions carrying a prohibition with nothing to signal when monty lifts it. Reuse the VFS and the pool when a session is replaced, so recovery skips the document scan. Mount metadata.json with a lambda rather than a factory. Cover open() in write mode. Fold the crash entry into the pydantic-monty bullet, because no release shipped the worker without it. --- CHANGELOG.md | 3 +- haiku_rag_slim/haiku/rag/sandbox/sandbox.py | 37 +++--- tests/sandbox/test_sandbox.py | 118 ++++++++++++++++++++ 3 files changed, 142 insertions(+), 16 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index c6dabef9..d42fc52a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,7 +8,7 @@ ### Changed - Require `pydantic-ai-slim>=2.18,<3`. -- `pydantic-monty` bumped to `>=0.0.19`; sandbox code now runs in a subprocess worker pool. +- `pydantic-monty` bumped to `>=0.0.19`; sandbox code now runs in a subprocess worker pool. A crashed or timed-out worker fails one call and the next call gets a replacement session. - The analysis sandbox gives Monty a duration budget of `analysis.code_timeout * analysis.max_executions` for the session, was `analysis.code_timeout`. - `enable_thinking` maps onto Pydantic AI's unified `thinking` setting for the `anthropic`, `gemini`, `groq` and `bedrock` providers. The Anthropic thinking budget is now Pydantic AI's default of 10000 tokens, was 4096, and `max_tokens` must exceed it on budget-based Claude models; `enable_thinking: false` disables Gemini thinking rather than only hiding thoughts; Groq maps reasoning effort rather than `groq_reasoning_format`; Bedrock Qwen with `enable_thinking: false` no longer sends `reasoning_config`, while Bedrock-served Claude keeps an explicit `thinking: disabled`. The `openai` and `ollama` providers still map to `openai_reasoning_effort`. - `provider: bedrock` with a proprietary OpenAI model such as `openai.o3-mini-v1:0` raises `UserError`: Bedrock Converse serves only the `gpt-oss` family. Use `provider: bedrock-mantle`. @@ -22,7 +22,6 @@ - `docs/installation.md` documents the `jina`, `s3` and `ingester` extras, names the extras the full package actually pulls, drops the removed MixedBread AI reranker, and no longer lists Anthropic as a built-in provider. - `docs/tuning.md` no longer points at the removed `claim_timeout_s` setting. - `analysis.code_timeout` is checked before each document read; code that reads in a loop no longer overruns it by the duration of the outstanding reads. -- A crashed Monty worker no longer poisons the analysis sandbox for the rest of the run: `execute` reports the crash and the next call checks out a replacement session. - `metadata.json` in the document VFS rejects writes, matching `content.txt`, `items.jsonl` and `toc.json`. ### Removed diff --git a/haiku_rag_slim/haiku/rag/sandbox/sandbox.py b/haiku_rag_slim/haiku/rag/sandbox/sandbox.py index a4256579..9e1cbfac 100644 --- a/haiku_rag_slim/haiku/rag/sandbox/sandbox.py +++ b/haiku_rag_slim/haiku/rag/sandbox/sandbox.py @@ -28,6 +28,10 @@ if TYPE_CHECKING: from haiku.rag.client import HaikuRAG +_REQUEST_TIMEOUT_MARGIN_S = 30.0 +"""Grace above ``code_timeout`` before the pool watchdog kills the worker.""" + + @dataclass class SandboxResult: """Result of executing code in the sandbox.""" @@ -417,20 +421,11 @@ class Sandbox: ) # A MemoryFile accepts writes, so mount metadata.json through the - # same read/deny pair as the rest. The content is already built, - # so the reader stays eager. - def _make_metadata_reader( - text: str, - ) -> Callable[["PurePosixPath"], str]: - def read_metadata(_path: "PurePosixPath") -> str: - return text - - return read_metadata - + # same read and deny pair as the rest. The content is already built. files.append( CallbackFile( f"{doc_dir}/metadata.json", - read=_make_metadata_reader(metadata), + read=lambda _path, text=metadata: text, write=_deny_write, ) ) @@ -476,6 +471,17 @@ class Sandbox: return OSAccess(files) + def _request_timeout(self) -> float: + """Hard per-call deadline for the pool watchdog. + + The read deadline refuses the next read at ``code_timeout`` and keeps the + session, so it wins for code that reads. This watchdog is the fallback + for code that computes without reading: it kills the worker, which loses + the variables the session held. Leave room for one read that is already + in flight, so the graceful refusal wins the race. + """ + return self._config.analysis.code_timeout + _REQUEST_TIMEOUT_MARGIN_S + def _session_limits(self) -> ResourceLimits: """Resource limits for the worker session. @@ -491,10 +497,13 @@ class Sandbox: async def _ensure_initialized(self) -> tuple[AsyncMontySession, OSAccess]: """Check out a worker session and build the VFS on first use.""" - if self._session is None: + if self._vfs is None: self._vfs = await self._build_vfs() - self._pool = AsyncMonty() - await self._pool.__aenter__() + if self._pool is None: + pool = AsyncMonty(request_timeout=self._request_timeout()) + await pool.__aenter__() + self._pool = pool + if self._session is None: session = self._pool.checkout(limits=self._session_limits()) await session.__aenter__() self._session = session diff --git a/tests/sandbox/test_sandbox.py b/tests/sandbox/test_sandbox.py index 107ae8bf..03d4f93f 100644 --- a/tests/sandbox/test_sandbox.py +++ b/tests/sandbox/test_sandbox.py @@ -471,6 +471,88 @@ class TestSandboxVFS: finally: await sb.close() + @pytest.mark.asyncio + async def test_open_for_writing_is_denied(self, temp_db_path): + """`open()` in write mode is refused, not only `Path.write_text`.""" + from docling_core.types.doc.document import DoclingDocument + from docling_core.types.doc.labels import DocItemLabel + + config = AppConfig() + 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://openwrite", + ) + + sb = Sandbox(db_path=temp_db_path, config=config, context=AnalysisContext()) + try: + result = await sb.execute( + "try:\n" + f" with open('/documents/{doc.id}/content.txt', 'w') as f:\n" + " f.write('nope')\n" + " print('WROTE')\n" + "except PermissionError:\n" + " print('DENIED')" + ) + assert result.success, result.stderr + assert "DENIED" in result.stdout + assert "WROTE" not in result.stdout + finally: + await sb.close() + + @pytest.mark.asyncio + async def test_open_file_objects_are_not_iterable(self, temp_db_path): + """Pins the limitation the instructions warn about: pydantic/monty#490. + + A failure here means Monty gained iteration support and the + `for line in f` prohibition in the analysis instructions is now wrong. + """ + from docling_core.types.doc.document import DoclingDocument + from docling_core.types.doc.labels import DocItemLabel + + config = AppConfig() + docling = DoclingDocument(name="d") + docling.add_text(label=DocItemLabel.TEXT, text="one\ntwo") + + async with HaikuRAG(temp_db_path, create=True) as client: + doc = await client.import_document( + docling, + [ + Chunk( + content="one\ntwo", + embedding=[0.1] * config.embeddings.model.vector_dim, + order=0, + ) + ], + uri="test://lines", + ) + + sb = Sandbox(db_path=temp_db_path, config=config, context=AnalysisContext()) + try: + result = await sb.execute( + f"for line in open('/documents/{doc.id}/content.txt'):\n print(line)" + ) + assert result.success is False + assert "not iterable" in result.stderr + + # The documented alternatives do work. + result = await sb.execute( + f"print(len(open('/documents/{doc.id}/content.txt').readlines()))" + ) + assert result.success, result.stderr + finally: + await sb.close() + @pytest.mark.asyncio @pytest.mark.vcr() async def test_context_filter_limits_vfs(self, temp_db_path): @@ -819,3 +901,39 @@ class TestSandboxWorkerCrash: assert lost.success is False finally: await sb.close() + + +class TestSandboxRequestTimeout: + """The pool watchdog bounds a call that never reads.""" + + def test_request_timeout_sits_above_the_read_deadline(self, temp_db_path): + """A read refusal keeps the session, so it must win the race.""" + config = AppConfig() + config.analysis.code_timeout = 5.0 + sb = Sandbox(db_path=temp_db_path, config=config, context=AnalysisContext()) + + assert sb._request_timeout() > config.analysis.code_timeout + + @pytest.mark.asyncio + async def test_runaway_compute_is_killed_and_the_next_call_recovers( + self, temp_db_path, monkeypatch + ): + """Code that never reads escapes the read deadline. The watchdog kills it.""" + config = AppConfig() + async with HaikuRAG(temp_db_path, create=True): + pass + + sb = Sandbox(db_path=temp_db_path, config=config, context=AnalysisContext()) + monkeypatch.setattr(sb, "_request_timeout", lambda: 1.0) + try: + runaway = await sb.execute( + "x = 0\nfor i in range(500000000):\n x += i\nprint(x)" + ) + assert runaway.success is False + assert "restarted" in runaway.stderr + + recovered = await sb.execute("print('alive')") + assert recovered.success, recovered.stderr + assert "alive" in recovered.stdout + finally: + await sb.close()