feat(code_mode): expose sandbox resource limits

pydantic_ai_harness/code_mode/README.md

@@ -244,6 +244,29 @@ Code runs inside [Monty](https://github.com/pydantic/monty), a sandboxed Python
 - Filesystem I/O needs an `os_access` handler or a `mount`; `os.getenv`/`os.environ` need an `os_access` handler
 - Tools requiring approval or with deferred execution are excluded from the sandbox
 
+## Resource limits
+
+Model-generated code is not always well behaved: an accidental `while True`, unbounded recursion,
+or a huge allocation can block the event loop or exhaust memory. `resource_limits` caps what a
+`run_code` execution may consume. The default backstop is 256 MiB of sandbox memory and 50 million
+allocations, with no time limit.
+
+```python
+CodeMode(resource_limits={'max_duration_secs': 30})
+```
+
+A partial mapping merges onto the backstop, so setting one key keeps the others. Pass
+`'unlimited'` to remove all limits. Unknown keys raise at construction. Available keys:
+`max_duration_secs`, `max_memory`, `max_allocations`, `max_recursion_depth`.
+
+There is no default duration cap because Monty checks the limit once per bytecode step: it
+measures time spent running sandbox code, not wall-clock time. While the code awaits a dispatched
+tool call it is suspended on the host, so slow tools do not count against the cap. A cap is worth
+setting when you want to bound pure-CPU loops, which the tool-call path never sees.
+
+When a limit trips, the sandbox raises and `run_code` surfaces the error to the model as a retry,
+the same path as any other runtime error in generated code.
+
 ## API
 
 ```python
@@ -252,6 +275,7 @@ CodeMode(
     max_retries: int = 3,               # retries on sandbox execution errors
     os_access: CodeModeOS | None = None,   # host handler for env vars, clock, and file I/O
     mount: CodeModeMount | None = None,    # host directories to share with the sandbox
+    resource_limits=None,               # sandbox caps; None = backstop, 'unlimited' = none
 )
 ```
 

pydantic_ai_harness/code_mode/__init__.py

@@ -1,6 +1,19 @@
 """Code mode capability: route tool calls through a sandboxed Python environment."""
 
 from pydantic_ai_harness.code_mode._capability import CodeMode
-from pydantic_ai_harness.code_mode._toolset import CodeModeMount, CodeModeOS, CodeModeOSCallback, CodeModeToolset
+from pydantic_ai_harness.code_mode._toolset import (
+    CodeModeMount,
+    CodeModeOS,
+    CodeModeOSCallback,
+    CodeModeResourceLimits,
+    CodeModeToolset,
+)
 
-__all__ = ['CodeMode', 'CodeModeMount', 'CodeModeOS', 'CodeModeOSCallback', 'CodeModeToolset']
+__all__ = [
+    'CodeMode',
+    'CodeModeMount',
+    'CodeModeOS',
+    'CodeModeOSCallback',
+    'CodeModeResourceLimits',
+    'CodeModeToolset',
+]

pydantic_ai_harness/code_mode/_capability.py

@@ -4,7 +4,7 @@
 
 from collections.abc import Sequence
 from dataclasses import KW_ONLY, dataclass, field, replace
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, Literal
 
 from pydantic import TypeAdapter, ValidationError
 from pydantic_ai import AbstractToolset
@@ -14,7 +14,12 @@
 from pydantic_ai.tools import AgentDepsT, RunContext, ToolDefinition, ToolSelector
 from typing_extensions import TypedDict
 
-from pydantic_ai_harness.code_mode._toolset import CodeModeMount, CodeModeOS, CodeModeToolset
+from pydantic_ai_harness.code_mode._toolset import (
+    CodeModeMount,
+    CodeModeOS,
+    CodeModeResourceLimits,
+    CodeModeToolset,
+)
 
 if TYPE_CHECKING:
     from pydantic_ai.capabilities.abstract import ValidatedToolArgs
@@ -89,6 +94,22 @@ class CodeMode(AbstractCapability[AgentDepsT]):
     mount: CodeModeMount | None = None
     """Host directories to expose to sandboxed `pathlib` code; each mount's `mode` controls whether writes reach the host."""
 
+    resource_limits: CodeModeResourceLimits | Literal['unlimited'] | None = None
+    """Sandbox limits on `run_code` executions: memory, allocations, recursion depth, and CPU time.
+
+    Model-generated code is not always well behaved -- an accidental `while True`, unbounded
+    recursion, or a huge allocation can block the event loop or exhaust memory. The sandbox
+    enforces these caps per execution.
+
+    - `None` (default): backstop limits apply (256 MiB memory, 50M allocations, no duration cap).
+    - A `CodeModeResourceLimits` mapping: merged onto the backstop, so setting one key does not
+      drop the others. Unknown keys raise at construction.
+    - `'unlimited'`: removes all limits, restoring the previous unguarded behavior.
+
+    `max_duration_secs` counts only sandbox bytecode execution; time spent awaiting dispatched
+    tool calls does not count against it.
+    """
+
     dynamic_catalog: bool = False
     """Keep the `run_code` tool definition cache-stable as the sandboxed toolset grows.
 
@@ -140,6 +161,7 @@ def get_wrapper_toolset(self, toolset: AbstractToolset[AgentDepsT]) -> AbstractT
             dynamic_catalog=self.dynamic_catalog,
             os_access=self.os_access,
             mount=self.mount,
+            resource_limits=self.resource_limits,
         )
 
     async def after_tool_execute(

pydantic_ai_harness/code_mode/_toolset.py

@@ -8,7 +8,7 @@
 import warnings
 from collections.abc import Callable, Sequence
 from dataclasses import dataclass, field, replace
-from typing import Annotated, Any
+from typing import Annotated, Any, Literal
 
 from pydantic import Field, TypeAdapter
 from pydantic_ai import AbstractToolset, RunContext, ToolDefinition, WrapperToolset
@@ -41,6 +41,7 @@
         MontyTypingError,
         MountDir,
         OsFunction,
+        ResourceLimits,
     )
 except ImportError as _import_error:  # pragma: no cover
     raise ImportError(
@@ -59,6 +60,63 @@
 CodeModeMount = MountDir | list[MountDir]
 
 
+class CodeModeResourceLimits(TypedDict, total=False):
+    """Caps on the sandbox resources available to `run_code` executions.
+
+    A harness-owned view of the sandbox limits the capability supports, so the public API does
+    not depend on the underlying sandbox's own types. Every field is optional; an omitted field
+    keeps its backstop value.
+    """
+
+    max_duration_secs: float
+    """Ceiling on the time the code spends executing sandbox bytecode. Monty checks it per
+    bytecode step, so time spent awaiting dispatched tool calls does not count against it --
+    during that wait the code is suspended on the host, not running sandbox code. There is no
+    default cap. Set one to bound a pure-CPU `while True` loop, which would otherwise burn a
+    core and block the event loop."""
+
+    max_memory: int
+    """Maximum sandbox memory, in bytes."""
+
+    max_allocations: int
+    """Maximum number of sandbox allocations."""
+
+    max_recursion_depth: int
+    """Maximum sandbox call-stack depth, bounding runaway recursion."""
+
+
+def _default_resource_limits() -> ResourceLimits:
+    """Backstop sandbox limits; no duration cap -- see `CodeModeResourceLimits.max_duration_secs`."""
+    return {
+        'max_memory': 256 * 1024 * 1024,
+        'max_allocations': 50_000_000,
+    }
+
+
+# The keys `CodeModeResourceLimits` accepts. A `total=False` TypedDict does not validate keys at
+# runtime, so a typo (e.g. `max_durations_secs`) would otherwise merge through and be silently
+# dropped, disabling the limit the caller thought they set. We reject unknowns.
+_RESOURCE_LIMIT_KEYS = frozenset(CodeModeResourceLimits.__annotations__)
+
+
+def _resolve_resource_limits(limits: CodeModeResourceLimits | Literal['unlimited'] | None) -> ResourceLimits:
+    """Resolve the public `resource_limits` value to the limits handed to the sandbox.
+
+    A partial mapping merges *onto* the backstop rather than replacing it, so `{'max_memory': ...}`
+    never drops the allocations backstop. Full semantics: `CodeMode.resource_limits`.
+    """
+    if limits is None:
+        return _default_resource_limits()
+    if limits == 'unlimited':
+        return {}
+    unknown = set(limits) - _RESOURCE_LIMIT_KEYS
+    if unknown:
+        raise UserError(
+            f'Unknown `resource_limits` key(s): {sorted(unknown)}. Valid keys are {sorted(_RESOURCE_LIMIT_KEYS)}.'
+        )
+    return {**_default_resource_limits(), **limits}
+
+
 class _RunCodeArguments(TypedDict):
     code: Annotated[str, Field(description='The Python code to execute in the sandbox.')]
     restart: NotRequired[
@@ -284,6 +342,11 @@ class CodeModeToolset(WrapperToolset[AgentDepsT]):
     so Tool Search discoveries don't bust the tool-definitions cache prefix.
     """
 
+    resource_limits: CodeModeResourceLimits | Literal['unlimited'] | None = None
+    """Sandbox limits guarding `run_code` executions. `None` applies the backstop defaults,
+    `'unlimited'` removes all limits; a `CodeModeResourceLimits` mapping is merged onto the
+    backstop. Full semantics: `CodeMode.resource_limits`."""
+
     # init=False so `replace()` in `for_run` produces a fresh instance with _repl=None,
     # giving each agent run isolated REPL state. Lazy-initialized on first call_tool.
     _repl: MontyRepl | None = field(default=None, init=False, repr=False)
@@ -296,6 +359,9 @@ class CodeModeToolset(WrapperToolset[AgentDepsT]):
     # logs every step. Reset on `for_run` because each run gets a fresh instance.
     _warned_deferred: set[str] = field(default_factory=set[str], init=False, repr=False)
 
+    def __post_init__(self) -> None:
+        _resolve_resource_limits(self.resource_limits)  # validate keys now, not at the first tool call
+
     async def for_run(self, ctx: RunContext[AgentDepsT]) -> AbstractToolset[AgentDepsT]:
         """Return a fresh toolset instance with isolated REPL state for this agent run."""
         wrapped = await self.wrapped.for_run(ctx)
@@ -536,7 +602,7 @@ async def dispatch_tool_call(sandbox_name: str, kwargs: dict[str, Any]) -> Any:
 
         # Create the REPL after type checking passes.
         if fresh_repl:
-            self._repl = MontyRepl()
+            self._repl = MontyRepl(limits=_resolve_resource_limits(self.resource_limits))
         assert self._repl is not None
 
         capture = PrintCapture()

tests/code_mode/test_code_mode.py

@@ -9,7 +9,7 @@
 from __future__ import annotations
 
 from pathlib import Path
-from typing import Any, TypeVar
+from typing import Any, Literal, TypeVar
 
 import pytest
 from pydantic_ai import (
@@ -19,7 +19,7 @@
     Tool,
     ToolDefinition,
 )
-from pydantic_ai.exceptions import ModelRetry
+from pydantic_ai.exceptions import ModelRetry, UserError
 from pydantic_ai.messages import ToolCallPart
 from pydantic_ai.models.test import TestModel
 from pydantic_ai.tool_manager import ParallelExecutionMode
@@ -32,7 +32,7 @@
 
 from pydantic_ai_harness import CodeMode
 from pydantic_ai_harness._monty_exec import PrintCapture
-from pydantic_ai_harness.code_mode import CodeModeToolset
+from pydantic_ai_harness.code_mode import CodeModeResourceLimits, CodeModeToolset
 from pydantic_ai_harness.code_mode._toolset import (  # pyright: ignore[reportPrivateUsage]
     _SEARCH_TOOLS_MODIFIER,
     _TOOL_SEARCH_ADDENDUM,
@@ -2656,3 +2656,54 @@ def sequential() -> ParallelExecutionMode:
 
         assert _global_mode_is_sequential(parallel) is False
         assert _global_mode_is_sequential(sequential) is True
+
+
+class TestResourceLimits:
+    """Sandbox resource limits on `run_code` executions."""
+
+    async def _run_code(self, code: str, limits: CodeModeResourceLimits | Literal['unlimited'] | None) -> Any:
+        toolset = _build_function_toolset(add)
+        wrapper = CodeMode[object](resource_limits=limits).get_wrapper_toolset(toolset)
+        assert isinstance(wrapper, CodeModeToolset)
+        ctx = await build_ctx(None, wrapper)
+        tools = await wrapper.get_tools(ctx)
+        return await wrapper.call_tool('run_code', {'code': code}, ctx, tools['run_code'])
+
+    async def test_runaway_loop_stopped_by_duration_cap(self) -> None:
+        with pytest.raises(ModelRetry, match='Runtime error'):
+            await self._run_code('while True:\n    x = 1', {'max_duration_secs': 0.2})
+
+    async def test_partial_mapping_is_applied_in_the_sandbox(self) -> None:
+        # Code comfortably under the default backstop trips a small explicit `max_memory`,
+        # proving a partial dict reaches the sandbox (merged, not dropped).
+        with pytest.raises(ModelRetry, match='Runtime error'):
+            await self._run_code("x = ['data'] * 100000\nlen(x)", {'max_memory': 4096})
+
+    async def test_recursion_depth_cap(self) -> None:
+        code = 'def f(n: int) -> int:\n    return f(n + 1)\nf(0)'
+        with pytest.raises(ModelRetry, match='Runtime error'):
+            await self._run_code(code, {'max_recursion_depth': 4})
+
+    async def test_default_backstop_allows_normal_code(self) -> None:
+        result = await self._run_code('print(await add(a=2, b=3))', None)
+        assert result.return_value == {'output': '5\n'}
+
+    async def test_unlimited_removes_limits(self) -> None:
+        # The same allocation that trips a small explicit cap runs to completion.
+        result = await self._run_code("x = ['data'] * 100000\nprint(len(x))", 'unlimited')
+        assert result.return_value == {'output': '100000\n'}
+
+    def test_unknown_key_raises_at_construction(self) -> None:
+        # A typo'd key (e.g. plural `max_durations_secs`) must not be silently dropped -- that
+        # would disable the cap it was meant to set. Validated eagerly, not at the first call.
+        toolset = _build_function_toolset(add)
+        with pytest.raises(UserError, match='Unknown `resource_limits` key'):
+            CodeMode[object](
+                resource_limits={'max_durations_secs': 5},  # pyright: ignore[reportArgumentType]
+            ).get_wrapper_toolset(toolset)
+
+    def test_capability_field_reaches_toolset(self) -> None:
+        toolset = _build_function_toolset(add)
+        wrapper = CodeMode[object](resource_limits='unlimited').get_wrapper_toolset(toolset)
+        assert isinstance(wrapper, CodeModeToolset)
+        assert wrapper.resource_limits == 'unlimited'