From 4a9dd9b49ad407a551356ee54d39ab2af2498d58 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Wed, 22 Apr 2026 13:30:38 +0300 Subject: [PATCH] persist sandbox variables across execute_code calls within one invocation --- CHANGELOG.md | 1 + .../haiku/rag/agents/analysis/sandbox.py | 97 +++++++++++-------- haiku_rag_slim/haiku/rag/skills/_deps.py | 21 ++++ haiku_rag_slim/haiku/rag/skills/_tools.py | 19 ++-- haiku_rag_slim/haiku/rag/skills/analysis.py | 4 +- .../haiku/rag/skills/rag-analysis/SKILL.md | 4 +- tests/skills/conftest.py | 22 ++++- tests/skills/test_analysis.py | 78 ++++++++++++--- 8 files changed, 172 insertions(+), 74 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 7bdb3a64..a441e78c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,7 @@ - **Skills share a single `HaikuRAG` client per invocation** via the new `haiku.skills>=0.15.0` `lifespan` hook. The skill's sub-agent opens one read-only client on entry, all tool calls reuse it, and it closes on exit — replacing the old pattern of open/close around every `search` / `list_documents` / `get_document` call. - **`max_searches` tracked on `RAGRunDeps.search_count`** instead of a module-level `ctx.run_id`-keyed dict. Eliminates a memory leak in long-running processes where old run ids were never evicted. +- **Analysis sandbox persists variables across `execute_code` calls within one invocation.** Re-enables the incremental-exploration workflow (search in one call, process results in the next). Each new skill invocation constructs a fresh `Sandbox` via the analysis lifespan, so there is no cross-invocation leak. ## [0.41.0] - 2026-04-20 diff --git a/haiku_rag_slim/haiku/rag/agents/analysis/sandbox.py b/haiku_rag_slim/haiku/rag/agents/analysis/sandbox.py index 88a61ea9..3232a03d 100644 --- a/haiku_rag_slim/haiku/rag/agents/analysis/sandbox.py +++ b/haiku_rag_slim/haiku/rag/agents/analysis/sandbox.py @@ -8,7 +8,7 @@ from pathlib import Path from typing import TYPE_CHECKING, Any, Literal import pydantic_monty -from pydantic_monty import CallbackFile, MemoryFile, OSAccess +from pydantic_monty import CallbackFile, MemoryFile, MontyRepl, OSAccess from haiku.rag.agents.analysis.dependencies import AnalysisContext from haiku.rag.config.models import AppConfig @@ -44,11 +44,12 @@ class Sandbox: and resolved asynchronously on the host. Documents are exposed via a virtual filesystem at ``/documents/{id}/``. - Each ``execute()`` call runs in a fresh interpreter — variables do not - persist between calls. + The interpreter uses a REPL session — variables persist across + ``execute()`` calls within the same Sandbox instance. sandbox = Sandbox(db_path, config, context) - result = await sandbox.execute("print('hello')") + result = await sandbox.execute("x = await search('query')") + result = await sandbox.execute("print(x[0]['content'])") # x persists """ _db_path: Path @@ -56,6 +57,8 @@ class Sandbox: _context: AnalysisContext _search_results: "list[SearchResult]" _items_cache: dict[str, str] | None + _repl: MontyRepl | None + _vfs: OSAccess | None def __init__( self, @@ -68,6 +71,8 @@ class Sandbox: self._context = context self._search_results = [] self._items_cache = None + self._repl = None + self._vfs = None def _build_external_functions(self) -> dict[str, Any]: """Build async external functions for the Monty interpreter.""" @@ -245,37 +250,46 @@ class Sandbox: return OSAccess(files) - async def execute(self, code: str) -> SandboxResult: - """Execute Python code in the Monty interpreter.""" - external_fns = self._build_external_functions() - vfs = await self._build_vfs() - - input_names: list[str] = [] - inputs: dict[str, Any] | None = None - if self._context.documents: - input_names.append("documents") - inputs = { - "documents": [ - { - "id": d.id, - "title": d.title, - "uri": d.uri, - "content": d.content, - } - for d in self._context.documents - ] - } - - try: - monty = pydantic_monty.Monty( - code, - inputs=input_names, + async def _ensure_initialized(self) -> tuple[MontyRepl, OSAccess]: + """Initialize the REPL session and VFS on first use.""" + if self._repl is None: + self._vfs = await self._build_vfs() + self._repl = MontyRepl( + limits={ + "max_duration_secs": self._config.analysis.code_timeout, + }, ) - except ( - pydantic_monty.MontySyntaxError, - pydantic_monty.MontyRuntimeError, - ) as e: - return SandboxResult(stdout="", stderr=str(e), success=False) + if self._context.documents: + await pydantic_monty.run_repl_async( + self._repl, + "pass", + inputs={ + "documents": [ + { + "id": d.id, + "title": d.title, + "uri": d.uri, + "content": d.content, + } + for d in self._context.documents + ] + }, + external_functions=self._build_external_functions(), + os=self._vfs, + ) + repl = self._repl + vfs = self._vfs + if repl is None or vfs is None: + raise RuntimeError("Sandbox initialization failed") + return repl, vfs + + async def execute(self, code: str) -> SandboxResult: + """Execute Python code in the Monty REPL. + + Variables persist across calls within the same Sandbox instance. + """ + repl, vfs = await self._ensure_initialized() + external_fns = self._build_external_functions() stdout_lines: list[str] = [] @@ -283,20 +297,19 @@ class Sandbox: stdout_lines.append(text) max_chars = self._config.analysis.max_output_chars - limits: pydantic_monty.ResourceLimits = { - "max_duration_secs": self._config.analysis.code_timeout, - } try: - output = await pydantic_monty.run_monty_async( - monty, - inputs=inputs, + output = await pydantic_monty.run_repl_async( + repl, + code, external_functions=external_fns, - limits=limits, print_callback=print_callback, os=vfs, ) - except pydantic_monty.MontyRuntimeError as e: + except ( + pydantic_monty.MontySyntaxError, + pydantic_monty.MontyRuntimeError, + ) as e: stdout = "".join(stdout_lines) if len(stdout) > max_chars: stdout = stdout[:max_chars] + "\n... (output truncated)" diff --git a/haiku_rag_slim/haiku/rag/skills/_deps.py b/haiku_rag_slim/haiku/rag/skills/_deps.py index eb20f6df..8f8795d3 100644 --- a/haiku_rag_slim/haiku/rag/skills/_deps.py +++ b/haiku_rag_slim/haiku/rag/skills/_deps.py @@ -34,3 +34,24 @@ def make_rag_lifespan(db_path: Path, config: AppConfig): yield return lifespan + + +def make_analysis_lifespan(db_path: Path, config: AppConfig): + @asynccontextmanager + async def lifespan(deps: AnalysisRunDeps) -> AsyncIterator[None]: + from haiku.rag.agents.analysis.dependencies import AnalysisContext + from haiku.rag.agents.analysis.sandbox import Sandbox + from haiku.rag.client import HaikuRAG + + doc_filter = getattr(deps.state, "document_filter", None) + async with HaikuRAG(db_path, config=config, read_only=True) as rag: + deps.rag = rag + deps.search_count = 0 + deps.sandbox = Sandbox( + db_path=db_path, + config=config, + context=AnalysisContext(filter=doc_filter), + ) + yield + + return lifespan diff --git a/haiku_rag_slim/haiku/rag/skills/_tools.py b/haiku_rag_slim/haiku/rag/skills/_tools.py index 909ec65d..0869c16c 100644 --- a/haiku_rag_slim/haiku/rag/skills/_tools.py +++ b/haiku_rag_slim/haiku/rag/skills/_tools.py @@ -227,27 +227,26 @@ def create_skill_tools( tools["get_document"] = get_document if "execute_code" in tool_names: + from haiku.rag.skills._deps import AnalysisRunDeps - async def execute_code(ctx: RunContext[RAGRunDeps], code: str) -> str: + async def execute_code(ctx: RunContext[AnalysisRunDeps], code: str) -> str: """Execute Python code in a sandboxed interpreter. The code has access to search(), list_documents(), llm() functions and a virtual filesystem at /documents/ with document content and structure (metadata.json, content.txt, items.jsonl per document). - Use print() to output results. Each call runs in a fresh - interpreter — variables do not persist between calls. + Use print() to output results. Variables persist between calls + within the same skill invocation. Args: code: Python code to execute. """ - from haiku.rag.agents.analysis.dependencies import AnalysisContext - from haiku.rag.agents.analysis.sandbox import Sandbox - - state = _get_state(ctx, state_type) - doc_filter = state.document_filter if state else None - context = AnalysisContext(filter=doc_filter) - sandbox = Sandbox(db_path=db_path, config=config, context=context) + if ctx.deps is None or ctx.deps.sandbox is None: + raise RuntimeError( + "AnalysisRunDeps.sandbox is not set — skill lifespan must run before execute_code." + ) + sandbox = ctx.deps.sandbox result = await sandbox.execute(code) state = _get_state(ctx, state_type) diff --git a/haiku_rag_slim/haiku/rag/skills/analysis.py b/haiku_rag_slim/haiku/rag/skills/analysis.py index c74637b3..dd05d236 100644 --- a/haiku_rag_slim/haiku/rag/skills/analysis.py +++ b/haiku_rag_slim/haiku/rag/skills/analysis.py @@ -60,7 +60,7 @@ def create_skill( config: haiku.rag AppConfig instance. If None, uses get_config(). """ from haiku.rag.config import get_config - from haiku.rag.skills._deps import AnalysisRunDeps, make_rag_lifespan + from haiku.rag.skills._deps import AnalysisRunDeps, make_analysis_lifespan from haiku.rag.skills._tools import create_skill_extras, create_skill_tools if config is None: @@ -95,5 +95,5 @@ def create_skill( state_type=STATE_TYPE, state_namespace=STATE_NAMESPACE, deps_type=AnalysisRunDeps, - lifespan=make_rag_lifespan(db_path, config), + lifespan=make_analysis_lifespan(db_path, config), ) diff --git a/haiku_rag_slim/haiku/rag/skills/rag-analysis/SKILL.md b/haiku_rag_slim/haiku/rag/skills/rag-analysis/SKILL.md index 80d9c099..65738896 100644 --- a/haiku_rag_slim/haiku/rag/skills/rag-analysis/SKILL.md +++ b/haiku_rag_slim/haiku/rag/skills/rag-analysis/SKILL.md @@ -15,7 +15,7 @@ You solve complex analytical questions by writing and executing Python code agai ## Tools ### execute_code -Execute Python code in a sandboxed interpreter. Each call runs in a fresh interpreter — write self-contained code. Use `print()` to output results. +Execute Python code in a sandboxed interpreter. Variables persist between calls — you can build state incrementally. Use `print()` to output results. Inside the code, these functions are available (use `await`): - `await search(query, limit=10)` → list of dicts with keys: chunk_id, content, document_id, document_title, document_uri, score, page_numbers, headings, doc_item_refs, labels @@ -93,7 +93,7 @@ Search results include `doc_item_refs` (e.g. `["#/texts/48", "#/tables/0"]`) tha ## Important -- Each `execute_code` call runs in a fresh interpreter — write self-contained code blocks +- Variables persist between `execute_code` calls — you can search in one call and process results in the next - Use `print()` to output results — the output is your only feedback - Always execute code to answer questions — don't just describe what code would do - Use `await` for all async functions inside execute_code (search, list_documents, llm) diff --git a/tests/skills/conftest.py b/tests/skills/conftest.py index 3bed1278..25a419b3 100644 --- a/tests/skills/conftest.py +++ b/tests/skills/conftest.py @@ -12,13 +12,13 @@ from haiku.rag.skills._deps import AnalysisRunDeps, RAGRunDeps VECTOR_DIM = 2560 -def _make_ctx(state=None, rag=None): +def _make_ctx(state=None, rag=None, sandbox=None): """Create a mock RunContext with RAGRunDeps (or AnalysisRunDeps when state is AnalysisState).""" from haiku.rag.skills.analysis import AnalysisState ctx = MagicMock(spec=RunContext) - if isinstance(state, AnalysisState): - ctx.deps = AnalysisRunDeps(state=state, rag=rag) + if isinstance(state, AnalysisState) or sandbox is not None: + ctx.deps = AnalysisRunDeps(state=state, rag=rag, sandbox=sandbox) else: ctx.deps = RAGRunDeps(state=state, rag=rag) return ctx @@ -80,3 +80,19 @@ async def rag_client(rag_db): """Yield an open read-only HaikuRAG client on the sample db.""" async with HaikuRAG(rag_db, read_only=True) as rag: yield rag + + +@pytest.fixture +def sandbox_factory(rag_db, test_app_config): + """Build Sandbox instances bound to the sample db, optionally with a doc filter.""" + from haiku.rag.agents.analysis.dependencies import AnalysisContext + from haiku.rag.agents.analysis.sandbox import Sandbox + + def _make(filter: str | None = None) -> Sandbox: + return Sandbox( + db_path=rag_db, + config=test_app_config, + context=AnalysisContext(filter=filter), + ) + + return _make diff --git a/tests/skills/test_analysis.py b/tests/skills/test_analysis.py index dffc5787..1119b4d8 100644 --- a/tests/skills/test_analysis.py +++ b/tests/skills/test_analysis.py @@ -113,73 +113,75 @@ class TestDomainPreambleInAnalysisSkillInstructions: class TestExecuteCodeTool: - async def test_execute_code_returns_output(self, rag_db): + async def test_execute_code_returns_output(self, rag_db, sandbox_factory): from haiku.rag.skills.analysis import create_skill skill = create_skill(db_path=rag_db) execute_code = _get_tool(skill, "execute_code") state = AnalysisState() - ctx = _make_ctx(state) + ctx = _make_ctx(state, sandbox=sandbox_factory()) result = await execute_code(ctx, code="print('hello')") assert "hello" in result - async def test_execute_code_updates_state(self, rag_db): + async def test_execute_code_updates_state(self, rag_db, sandbox_factory): from haiku.rag.skills.analysis import create_skill skill = create_skill(db_path=rag_db) execute_code = _get_tool(skill, "execute_code") state = AnalysisState() - ctx = _make_ctx(state) + ctx = _make_ctx(state, sandbox=sandbox_factory()) await execute_code(ctx, code="print('hello')") assert len(state.executions) == 1 assert state.executions[0].code == "print('hello')" assert state.executions[0].success is True assert "hello" in state.executions[0].stdout - async def test_execute_code_reports_errors(self, rag_db): + async def test_execute_code_reports_errors(self, rag_db, sandbox_factory): from haiku.rag.skills.analysis import create_skill skill = create_skill(db_path=rag_db) execute_code = _get_tool(skill, "execute_code") state = AnalysisState() - ctx = _make_ctx(state) + ctx = _make_ctx(state, sandbox=sandbox_factory()) result = await execute_code(ctx, code="x = 1/0") assert "Error" in result assert "ZeroDivisionError" in result assert state.executions[0].success is False - async def test_execute_code_applies_document_filter(self, rag_db): + async def test_execute_code_applies_document_filter(self, rag_db, sandbox_factory): from haiku.rag.skills.analysis import create_skill skill = create_skill(db_path=rag_db) execute_code = _get_tool(skill, "execute_code") state = AnalysisState(document_filter="title = 'AI Overview'") - ctx = _make_ctx(state) + ctx = _make_ctx(state, sandbox=sandbox_factory(filter=state.document_filter)) result = await execute_code( ctx, code="docs = await list_documents()\nprint(len(docs))" ) assert "1" in result - async def test_execute_code_accumulates_search_results(self, rag_db): + async def test_execute_code_accumulates_search_results( + self, rag_db, sandbox_factory + ): from haiku.rag.skills.analysis import create_skill skill = create_skill(db_path=rag_db) execute_code = _get_tool(skill, "execute_code") state = AnalysisState() - ctx = _make_ctx(state) + ctx = _make_ctx(state, sandbox=sandbox_factory()) await execute_code( ctx, code="results = await search('intelligence')\nprint(len(results))" ) assert "_sandbox" in state.searches assert len(state.searches["_sandbox"]) > 0 - async def test_execute_code_vfs_write_denied(self, rag_db): + async def test_execute_code_vfs_write_denied(self, rag_db, sandbox_factory): from haiku.rag.skills.analysis import create_skill skill = create_skill(db_path=rag_db) execute_code = _get_tool(skill, "execute_code") state = AnalysisState() - ctx = _make_ctx(state) + ctx = _make_ctx(state, sandbox=sandbox_factory()) result = await execute_code( ctx, code=( @@ -193,21 +195,67 @@ class TestExecuteCodeTool: assert "Error" in result assert "read-only" in result + async def test_execute_code_variables_persist_within_invocation( + self, rag_db, sandbox_factory + ): + """Same sandbox across two calls → vars persist (one skill invocation).""" + from haiku.rag.skills.analysis import create_skill + + skill = create_skill(db_path=rag_db) + execute_code = _get_tool(skill, "execute_code") + state = AnalysisState() + ctx = _make_ctx(state, sandbox=sandbox_factory()) + + await execute_code(ctx, code="x = 42") + result = await execute_code(ctx, code="print(x * 2)") + assert "84" in result + + async def test_execute_code_isolated_across_invocations( + self, rag_db, sandbox_factory + ): + """Different Sandbox instances → no cross-invocation leak.""" + from haiku.rag.skills.analysis import create_skill + + skill = create_skill(db_path=rag_db) + execute_code = _get_tool(skill, "execute_code") + + ctx1 = _make_ctx(AnalysisState(), sandbox=sandbox_factory()) + await execute_code(ctx1, code="secret = 'do not leak'") + + ctx2 = _make_ctx(AnalysisState(), sandbox=sandbox_factory()) + result = await execute_code(ctx2, code="print(secret)") + assert not result.startswith("do not leak") + assert "Error" in result or "NameError" in result + class TestAnalysisLifespan: - async def test_opens_one_client_per_invocation(self, rag_db): - from haiku.rag.skills._deps import AnalysisRunDeps, make_rag_lifespan + async def test_opens_client_and_sandbox_per_invocation(self, rag_db): + from haiku.rag.agents.analysis.sandbox import Sandbox + from haiku.rag.skills._deps import AnalysisRunDeps, make_analysis_lifespan config = AppConfig() - lifespan = make_rag_lifespan(rag_db, config) + lifespan = make_analysis_lifespan(rag_db, config) deps = AnalysisRunDeps() async with lifespan(deps): assert deps.rag is not None assert deps.rag.is_read_only assert deps.search_count == 0 + assert isinstance(deps.sandbox, Sandbox) docs = await deps.rag.list_documents() assert len(docs) == 2 + async def test_lifespan_reads_document_filter_from_state(self, rag_db): + from haiku.rag.skills._deps import AnalysisRunDeps, make_analysis_lifespan + from haiku.rag.skills.analysis import AnalysisState + + config = AppConfig() + lifespan = make_analysis_lifespan(rag_db, config) + state = AnalysisState(document_filter="title = 'AI Overview'") + deps = AnalysisRunDeps(state=state) + async with lifespan(deps): + assert deps.sandbox is not None + assert deps.sandbox._context.filter == "title = 'AI Overview'" + async def test_skill_has_lifespan_and_deps_type( self, test_app_config, temp_db_path ):