persist sandbox variables across execute_code calls within one invocation

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
 

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
+    The interpreter uses a REPL session — variables persist across
-    persist between calls.
+    ``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,16 +250,20 @@ class Sandbox:
 
         return OSAccess(files)
 
-    async def execute(self, code: str) -> SandboxResult:
+    async def _ensure_initialized(self) -> tuple[MontyRepl, OSAccess]:
-        """Execute Python code in the Monty interpreter."""
+        """Initialize the REPL session and VFS on first use."""
-        external_fns = self._build_external_functions()
+        if self._repl is None:
-        vfs = await self._build_vfs()
+            self._vfs = await self._build_vfs()
-
+            self._repl = MontyRepl(
-        input_names: list[str] = []
+                limits={
-        inputs: dict[str, Any] | None = None
+                    "max_duration_secs": self._config.analysis.code_timeout,
+                },
+            )
             if self._context.documents:
-            input_names.append("documents")
+                await pydantic_monty.run_repl_async(
-            inputs = {
+                    self._repl,
+                    "pass",
+                    inputs={
                         "documents": [
                             {
                                 "id": d.id,
@@ -264,18 +273,23 @@ class Sandbox:
                             }
                             for d in self._context.documents
                         ]
-            }
+                    },
-
+                    external_functions=self._build_external_functions(),
-        try:
+                    os=self._vfs,
-            monty = pydantic_monty.Monty(
-                code,
-                inputs=input_names,
                 )
-        except (
+        repl = self._repl
-            pydantic_monty.MontySyntaxError,
+        vfs = self._vfs
-            pydantic_monty.MontyRuntimeError,
+        if repl is None or vfs is None:
-        ) as e:
+            raise RuntimeError("Sandbox initialization failed")
-            return SandboxResult(stdout="", stderr=str(e), success=False)
+        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(
+            output = await pydantic_monty.run_repl_async(
-                monty,
+                repl,
-                inputs=inputs,
+                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)"

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

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
+            Use print() to output results. Variables persist between calls
-            interpreter — variables do not persist between calls.
+            within the same skill invocation.
 
             Args:
                 code: Python code to execute.
             """
-            from haiku.rag.agents.analysis.dependencies import AnalysisContext
+            if ctx.deps is None or ctx.deps.sandbox is None:
-            from haiku.rag.agents.analysis.sandbox import Sandbox
+                raise RuntimeError(
-
+                    "AnalysisRunDeps.sandbox is not set — skill lifespan must run before execute_code."
-            state = _get_state(ctx, state_type)
+                )
-            doc_filter = state.document_filter if state else None
+            sandbox = ctx.deps.sandbox
-            context = AnalysisContext(filter=doc_filter)
-            sandbox = Sandbox(db_path=db_path, config=config, context=context)
             result = await sandbox.execute(code)
 
             state = _get_state(ctx, state_type)

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),
     )

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)

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):
+    if isinstance(state, AnalysisState) or sandbox is not None:
-        ctx.deps = AnalysisRunDeps(state=state, rag=rag)
+        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

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):
+    async def test_opens_client_and_sandbox_per_invocation(self, rag_db):
-        from haiku.rag.skills._deps import AnalysisRunDeps, make_rag_lifespan
+        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
     ):