furyhawk/pydantic-ai-harness
1
0
Fork 0
mirror of https://github.com/pydantic/pydantic-ai-harness.git synced 2026-07-21 10:55:35 +00:00
Code Issues Packages Projects Releases Wiki Activity

entering and keeping a pool for Monty, skeptical about where the toolset is entered

Browse Source
This commit is contained in:
Aditya Vardhan
2026-07-13 20:15:08 +05:30
parent 1571f4af36
commit c7895b7556
3 changed files with 42 additions and 328 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pydantic_ai_harness/code_mode/README.md
-40
View File
@@ -154,46 +154,6 @@ The last expression in the code snippet is automatically captured as the return
State persists between `run_code` calls within the same agent run -- variables, imports, and function definitions carry over. Pass `restart: true` in the tool call to reset state.
## Temporal durable execution
Code mode works with Pydantic AI's Temporal integration, but the Temporal worker must pass `pydantic_monty`
through its workflow import sandbox. Monty's native module manages the subprocess pool used to execute code. Passing
the module through keeps that pool outside the per-workflow import sandbox without disabling sandboxing for the rest
of the workflow.
Add a custom workflow runner to the `Worker` that runs the durable agent:
```python
from pydantic_ai.durable_exec.temporal import PydanticAIPlugin
from temporalio.client import Client
from temporalio.worker import Worker
from temporalio.worker.workflow_sandbox import SandboxedWorkflowRunner, SandboxRestrictions
workflow_runner = SandboxedWorkflowRunner(
restrictions=SandboxRestrictions.default.with_passthrough_modules('pydantic_monty')
)
client = await Client.connect(
'localhost:7233',
plugins=[PydanticAIPlugin()],
)
async with Worker(
client,
task_queue='code-mode',
workflows=[CodeModeWorkflow],
workflow_runner=workflow_runner,
):
...
```
`PydanticAIPlugin` merges its normal passthrough modules into the runner's restrictions. The workflow remains
sandboxed, and the same runner can be passed to Temporal's `Replayer`.
Monty code runs again when Temporal replays the workflow. Keep external state access out of the code sandbox: do not
connect `os_access` or `mount` to changing host state. Put filesystem, network, clock, and environment access in tools
so `TemporalAgent` can execute those calls as activities and record their results in workflow history.
## Observability
Nested tool calls inside `run_code` produce their own spans when instrumented with [Logfire](https://pydantic.dev/logfire) or any OpenTelemetry backend. The `run_code` tool return includes metadata with all nested calls:
pydantic_ai_harness/code_mode/_toolset.py
+42 -3
View File
@@ -6,7 +6,8 @@ import inspect
import keyword
import re
import warnings
from collections.abc import Callable, Sequence
from collections.abc import Callable, Generator, Sequence
from contextlib import contextmanager
from dataclasses import dataclass, field, replace
from typing import Annotated, Any
@@ -45,7 +46,7 @@ except ImportError as _import_error: # pragma: no cover
raise ImportError(
'pydantic-monty is required for CodeMode. Install it with: pip install "pydantic-ai-harness[code-mode]"'
) from _import_error
from typing_extensions import NotRequired, TypedDict
from typing_extensions import NotRequired, Self, TypedDict
from pydantic_ai_harness._monty_exec import MontyExecutor, PrintCapture, is_sandbox_panic
@@ -289,6 +290,13 @@ class CodeModeToolset(WrapperToolset[AgentDepsT]):
# (`session.dump()`), reloaded into a fresh session on the next call.
_repl_state: bytes | None = field(default=None, init=False, repr=False)
# The Monty worker pool, created in `__aenter__` and reused by every `run_code` call for
# the lifetime of the entered toolset (one pool per agent run, rather than one per call).
# `None` when the toolset was not entered through its lifecycle -- `call_tool` then falls
# back to a per-call pool. init=False so `for_run` copies start poolless and each entered
# copy creates its own; `for_run_step` copies borrow the entered instance's pool by reference.
_monty_pool: Monty | None = field(default=None, init=False, repr=False, compare=False)
# Catalog string stashed during `get_tools` (when `dynamic_catalog`) and read back by
# `get_instructions` in the same step. Empty when there's nothing to surface.
_last_catalog: str = field(default='', init=False, repr=False)
@@ -309,10 +317,41 @@ class CodeModeToolset(WrapperToolset[AgentDepsT]):
return self
new_self = replace(self, wrapped=new_wrapped)
new_self._repl_state = self._repl_state
new_self._monty_pool = self._monty_pool
new_self._warned_deferred = self._warned_deferred
new_self._last_catalog = self._last_catalog
return new_self
async def __aenter__(self) -> Self:
"""Create the Monty worker pool for this toolset's lifetime, then enter the wrapped toolset.
The pool is created before entering the wrapped toolset so that a failure to spawn workers
(e.g. inside a durable-execution sandbox that forbids subprocesses) fails fast without
leaving the wrapped toolset half-entered.
"""
self._monty_pool = Monty().__enter__()
await self.wrapped.__aenter__()
return self
async def __aexit__(self, *args: Any) -> bool | None:
"""Exit the wrapped toolset, then tear down the worker pool."""
assert self._monty_pool is not None
monty_pool = self._monty_pool
self._monty_pool = None
try:
return await self.wrapped.__aexit__(*args)
finally:
monty_pool.__exit__(*args)
@contextmanager
def _acquire_pool(self) -> Generator[Monty]:
"""Yield the pool created in `__aenter__`; if the toolset was not entered, spin up a per-call pool."""
if self._monty_pool is not None:
yield self._monty_pool
else:
with Monty() as pool:
yield pool
async def get_instructions(
self, ctx: RunContext[AgentDepsT]
) -> str | InstructionPart | Sequence[str | InstructionPart] | None:
@@ -546,7 +585,7 @@ class CodeModeToolset(WrapperToolset[AgentDepsT]):
capture = PrintCapture()
try:
with Monty() as monty_pool:
with self._acquire_pool() as monty_pool:
with monty_pool.checkout(type_check=type_check, type_check_stubs=type_check_stubs) as session:
if self._repl_state is not None:
session.load(self._repl_state)
tests/code_mode/test_temporal.py
-285
View File
@@ -1,285 +0,0 @@
"""Temporal integration tests for CodeMode.
Verifies that the snapshot-based execution loop (`feed_start`/`resume`)
works with Temporal's workflow sandbox and history replay.
Monty executes snippets in subprocess workers. Passing `pydantic_monty`
through Temporal's import sandbox keeps its native module and subprocess pool
outside the per-workflow sandbox while the rest of the workflow remains
sandboxed.
These tests start a local Temporal dev server via
`WorkflowEnvironment.start_local()` -- the Temporal SDK downloads and
runs `temporalite` automatically.
"""
from __future__ import annotations
import json
import subprocess
import sys
from collections.abc import AsyncIterator
from datetime import timedelta
from typing import Any
import pytest
try:
from pydantic_ai.durable_exec.temporal import (
AgentPlugin,
PydanticAIPlugin,
TemporalAgent,
)
from temporalio import workflow
from temporalio.client import Client
from temporalio.common import RetryPolicy
from temporalio.testing import WorkflowEnvironment
from temporalio.worker import Replayer, Worker
from temporalio.worker.workflow_sandbox import (
RestrictedWorkflowAccessError,
SandboxedWorkflowRunner,
SandboxRestrictions,
)
from temporalio.workflow import ActivityConfig
except ImportError: # pragma: lax no cover
pytest.skip('temporalio not installed', allow_module_level=True)
from pydantic_ai import Agent, ToolDefinition
from pydantic_ai.messages import ModelRequest, ModelResponse, TextPart, ToolCallPart, ToolReturnPart
from pydantic_ai.models.function import AgentInfo, FunctionModel
from pydantic_ai.toolsets.function import FunctionToolset
from pydantic_ai_harness import CodeMode
pytestmark = pytest.mark.anyio
TEMPORAL_PORT = 7244 # avoid conflict with other test suites
TASK_QUEUE = 'pydantic-ai-harness-code-mode-queue'
BASE_ACTIVITY_CONFIG = ActivityConfig(
start_to_close_timeout=timedelta(seconds=60),
retry_policy=RetryPolicy(maximum_attempts=1),
)
def _workflow_runner() -> SandboxedWorkflowRunner:
return SandboxedWorkflowRunner(
restrictions=SandboxRestrictions.default.with_passthrough_modules(
'pydantic_monty',
# Coverage imports parser modules lazily while tracing workflow code.
'coverage',
)
)
# ---------------------------------------------------------------------------
# Fixtures
# ---------------------------------------------------------------------------
@pytest.fixture(scope='module')
def anyio_backend() -> str:
"""Temporal's Python SDK runs on asyncio."""
return 'asyncio'
@pytest.fixture(scope='module')
async def temporal_env() -> AsyncIterator[WorkflowEnvironment]:
async with await WorkflowEnvironment.start_local( # pyright: ignore[reportUnknownMemberType]
port=TEMPORAL_PORT,
dev_server_extra_args=[
'--dynamic-config-value',
'frontend.enableServerVersionCheck=false',
],
) as env:
yield env
@pytest.fixture
async def client(temporal_env: WorkflowEnvironment) -> Client:
return await Client.connect(
f'localhost:{TEMPORAL_PORT}',
plugins=[PydanticAIPlugin()],
)
# ---------------------------------------------------------------------------
# Tools and agents (module-level -- Temporal requirement)
# ---------------------------------------------------------------------------
def add(a: int, b: int) -> int:
"""Add two numbers."""
return a + b
_captured_tool_defs: list[list[ToolDefinition]] = []
# FunctionModel that emits a run_code tool call for the given code snippet.
def _code_mode_model(messages: list[ModelRequest | ModelResponse], info: AgentInfo) -> ModelResponse:
"""Model that generates a run_code call on the first request, then returns the result as text."""
_captured_tool_defs.append(info.function_tools)
# Check if we already got a tool result back.
for msg in messages:
if isinstance(msg, ModelResponse):
continue
for part in msg.parts:
if isinstance(part, ToolReturnPart) and part.tool_name == 'run_code':
return ModelResponse(parts=[TextPart(content=f'done: {part.content}')])
# First call -- emit run_code.
return ModelResponse(
parts=[
ToolCallPart(
tool_name='run_code',
args={'code': 'result = await add(a=3, b=4)\nresult'},
tool_call_id='test_tc_1',
)
]
)
code_mode_agent = Agent(
FunctionModel(_code_mode_model),
name='code_mode_temporal_agent',
toolsets=[FunctionToolset(tools=[add], id='math')],
capabilities=[CodeMode()],
)
temporal_code_mode_agent = TemporalAgent(
code_mode_agent,
activity_config=BASE_ACTIVITY_CONFIG,
)
@workflow.defn
class CodeModeWorkflow:
@workflow.run
async def run(self, prompt: str) -> dict[str, Any]:
result = await temporal_code_mode_agent.run(prompt)
return {
'output': str(result.output),
'messages': result.all_messages_json().decode(),
}
@workflow.defn
class SandboxRestrictionWorkflow:
"""Probe that passing Monty through does not allow Python subprocess calls."""
@workflow.run
async def run(self) -> str:
try:
subprocess.run([sys.executable, '-c', 'pass'], check=True)
except RestrictedWorkflowAccessError as e:
return e.qualified_name
return 'subprocess was allowed'
# ---------------------------------------------------------------------------
# Tests
# ---------------------------------------------------------------------------
async def test_code_mode_runs_in_temporal_workflow(client: Client) -> None:
"""CodeMode's snapshot-based execution loop works inside a Temporal workflow.
This is the core regression test for the `call_soon_threadsafe` issue:
the old `feed_run_async` approach hung because Temporal's sandboxed
event loop doesn't implement `call_soon_threadsafe`. The snapshot
approach (`feed_start`/`resume`) avoids threads entirely. The worker passes
`pydantic_monty` through the import sandbox so Monty's subprocess pool can
run without disabling Temporal's sandbox for the workflow.
"""
_captured_tool_defs.clear()
workflow_id = 'test_code_mode_temporal_1'
async with Worker(
client,
task_queue=TASK_QUEUE,
workflows=[CodeModeWorkflow, SandboxRestrictionWorkflow],
plugins=[AgentPlugin(temporal_code_mode_agent)],
workflow_runner=_workflow_runner(),
):
result = await client.execute_workflow(
CodeModeWorkflow.run,
args=['Calculate 3 + 4'],
id=workflow_id,
task_queue=TASK_QUEUE,
)
sandbox_result = await client.execute_workflow(
SandboxRestrictionWorkflow.run,
id='test_code_mode_temporal_sandbox_restrictions',
task_queue=TASK_QUEUE,
)
assert result['output'] == 'done: 7'
assert sandbox_result == 'subprocess.run.__call__'
messages = json.loads(result['messages'])
assert len(messages) == 4
# 1. User prompt
assert messages[0]['kind'] == 'request'
assert messages[0]['parts'][0]['part_kind'] == 'user-prompt'
assert messages[0]['parts'][0]['content'] == 'Calculate 3 + 4'
# 2. Model response -- run_code tool call
assert messages[1]['kind'] == 'response'
tc = messages[1]['parts'][0]
assert tc['part_kind'] == 'tool-call'
assert tc['tool_name'] == 'run_code'
assert tc['args'] == {'code': 'result = await add(a=3, b=4)\nresult'}
assert tc['tool_call_id'] == 'test_tc_1'
# 3. Tool return with nested tool call metadata
assert messages[2]['kind'] == 'request'
tr = messages[2]['parts'][0]
assert tr['part_kind'] == 'tool-return'
assert tr['tool_name'] == 'run_code'
assert tr['content'] == 7
assert tr['tool_call_id'] == 'test_tc_1'
# Verify nested tool call/return metadata
metadata = tr['metadata']
assert metadata is not None
assert metadata['code_mode'] is True
nested_calls = metadata['tool_calls']
nested_returns = metadata['tool_returns']
assert len(nested_calls) == 1
assert len(nested_returns) == 1
nested_call = next(iter(nested_calls.values()))
assert nested_call['tool_name'] == 'add'
assert nested_call['args'] == {'a': 3, 'b': 4}
nested_return = next(iter(nested_returns.values()))
assert nested_return['tool_name'] == 'add'
assert nested_return['content'] == 7
assert nested_return['tool_call_id'] == nested_call['tool_call_id']
# 4. Final text response
assert messages[3]['kind'] == 'response'
assert messages[3]['parts'][0]['part_kind'] == 'text'
assert messages[3]['parts'][0]['content'] == 'done: 7'
# 5. Verify tool definitions sent to the model
assert len(_captured_tool_defs) == 2
for tool_defs in _captured_tool_defs:
tool_names = [td.name for td in tool_defs]
# CodeMode wraps `add` into `run_code` -- the model should only see `run_code`
assert 'run_code' in tool_names
assert 'add' not in tool_names
run_code_td = next(td for td in tool_defs if td.name == 'run_code')
assert run_code_td.description is not None
assert 'async def add' in run_code_td.description
assert run_code_td.parameters_json_schema['properties']['code']['type'] == 'string'
history = await client.get_workflow_handle(workflow_id).fetch_history()
replay_result = await Replayer(
workflows=[CodeModeWorkflow],
plugins=[PydanticAIPlugin()],
workflow_runner=_workflow_runner(),
).replay_workflow(history)
assert replay_result.replay_failure is None