RLM sandbox environment

haiku_rag_slim/haiku/rag/agents/rlm/__init__.py

@@ -0,0 +1,10 @@
+from haiku.rag.agents.rlm.dependencies import RLMConfig, RLMContext, RLMDeps
+from haiku.rag.agents.rlm.sandbox import REPLEnvironment, REPLResult
+
+__all__ = [
+    "RLMConfig",
+    "RLMContext",
+    "RLMDeps",
+    "REPLEnvironment",
+    "REPLResult",
+]

haiku_rag_slim/haiku/rag/agents/rlm/dependencies.py

@@ -0,0 +1,37 @@
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING
+
+from pydantic import BaseModel
+
+from haiku.rag.store.models import Document, SearchResult
+
+if TYPE_CHECKING:
+    from haiku.rag.client import HaikuRAG
+    from haiku.rag.config.models import AppConfig
+
+
+class RLMConfig(BaseModel):
+    """Configuration for RLM agent sandbox execution."""
+
+    code_timeout: float = 60.0
+    max_output_chars: int = 50_000
+    max_tool_calls: int = 20
+
+
+@dataclass
+class RLMContext:
+    """Mutable context accumulating data during RLM execution."""
+
+    documents: list[Document] | None = None
+    search_results: list[SearchResult] = field(default_factory=list)
+    code_executions: list[dict] = field(default_factory=list)
+
+
+@dataclass
+class RLMDeps:
+    """Dependencies for RLM agent."""
+
+    client: "HaikuRAG"
+    config: "AppConfig"
+    rlm_config: RLMConfig = field(default_factory=RLMConfig)
+    context: RLMContext = field(default_factory=RLMContext)

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

@@ -0,0 +1,413 @@
+import ast
+import asyncio
+import concurrent.futures
+import sys
+import traceback
+from io import StringIO
+from typing import TYPE_CHECKING, Any
+
+from haiku.rag.agents.rlm.dependencies import RLMConfig, RLMContext
+
+if TYPE_CHECKING:
+    from haiku.rag.client import HaikuRAG
+
+
+class REPLResult:
+    """Result of executing code in the REPL environment."""
+
+    def __init__(
+        self, stdout: str, stderr: str, success: bool, locals_: dict | None = None
+    ):
+        self.stdout = stdout
+        self.stderr = stderr
+        self.success = success
+        self.locals = locals_ or {}
+
+    def __repr__(self) -> str:
+        return f"REPLResult(success={self.success}, stdout={self.stdout!r}, stderr={self.stderr!r})"
+
+
+class REPLEnvironment:
+    """Sandboxed Python execution environment with haiku.rag access."""
+
+    SAFE_BUILTINS: dict[str, Any] = {
+        "True": True,
+        "False": False,
+        "None": None,
+        "__build_class__": __builtins__["__build_class__"]
+        if isinstance(__builtins__, dict)
+        else getattr(__builtins__, "__build_class__"),
+        "abs": abs,
+        "all": all,
+        "any": any,
+        "ascii": ascii,
+        "bin": bin,
+        "bool": bool,
+        "bytearray": bytearray,
+        "bytes": bytes,
+        "callable": callable,
+        "chr": chr,
+        "complex": complex,
+        "dict": dict,
+        "divmod": divmod,
+        "enumerate": enumerate,
+        "filter": filter,
+        "float": float,
+        "format": format,
+        "frozenset": frozenset,
+        "hash": hash,
+        "hex": hex,
+        "id": id,
+        "int": int,
+        "isinstance": isinstance,
+        "issubclass": issubclass,
+        "iter": iter,
+        "len": len,
+        "list": list,
+        "map": map,
+        "max": max,
+        "min": min,
+        "next": next,
+        "object": object,
+        "oct": oct,
+        "ord": ord,
+        "pow": pow,
+        "print": print,
+        "range": range,
+        "repr": repr,
+        "reversed": reversed,
+        "round": round,
+        "set": set,
+        "slice": slice,
+        "sorted": sorted,
+        "str": str,
+        "sum": sum,
+        "tuple": tuple,
+        "type": type,
+        "zip": zip,
+        "Exception": Exception,
+        "ValueError": ValueError,
+        "TypeError": TypeError,
+        "KeyError": KeyError,
+        "IndexError": IndexError,
+        "AttributeError": AttributeError,
+        "RuntimeError": RuntimeError,
+        "StopIteration": StopIteration,
+        "ZeroDivisionError": ZeroDivisionError,
+        "AssertionError": AssertionError,
+    }
+
+    ALLOWED_IMPORTS = {
+        "json",
+        "re",
+        "collections",
+        "math",
+        "statistics",
+        "itertools",
+        "functools",
+        "datetime",
+        "typing",
+    }
+
+    def __init__(
+        self,
+        client: "HaikuRAG",
+        config: RLMConfig,
+        context: RLMContext,
+        event_loop: asyncio.AbstractEventLoop | None = None,
+    ):
+        self.client = client
+        self.config = config
+        self.context = context
+        self._event_loop = event_loop
+        self._setup_namespace()
+
+    def _run_async_from_thread(self, coro):
+        """Run async coroutine from a worker thread using run_coroutine_threadsafe."""
+        if self._event_loop is None:
+            raise RuntimeError("Event loop not set. Cannot call async functions.")
+        future = asyncio.run_coroutine_threadsafe(coro, self._event_loop)
+        return future.result(timeout=self.config.code_timeout)
+
+    def _setup_namespace(self) -> None:
+        """Build execution namespace with haiku.rag functions."""
+        self.globals: dict[str, Any] = {
+            "__builtins__": dict(self.SAFE_BUILTINS),
+            "__name__": "__sandbox__",
+            "search": self._make_search(),
+            "list_documents": self._make_list_documents(),
+            "get_document": self._make_get_document(),
+            "get_docling_document": self._make_get_docling_document(),
+            "ask": self._make_ask(),
+        }
+        self.locals: dict[str, Any] = {}
+
+        if self.context.documents:
+            self.globals["documents"] = [
+                {"id": d.id, "title": d.title, "uri": d.uri, "content": d.content}
+                for d in self.context.documents
+            ]
+
+    def _make_search(self):
+        """Create sync search function that bridges to async client."""
+
+        def search(
+            query: str, limit: int = 10, filter: str | None = None
+        ) -> list[dict]:
+            async def _search():
+                return await self.client.search(query, limit=limit, filter=filter)
+
+            results = self._run_async_from_thread(_search())
+            self.context.search_results.extend(results)
+            return [
+                {
+                    "chunk_id": r.chunk_id,
+                    "content": r.content,
+                    "document_id": r.document_id,
+                    "document_title": r.document_title,
+                    "document_uri": r.document_uri,
+                    "score": r.score,
+                    "page_numbers": r.page_numbers,
+                    "headings": r.headings,
+                }
+                for r in results
+            ]
+
+        return search
+
+    def _make_list_documents(self):
+        """Create sync list_documents function."""
+
+        def list_documents(
+            limit: int = 10, offset: int = 0, filter: str | None = None
+        ) -> list[dict]:
+            async def _list():
+                return await self.client.list_documents(
+                    limit=limit, offset=offset, filter=filter
+                )
+
+            docs = self._run_async_from_thread(_list())
+            return [
+                {
+                    "id": d.id,
+                    "title": d.title,
+                    "uri": d.uri,
+                    "created_at": str(d.created_at),
+                }
+                for d in docs
+            ]
+
+        return list_documents
+
+    def _make_get_document(self):
+        """Create sync get_document function that returns text content."""
+
+        def get_document(id_or_title: str) -> str | None:
+            async def _get():
+                doc = await self.client.get_document_by_id(id_or_title)
+                if doc:
+                    return doc.content
+                docs = await self.client.list_documents(
+                    filter=f"title = '{id_or_title}'"
+                )
+                if docs and docs[0].id:
+                    full_doc = await self.client.get_document_by_id(docs[0].id)
+                    return full_doc.content if full_doc else None
+                docs = await self.client.list_documents(filter=f"uri = '{id_or_title}'")
+                if docs and docs[0].id:
+                    full_doc = await self.client.get_document_by_id(docs[0].id)
+                    return full_doc.content if full_doc else None
+                return None
+
+            return self._run_async_from_thread(_get())
+
+        return get_document
+
+    def _make_get_docling_document(self):
+        """Create sync get_docling_document function that returns DoclingDocument."""
+
+        def get_docling_document(id_or_title: str):
+            async def _get():
+                doc = await self.client.get_document_by_id(id_or_title)
+                if doc:
+                    return doc.docling_document
+                docs = await self.client.list_documents(
+                    filter=f"title = '{id_or_title}'"
+                )
+                if docs and docs[0].id:
+                    full_doc = await self.client.get_document_by_id(docs[0].id)
+                    return full_doc.docling_document if full_doc else None
+                docs = await self.client.list_documents(filter=f"uri = '{id_or_title}'")
+                if docs and docs[0].id:
+                    full_doc = await self.client.get_document_by_id(docs[0].id)
+                    return full_doc.docling_document if full_doc else None
+                return None
+
+            return self._run_async_from_thread(_get())
+
+        return get_docling_document
+
+    def _make_ask(self):
+        """Create sync ask function that uses QA agent."""
+
+        def ask(question: str, filter: str | None = None) -> str:
+            async def _ask():
+                answer, citations = await self.client.ask(question, filter=filter)
+                for c in citations:
+                    for sr in self.context.search_results:
+                        if sr.chunk_id == c.chunk_id:
+                            break
+                    else:
+                        from haiku.rag.store.models import SearchResult
+
+                        self.context.search_results.append(
+                            SearchResult(
+                                chunk_id=c.chunk_id,
+                                document_id=c.document_id,
+                                document_title=c.document_title or "",
+                                document_uri=c.document_uri,
+                                content=c.content,
+                                score=1.0,
+                                page_numbers=c.page_numbers,
+                                headings=c.headings or [],
+                            )
+                        )
+                return answer
+
+            return self._run_async_from_thread(_ask())
+
+        return ask
+
+    def _safe_import(
+        self,
+        name: str,
+        globals: dict | None = None,
+        locals: dict | None = None,
+        fromlist: tuple = (),
+        level: int = 0,
+    ):
+        """Import hook that only allows safe modules."""
+        base_module = name.split(".")[0]
+        if base_module not in self.ALLOWED_IMPORTS:
+            raise ImportError(f"Import of '{name}' is not allowed in sandbox")
+
+        import importlib
+
+        module = importlib.import_module(name)
+        if fromlist:
+            for attr in fromlist:
+                if not hasattr(module, attr):
+                    raise ImportError(f"cannot import name '{attr}' from '{name}'")
+            return module
+        return module
+
+    def _validate_code(self, code: str) -> None:
+        """Validate code AST for security issues."""
+        tree = ast.parse(code)
+
+        for node in ast.walk(tree):
+            if isinstance(node, ast.Attribute):
+                if node.attr.startswith("_") and node.attr not in (
+                    "__init__",
+                    "__str__",
+                    "__repr__",
+                    "__class__",
+                    "__name__",
+                    "__doc__",
+                    "__dict__",
+                ):
+                    raise SecurityError(
+                        f"Access to private/dunder attribute '{node.attr}' is not allowed"
+                    )
+
+    def _execute_sync(self, code: str) -> REPLResult:
+        """Internal synchronous execution - must be called from executor thread."""
+        stdout_capture = StringIO()
+        stderr_capture = StringIO()
+
+        original_stdout = sys.stdout
+        original_stderr = sys.stderr
+
+        try:
+            self._validate_code(code)
+        except SyntaxError as e:
+            return REPLResult(
+                stdout="",
+                stderr=f"SyntaxError: {e}",
+                success=False,
+            )
+        except SecurityError as e:
+            return REPLResult(
+                stdout="",
+                stderr=str(e),
+                success=False,
+            )
+
+        exec_globals = dict(self.globals)
+        exec_globals["__builtins__"] = dict(self.SAFE_BUILTINS)
+        exec_globals["__builtins__"]["__import__"] = self._safe_import
+
+        try:
+            sys.stdout = stdout_capture
+            sys.stderr = stderr_capture
+
+            exec(code, exec_globals, self.locals)
+
+            for key, value in self.locals.items():
+                if not key.startswith("_"):
+                    self.globals[key] = value
+
+            stdout = stdout_capture.getvalue()
+            if len(stdout) > self.config.max_output_chars:
+                stdout = (
+                    stdout[: self.config.max_output_chars] + "\n... (output truncated)"
+                )
+
+            return REPLResult(
+                stdout=stdout,
+                stderr=stderr_capture.getvalue(),
+                success=True,
+                locals_=dict(self.locals),
+            )
+
+        except Exception:
+            tb = traceback.format_exc()
+            return REPLResult(
+                stdout=stdout_capture.getvalue(),
+                stderr=tb,
+                success=False,
+            )
+
+        finally:
+            sys.stdout = original_stdout
+            sys.stderr = original_stderr
+
+    def execute(self, code: str) -> REPLResult:
+        """Execute code in sandbox synchronously.
+
+        This method runs code directly in the current thread.
+        For async contexts, use execute_async() instead.
+        """
+        return self._execute_sync(code)
+
+    async def execute_async(self, code: str) -> REPLResult:
+        """Execute code in sandbox from async context.
+
+        Runs the synchronous code in a thread executor, allowing
+        sandbox functions to call back to async client methods.
+        """
+        loop = asyncio.get_running_loop()
+        self._event_loop = loop
+
+        with concurrent.futures.ThreadPoolExecutor(max_workers=1) as executor:
+            result = await asyncio.wait_for(
+                loop.run_in_executor(executor, self._execute_sync, code),
+                timeout=self.config.code_timeout,
+            )
+        return result
+
+
+class SecurityError(Exception):
+    """Raised when sandbox security is violated."""
+
+    pass

tests/agents/rlm/conftest.py

@@ -0,0 +1,20 @@
+import pytest
+
+from haiku.rag.agents.rlm.dependencies import RLMConfig, RLMContext
+from haiku.rag.agents.rlm.sandbox import REPLEnvironment
+from haiku.rag.client import HaikuRAG
+
+
+@pytest.fixture
+async def empty_client(temp_db_path):
+    """Create an empty HaikuRAG client without documents."""
+    async with HaikuRAG(temp_db_path, create=True) as client:
+        yield client
+
+
+@pytest.fixture
+async def repl_env_empty(empty_client):
+    """Create a REPL environment without documents."""
+    config = RLMConfig()
+    context = RLMContext()
+    return REPLEnvironment(client=empty_client, config=config, context=context)

tests/agents/rlm/test_sandbox.py

@@ -0,0 +1,490 @@
+import pytest
+
+
+class TestSafeBuiltins:
+    """Test that safe builtins are available."""
+
+    @pytest.mark.asyncio
+    async def test_print_available(self, repl_env_empty):
+        result = await repl_env_empty.execute_async("print('hello')")
+        assert result.success
+        assert "hello" in result.stdout
+
+    @pytest.mark.asyncio
+    async def test_len_available(self, repl_env_empty):
+        result = await repl_env_empty.execute_async("print(len([1, 2, 3]))")
+        assert result.success
+        assert "3" in result.stdout
+
+    @pytest.mark.asyncio
+    async def test_range_available(self, repl_env_empty):
+        result = await repl_env_empty.execute_async("print(list(range(3)))")
+        assert result.success
+        assert "[0, 1, 2]" in result.stdout
+
+    @pytest.mark.asyncio
+    async def test_enumerate_available(self, repl_env_empty):
+        result = await repl_env_empty.execute_async(
+            "print(list(enumerate(['a', 'b'])))"
+        )
+        assert result.success
+        assert "[(0, 'a'), (1, 'b')]" in result.stdout
+
+    @pytest.mark.asyncio
+    async def test_sorted_available(self, repl_env_empty):
+        result = await repl_env_empty.execute_async("print(sorted([3, 1, 2]))")
+        assert result.success
+        assert "[1, 2, 3]" in result.stdout
+
+    @pytest.mark.asyncio
+    async def test_sum_available(self, repl_env_empty):
+        result = await repl_env_empty.execute_async("print(sum([1, 2, 3]))")
+        assert result.success
+        assert "6" in result.stdout
+
+    @pytest.mark.asyncio
+    async def test_min_max_available(self, repl_env_empty):
+        result = await repl_env_empty.execute_async(
+            "print(min([3, 1, 2]), max([3, 1, 2]))"
+        )
+        assert result.success
+        assert "1 3" in result.stdout
+
+    @pytest.mark.asyncio
+    async def test_all_any_available(self, repl_env_empty):
+        result = await repl_env_empty.execute_async(
+            "print(all([True, True]), any([False, True]))"
+        )
+        assert result.success
+        assert "True True" in result.stdout
+
+    @pytest.mark.asyncio
+    async def test_dict_list_set_tuple_available(self, repl_env_empty):
+        result = await repl_env_empty.execute_async(
+            "print(dict(a=1), list((1,2)), set([1,2,1]), tuple([1,2]))"
+        )
+        assert result.success
+        assert "{'a': 1}" in result.stdout
+
+    @pytest.mark.asyncio
+    async def test_str_int_float_bool_available(self, repl_env_empty):
+        result = await repl_env_empty.execute_async(
+            "print(str(1), int('2'), float('3.0'), bool(1))"
+        )
+        assert result.success
+        assert "1 2 3.0 True" in result.stdout
+
+    @pytest.mark.asyncio
+    async def test_zip_map_filter_available(self, repl_env_empty):
+        result = await repl_env_empty.execute_async(
+            "print(list(zip([1,2], ['a','b'])), "
+            "list(map(str, [1,2])), "
+            "list(filter(lambda x: x > 1, [1,2,3])))"
+        )
+        assert result.success
+        assert "[(1, 'a'), (2, 'b')]" in result.stdout
+
+    @pytest.mark.asyncio
+    async def test_isinstance_type_available(self, repl_env_empty):
+        result = await repl_env_empty.execute_async(
+            "print(isinstance(1, int), type([]))"
+        )
+        assert result.success
+        assert "True" in result.stdout
+
+
+class TestDangerousBuiltinsBlocked:
+    """Test that dangerous builtins are blocked."""
+
+    @pytest.mark.asyncio
+    async def test_eval_blocked(self, repl_env_empty):
+        result = await repl_env_empty.execute_async("eval('1+1')")
+        assert not result.success
+        assert "eval" in result.stderr.lower() or "not defined" in result.stderr.lower()
+
+    @pytest.mark.asyncio
+    async def test_exec_blocked(self, repl_env_empty):
+        result = await repl_env_empty.execute_async("exec('x = 1')")
+        assert not result.success
+        assert "exec" in result.stderr.lower() or "not defined" in result.stderr.lower()
+
+    @pytest.mark.asyncio
+    async def test_compile_blocked(self, repl_env_empty):
+        result = await repl_env_empty.execute_async(
+            "compile('1+1', '<string>', 'eval')"
+        )
+        assert not result.success
+        assert (
+            "compile" in result.stderr.lower() or "not defined" in result.stderr.lower()
+        )
+
+    @pytest.mark.asyncio
+    async def test_open_blocked(self, repl_env_empty):
+        result = await repl_env_empty.execute_async("open('/etc/passwd')")
+        assert not result.success
+        assert "open" in result.stderr.lower() or "not defined" in result.stderr.lower()
+
+    @pytest.mark.asyncio
+    async def test_input_blocked(self, repl_env_empty):
+        result = await repl_env_empty.execute_async("input('Enter: ')")
+        assert not result.success
+        assert (
+            "input" in result.stderr.lower() or "not defined" in result.stderr.lower()
+        )
+
+    @pytest.mark.asyncio
+    async def test___import___blocked(self, repl_env_empty):
+        result = await repl_env_empty.execute_async("__import__('os')")
+        assert not result.success
+
+    @pytest.mark.asyncio
+    async def test_globals_blocked(self, repl_env_empty):
+        result = await repl_env_empty.execute_async("globals()")
+        assert not result.success
+
+    @pytest.mark.asyncio
+    async def test_locals_blocked(self, repl_env_empty):
+        result = await repl_env_empty.execute_async("locals()")
+        assert not result.success
+
+    @pytest.mark.asyncio
+    async def test_breakpoint_blocked(self, repl_env_empty):
+        result = await repl_env_empty.execute_async("breakpoint()")
+        assert not result.success
+
+    @pytest.mark.asyncio
+    async def test_getattr_setattr_blocked(self, repl_env_empty):
+        result = await repl_env_empty.execute_async("getattr(object, '__class__')")
+        assert not result.success
+
+    @pytest.mark.asyncio
+    async def test_delattr_blocked(self, repl_env_empty):
+        result = await repl_env_empty.execute_async("delattr(object, 'x')")
+        assert not result.success
+
+
+class TestAllowedImports:
+    """Test that allowed imports work."""
+
+    @pytest.mark.asyncio
+    async def test_json_import(self, repl_env_empty):
+        result = await repl_env_empty.execute_async(
+            "import json\nprint(json.dumps({'a': 1}))"
+        )
+        assert result.success
+        assert '{"a": 1}' in result.stdout
+
+    @pytest.mark.asyncio
+    async def test_re_import(self, repl_env_empty):
+        result = await repl_env_empty.execute_async(
+            "import re\nprint(re.match(r'\\d+', '123').group())"
+        )
+        assert result.success
+        assert "123" in result.stdout
+
+    @pytest.mark.asyncio
+    async def test_math_import(self, repl_env_empty):
+        result = await repl_env_empty.execute_async("import math\nprint(math.sqrt(4))")
+        assert result.success
+        assert "2.0" in result.stdout
+
+    @pytest.mark.asyncio
+    async def test_statistics_import(self, repl_env_empty):
+        result = await repl_env_empty.execute_async(
+            "import statistics\nprint(statistics.mean([1, 2, 3]))"
+        )
+        assert result.success
+        assert "2" in result.stdout
+
+    @pytest.mark.asyncio
+    async def test_collections_import(self, repl_env_empty):
+        result = await repl_env_empty.execute_async(
+            "from collections import Counter\nprint(Counter(['a', 'b', 'a']))"
+        )
+        assert result.success
+        assert "'a': 2" in result.stdout
+
+    @pytest.mark.asyncio
+    async def test_itertools_import(self, repl_env_empty):
+        result = await repl_env_empty.execute_async(
+            "from itertools import chain\nprint(list(chain([1], [2])))"
+        )
+        assert result.success
+        assert "[1, 2]" in result.stdout
+
+    @pytest.mark.asyncio
+    async def test_functools_import(self, repl_env_empty):
+        result = await repl_env_empty.execute_async(
+            "from functools import reduce\nprint(reduce(lambda a, b: a+b, [1,2,3]))"
+        )
+        assert result.success
+        assert "6" in result.stdout
+
+    @pytest.mark.asyncio
+    async def test_datetime_import(self, repl_env_empty):
+        result = await repl_env_empty.execute_async(
+            "from datetime import date\nprint(date(2025, 1, 1))"
+        )
+        assert result.success
+        assert "2025-01-01" in result.stdout
+
+
+class TestDangerousImportsBlocked:
+    """Test that dangerous imports are blocked."""
+
+    @pytest.mark.asyncio
+    async def test_os_blocked(self, repl_env_empty):
+        result = await repl_env_empty.execute_async("import os")
+        assert not result.success
+        assert (
+            "not allowed" in result.stderr.lower() or "error" in result.stderr.lower()
+        )
+
+    @pytest.mark.asyncio
+    async def test_sys_blocked(self, repl_env_empty):
+        result = await repl_env_empty.execute_async("import sys")
+        assert not result.success
+
+    @pytest.mark.asyncio
+    async def test_subprocess_blocked(self, repl_env_empty):
+        result = await repl_env_empty.execute_async("import subprocess")
+        assert not result.success
+
+    @pytest.mark.asyncio
+    async def test_shutil_blocked(self, repl_env_empty):
+        result = await repl_env_empty.execute_async("import shutil")
+        assert not result.success
+
+    @pytest.mark.asyncio
+    async def test_socket_blocked(self, repl_env_empty):
+        result = await repl_env_empty.execute_async("import socket")
+        assert not result.success
+
+    @pytest.mark.asyncio
+    async def test_requests_blocked(self, repl_env_empty):
+        result = await repl_env_empty.execute_async("import requests")
+        assert not result.success
+
+    @pytest.mark.asyncio
+    async def test_builtins_module_blocked(self, repl_env_empty):
+        result = await repl_env_empty.execute_async("import builtins")
+        assert not result.success
+
+
+class TestHaikuRAGBridgeFunctions:
+    """Test haiku.rag bridge functions in sandbox."""
+
+    @pytest.mark.asyncio
+    async def test_search(self, repl_env_empty):
+        """Test search function calls client with correct args."""
+        from unittest.mock import AsyncMock
+
+        from haiku.rag.store.models import SearchResult
+
+        mock_results = [
+            SearchResult(
+                chunk_id="chunk-1",
+                document_id="doc-1",
+                document_title="Test Doc",
+                document_uri="test://doc",
+                content="Test content about foxes",
+                score=0.9,
+                page_numbers=[1],
+                headings=["Heading"],
+            )
+        ]
+        repl_env_empty.client.search = AsyncMock(return_value=mock_results)
+
+        result = await repl_env_empty.execute_async(
+            "results = search('fox', limit=5)\n"
+            "print(len(results), results[0]['chunk_id'], 'fox' in results[0]['content'].lower())"
+        )
+        assert result.success
+        assert "1 chunk-1 True" in result.stdout
+        repl_env_empty.client.search.assert_called_once_with(
+            "fox", limit=5, filter=None
+        )
+
+    @pytest.mark.asyncio
+    async def test_list_documents(self, repl_env_empty):
+        """Test list_documents returns list structure."""
+        result = await repl_env_empty.execute_async(
+            "docs = list_documents()\nprint(type(docs).__name__, len(docs))"
+        )
+        assert result.success
+        assert "list 0" in result.stdout
+
+    @pytest.mark.asyncio
+    async def test_get_document(self, repl_env_empty):
+        """Test get_document calls client correctly."""
+        from unittest.mock import AsyncMock
+
+        from haiku.rag.store.models import Document
+
+        mock_doc = Document(
+            id="doc-1",
+            uri="test://doc",
+            title="Test Doc",
+            content="The quick brown fox",
+        )
+        repl_env_empty.client.get_document_by_id = AsyncMock(return_value=mock_doc)
+
+        result = await repl_env_empty.execute_async(
+            "doc = get_document('doc-1')\nprint('fox' in doc.lower())"
+        )
+        assert result.success
+        assert "True" in result.stdout
+
+    @pytest.mark.asyncio
+    async def test_get_document_missing(self, repl_env_empty):
+        """Test get_document returns None for missing document."""
+        result = await repl_env_empty.execute_async(
+            "doc = get_document('Nonexistent')\nprint(doc is None)"
+        )
+        assert result.success
+        assert "True" in result.stdout
+
+    @pytest.mark.asyncio
+    async def test_ask(self, repl_env_empty):
+        """Test ask function calls client with correct args."""
+        from unittest.mock import AsyncMock
+
+        repl_env_empty.client.ask = AsyncMock(return_value=("The fox is brown.", []))
+
+        result = await repl_env_empty.execute_async(
+            "answer = ask('What color is the fox?')\nprint('fox' in answer.lower())"
+        )
+        assert result.success
+        assert "True" in result.stdout
+        repl_env_empty.client.ask.assert_called_once_with(
+            "What color is the fox?", filter=None
+        )
+
+
+class TestSandboxExecution:
+    """Test general sandbox execution behavior."""
+
+    @pytest.mark.asyncio
+    async def test_variable_persistence(self, repl_env_empty):
+        """Variables persist across executions."""
+        await repl_env_empty.execute_async("x = 42")
+        result = await repl_env_empty.execute_async("print(x)")
+        assert result.success
+        assert "42" in result.stdout
+
+    @pytest.mark.asyncio
+    async def test_function_definition(self, repl_env_empty):
+        """Can define and call functions."""
+        result = await repl_env_empty.execute_async(
+            "def add(a, b):\n    return a + b\nprint(add(1, 2))"
+        )
+        assert result.success
+        assert "3" in result.stdout
+
+    @pytest.mark.asyncio
+    async def test_class_definition(self, repl_env_empty):
+        """Can define and use classes."""
+        result = await repl_env_empty.execute_async(
+            "class Point:\n"
+            "    def __init__(self, x, y):\n"
+            "        self.x = x\n"
+            "        self.y = y\n"
+            "p = Point(1, 2)\n"
+            "print(p.x, p.y)"
+        )
+        assert result.success
+        assert "1 2" in result.stdout
+
+    @pytest.mark.asyncio
+    async def test_list_comprehension(self, repl_env_empty):
+        """List comprehensions work."""
+        result = await repl_env_empty.execute_async("print([x**2 for x in range(5)])")
+        assert result.success
+        assert "[0, 1, 4, 9, 16]" in result.stdout
+
+    @pytest.mark.asyncio
+    async def test_dict_comprehension(self, repl_env_empty):
+        """Dict comprehensions work."""
+        result = await repl_env_empty.execute_async(
+            "print({x: x**2 for x in range(3)})"
+        )
+        assert result.success
+        assert "{0: 0, 1: 1, 2: 4}" in result.stdout
+
+    @pytest.mark.asyncio
+    async def test_exception_handling(self, repl_env_empty):
+        """Can catch and handle exceptions."""
+        result = await repl_env_empty.execute_async(
+            "try:\n    x = 1/0\nexcept ZeroDivisionError:\n    print('caught')"
+        )
+        assert result.success
+        assert "caught" in result.stdout
+
+    @pytest.mark.asyncio
+    async def test_uncaught_exception_reports_error(self, repl_env_empty):
+        """Uncaught exceptions are reported."""
+        result = await repl_env_empty.execute_async("x = 1/0")
+        assert not result.success
+        assert "ZeroDivisionError" in result.stderr
+
+    @pytest.mark.asyncio
+    async def test_syntax_error_reports_error(self, repl_env_empty):
+        """Syntax errors are reported."""
+        result = await repl_env_empty.execute_async("def foo(")
+        assert not result.success
+        assert "SyntaxError" in result.stderr
+
+    @pytest.mark.asyncio
+    async def test_output_truncation(self, repl_env_empty):
+        """Output is truncated if too long."""
+        repl_env_empty.config.max_output_chars = 100
+        result = await repl_env_empty.execute_async("print('x' * 1000)")
+        assert result.success
+        assert (
+            len(result.stdout) <= 100 + 50
+        )  # Allow some margin for truncation message
+
+
+class TestSecurityEscapes:
+    """Test that common security escape attempts are blocked."""
+
+    @pytest.mark.asyncio
+    async def test_eval_via_builtins_dict(self, repl_env_empty):
+        """Cannot access eval through __builtins__."""
+        result = await repl_env_empty.execute_async("__builtins__['eval']('1+1')")
+        assert not result.success
+
+    @pytest.mark.asyncio
+    async def test_import_via_builtins(self, repl_env_empty):
+        """Cannot import os through builtins trickery."""
+        result = await repl_env_empty.execute_async("__builtins__.__import__('os')")
+        assert not result.success
+
+    @pytest.mark.asyncio
+    async def test_class_bases_escape(self, repl_env_empty):
+        """Cannot escape through __class__.__bases__."""
+        result = await repl_env_empty.execute_async(
+            "().__class__.__bases__[0].__subclasses__()"
+        )
+        assert not result.success
+
+    @pytest.mark.asyncio
+    async def test_code_object_escape(self, repl_env_empty):
+        """Cannot create code objects."""
+        result = await repl_env_empty.execute_async(
+            "def f(): pass\n"
+            "type(f.__code__)(0, 0, 0, 0, 0, 0, b'', (), (), (), '', '', 0, b'')"
+        )
+        assert not result.success
+
+    @pytest.mark.asyncio
+    async def test_import_system_escape(self, repl_env_empty):
+        """Cannot escape through importlib."""
+        result = await repl_env_empty.execute_async("import importlib")
+        assert not result.success
+
+    @pytest.mark.asyncio
+    async def test_pickle_escape(self, repl_env_empty):
+        """Cannot use pickle for code execution."""
+        result = await repl_env_empty.execute_async("import pickle")
+        assert not result.success