diff --git a/CHANGELOG.md b/CHANGELOG.md index f40fd167..807fd72d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,6 +9,7 @@ - Require `pydantic-ai-slim>=2.18,<3`. - `pydantic-monty` bumped to `>=0.0.19`; sandbox code now runs in a subprocess worker pool. +- 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`. - 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. @@ -20,6 +21,7 @@ - `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. - `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. ### Removed diff --git a/docs/configuration/qa.md b/docs/configuration/qa.md index 601a0165..afc4ea05 100644 --- a/docs/configuration/qa.md +++ b/docs/configuration/qa.md @@ -50,13 +50,13 @@ analysis: provider: anthropic name: claude-sonnet-4-20250514 temperature: 0.0 # Default: 0.0 (deterministic for code generation) - code_timeout: 60.0 # Max seconds for code execution + code_timeout: 60.0 # Max seconds a call may spend reading documents max_output_chars: 50000 # Truncate output after this many chars max_executions: 15 # Max execute_code calls per question ``` - **model**: LLM configuration (see [Providers](providers.md#model-settings)). When unset, falls back to `qa.model`. -- **code_timeout**: Maximum seconds for each code execution (default: 60) +- **code_timeout**: Seconds a single `execute_code` call may spend reading documents (default: 60). The sandbox refuses further reads past this point. Code that computes without reading is bounded instead by the session budget of `code_timeout * max_executions`. - **max_output_chars**: Truncate code output after this many characters (default: 50000) - **max_executions**: Maximum `execute_code` calls per question before the capability is told to answer from what it has (default: 15) diff --git a/haiku_rag_slim/haiku/rag/sandbox/sandbox.py b/haiku_rag_slim/haiku/rag/sandbox/sandbox.py index a32b4973..e788ea7a 100644 --- a/haiku_rag_slim/haiku/rag/sandbox/sandbox.py +++ b/haiku_rag_slim/haiku/rag/sandbox/sandbox.py @@ -14,6 +14,7 @@ from pydantic_monty import ( CallbackFile, MemoryFile, OSAccess, + ResourceLimits, ) from haiku.rag.config.models import AppConfig @@ -151,6 +152,7 @@ class Sandbox: _session: AsyncMontySession | None _vfs: OSAccess | None _loop: asyncio.AbstractEventLoop | None + _deadline: float | None def __init__( self, @@ -174,6 +176,7 @@ class Sandbox: self._session = None self._vfs = None self._loop = None + self._deadline = None @asynccontextmanager async def _connection(self) -> "AsyncIterator[HaikuRAG]": @@ -197,10 +200,24 @@ class Sandbox: Called off the loop while ``feed_run`` is awaited, so scheduling onto it and blocking for the result is safe. + + Blocking here suspends the worker, and Monty checks its duration budget + between interpreter steps, so it cannot check while a read is in flight. + Enforce the budget before starting another read, or code that reads in a + 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. """ 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() async def close(self) -> None: @@ -428,15 +445,26 @@ class Sandbox: return OSAccess(files) + def _session_limits(self) -> ResourceLimits: + """Resource limits for the worker session. + + Monty spends ``max_duration_secs`` across the session's whole life, and + the session 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 all of it. + """ + analysis = self._config.analysis + return {"max_duration_secs": analysis.code_timeout * analysis.max_executions} + 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: self._vfs = await self._build_vfs() self._pool = AsyncMonty() await self._pool.__aenter__() - session = self._pool.checkout( - limits={"max_duration_secs": self._config.analysis.code_timeout}, - ) + session = self._pool.checkout(limits=self._session_limits()) await session.__aenter__() self._session = session assert self._session is not None and self._vfs is not None @@ -449,6 +477,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 session, 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 0a19328c..a2a5fac3 100644 --- a/tests/sandbox/test_sandbox.py +++ b/tests/sandbox/test_sandbox.py @@ -733,3 +733,77 @@ class TestSandboxHeldConnection: assert not any(t.name == "sandbox-vfs" for t in threading.enumerate()) await sb.close() await sb.close() + + +class TestSandboxReadDeadline: + """The VFS bridge suspends the worker for the length of a read, so Monty + cannot check its duration budget while one 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 + + def test_session_budget_covers_every_permitted_execution(self, temp_db_path): + """Monty spends its duration budget across the session's whole life, so a + per-call value would let the first call starve the rest.""" + config = AppConfig() + config.analysis.code_timeout = 5.0 + config.analysis.max_executions = 3 + + sb = Sandbox(db_path=temp_db_path, config=config, context=AnalysisContext()) + + assert sb._session_limits() == {"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()) + try: + 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 + finally: + await sb.close()