diff --git a/CHANGELOG.md b/CHANGELOG.md index 70ecb1d7..5235be87 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,12 +7,14 @@ - `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`. - Tool failures raise `pydantic_ai.ToolFailed` instead of returning failure text: search and code-execution limits, sandbox execution errors, and `get_document`/`summarize_document` misses. +- The analysis sandbox gives Monty a duration budget of `analysis.code_timeout * analysis.max_executions` for the session, was `analysis.code_timeout`. ### Fixed - `Store.set_haiku_version` stamps the store's own config into a recreated settings row instead of the process-global `Config`. - `check_source_accessible` returns `False` for a URI it cannot resolve (unparseable host, unreadable path) instead of raising and aborting a full rebuild. - `evaluations run` opens the database read-only outside the population phase, so an embedder identity differing from the stored one warns instead of aborting the run. +- `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. ### Removed diff --git a/haiku_rag_slim/haiku/rag/sandbox/sandbox.py b/haiku_rag_slim/haiku/rag/sandbox/sandbox.py index 76486690..d5417d94 100644 --- a/haiku_rag_slim/haiku/rag/sandbox/sandbox.py +++ b/haiku_rag_slim/haiku/rag/sandbox/sandbox.py @@ -142,6 +142,7 @@ class Sandbox: _repl: MontyRepl | None _vfs: OSAccess | None _loop: asyncio.AbstractEventLoop | None + _deadline: float | None def __init__( self, @@ -164,6 +165,7 @@ class Sandbox: self._repl = None self._vfs = None self._loop = None + self._deadline = None @asynccontextmanager async def _connection(self) -> "AsyncIterator[HaikuRAG]": @@ -187,10 +189,21 @@ class Sandbox: Called from Monty's worker thread while ``feed_run_async`` leaves the loop free, so scheduling onto it and blocking for the result is safe. + + Blocking here parks the worker thread, so Monty cannot check its own + duration budget until the read returns. Enforce the budget before + starting another one, or code that reads in a loop overruns it by + however long the outstanding reads take. """ 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: + coro.close() + raise TimeoutError( + "time limit exceeded: no further document reads after " + f"{self._config.analysis.code_timeout}s" + ) return asyncio.run_coroutine_threadsafe(coro, self._loop).result() def close(self) -> None: @@ -417,9 +430,18 @@ class Sandbox: """Initialize the REPL session and VFS on first use.""" if self._repl is None: self._vfs = await self._build_vfs() + # Monty spends `max_duration_secs` across the REPL's whole life, and + # the REPL is reused so variables persist between calls. Budget it for + # the run rather than for one call, or the first slow call starves + # every later one. `code_timeout` is enforced per call by the read + # deadline in `_run_on_loop`; this is the backstop for code that + # computes without reading, and one such call can spend it all. self._repl = MontyRepl( limits={ - "max_duration_secs": self._config.analysis.code_timeout, + "max_duration_secs": ( + self._config.analysis.code_timeout + * self._config.analysis.max_executions + ), }, ) assert self._repl is not None and self._vfs is not None @@ -432,6 +454,7 @@ class Sandbox: """ # Monty's synchronous file callbacks bridge DB reads back to this loop. self._loop = asyncio.get_running_loop() + self._deadline = self._loop.time() + self._config.analysis.code_timeout repl, vfs = await self._ensure_initialized() external_fns = self._build_external_functions() diff --git a/tests/sandbox/test_sandbox.py b/tests/sandbox/test_sandbox.py index b996fe9d..7fa8266b 100644 --- a/tests/sandbox/test_sandbox.py +++ b/tests/sandbox/test_sandbox.py @@ -7,6 +7,7 @@ import pytest from haiku.rag.client import HaikuRAG from haiku.rag.config.models import AppConfig from haiku.rag.sandbox import AnalysisContext, Sandbox, SandboxResult +from haiku.rag.store.models.chunk import Chunk @pytest.fixture(scope="module") @@ -618,3 +619,93 @@ class TestSandboxHeldConnection: assert not any(t.name == "sandbox-vfs" for t in threading.enumerate()) sb.close() sb.close() + + +class TestSandboxReadDeadline: + """The VFS bridge blocks Monty's worker thread, so its duration budget + cannot be checked while a read is in flight. The sandbox enforces the + budget itself, before each read.""" + + @pytest.mark.asyncio + async def test_read_after_deadline_raises_without_scheduling(self, sandbox): + """A read attempted past the deadline fails instead of querying.""" + scheduled = False + + async def _never_runs(): + nonlocal scheduled + scheduled = True + + sandbox._loop = asyncio.get_running_loop() + sandbox._deadline = sandbox._loop.time() - 1.0 + + coro = _never_runs() + with pytest.raises(TimeoutError, match="time limit exceeded"): + sandbox._run_on_loop(coro) + + coro.close() + assert scheduled is False + + @pytest.mark.asyncio + async def test_session_budget_covers_every_permitted_execution( + self, temp_db_path, monkeypatch + ): + """Monty spends its duration budget across the REPL's whole life, so a + per-call value would let the first call starve the rest.""" + import haiku.rag.sandbox.sandbox as sandbox_mod + + captured: dict[str, float] = {} + real_repl = sandbox_mod.MontyRepl + + def _capture(*args, **kwargs): + captured.update(kwargs.get("limits", {})) + return real_repl(*args, **kwargs) + + monkeypatch.setattr(sandbox_mod, "MontyRepl", _capture) + + config = AppConfig() + config.analysis.code_timeout = 5.0 + config.analysis.max_executions = 3 + + async with HaikuRAG(temp_db_path, create=True): + pass + sb = Sandbox(db_path=temp_db_path, config=config, context=AnalysisContext()) + result = await sb.execute("x = 1") + + assert result.success + assert captured["max_duration_secs"] == 15.0 + + @pytest.mark.asyncio + async def test_document_read_past_the_deadline_fails_the_execution( + self, temp_db_path + ): + """The overrun surfaces as a failed result, not a raised exception.""" + from docling_core.types.doc.document import DoclingDocument + from docling_core.types.doc.labels import DocItemLabel + + config = AppConfig() + config.analysis.code_timeout = 0.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", + ) + + sb = Sandbox(db_path=temp_db_path, config=config, context=AnalysisContext()) + result = await sb.execute( + "from pathlib import Path\n" + f"print(Path('/documents/{doc.id}/content.txt').read_text())" + ) + + assert result.success is False + assert "no further document reads" in result.stderr