use MontyRepl for persistent variables across execute_code calls

haiku_rag_slim/haiku/rag/agents/analysis/sandbox.py

@@ -3,10 +3,11 @@ import concurrent.futures
 import json
 from collections.abc import Callable
 from dataclasses import dataclass
+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
@@ -15,8 +16,6 @@ from haiku.rag.store.models.chunk import SearchResult
 if TYPE_CHECKING:
     from pathlib import PurePosixPath
 
-    from haiku.rag.client import HaikuRAG
-
 
 @dataclass
 class SandboxResult:
@@ -41,35 +40,46 @@ class Sandbox:
     and resolved asynchronously on the host.
     Documents are exposed via a virtual filesystem at ``/documents/{id}/``.
 
-        sandbox = Sandbox(client, config, context)
+    The interpreter uses a REPL session — variables persist across
-        result = await sandbox.execute("print('hello')")
+    ``execute()`` calls within the same Sandbox instance.
+
+        sandbox = Sandbox(db_path, config, context)
+        result = await sandbox.execute("x = await search('query')")
+        result = await sandbox.execute("print(x[0]['content'])")  # x persists
     """
 
-    _client: "HaikuRAG"
+    _db_path: Path
     _config: AppConfig
     _context: AnalysisContext
     _search_results: "list[SearchResult]"
+    _repl: MontyRepl | None
+    _vfs: OSAccess | None
 
     def __init__(
         self,
-        client: "HaikuRAG",
+        db_path: Path,
         config: AppConfig,
         context: AnalysisContext,
     ):
-        self._client = client
+        self._db_path = db_path
         self._config = config
         self._context = context
         self._search_results = []
+        self._repl = None
+        self._vfs = None
 
     def _build_external_functions(self) -> dict[str, Any]:
         """Build async external functions for the Monty interpreter."""
-        client = self._client
+        db_path = self._db_path
         config = self._config
         context = self._context
 
         async def search(query: str, limit: int = 10) -> list[dict[str, Any]]:
-            results = await client.search(query, limit=limit, filter=context.filter)
+            from haiku.rag.client import HaikuRAG
-            expanded = await client.expand_context(results)
+
+            async with HaikuRAG(db_path, config=config, read_only=True) as rag:
+                results = await rag.search(query, limit=limit, filter=context.filter)
+                expanded = await rag.expand_context(results)
             self._search_results.extend(expanded)
             return [
                 {
@@ -88,7 +98,10 @@ class Sandbox:
             ]
 
         async def list_documents() -> list[dict[str, Any]]:
-            docs = await client.list_documents(filter=context.filter)
+            from haiku.rag.client import HaikuRAG
+
+            async with HaikuRAG(db_path, config=config, read_only=True) as rag:
+                docs = await rag.list_documents(filter=context.filter)
             return [
                 {
                     "id": d.id,
@@ -123,10 +136,14 @@ class Sandbox:
         - content.txt: CallbackFile (lazy, can be large)
         - items.jsonl: CallbackFile (lazy, can be large)
         """
-        client = self._client
+        from haiku.rag.client import HaikuRAG
+
+        db_path = self._db_path
+        config = self._config
         files: list[MemoryFile | CallbackFile] = []
 
-        docs = await client.list_documents(filter=self._context.filter)
+        async with HaikuRAG(db_path, config=config, read_only=True) as rag:
+            docs = await rag.list_documents(filter=self._context.filter)
 
         for doc in docs:
             if not doc.id:
@@ -150,11 +167,15 @@ class Sandbox:
             ) -> Callable[["PurePosixPath"], str]:
                 def read_content(_path: "PurePosixPath") -> str:
                     async def _fetch() -> str:
+                        from haiku.rag.client import HaikuRAG
                         from haiku.rag.utils import escape_sql_string
 
+                        async with HaikuRAG(
+                            db_path, config=config, read_only=True
+                        ) as rag:
                             safe_id = escape_sql_string(did)
                             rows = list(
-                            client.store.documents_table.search()
+                                rag.store.documents_table.search()
                                 .select(["content"])
                                 .where(f"id = '{safe_id}'")
                                 .limit(1)
@@ -171,8 +192,13 @@ class Sandbox:
             ) -> Callable[["PurePosixPath"], str]:
                 def read_items(_path: "PurePosixPath") -> str:
                     async def _fetch() -> str:
+                        from haiku.rag.client import HaikuRAG
+
+                        async with HaikuRAG(
+                            db_path, config=config, read_only=True
+                        ) as rag:
                             items = (
-                            await client.document_item_repository.get_items_in_range(
+                                await rag.document_item_repository.get_items_in_range(
                                     did, 0, 999999
                                 )
                             )
@@ -213,15 +239,19 @@ class Sandbox:
 
         return OSAccess(files)
 
-    async def execute(self, code: str) -> SandboxResult:
+    async def _ensure_initialized(self) -> None:
-        """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(
+                    self._repl,
+                    "pass",
                     inputs={
                         "documents": [
                             {
@@ -232,18 +262,21 @@ 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 (
+
-            pydantic_monty.MontySyntaxError,
+    async def execute(self, code: str) -> SandboxResult:
-            pydantic_monty.MontyRuntimeError,
+        """Execute Python code in the Monty REPL.
-        ) as e:
+
-            return SandboxResult(stdout="", stderr=str(e), success=False)
+        Variables persist across calls within the same Sandbox instance.
+        """
+        await self._ensure_initialized()
+        assert self._repl is not None
+        assert self._vfs is not None
+
+        external_fns = self._build_external_functions()
 
         stdout_lines: list[str] = []
 
@@ -251,20 +284,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,
+                self._repl,
-                inputs=inputs,
+                code,
                 external_functions=external_fns,
-                limits=limits,
                 print_callback=print_callback,
-                os=vfs,
+                os=self._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/client.py

@@ -1230,7 +1230,7 @@ class HaikuRAG:
             context.documents = loaded_docs if loaded_docs else None
 
         sandbox = Sandbox(
-            client=self,
+            db_path=self.store.db_path,
             config=self._config,
             context=context,
         )

haiku_rag_slim/haiku/rag/skills/_tools.py

@@ -240,6 +240,7 @@ def create_skill_tools(
         tools["get_document"] = get_document
 
     if "execute_code" in tool_names:
+        _sandbox: list[Any] = []  # mutable container for closure; holds [Sandbox] or []
 
         async def execute_code(ctx: RunContext[SkillRunDeps], code: str) -> str:
             """Execute Python code in a sandboxed interpreter.
@@ -248,24 +249,26 @@ def create_skill_tools(
             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.
 
             Args:
                 code: Python code to execute.
             """
             from haiku.rag.agents.analysis.dependencies import AnalysisContext
             from haiku.rag.agents.analysis.sandbox import Sandbox
-            from haiku.rag.client import HaikuRAG
 
+            if not _sandbox:
                 state = _get_state(ctx, state_type)
                 doc_filter = state.document_filter if state else None
                 context = AnalysisContext(filter=doc_filter)
+                _sandbox.append(
+                    Sandbox(db_path=db_path, config=config, context=context)
+                )
 
-            async with HaikuRAG(db_path, config=config, read_only=True) as rag:
+            sandbox = _sandbox[0]
-                sandbox = Sandbox(client=rag, config=config, context=context)
             result = await sandbox.execute(code)
 
+            state = _get_state(ctx, state_type)
             if state and sandbox._search_results:
                 existing = state.searches.get("_sandbox", [])
                 seen = {r.chunk_id for r in existing}

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 — variables do not persist between calls. 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
@@ -71,7 +71,7 @@ Search results include `doc_item_refs` (e.g. `["#/texts/48", "#/tables/0"]`) tha
 
 ## Important
 
-- Each `execute_code` call runs in a fresh interpreter (no persistent variables between calls)
+- 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/agents/analysis/conftest.py

@@ -14,8 +14,9 @@ async def empty_client(temp_db_path):
 
 
 @pytest.fixture
-async def sandbox(empty_client):
+async def sandbox(temp_db_path):
     """Create a Monty sandbox for testing."""
+    async with HaikuRAG(temp_db_path, create=True):
         config = AppConfig()
         context = AnalysisContext()
-    return Sandbox(client=empty_client, config=config, context=context)
+        return Sandbox(db_path=temp_db_path, config=config, context=context)

tests/agents/analysis/test_sandbox.py

@@ -99,7 +99,7 @@ class TestSandboxListDocuments:
             )
 
             context = AnalysisContext()
-            sb = Sandbox(client=client, config=config, context=context)
+            sb = Sandbox(db_path=temp_db_path, config=config, context=context)
             result = await sb.execute(
                 "docs = await list_documents()\n"
                 "print(len(docs))\n"
@@ -126,7 +126,7 @@ class TestSandboxSearch:
             )
 
             context = AnalysisContext()
-            sb = Sandbox(client=client, config=config, context=context)
+            sb = Sandbox(db_path=temp_db_path, config=config, context=context)
             result = await sb.execute(
                 "results = await search('fox', limit=5)\n"
                 "print(len(results))\n"
@@ -149,7 +149,7 @@ class TestSandboxSearch:
             )
 
             context = AnalysisContext()
-            sb = Sandbox(client=client, config=config, context=context)
+            sb = Sandbox(db_path=temp_db_path, config=config, context=context)
             result = await sb.execute(
                 "results = await search('fox', limit=1)\n"
                 "r = results[0]\n"
@@ -175,7 +175,7 @@ class TestSandboxSearch:
             )
 
             context = AnalysisContext()
-            sb = Sandbox(client=client, config=config, context=context)
+            sb = Sandbox(db_path=temp_db_path, config=config, context=context)
             result = await sb.execute(
                 "results = await search('fox', limit=1)\n"
                 "print(type(results[0]['content']).__name__)\n"
@@ -242,12 +242,13 @@ class TestSandboxOutputTruncation:
     """Test output truncation behavior."""
 
     @pytest.mark.asyncio
-    async def test_truncate_stdout_on_runtime_error(self, empty_client):
+    async def test_truncate_stdout_on_runtime_error(self, temp_db_path):
         """Test stdout is truncated when a runtime error occurs after large output."""
+        async with HaikuRAG(temp_db_path, create=True):
             config = AppConfig()
             config.analysis.max_output_chars = 20
             context = AnalysisContext()
-        sb = Sandbox(client=empty_client, config=config, context=context)
+            sb = Sandbox(db_path=temp_db_path, config=config, context=context)
             result = await sb.execute("print('a' * 100)\nx = 1/0")
             assert not result.success
             assert "ZeroDivisionError" in result.stderr
@@ -255,12 +256,13 @@ class TestSandboxOutputTruncation:
             assert len(result.stdout) < 100
 
     @pytest.mark.asyncio
-    async def test_truncate_successful_output(self, empty_client):
+    async def test_truncate_successful_output(self, temp_db_path):
         """Test output is truncated on successful execution with large output."""
+        async with HaikuRAG(temp_db_path, create=True):
             config = AppConfig()
             config.analysis.max_output_chars = 20
             context = AnalysisContext()
-        sb = Sandbox(client=empty_client, config=config, context=context)
+            sb = Sandbox(db_path=temp_db_path, config=config, context=context)
             result = await sb.execute("print('b' * 100)")
             assert result.success
             assert result.stdout.endswith("... (output truncated)")
@@ -293,7 +295,7 @@ class TestSandboxVFS:
             )
 
             context = AnalysisContext()
-            sb = Sandbox(client=client, config=config, context=context)
+            sb = Sandbox(db_path=temp_db_path, config=config, context=context)
             result = await sb.execute(
                 "from pathlib import Path\n"
                 "dirs = list(Path('/documents').iterdir())\n"
@@ -317,7 +319,7 @@ class TestSandboxVFS:
             )
 
             context = AnalysisContext()
-            sb = Sandbox(client=client, config=config, context=context)
+            sb = Sandbox(db_path=temp_db_path, config=config, context=context)
             result = await sb.execute(
                 "from pathlib import Path\n"
                 "import json\n"
@@ -342,7 +344,7 @@ class TestSandboxVFS:
             )
 
             context = AnalysisContext()
-            sb = Sandbox(client=client, config=config, context=context)
+            sb = Sandbox(db_path=temp_db_path, config=config, context=context)
             result = await sb.execute(
                 "from pathlib import Path\n"
                 f"content = Path('/documents/{doc.id}/content.txt').read_text()\n"
@@ -364,7 +366,7 @@ class TestSandboxVFS:
             )
 
             context = AnalysisContext()
-            sb = Sandbox(client=client, config=config, context=context)
+            sb = Sandbox(db_path=temp_db_path, config=config, context=context)
             result = await sb.execute(
                 "from pathlib import Path\n"
                 "import json\n"
@@ -399,7 +401,7 @@ class TestSandboxVFS:
             )
 
             context = AnalysisContext(filter="uri LIKE 'public://%'")
-            sb = Sandbox(client=client, config=config, context=context)
+            sb = Sandbox(db_path=temp_db_path, config=config, context=context)
             result = await sb.execute(
                 "from pathlib import Path\n"
                 "import json\n"
@@ -425,15 +427,16 @@ class TestSandboxPreloadedDocuments:
         assert "NameError" in result.stderr
 
     @pytest.mark.asyncio
-    async def test_documents_variable_available_with_preload(self, empty_client):
+    async def test_documents_variable_available_with_preload(self, temp_db_path):
         """documents variable is available when context.documents is set."""
+        async with HaikuRAG(temp_db_path, create=True):
             config = AppConfig()
             docs = [
                 Document(id="1", content="Content A", title="Doc A", uri="a://1"),
                 Document(id="2", content="Content B", title="Doc B", uri="b://2"),
             ]
             context = AnalysisContext(documents=docs)
-        sb = Sandbox(client=empty_client, config=config, context=context)
+            sb = Sandbox(db_path=temp_db_path, config=config, context=context)
             result = await sb.execute(
                 "print(len(documents))\n"
                 "print(documents[0]['title'])\n"
@@ -450,11 +453,12 @@ class TestSandboxLLM:
 
     @pytest.mark.asyncio
     @pytest.mark.vcr()
-    async def test_llm_function(self, allow_model_requests, empty_client):
+    async def test_llm_function(self, allow_model_requests, temp_db_path):
         """Test llm() calls the model and returns a string."""
+        async with HaikuRAG(temp_db_path, create=True):
             config = AppConfig()
             context = AnalysisContext()
-        sb = Sandbox(client=empty_client, config=config, context=context)
+            sb = Sandbox(db_path=temp_db_path, config=config, context=context)
             result = await sb.execute(
                 "answer = await llm('What is 2 + 2? Reply with just the number.')\n"
                 "print(answer)"