revert REPL persistence: fresh sandbox per execute_code call

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

@@ -1,4 +1,5 @@
 import asyncio
+import atexit
 import concurrent.futures
 import json
 from collections.abc import Callable
@@ -7,7 +8,7 @@ from pathlib import Path
 from typing import TYPE_CHECKING, Any, Literal
 
 import pydantic_monty
-from pydantic_monty import CallbackFile, MemoryFile, MontyRepl, OSAccess
+from pydantic_monty import CallbackFile, MemoryFile, OSAccess
 
 from haiku.rag.agents.analysis.dependencies import AnalysisContext
 from haiku.rag.config.models import AppConfig
@@ -27,6 +28,7 @@ class SandboxResult:
 
 
 _executor = concurrent.futures.ThreadPoolExecutor(max_workers=1)
+atexit.register(_executor.shutdown, wait=False)
 
 
 def _run_async(coro: Any) -> Any:
@@ -42,20 +44,17 @@ class Sandbox:
     and resolved asynchronously on the host.
     Documents are exposed via a virtual filesystem at ``/documents/{id}/``.
 
-    The interpreter uses a REPL session — variables persist across
+    Each ``execute()`` call runs in a fresh interpreter — variables do not
-    ``execute()`` calls within the same Sandbox instance.
+    persist between calls.
 
         sandbox = Sandbox(db_path, config, context)
-        result = await sandbox.execute("x = await search('query')")
+        result = await sandbox.execute("print('hello')")
-        result = await sandbox.execute("print(x[0]['content'])")  # x persists
     """
 
     _db_path: Path
     _config: AppConfig
     _context: AnalysisContext
     _search_results: "list[SearchResult]"
-    _repl: MontyRepl | None
-    _vfs: OSAccess | None
 
     def __init__(
         self,
@@ -67,8 +66,6 @@ class Sandbox:
         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."""
@@ -234,47 +231,37 @@ class Sandbox:
 
         return OSAccess(files)
 
-    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,
-                },
-            )
-            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,
-                )
-        # Both are guaranteed non-None after initialization
-        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.
+        """Execute Python code in the Monty interpreter."""
-
-        Variables persist across calls within the same Sandbox instance.
-        """
-        repl, vfs = await self._ensure_initialized()
         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,
+            )
+        except (
+            pydantic_monty.MontySyntaxError,
+            pydantic_monty.MontyRuntimeError,
+        ) as e:
+            return SandboxResult(stdout="", stderr=str(e), success=False)
 
         stdout_lines: list[str] = []
 
@@ -282,19 +269,20 @@ 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_repl_async(
+            output = await pydantic_monty.run_monty_async(
-                repl,
+                monty,
-                code,
+                inputs=inputs,
                 external_functions=external_fns,
+                limits=limits,
                 print_callback=print_callback,
                 os=vfs,
             )
-        except (
+        except pydantic_monty.MontyRuntimeError as e:
-            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/_tools.py

@@ -240,7 +240,6 @@ def create_skill_tools(
         tools["get_document"] = get_document
 
     if "execute_code" in tool_names:
-        _sandbox_state: dict[str, Any] = {}  # {run_id, sandbox} — reset per invocation
 
         async def execute_code(ctx: RunContext[SkillRunDeps], code: str) -> str:
             """Execute Python code in a sandboxed interpreter.
@@ -249,8 +248,8 @@ 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. Variables persist between calls
+            Use print() to output results. Each call runs in a fresh
-            within the same skill invocation.
+            interpreter — variables do not persist between calls.
 
             Args:
                 code: Python code to execute.
@@ -258,17 +257,10 @@ def create_skill_tools(
             from haiku.rag.agents.analysis.dependencies import AnalysisContext
             from haiku.rag.agents.analysis.sandbox import Sandbox
 
-            rid = ctx.run_id or ""
+            state = _get_state(ctx, state_type)
-            if _sandbox_state.get("run_id") != rid:
+            doc_filter = state.document_filter if state else None
-                state = _get_state(ctx, state_type)
+            context = AnalysisContext(filter=doc_filter)
-                doc_filter = state.document_filter if state else None
+            sandbox = Sandbox(db_path=db_path, config=config, context=context)
-                context = AnalysisContext(filter=doc_filter)
-                _sandbox_state["run_id"] = rid
-                _sandbox_state["sandbox"] = Sandbox(
-                    db_path=db_path, config=config, context=context
-                )
-
-            sandbox = _sandbox_state["sandbox"]
             result = await sandbox.execute(code)
 
             state = _get_state(ctx, state_type)

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. Variables persist between calls — you can build state incrementally. Use `print()` to output results.
+Execute Python code in a sandboxed interpreter. Each call runs in a fresh interpreter — write self-contained code. 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
 
-- Variables persist between `execute_code` calls — you can search in one call and process results in the next
+- Each `execute_code` call runs in a fresh interpreter — write self-contained code blocks
 - 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)