"""Lightweight sandbox for executing LLM-generated Python code.""" from __future__ import annotations import collections import copy import io import json import logging import platform import re import math import signal import threading import time as _time_mod from concurrent.futures import ThreadPoolExecutor from contextlib import redirect_stdout, redirect_stderr from dataclasses import dataclass from datetime import datetime, timedelta, date, time, timezone from typing import Any, Callable, Dict, List, Optional logger = logging.getLogger(__name__) class ExecutionTimeoutError(Exception): """Raised when code execution exceeds the timeout.""" pass @dataclass class ExecutionResult: """Result of executing code in the sandbox. Attributes: stdout: Captured standard output stderr: Captured standard error final_value: Value passed to FINAL() if called, otherwise None exception: Exception that occurred during execution, if any """ stdout: str = "" stderr: str = "" final_value: Any = None exception: Optional[Exception] = None @property def success(self) -> bool: """Return True if execution completed without errors.""" return self.exception is None class TraceSandbox: """Lightweight sandbox using exec() with restricted builtins. This sandbox restricts builtins but is NOT secure against determined escape attempts. Security relies on trusting the LLM not to generate malicious code. Do not use this sandbox to execute untrusted or user-provided code. Restrictions (defense-in-depth, not security guarantees): - No file access: `open` and `__import__` are blocked - No code injection: `eval`, `exec`, `compile` are blocked - Read-only trace: trace data is injected as-is - Timeout protection: Configurable per-execution timeout (Unix only) - Worst case: bad code fails -> fallback to simple reflector Example: >>> sandbox = TraceSandbox(trace=trace, llm_query_fn=llm_query) >>> result = sandbox.execute("print(len(trace.steps))", timeout=30.0) >>> print(result.stdout) 5 """ # Safe builtins that don't allow file/network access or code injection SAFE_BUILTINS: Dict[str, Any] = { # Core types "print": print, "len": len, "str": str, "int": int, "float": float, "list": list, "dict": dict, "set": set, "tuple": tuple, "bool": bool, "type": type, "isinstance": isinstance, "issubclass": issubclass, "range": range, "bytes": bytes, "bytearray": bytearray, # Iteration "enumerate": enumerate, "zip": zip, "map": map, "filter": filter, "sorted": sorted, "reversed": reversed, "iter": iter, "next": next, # Math "min": min, "max": max, "sum": sum, "abs": abs, "round": round, "pow": pow, "divmod": divmod, # Logic "any": any, "all": all, "not": lambda x: not x, # String/Formatting "chr": chr, "ord": ord, "repr": repr, "format": format, "ascii": ascii, "bin": bin, "hex": hex, "oct": oct, # Object inspection (getattr blocks dunder access — see __init__) "hasattr": hasattr, "getattr": None, # Replaced with safe_getattr in __init__ "dir": dir, "vars": lambda obj=None: {} if obj is None else vars(obj), "id": id, "hash": hash, "callable": callable, # Exceptions (for try/except in generated code) "Exception": Exception, "BaseException": BaseException, "ValueError": ValueError, "KeyError": KeyError, "IndexError": IndexError, "TypeError": TypeError, "AttributeError": AttributeError, "RuntimeError": RuntimeError, "StopIteration": StopIteration, "AssertionError": AssertionError, "LookupError": LookupError, "ZeroDivisionError": ZeroDivisionError, "NameError": NameError, "OverflowError": OverflowError, "FloatingPointError": FloatingPointError, "ArithmeticError": ArithmeticError, "SyntaxError": SyntaxError, "IndentationError": IndentationError, "TabError": TabError, "UnicodeError": UnicodeError, "UnicodeDecodeError": UnicodeDecodeError, "UnicodeEncodeError": UnicodeEncodeError, "NotImplementedError": NotImplementedError, "RecursionError": RecursionError, # Constants "True": True, "False": False, "None": None, # BLOCKED - security sensitive (raise clear errors, not NoneType) "open": None, "__import__": None, # replaced with _safe_import in __init__ "eval": None, "exec": None, "compile": None, "input": None, "globals": None, "locals": None, "breakpoint": None, "memoryview": None, } def __init__( self, trace: Optional[str] = None, llm_query_fn: Optional[Callable[[str], str]] = None, additional_globals: Optional[Dict[str, Any]] = None, *, parallel_max_concurrency: int = 10, parallel_max_retries: int = 3, parallel_retry_delay: float = 1.0, parallel_timeout: Optional[float] = None, ) -> None: """Initialize the sandbox with trace and optional LLM query function. Args: trace: Trace string for exploration (can be None). Non-string values are coerced to str; None is left as-is. llm_query_fn: Function to call for sub-LLM queries additional_globals: Extra variables to inject into the namespace parallel_max_concurrency: Max concurrent workers for parallel_map parallel_max_retries: Max retries per item in parallel_map parallel_retry_delay: Base delay (seconds) for exponential backoff parallel_timeout: Per-item timeout in seconds (None = no timeout) """ self._final_value: Any = None self._final_called = False # parallel_map configuration (infrastructure-side only) self._parallel_max_concurrency = parallel_max_concurrency self._parallel_max_retries = parallel_max_retries self._parallel_retry_delay = parallel_retry_delay self._parallel_timeout = parallel_timeout # Sanitize trace: coerce to str if provided if trace is not None and not isinstance(trace, str): trace = str(trace) # Build the namespace self.namespace: Dict[str, Any] = { "__builtins__": self.SAFE_BUILTINS.copy(), # Core analysis objects "trace": trace, "FINAL": self._final, "FINAL_VAR": self._final_var, "SHOW_VARS": self._show_vars, "helper_registry": {}, "register_helper": self._register_helper, "list_helpers": self._list_helpers, "run_helper": self._run_helper, "get_batch_item": self._get_batch_item, "get_item_payload": self._get_item_payload, "get_item_messages": self._get_item_messages, "get_item_question": self._get_item_question, "get_item_feedback": self._get_item_feedback, "get_item_id": self._get_item_id, "get_message_text": self._get_message_text, "preview_item": self._preview_item, "parallel_map": self._parallel_map, # Safe stdlib modules "json": json, "re": re, "math": math, "collections": collections, # datetime module and commonly used classes "datetime": datetime, "timedelta": timedelta, "date": date, "time": time, "timezone": timezone, } # Safe getattr that blocks dunder access — override in both # builtins (so bare getattr() works) and namespace (for direct ref) def safe_getattr(obj, name, *default): if name.startswith("_"): raise AttributeError(f"Access to '{name}' blocked") return getattr(obj, name, *default) if default else getattr(obj, name) self.namespace["__builtins__"]["getattr"] = safe_getattr self.namespace["safe_getattr"] = safe_getattr # Safe import — allows pre-loaded modules, blocks everything else. # LLMs often write `import json` even when json is already available. _allowed_modules = { "json": json, "re": re, "math": math, "collections": collections, "datetime": __import__("datetime"), } def _safe_import(name: str, *args: Any, **kwargs: Any) -> Any: if name in _allowed_modules: return _allowed_modules[name] raise ImportError( f"import {name!r} is blocked in sandbox. " f"Pre-loaded modules ({', '.join(sorted(_allowed_modules))}) " f"are already available — use them directly." ) self.namespace["__builtins__"]["__import__"] = _safe_import # Add llm_query if provided if llm_query_fn is not None: self.namespace["llm_query"] = llm_query_fn else: # Provide a stub that explains the feature is disabled self.namespace["llm_query"] = lambda _prompt: ( "(llm_query disabled - analyze with available data)" ) # Add any additional globals if additional_globals: self.namespace.update(additional_globals) def _final(self, value: Any) -> None: """Called by LLM code to output the final result. Args: value: The final analysis result (should be a dict matching ReflectorOutput) Raises: StopIteration: Always raised to signal completion """ self._final_value = value self._final_called = True raise StopIteration("FINAL called - analysis complete") def _final_var(self, var_name: str) -> None: """Called by LLM code to output a variable as the final result. Convenience function to finalize with a pre-built result stored in a variable. Useful when the analysis result is built up across multiple code blocks. Args: var_name: Name of the variable in the namespace to use as the result Raises: ValueError: If the variable doesn't exist StopIteration: Always raised to signal completion """ if var_name not in self.namespace: available = [k for k in self.namespace.keys() if not k.startswith("_")] raise ValueError( f"Variable '{var_name}' not found. Available: {available[:20]}" ) self._final(self.namespace[var_name]) def _show_vars(self) -> None: """Print available variables in the namespace for debugging. Prints a list of user-accessible variables (excludes internal/dunder names). """ user_vars = [k for k in self.namespace.keys() if not k.startswith("_")] # Exclude builtins and modules for cleaner output excluded = { "__builtins__", "json", "re", "collections", "datetime", "timedelta", "date", "time", "timezone", "safe_getattr", } user_vars = [k for k in user_vars if k not in excluded] logger.debug("Available variables: %s", sorted(user_vars)) def _register_helper( self, name: str, source: str, description: str = "", ) -> str: """Register reusable helper code in the sandbox. The helper source is executed immediately and stored so that future sandbox snapshots can recreate the same helper definitions for sub-agents. """ if not name.isidentifier(): raise ValueError(f"Invalid helper name: {name!r}") if not source.strip(): raise ValueError("Helper source cannot be empty") exec(source, self.namespace, self.namespace) helper = self.namespace.get(name) if not callable(helper): raise ValueError(f"Helper source must define a callable named {name!r}") registry = self.namespace.setdefault("helper_registry", {}) registry[name] = { "description": description, "source": source, } return f"Registered helper {name}" def _list_helpers(self) -> list[dict[str, str]]: """Return metadata for registered helpers.""" registry = self.namespace.get("helper_registry", {}) if not isinstance(registry, dict): return [] helpers: list[dict[str, str]] = [] for name, meta in registry.items(): if not isinstance(meta, dict): continue helpers.append( { "name": str(name), "description": str(meta.get("description", "")), } ) return helpers def _run_helper(self, name: str, *args: Any, **kwargs: Any) -> Any: """Invoke a registered helper by name.""" helper = self.namespace.get(name) if not callable(helper): raise KeyError(f"Helper {name!r} is not registered") return helper(*args, **kwargs) def _get_batch_item(self, index: int) -> Any: """Return a batch item by index when batch helpers are available.""" batch_items = self.namespace.get("batch_items") if not isinstance(batch_items, list): raise RuntimeError("batch_items is not available in this sandbox") return batch_items[index] def _resolve_batch_item(self, item_or_index: Any) -> Any: """Resolve a batch helper argument to the underlying item.""" if isinstance(item_or_index, int): return self._get_batch_item(item_or_index) return item_or_index def _get_item_payload(self, item_or_index: Any) -> Any: """Return the payload for a batch item or index without rewriting it.""" item = self._resolve_batch_item(item_or_index) if ( isinstance(item, dict) and item.get("role") == "conversation" and isinstance(item.get("content"), dict) ): return item["content"] return item def _get_item_messages(self, item_or_index: Any) -> list[Any]: """Return a best-effort message list for a batch item or index.""" payload = self._get_item_payload(item_or_index) if isinstance(payload, list): return payload if isinstance(payload, dict): trace_value = payload.get("trace") if isinstance(trace_value, list): return trace_value if isinstance(trace_value, dict): for key in ("messages", "steps", "trace"): nested = trace_value.get(key) if isinstance(nested, list): return nested for key in ("messages", "steps"): value = payload.get(key) if isinstance(value, list): return value return [] def _get_item_field(self, item_or_index: Any, field: str) -> str: """Extract a string field from a batch item payload when present.""" payload = self._get_item_payload(item_or_index) if isinstance(payload, dict): value = payload.get(field) if value is not None: return str(value) return "" def _get_item_question(self, item_or_index: Any) -> str: """Return the question field for a batch item or index.""" return self._get_item_field(item_or_index, "question") def _get_item_feedback(self, item_or_index: Any) -> str: """Return the feedback field for a batch item or index.""" return self._get_item_field(item_or_index, "feedback") def _get_item_id(self, item_or_index: Any) -> str: """Return a stable identifier for a batch item or index.""" item = self._resolve_batch_item(item_or_index) payload = self._get_item_payload(item) if isinstance(item_or_index, int): item_ids = self.namespace.get("item_ids") if isinstance(item_ids, list) and 0 <= item_or_index < len(item_ids): return str(item_ids[item_or_index]) if isinstance(item, dict): for key in ("item_id", "task_id", "id"): value = item.get(key) if value is not None: return str(value) if isinstance(payload, dict): for key in ("item_id", "task_id", "id"): value = payload.get(key) if value is not None: return str(value) return "unknown_item" def _get_message_text(self, message: Any) -> str: """Return a readable text summary for a message-like object.""" if isinstance(message, dict): content = message.get("content") if content not in (None, ""): if isinstance(content, str): return content try: return json.dumps(content, default=str) except Exception: return str(content) tool_calls = message.get("tool_calls") if tool_calls: return f"tool_calls={json.dumps(tool_calls, default=str)}" tool_results = message.get("tool_results") if tool_results: return f"tool_results={json.dumps(tool_results, default=str)}" for key in ("reasoning", "answer", "text"): value = message.get(key) if value not in (None, ""): return str(value) try: return json.dumps(message, default=str) except Exception: return str(message) return str(message) def _preview_item(self, item_or_index: Any) -> dict[str, Any]: """Return a compact preview for a batch item or index.""" messages = self._get_item_messages(item_or_index) first_message = self._get_message_text(messages[0]) if messages else "" payload = self._get_item_payload(item_or_index) return { "item_id": self._get_item_id(item_or_index), "question_preview": self._get_item_question(item_or_index)[:120], "feedback_preview": self._get_item_feedback(item_or_index)[:120], "message_count": len(messages), "first_message_preview": first_message[:120], "payload_type": type(payload).__name__, } def _parallel_map( self, fn: Callable[[Any], Any], inputs: list, *, return_exceptions: bool = False ) -> List[Any]: """Execute fn over inputs in parallel using a thread pool. Concurrency, retries, backoff, and timeout are controlled by the sandbox configuration — the agent cannot override them. Args: fn: A callable to apply to each input inputs: Ordered list of inputs return_exceptions: If True, failed items appear as exceptions in the results list instead of raising immediately Returns: Ordered list of results (same length/order as inputs) Raises: Exception: Re-raises the first worker exception when return_exceptions is False """ if not inputs: return [] max_concurrency = self._parallel_max_concurrency max_retries = self._parallel_max_retries retry_delay = self._parallel_retry_delay timeout = self._parallel_timeout def _worker(item: Any) -> Any: last_exc: Optional[Exception] = None for attempt in range(max_retries + 1): try: return fn(item) except Exception as exc: last_exc = exc if attempt < max_retries: backoff = retry_delay * (2**attempt) _time_mod.sleep(backoff) raise last_exc # type: ignore[misc] pool_size = min(len(inputs), max_concurrency) results: List[Any] = [None] * len(inputs) first_exc: Optional[Exception] = None first_exc_idx: Optional[int] = None with ThreadPoolExecutor(max_workers=pool_size) as pool: futures = { pool.submit(_worker, item): idx for idx, item in enumerate(inputs) } for future in futures: idx = futures[future] try: results[idx] = future.result(timeout=timeout) except Exception as exc: if return_exceptions: results[idx] = exc else: if first_exc_idx is None or idx < first_exc_idx: first_exc = exc first_exc_idx = idx if first_exc is not None and not return_exceptions: raise first_exc return results @property def final_value(self) -> Any: """Return the value passed to FINAL(), or None if not called.""" return self._final_value @property def final_called(self) -> bool: """Return True if FINAL() was called.""" return self._final_called def inject(self, name: str, value: Any) -> None: """Inject a variable into the sandbox namespace. Args: name: Variable name value: Variable value """ self.namespace[name] = value def execute(self, code: str, timeout: float = 30.0) -> ExecutionResult: """Execute code in the sandbox and capture output. Args: code: Python code to execute timeout: Maximum execution time in seconds (default: 30.0). - Unix: uses signal.SIGALRM - Windows: not enforced (in-process execution) Returns: ExecutionResult with stdout, stderr, final_value, and exception """ if platform.system() == "Windows": return self._execute_no_timeout(code) elif threading.current_thread() is not threading.main_thread(): return self._execute_no_timeout(code) else: return self._execute_unix(code, timeout) def _execute_unix(self, code: str, timeout: float) -> ExecutionResult: """Execute code using signal-based timeout (Unix only). Args: code: Python code to execute timeout: Maximum execution time in seconds Returns: ExecutionResult with stdout, stderr, final_value, and exception """ stdout_buf = io.StringIO() stderr_buf = io.StringIO() # Set up timeout handler (Unix only) use_timeout = timeout > 0 old_handler = None def timeout_handler(_signum: int, _frame: Any) -> None: raise ExecutionTimeoutError(f"Execution exceeded {timeout}s timeout") if use_timeout: old_handler = signal.signal(signal.SIGALRM, timeout_handler) signal.alarm(math.ceil(timeout)) try: with redirect_stdout(stdout_buf), redirect_stderr(stderr_buf): exec(code, self.namespace, self.namespace) except StopIteration: # FINAL() was called - this is expected pass except ExecutionTimeoutError as e: stderr_buf.write(f"\nExecutionTimeoutError: {e}") return ExecutionResult( stdout=stdout_buf.getvalue(), stderr=stderr_buf.getvalue(), final_value=self._final_value, exception=e, ) except Exception as e: # Capture the exception info stderr_buf.write(f"\n{type(e).__name__}: {e}") return ExecutionResult( stdout=stdout_buf.getvalue(), stderr=stderr_buf.getvalue(), final_value=self._final_value, exception=e, ) finally: if use_timeout: signal.alarm(0) # Cancel the alarm signal.signal(signal.SIGALRM, old_handler) return ExecutionResult( stdout=stdout_buf.getvalue(), stderr=stderr_buf.getvalue(), final_value=self._final_value, exception=None, ) def _execute_windows(self, code: str, timeout: float) -> ExecutionResult: """Execute code on Windows without timeout enforcement. Windows multiprocessing uses 'spawn' which cannot pass functions, trace objects, or injected variables to subprocesses. Instead, execute in-process for full feature support (no timeout enforcement). Args: code: Python code to execute timeout: Ignored on Windows (logged as warning) Returns: ExecutionResult with stdout, stderr, final_value, and exception """ logger.debug("Windows: executing in-process (timeout not enforced)") return self._execute_no_timeout(code) def _execute_no_timeout(self, code: str) -> ExecutionResult: """Execute code without timeout enforcement. Fallback when multiprocessing is unavailable or fails. Args: code: Python code to execute Returns: ExecutionResult with stdout, stderr, final_value, and exception """ stdout_buf = io.StringIO() stderr_buf = io.StringIO() try: with redirect_stdout(stdout_buf), redirect_stderr(stderr_buf): exec(code, self.namespace, self.namespace) except StopIteration: # FINAL() was called - this is expected pass except Exception as e: stderr_buf.write(f"\n{type(e).__name__}: {e}") return ExecutionResult( stdout=stdout_buf.getvalue(), stderr=stderr_buf.getvalue(), final_value=self._final_value, exception=e, ) return ExecutionResult( stdout=stdout_buf.getvalue(), stderr=stderr_buf.getvalue(), final_value=self._final_value, exception=None, ) def reset(self) -> None: """Reset the sandbox state for a new execution.""" self._final_value = None self._final_called = False def create_readonly_sandbox(parent: TraceSandbox) -> TraceSandbox: """Create an isolated sandbox snapshot for sub-agent use. Deep-copies data variables from the parent sandbox so the sub-agent can explore trace data via ``execute_code`` without affecting the parent's state. Safe for parallel use — each snapshot is independent. Args: parent: The parent sandbox to snapshot. Returns: A new TraceSandbox with deep-copied data variables. """ sandbox = TraceSandbox( trace=None, llm_query_fn=None, parallel_max_concurrency=parent._parallel_max_concurrency, parallel_max_retries=parent._parallel_max_retries, parallel_retry_delay=parent._parallel_retry_delay, parallel_timeout=parent._parallel_timeout, ) # Keys already set up by TraceSandbox.__init__ — skip them infrastructure = { "__builtins__", "FINAL", "FINAL_VAR", "SHOW_VARS", "parallel_map", "llm_query", "safe_getattr", "trace", "register_helper", "list_helpers", "run_helper", "get_batch_item", "get_item_payload", "get_item_messages", "get_item_question", "get_item_feedback", "get_item_id", "get_message_text", "preview_item", "json", "re", "math", "collections", "datetime", "timedelta", "date", "time", "timezone", } for key, value in parent.namespace.items(): if key in infrastructure or key.startswith("_"): continue try: sandbox.namespace[key] = copy.deepcopy(value) except (TypeError, copy.Error): # Modules, functions, etc. — share by reference sandbox.namespace[key] = value registry = sandbox.namespace.get("helper_registry", {}) if isinstance(registry, dict): for name, meta in registry.items(): if not isinstance(meta, dict): continue source = meta.get("source") if not isinstance(source, str) or not source.strip(): continue try: exec(source, sandbox.namespace, sandbox.namespace) except Exception as exc: logger.warning("Failed to restore helper %s in snapshot: %s", name, exc) return sandbox