feat(code_mode): expose sandbox resource limits#352
Conversation
Model-generated code can loop forever, recurse without bound, or allocate until the host suffers; the sandbox supported limits but CodeMode gave callers no way to set them. Reuses the normalized-limits pattern from DynamicWorkflow (None backstop / partial merge / 'unlimited') so the two capabilities configure their sandboxes the same way. Fixes pydantic#310
📝 WalkthroughWalkthroughCodeMode now accepts 🚥 Pre-merge checks | ✅ 4✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@pydantic_ai_harness/code_mode/_toolset.py`:
- Around line 112-117: Update the resource-limits validation flow around
_RESOURCE_LIMIT_KEYS and _default_resource_limits() to reject any None values in
the supplied limits mapping before merging. Raise UserError identifying the
invalid key or keys, preserve valid override behavior and defaults, and add a
regression test covering {'max_memory': None} so the 256 MiB safeguard remains
enforced.
In `@tests/code_mode/test_code_mode.py`:
- Around line 2661-2709: The TestResourceLimits end-to-end tests currently
exercise CodeModeToolset directly instead of the public Agent capability API.
Rework _run_code and the execution-focused tests to create an Agent with
capabilities=[CodeMode(...)] and run it using pydantic_ai.models.TestModel;
retain direct toolset construction only for construction validation and
capability-field propagation that cannot be covered through Agent.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 91bf3b77-7253-46a3-bace-04a19c6203a3
📒 Files selected for processing (5)
pydantic_ai_harness/code_mode/README.mdpydantic_ai_harness/code_mode/__init__.pypydantic_ai_harness/code_mode/_capability.pypydantic_ai_harness/code_mode/_toolset.pytests/code_mode/test_code_mode.py
🔗 Linked repositories identified
CodeRabbit considers these linked repositories for cross-repo context during reviews:
pydantic/logfire(manual)pydantic/pydantic-ai(manual)pydantic/monty(auto-detected)pydantic/pydantic(auto-detected)
TypedDict annotations are advisory: a config-driven caller can pass None or a bool, and None would overwrite the backstop in the merge, disabling the guard the caller thought was in place. Same rationale as the FileSystem field validation.
There was a problem hiding this comment.
🧹 Nitpick comments (1)
tests/code_mode/test_code_mode.py (1)
2716-2716: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAvoid
Anyin the invalid-input matrix.
list[Any]weakens strict Pyright at the boundary this test is meant to validate. Uselist[object](or a raw mapping type) and apply a narrowly scoped# pyright: ignore[reportArgumentType]only at the intentional invalid call.As per coding guidelines, Python files must use strict Pyright, avoid
Any, and provide complete type annotations.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/code_mode/test_code_mode.py` at line 2716, Replace the bad_values annotation in the invalid-input matrix with list[object] or an appropriately typed raw mapping, removing Any. Add a narrowly scoped # pyright: ignore[reportArgumentType] only on the intentional invalid call, while preserving strict typing and complete annotations elsewhere in the test.Source: Coding guidelines
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@tests/code_mode/test_code_mode.py`:
- Line 2716: Replace the bad_values annotation in the invalid-input matrix with
list[object] or an appropriately typed raw mapping, removing Any. Add a narrowly
scoped # pyright: ignore[reportArgumentType] only on the intentional invalid
call, while preserving strict typing and complete annotations elsewhere in the
test.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 26c3f971-229c-454e-9ff2-fcfe629bb63a
📒 Files selected for processing (2)
pydantic_ai_harness/code_mode/_toolset.pytests/code_mode/test_code_mode.py
🔗 Linked repositories identified
CodeRabbit considers these linked repositories for cross-repo context during reviews:
pydantic/logfire(manual)pydantic/pydantic-ai(manual)pydantic/monty(auto-detected)pydantic/pydantic(auto-detected)
🚧 Files skipped from review as they are similar to previous changes (1)
- pydantic_ai_harness/code_mode/_toolset.py
1 participant
run_code executes model-generated Python, and that code is not always well behaved: an accidental
while True, unbounded recursion, or a huge allocation can block the event loop or exhaust memory. Monty already supports ResourceLimits but CodeMode had no way to configure them, so a single bad snippet could hang the whole agent run.This adds
CodeMode(resource_limits=...)following the normalized-limits pattern DynamicWorkflow already ships:Noneapplies a backstop (256 MiB, 50M allocations, no duration cap), a partial mapping merges onto that backstop,'unlimited'restores the old unguarded behavior, and unknown keys fail at construction rather than being dropped. Keeping the semantics identical between the two capabilities seemed better than inventing a second dialect, and it matches the direction suggested on the issue. The only plumbing point is the MontyRepl construction; when a limit trips, the error reaches the model through the existing ModelRetry path like any other runtime error.Tested with the same style of checks DynamicWorkflow uses for its limits: a runaway loop stopped by a duration cap, a small memory cap tripping code that fits under the backstop, a recursion cap,
'unlimited'letting that same code run, and eager rejection of typo'd keys. Lint, pyright strict and the coverage gate all pass.Fixes #310