diff --git a/haiku_rag_slim/haiku/rag/agents/rlm/__init__.py b/haiku_rag_slim/haiku/rag/agents/rlm/__init__.py new file mode 100644 index 00000000..16194705 --- /dev/null +++ b/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", +] diff --git a/haiku_rag_slim/haiku/rag/agents/rlm/dependencies.py b/haiku_rag_slim/haiku/rag/agents/rlm/dependencies.py new file mode 100644 index 00000000..c5b968c7 --- /dev/null +++ b/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) diff --git a/haiku_rag_slim/haiku/rag/agents/rlm/sandbox.py b/haiku_rag_slim/haiku/rag/agents/rlm/sandbox.py new file mode 100644 index 00000000..31f19d44 --- /dev/null +++ b/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 diff --git a/tests/agents/rlm/__init__.py b/tests/agents/rlm/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/agents/rlm/conftest.py b/tests/agents/rlm/conftest.py new file mode 100644 index 00000000..580671c5 --- /dev/null +++ b/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) diff --git a/tests/agents/rlm/test_sandbox.py b/tests/agents/rlm/test_sandbox.py new file mode 100644 index 00000000..c158a80e --- /dev/null +++ b/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', '', '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