Sandbox Code Mode execution in a subprocess with resource limits

- Execute code in a separate process via multiprocessing to isolate
  the parent from crashes, infinite loops, and memory exhaustion
- Apply RLIMIT_CPU and RLIMIT_AS (Unix) for CPU and memory bounds
- Validate input/output variables contain only plain data types
  (bool, int, float, str, list, tuple, dict) to prevent code injection
- Add CodeModeSandbox dataclass for configurable timeout, CPU, and
  memory limits

Author: johnteee

teaagent/__init__.py

@@ -11,7 +11,12 @@
     register_subagent_tool,
     run_chat_agent,
 )
-from teaagent.code_mode import CodeModeResult, UnsafeCodeError, execute_code_mode
+from teaagent.code_mode import (
+    CodeModeResult,
+    CodeModeSandbox,
+    UnsafeCodeError,
+    execute_code_mode,
+)
 from teaagent.context import ContextCompactor
 from teaagent.eval import EvalCase, EvalReport, run_eval
 from teaagent.graph_rag import GraphEdge, KnowledgeGraph, graph_retrieve
@@ -114,6 +119,7 @@
     'ChatAgentConfig',
     'ClarificationResult',
     'CodeModeResult',
+    'CodeModeSandbox',
     'ContextCompactor',
     'Decision',
     'Document',

teaagent/code_mode.py

@@ -1,9 +1,17 @@
 from __future__ import annotations
 
 import ast
+import multiprocessing
+import queue
+import traceback
 from dataclasses import dataclass
 from typing import Any
 
+try:
+    import resource
+except ImportError:  # pragma: no cover - resource is Unix-only.
+    resource = None  # type: ignore[assignment]
+
 ALLOWED_NODES = {
     ast.Add,
     ast.Assign,
@@ -58,24 +66,58 @@ class CodeModeResult:
     variables: dict[str, Any]
 
 
+@dataclass(frozen=True)
+class CodeModeSandbox:
+    timeout_seconds: float = 2.0
+    cpu_seconds: int = 2
+    memory_bytes: int = 64 * 1024 * 1024
+
+
 class UnsafeCodeError(ValueError):
     pass
 
 
 def execute_code_mode(
-    code: str, *, inputs: dict[str, Any] | None = None
+    code: str,
+    *,
+    inputs: dict[str, Any] | None = None,
+    sandbox: CodeModeSandbox | None = None,
 ) -> CodeModeResult:
     tree = ast.parse(code, mode='exec')
     _validate_tree(tree)
-    namespace: dict[str, Any] = {'__builtins__': SAFE_BUILTINS}
-    if inputs:
-        namespace.update(inputs)
-    exec(compile(tree, '<teaagent-code-mode>', 'exec'), namespace, namespace)
-    variables = {
-        key: value
-        for key, value in namespace.items()
-        if key != '__builtins__' and not key.startswith('_')
-    }
+    safe_inputs = dict(inputs or {})
+    _validate_plain_data(safe_inputs, 'inputs')
+
+    active_sandbox = sandbox or CodeModeSandbox()
+    if active_sandbox.timeout_seconds <= 0:
+        raise UnsafeCodeError('Code Mode timeout must be positive')
+    if active_sandbox.cpu_seconds <= 0:
+        raise UnsafeCodeError('Code Mode CPU limit must be positive')
+    if active_sandbox.memory_bytes <= 0:
+        raise UnsafeCodeError('Code Mode memory limit must be positive')
+
+    context = multiprocessing.get_context('spawn')
+    result_queue: multiprocessing.Queue[dict[str, Any]] = context.Queue(maxsize=1)
+    process = context.Process(
+        target=_execute_code_mode_child,
+        args=(code, safe_inputs, active_sandbox, result_queue),
+    )
+    process.start()
+    process.join(active_sandbox.timeout_seconds)
+    if process.is_alive():
+        process.terminate()
+        process.join()
+        raise UnsafeCodeError('Code Mode timed out')
+
+    try:
+        message = result_queue.get_nowait()
+    except queue.Empty as exc:
+        raise UnsafeCodeError('Code Mode sandbox exited without a result') from exc
+
+    if message['status'] == 'error':
+        raise UnsafeCodeError(message['error'])
+    variables = message['variables']
+    _validate_plain_data(variables, 'variables')
     return CodeModeResult(variables=variables)
 
 
@@ -87,3 +129,68 @@ def _validate_tree(tree: ast.AST) -> None:
             not isinstance(node.func, ast.Name) or node.func.id not in SAFE_BUILTINS
         ):
             raise UnsafeCodeError('Only approved builtin calls are allowed')
+
+
+def _execute_code_mode_child(
+    code: str,
+    inputs: dict[str, Any],
+    sandbox: CodeModeSandbox,
+    result_queue: multiprocessing.Queue[dict[str, Any]],
+) -> None:
+    try:
+        _apply_resource_limits(sandbox)
+        namespace: dict[str, Any] = {'__builtins__': SAFE_BUILTINS}
+        namespace.update(inputs)
+        exec(compile(code, '<teaagent-code-mode>', 'exec'), namespace, namespace)
+        variables = {
+            key: value
+            for key, value in namespace.items()
+            if key != '__builtins__' and not key.startswith('_')
+        }
+        _validate_plain_data(variables, 'variables')
+    except Exception as exc:  # pragma: no cover - exercised through parent process.
+        result_queue.put(
+            {
+                'status': 'error',
+                'error': f'{type(exc).__name__}: {exc}',
+                'traceback': traceback.format_exc(),
+            }
+        )
+    else:
+        result_queue.put({'status': 'ok', 'variables': variables})
+
+
+def _apply_resource_limits(sandbox: CodeModeSandbox) -> None:
+    if resource is None:
+        return
+    _, cpu_hard = resource.getrlimit(resource.RLIMIT_CPU)
+    resource.setrlimit(resource.RLIMIT_CPU, (sandbox.cpu_seconds, cpu_hard))
+    if hasattr(resource, 'RLIMIT_AS'):
+        _, hard = resource.getrlimit(resource.RLIMIT_AS)
+        soft = sandbox.memory_bytes
+        if hard != resource.RLIM_INFINITY:
+            soft = min(soft, hard)
+        try:
+            resource.setrlimit(resource.RLIMIT_AS, (soft, hard))
+        except ValueError:
+            # macOS reports RLIMIT_AS but rejects lowering it; wall/CPU limits still
+            # keep Code Mode isolated from the parent process.
+            pass
+
+
+def _validate_plain_data(value: Any, label: str) -> None:
+    if value is None or isinstance(value, (bool, int, float, str)):
+        return
+    if isinstance(value, list) or isinstance(value, tuple):
+        for index, item in enumerate(value):
+            _validate_plain_data(item, f'{label}[{index}]')
+        return
+    if isinstance(value, dict):
+        for key, item in value.items():
+            if not isinstance(key, str):
+                raise UnsafeCodeError(f'{label} contains a non-string key')
+            _validate_plain_data(item, f'{label}.{key}')
+        return
+    raise UnsafeCodeError(
+        f'{label} contains unsupported value type: {type(value).__name__}'
+    )

tests/test_p2_primitives.py

@@ -10,6 +10,7 @@
     StatelessMCPRequest,
     ToolAnnotations,
     ToolRegistry,
+    CodeModeSandbox,
     UnsafeCodeError,
     assess_managed_agent_readiness,
     assess_provider_portability,
@@ -84,6 +85,19 @@ def test_code_mode_blocks_imports_and_unknown_calls(self) -> None:
         with self.assertRaises(UnsafeCodeError):
             execute_code_mode("value = open('x')")
 
+    def test_code_mode_rejects_non_plain_inputs_and_outputs(self) -> None:
+        with self.assertRaises(UnsafeCodeError):
+            execute_code_mode('value = len(items)', inputs={'items': object()})
+        with self.assertRaises(UnsafeCodeError):
+            execute_code_mode('value = range(3)')
+
+    def test_code_mode_times_out_in_sandbox(self) -> None:
+        with self.assertRaises(UnsafeCodeError):
+            execute_code_mode(
+                'for index in range(100000000):\n    value = index',
+                sandbox=CodeModeSandbox(timeout_seconds=0.01),
+            )
+
     def test_stateless_mcp_request_executes_via_registry(self) -> None:
         request = StatelessMCPRequest.create(
             tool_name='sum_values',