Feature: Code Mode Capability with Monty (#167)

Co-authored-by: Aditya Vardhan <adtyavrdhn@gmail.com>
Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

.github/workflows/main.yml

@@ -18,16 +18,17 @@ jobs:
   lint:
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
+      - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
         with:
           persist-credentials: false
 
       - uses: astral-sh/setup-uv@cec208311dfd045dd5311c1add060b2062131d57 # v8.0.0
         with:
-          python-version: ${{ matrix.python-version }}
+          python-version: '3.13'
           enable-cache: true # zizmor: ignore[cache-poisoning] -- Job does not produce release artifacts and does not have sensitive permissions
+          cache-suffix: lint
 
-      - run: uv sync --frozen --all-groups
+      - run: uv sync --frozen --all-groups --all-extras
 
       - uses: pre-commit/action@646c83fcd040023954eafda54b4db0192ce70507 # v3.0.0
         with:
@@ -43,8 +44,16 @@ jobs:
       fail-fast: false
       matrix:
         python-version: ['3.10', '3.12', '3.13']
+        install:
+          - name: slim
+            extras: ''
+          - name: all-extras
+            extras: '--all-extras'
+    name: test on ${{ matrix.python-version }} (${{ matrix.install.name }})
+    env:
+      COVERAGE_FILE: coverage-data/.coverage.${{ matrix.python-version }}-${{ matrix.install.name }}
     steps:
-      - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
+      - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
         with:
           persist-credentials: false
 
@@ -53,12 +62,41 @@ jobs:
           python-version: ${{ matrix.python-version }}
           enable-cache: true # zizmor: ignore[cache-poisoning] -- Job does not produce release artifacts and does not have sensitive permissions
 
-      - run: uv sync --frozen --all-groups
-      - run: make testcov
+      - run: uv sync --frozen --group dev ${{ matrix.install.extras }}
+      - run: mkdir -p coverage-data
+      - run: uv run coverage run -m pytest
+      - uses: actions/upload-artifact@ea165f8d65b6e75b540449e92b4886f43607fa02 # v4.6.2
+        with:
+          name: coverage-${{ matrix.python-version }}-${{ matrix.install.name }}
+          path: coverage-data/.coverage.*
+          include-hidden-files: true
+
+  coverage:
+    needs: [test]
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
+        with:
+          persist-credentials: false
+
+      - uses: astral-sh/setup-uv@cec208311dfd045dd5311c1add060b2062131d57 # v8.0.0
+        with:
+          python-version: '3.13'
+          enable-cache: true # zizmor: ignore[cache-poisoning] -- Job does not produce release artifacts and does not have sensitive permissions
+
+      - uses: actions/download-artifact@d3f86a106a0bac45b974a628896c90dbdf5c8093 # v4.3.0
+        with:
+          pattern: coverage-*
+          path: coverage-data
+          merge-multiple: true
+
+      - run: uv sync --frozen --group dev --all-extras
+      - run: uv run coverage combine coverage-data
+      - run: uv run coverage report
 
   check:
     if: always()
-    needs: [lint, test]
+    needs: [lint, test, coverage]
     runs-on: ubuntu-latest
     steps:
       - uses: re-actors/alls-green@05ac9388f0aebcb5727afa17fcccfecd6f8ec5fe # v1.2.2
@@ -74,7 +112,7 @@ jobs:
       contents: read
       id-token: write
     steps:
-      - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
+      - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
         with:
           persist-credentials: false
 
@@ -84,4 +122,4 @@ jobs:
           enable-cache: false
 
       - run: uv build
-      - uses: pypa/gh-action-pypi-publish@ed0c53931b1dc9bd32cbe73a98c7f6766f8a527e # v1.13.0
+      - uses: pypa/gh-action-pypi-publish@ed0c53931b1dc9bd32cbe73a98c7f6766f8a527e # v1.13.0

AGENTS.md

@@ -71,15 +71,18 @@ Always run `make lint && make typecheck && make test` before committing.
 ## File structure
 
 ```
-src/pydantic_harness/
+pydantic_harness/
   __init__.py          # public API re-exports
-  guardrails.py        # guardrail capabilities (InputGuardrail, OutputGuardrail, CostGuard, ToolGuard, AsyncGuardrail)
-  <capability>.py      # each capability gets its own module
+  <capability>/        # each capability gets its own package
+    __init__.py        # public exports for the capability
+    _capability.py     # capability class (AbstractCapability subclass)
+    _toolset.py        # toolset implementation
 tests/
   conftest.py          # shared fixtures (TestModel, test_agent)
-  test_<capability>.py # tests mirror source modules
-examples/
-  <capability>/        # runnable examples per capability
+  _<capability>/       # tests mirror source packages
+    test_<capability>.py
+<capability>/
+  README.md            # standalone docs for the capability
 ```
 
 ## Testing patterns

Makefile

@@ -1,3 +1,5 @@
+.DEFAULT_GOAL := all
+
 .PHONY: .uv .prek install format lint typecheck test testcov all
 
 .uv:
@@ -7,12 +9,12 @@
 	@prek --version || echo 'Please install prek: https://github.com/j178/pre-commit-rs'
 
 install: .uv .prek
-	uv sync --frozen --all-groups
+	uv sync --frozen --all-extras --group lint
 	prek install --install-hooks
 
 format:
 	uv run ruff format
-	uv run ruff check --fix
+	uv run ruff check --fix --fix-only
 
 lint:
 	uv run ruff format --check

README.md

@@ -29,7 +29,7 @@ Extras for specific capabilities:
 uv add "pydantic-harness[code-mode]"   # CodeMode (adds the Monty sandbox)
 ```
 
-Requires Python 3.10+ and `pydantic-ai-slim>=1.76.0`.
+Requires Python 3.10+ and `pydantic-ai-slim>=1.78.0`.
 
 ## Quick start
 

demo_code_mode.py

@@ -0,0 +1,123 @@
+"""End-to-end CodeMode demo against Anthropic Claude Sonnet 4.6.
+
+Run with:
+
+    uv pip install anthropic logfire
+    ANTHROPIC_API_KEY=... LOGFIRE_TOKEN=... uv run python demo_code_mode.py
+
+Both `anthropic` and `logfire` are optional. Without `LOGFIRE_TOKEN` Logfire
+prints structured spans to stdout instead of streaming to the UI.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import os
+import sys
+
+from pydantic_ai import Agent
+from pydantic_ai.messages import ToolCallPart
+
+from pydantic_harness import CodeMode
+
+try:
+    import logfire
+except ImportError:  # pragma: no cover - logfire is optional
+    logfire = None  # type: ignore[assignment]
+
+
+def _configure_logfire() -> None:
+    """Configure Logfire if available, otherwise do nothing."""
+    if logfire is None:
+        return
+    logfire.configure(
+        service_name='pydantic-harness-demo',
+        send_to_logfire='if-token-present',
+        console=logfire.ConsoleOptions(verbose=True),
+    )
+    logfire.instrument_pydantic_ai()
+    try:
+        logfire.instrument_anthropic()  # pyright: ignore[reportUnknownMemberType]
+    except Exception:  # pragma: no cover - anthropic instrumentation is optional
+        pass
+
+
+PRODUCTS: dict[str, dict[str, float | int]] = {
+    'apple': {'price': 1.20, 'stock': 50},
+    'banana': {'price': 0.50, 'stock': 200},
+    'cherry': {'price': 3.00, 'stock': 12},
+    'date': {'price': 5.50, 'stock': 8},
+    'elderberry': {'price': 7.25, 'stock': 4},
+}
+
+
+def get_price(item: str) -> float:
+    """Look up the unit price of an item in dollars."""
+    if item not in PRODUCTS:
+        raise ValueError(f'Unknown item: {item!r}. Available: {sorted(PRODUCTS)}')
+    return float(PRODUCTS[item]['price'])
+
+
+def get_stock(item: str) -> int:
+    """Return the current available stock for an item."""
+    if item not in PRODUCTS:
+        raise ValueError(f'Unknown item: {item!r}. Available: {sorted(PRODUCTS)}')
+    return int(PRODUCTS[item]['stock'])
+
+
+def apply_discount(amount: float, percent: int) -> float:
+    """Apply a percentage discount to a dollar amount and return the new total."""
+    return round(amount * (1 - percent / 100), 2)
+
+
+async def main() -> int:
+    """Run the CodeMode demo against Claude Sonnet 4.6."""
+    if not os.environ.get('ANTHROPIC_API_KEY'):
+        print('ANTHROPIC_API_KEY not set', file=sys.stderr)
+        return 1
+
+    _configure_logfire()
+
+    agent: Agent[None, str] = Agent(
+        'anthropic:claude-sonnet-4-6',
+        capabilities=[CodeMode[None]()],
+        instructions=(
+            'You are a shopping assistant. When the user asks a question that requires '
+            'multiple tool calls or arithmetic, use the `run_code` tool to write a single '
+            'Python snippet that calls the available functions (remember to `await` them) '
+            'and prints the final answer. Then summarise the result for the user.'
+        ),
+    )
+
+    agent.tool_plain(get_price)
+    agent.tool_plain(get_stock)
+    agent.tool_plain(apply_discount)
+
+    question = (
+        "I'd like to buy 4 apples, 3 bananas, 2 cherries, and 1 date. "
+        'For any item where I am ordering more than 25% of available stock, '
+        'apply a 15% discount on that item only. '
+        'What is my total in dollars?'
+    )
+
+    if logfire is not None:
+        with logfire.span('code_mode_demo', question=question):
+            result = await agent.run(question)
+    else:
+        result = await agent.run(question)
+
+    for msg in result.all_messages():
+        for part in msg.parts:
+            if isinstance(part, ToolCallPart) and part.tool_name == 'run_code':
+                args = part.args if isinstance(part.args, dict) else None
+                if args and 'code' in args:
+                    print(args['code'])
+                    print()
+
+    print(result.output)
+
+    return 0
+
+
+if __name__ == '__main__':
+    sys.exit(asyncio.run(main()))

pydantic_harness/__init__.py

@@ -0,0 +1,16 @@
+"""Agent harness for composable, reusable AI agent capabilities, for Pydantic AI."""
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from .code_mode import CodeMode
+
+__all__ = ['CodeMode']
+
+
+def __getattr__(name: str) -> object:
+    if name == 'CodeMode':
+        from .code_mode import CodeMode
+
+        return CodeMode
+    raise AttributeError(f'module {__name__!r} has no attribute {name!r}')

pydantic_harness/code_mode/__init__.py

@@ -0,0 +1,6 @@
+"""Code mode capability: route tool calls through a sandboxed Python environment."""
+
+from pydantic_harness.code_mode._capability import CodeMode
+from pydantic_harness.code_mode._toolset import CodeModeToolset
+
+__all__ = ['CodeMode', 'CodeModeToolset']

pydantic_harness/code_mode/_capability.py

@@ -0,0 +1,57 @@
+"""Code mode capability that routes selected tools through a Monty sandbox."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+
+from pydantic_ai import AbstractToolset
+from pydantic_ai.capabilities import AbstractCapability, CapabilityOrdering
+from pydantic_ai.capabilities._tool_search import ToolSearch as _ToolSearch
+from pydantic_ai.tools import AgentDepsT, ToolSelector
+
+from pydantic_harness.code_mode._toolset import CodeModeToolset
+
+
+@dataclass
+class CodeMode(AbstractCapability[AgentDepsT]):
+    """Capability that exposes selected tools as callables inside a `run_code` sandbox.
+
+    By default (`tools='all'`) every tool the agent has is wrapped behind a single
+    `run_code` tool -- the model writes Python that calls them as functions instead
+    of issuing tool calls directly.
+
+    Pass a list of tool names or a callable predicate to `tools` to split the
+    toolset: matching tools become callables inside the sandbox, and the rest
+    stay visible to the model as normal tool calls.
+
+    ```python
+    from pydantic_ai import Agent
+    from pydantic_harness import CodeMode
+
+    # Sandbox all tools
+    agent = Agent('openai:gpt-5', capabilities=[CodeMode()])
+
+    # Sandbox only specific tools
+    agent = Agent('openai:gpt-5', capabilities=[CodeMode(tools=['search', 'fetch'])])
+    ```
+    """
+
+    tools: ToolSelector[AgentDepsT] = field(default='all')
+    """Which wrapped tools should be sandboxed inside `run_code`.
+
+    - `'all'` (default): every tool the agent has is sandboxed.
+    - `Sequence[str]`: only tools whose names are listed are sandboxed.
+    - Callable `(ctx, tool_def) -> bool | Awaitable[bool]`: tools where the
+      callable returns `True` are sandboxed; the rest stay as native tool calls.
+    """
+
+    max_retries: int = 3
+    """Maximum number of retries for the `run_code` tool (syntax errors count as retries)."""
+
+    def get_ordering(self) -> CapabilityOrdering:
+        """CodeMode wraps around ToolSearch so that search_tools stays native."""
+        return CapabilityOrdering(position='outermost', wraps=[_ToolSearch])
+
+    def get_wrapper_toolset(self, toolset: AbstractToolset[AgentDepsT]) -> AbstractToolset[AgentDepsT] | None:
+        """Wrap the agent's assembled toolset, splitting it into native + sandboxed subsets if needed."""
+        return CodeModeToolset(wrapped=toolset, tool_selector=self.tools, max_retries=self.max_retries)

pydantic_harness/code_mode/_toolset.py

@@ -0,0 +1,781 @@
+"""Code mode toolset that runs LLM-generated Python in a Monty sandbox."""
+
+from __future__ import annotations
+
+import asyncio
+import keyword
+import re
+import warnings
+from collections.abc import Callable, Coroutine
+from dataclasses import dataclass, field, replace
+from typing import Annotated, Any
+
+from pydantic import Field, TypeAdapter
+from pydantic_ai import AbstractToolset, RunContext, ToolDefinition, WrapperToolset
+from pydantic_ai.exceptions import ApprovalRequired, CallDeferred, ModelRetry, UserError
+from pydantic_ai.function_signature import FunctionSignature
+from pydantic_ai.messages import ToolCallPart, ToolReturn, ToolReturnContent, ToolReturnPart, is_multi_modal_content
+from pydantic_ai.tool_manager import ToolManager
+from pydantic_ai.tools import AgentDepsT, ToolSelector, matches_tool_selector
+from pydantic_ai.toolsets.abstract import SchemaValidatorProt, ToolsetTool
+
+try:
+    from pydantic_ai.toolsets._tool_search import _SEARCH_TOOLS_NAME  # pyright: ignore[reportPrivateUsage]
+except ImportError:  # pragma: no cover
+    _SEARCH_TOOLS_NAME = 'search_tools'  # pyright: ignore[reportConstantRedefinition]
+
+try:
+    from pydantic_monty import (
+        ExternalException,
+        ExternalResult,
+        ExternalReturnValue,
+        FunctionSnapshot,
+        FutureSnapshot,
+        Monty,
+        MontyComplete,
+        MontyRepl,
+        MontyRuntimeError,
+        MontySyntaxError,
+        MontyTypingError,
+        NameLookupSnapshot,
+    )
+except ImportError as _import_error:  # pragma: no cover
+    raise ImportError(
+        'pydantic-monty is required for CodeMode. Install it with: pip install "pydantic-harness[code-mode]"'
+    ) from _import_error
+from typing_extensions import NotRequired, TypedDict
+
+# Type alias for the dispatch callback passed to _execution_loop.
+_DispatchFn = Callable[[str, dict[str, Any]], Coroutine[Any, Any, Any]]
+
+
+class _RunCodeArguments(TypedDict):
+    code: Annotated[str, Field(description='The Python code to execute in the sandbox.')]
+    restart: NotRequired[
+        Annotated[
+            bool,
+            Field(
+                description='Set to true to reset REPL state. When false (default), state is preserved between calls.'
+            ),
+        ]
+    ]
+
+
+_RUN_CODE_TOOL_NAME = 'run_code'
+_RUN_CODE_ADAPTER = TypeAdapter(_RunCodeArguments)
+_RUN_CODE_JSON_SCHEMA = _RUN_CODE_ADAPTER.json_schema()
+_RUN_CODE_ARGS_VALIDATOR: SchemaValidatorProt = _RUN_CODE_ADAPTER.validator  # pyright: ignore[reportAssignmentType]
+# Used to serialize tool return values before sending into Monty (dump_python)
+# and to reconstruct multimodal types (e.g. BinaryContent) from Monty results (validate_python).
+_TOOL_RETURN_CONTENT_TA: TypeAdapter[Any] = TypeAdapter(ToolReturnContent)
+
+_RUN_CODE_BASE_DESCRIPTION = """\
+Write and run Python code in a sandboxed environment.
+
+The sandbox uses Monty, a subset of Python. Key restrictions:
+- **No classes**: class definitions are not supported
+- **No third-party libraries**: only the standard library modules listed below can be used
+- **Importable standard library modules**: `sys`, `typing`, `asyncio`, `math`, `json`, `re`, `datetime`, `os`, `pathlib`. These must be imported at the top of your snippet before use, just like in regular Python. For example: `import asyncio` then `results = await asyncio.gather(tool_one(...), tool_two(...))`.
+- **No `import *`**: wildcard imports are not supported
+
+State is preserved between calls (REPL-style). Set `restart: true` to reset state.
+
+The last expression's value is automatically captured as the return value — you do **not** need to \
+`print()` it. Avoid `print()` for return values as it produces Python string representations, not \
+structured data. Use `print()` only for supplementary logging or debug output.
+
+Returns the last expression's value directly. If `print()` was also called, returns \
+`{"output": "<printed text>", "result": <last expression>}`.\
+"""
+
+
+def _functions_header(*, has_sync: bool, has_async: bool) -> str:
+    """Build the functions-header paragraph for the `run_code` tool description."""
+    base = (
+        '\nThe following functions are available inside the sandbox. Call them directly '
+        '(do **not** redefine or import them). All parameters are keyword-only.'
+    )
+    if has_async and not has_sync:
+        return base + (
+            ' All tool functions are async: invoke them with `await`,'
+            ' e.g. `result = await tool_name(arg=value)`.'
+            ' Calling without `await` returns an unresolved future, not the value.'
+        )
+    if has_sync and not has_async:
+        return base + (' All tool functions are synchronous: call them directly, e.g. `result = tool_name(arg=value)`.')
+    return base + (
+        ' Async functions (`async def`) must be invoked with `await`,'
+        ' e.g. `result = await tool_name(arg=value)`.'
+        ' Sync functions (`def`) are called directly, e.g. `result = tool_name(arg=value)`.'
+    )
+
+
+_SEARCH_TOOLS_MODIFIER = (
+    ' Note: discovered tools become callable as functions inside the run_code sandbox in subsequent invocations.'
+)
+
+_TOOL_SEARCH_ADDENDUM = (
+    f'\n\nNot all functions may be available initially.'
+    f' Use the `{_SEARCH_TOOLS_NAME}` tool to discover additional functions'
+    f' that will become callable in subsequent `run_code` invocations.'
+)
+
+_INVALID_IDENT_CHARS = re.compile(r'[^a-zA-Z0-9_]')
+
+
+def _sanitize_tool_name(name: str) -> str:
+    """Turn a tool name into a valid Python identifier.
+
+    Replaces hyphens, dots, and other non-identifier characters with underscores,
+    prepends `_` if the result starts with a digit, appends `_` if it is a Python keyword.
+    """
+    sanitized = _INVALID_IDENT_CHARS.sub('_', name)
+    if sanitized and sanitized[0].isdigit():
+        sanitized = f'_{sanitized}'
+    if keyword.iskeyword(sanitized):
+        sanitized = f'{sanitized}_'
+    return sanitized or '_'
+
+
+@dataclass(kw_only=True)
+class _RunCodeTool(ToolsetTool[AgentDepsT]):
+    """ToolsetTool subclass that caches data computed during `get_tools`.
+
+    Avoids a redundant `get_tools` call in `call_tool` by storing the
+    callable tool definitions and name mapping on the tool instance itself.
+    Follows the same pattern as `_SearchTool` in pydantic-ai's
+    `ToolSearchToolset`.
+    """
+
+    callable_defs: dict[str, ToolDefinition]
+    """Tool definitions callable from inside the sandbox, keyed by (possibly sanitized) name."""
+
+    sanitized_to_original: dict[str, str]
+    """Maps sanitized Python-safe names back to original tool names (only for renamed tools)."""
+
+    wrapped_tools: dict[str, ToolsetTool[AgentDepsT]]
+    """The wrapped toolset's tools, keyed by original name."""
+
+
+@dataclass
+class CodeModeToolset(WrapperToolset[AgentDepsT]):
+    """Implementation toolset for the `CodeMode` capability.
+
+    Exposes a single `run_code` tool alongside any native (non-sandboxed) tools.
+    Tools selected by `tool_selector` are presented to the model as Python
+    function signatures inside the `run_code` tool description and become
+    callable from the sandbox at runtime. Non-selected tools remain visible
+    to the model as normal tool calls.
+
+    Tools that require deferred execution (kind `external`/`unapproved`) or
+    deferred loading (`defer_loading=True`) cannot be called from inside the
+    sandbox and are dropped with a one-time `UserWarning`.
+    """
+
+    tool_selector: ToolSelector[AgentDepsT] = 'all'
+    """Which wrapped tools to sandbox inside `run_code`. Non-matching tools
+    are exposed as native tools."""
+
+    max_retries: int = 3
+    """Maximum number of retries for the `run_code` tool (syntax errors count as retries)."""
+
+    # 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)
+
+    # Tracks deferred-tool names we've already warned about so we don't spam the
+    # 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)
+
+    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)
+        return replace(self, wrapped=wrapped)
+
+    async def for_run_step(self, ctx: RunContext[AgentDepsT]) -> AbstractToolset[AgentDepsT]:
+        """Update the wrapped toolset for this step while preserving REPL state."""
+        new_wrapped = await self.wrapped.for_run_step(ctx)
+        if new_wrapped is self.wrapped:
+            return self
+        new_self = replace(self, wrapped=new_wrapped)
+        new_self._repl = self._repl
+        new_self._warned_deferred = self._warned_deferred
+        return new_self
+
+    async def get_tools(self, ctx: RunContext[AgentDepsT]) -> dict[str, ToolsetTool[AgentDepsT]]:
+        """Return the `run_code` tool plus any native (non-sandboxed) tools."""
+        wrapped_tools = await self.wrapped.get_tools(ctx)
+
+        # Split tools into sandboxed vs native based on the selector.
+        # The search_tools tool (from ToolSearchToolset) is always kept native
+        # so the model can discover deferred tools alongside run_code.
+        sandboxed_tools: dict[str, ToolsetTool[AgentDepsT]] = {}
+        native_tools: dict[str, ToolsetTool[AgentDepsT]] = {}
+        for name, tool in wrapped_tools.items():
+            if name == _SEARCH_TOOLS_NAME:
+                native_tools[name] = tool
+            elif await matches_tool_selector(self.tool_selector, ctx, tool.tool_def):
+                sandboxed_tools[name] = tool
+            else:
+                native_tools[name] = tool
+
+        callable_defs, sanitized_to_original, native_fallbacks = self._partition_callable_tools(sandboxed_tools)
+
+        # Tools that matched the selector but can't run in the sandbox (deferred
+        # execution, deferred loading) are promoted back to native tools so
+        # they remain visible to the model.
+        for name in native_fallbacks:
+            native_tools[name] = sandboxed_tools[name]
+
+        description = self._build_description(callable_defs)
+
+        if _RUN_CODE_TOOL_NAME in native_tools:
+            raise UserError(
+                f"Tool name '{_RUN_CODE_TOOL_NAME}' is reserved for code mode. Rename your tool to avoid conflicts."
+            )
+
+        # When search_tools is present, append context about run_code to its
+        # description and add a discovery note to the run_code description.
+        has_search_tools = _SEARCH_TOOLS_NAME in native_tools
+        if has_search_tools:
+            search_tool = native_tools[_SEARCH_TOOLS_NAME]
+            native_tools[_SEARCH_TOOLS_NAME] = replace(
+                search_tool,
+                tool_def=replace(
+                    search_tool.tool_def,
+                    description=(search_tool.tool_def.description or '') + _SEARCH_TOOLS_MODIFIER,
+                ),
+            )
+            description += _TOOL_SEARCH_ADDENDUM
+
+        result: dict[str, ToolsetTool[AgentDepsT]] = dict(native_tools)
+        result[_RUN_CODE_TOOL_NAME] = _RunCodeTool(
+            toolset=self,
+            tool_def=ToolDefinition(
+                name=_RUN_CODE_TOOL_NAME,
+                description=description,
+                parameters_json_schema=_RUN_CODE_JSON_SCHEMA,
+                metadata={'code_arg_name': 'code', 'code_arg_language': 'python'},
+                sequential=True,
+            ),
+            max_retries=self.max_retries,
+            args_validator=_RUN_CODE_ARGS_VALIDATOR,
+            callable_defs=callable_defs,
+            sanitized_to_original=sanitized_to_original,
+            wrapped_tools=wrapped_tools,
+        )
+        return result
+
+    async def call_tool(
+        self, name: str, tool_args: dict[str, Any], ctx: RunContext[AgentDepsT], tool: ToolsetTool[AgentDepsT]
+    ) -> Any:
+        """Execute Python code in the sandbox, or pass through to a native tool."""
+        if not isinstance(tool, _RunCodeTool):
+            # Native (non-sandboxed) tool — pass through to the wrapped toolset.
+            return await self.wrapped.call_tool(name, tool_args, ctx, tool)
+
+        code = tool_args['code']
+        restart = tool_args.get('restart', False)
+
+        # Clear the REPL on restart so that if type checking fails, the
+        # next retry still gets fresh_repl=True and is type-checked again.
+        if restart:
+            self._repl = None
+        fresh_repl = self._repl is None
+
+        callable_defs = tool.callable_defs
+        sanitized_to_original = tool.sanitized_to_original
+
+        # Build a ToolManager for the sandbox's inner tools so that sandboxed
+        # tool calls go through the standard validation/execution path. We
+        # inherit `root_capability` from the agent's ToolManager (for capability
+        # hooks) but use the *wrapped* toolset and its tools.
+        # See https://github.com/pydantic/pydantic-ai/pull/4307
+        parent_tm = ctx.tool_manager
+        assert parent_tm is not None, 'CodeModeToolset requires ctx.tool_manager to be set'
+        tool_manager = ToolManager(
+            toolset=self.wrapped,
+            root_capability=parent_tm.root_capability,
+            ctx=ctx,
+            tools=tool.wrapped_tools,
+        )
+
+        # Determine execution mode for sandbox tool calls:
+        # - global_sequential: forced by durable execution engines (DBOS/Temporal)
+        #   via the parallel execution mode context var. Checked with empty calls
+        #   to isolate the context var from per-tool flags.
+        # - sequential_tools: per-tool `sequential` flags on ToolDefinition.
+        #   These tools are rendered as `def` (sync) and resolved inline.
+        global_sequential = tool_manager.get_parallel_execution_mode([]) != 'parallel'
+        sequential_tools = {name for name, td in callable_defs.items() if td.sequential}
+
+        # Collect nested tool calls and returns keyed by tool_call_id so they
+        # can be attached as metadata on the run_code ToolReturnPart.
+        nested_calls: dict[str, ToolCallPart] = {}
+        nested_returns: dict[str, ToolReturnPart] = {}
+        call_counter = 0
+
+        async def dispatch_tool_call(original_name: str, kwargs: dict[str, Any]) -> Any:
+            """Dispatch a single tool call from inside the sandbox.
+
+            Returns the serialized tool result on success. On failure, the
+            exception propagates — the execution loop passes it back into
+            Monty via `ExternalException` so the sandbox sees it at the
+            `await` site.
+            """
+            nonlocal call_counter
+            call_counter += 1
+            parent_id = ctx.tool_call_id or 'pyd_ai_code_mode'
+            tool_call_id = f'{parent_id}__{call_counter}'
+            call_part = ToolCallPart(tool_name=original_name, args=kwargs, tool_call_id=tool_call_id)
+            nested_calls[tool_call_id] = call_part
+
+            try:
+                result = await tool_manager.handle_call(call_part, wrap_validation_errors=False)
+            except (CallDeferred, ApprovalRequired) as e:
+                # Approval/deferral require a round-trip back to the caller,
+                # which the sandbox cannot do. Raise UserError so the execution
+                # loop passes it into Monty as an ExternalException; Monty
+                # re-raises it as MontyRuntimeError, which we catch and convert
+                # to ModelRetry. The error message is preserved through the chain.
+                raise UserError(
+                    'Tool approval and deferral are not supported in code mode. '
+                    f'Tool {original_name!r} raised {type(e).__name__}; ensure wrapped '
+                    'tools do not use approval or deferral when used with CodeMode.'
+                ) from e
+
+            # Unwrap ToolReturn to get the plain value for the sandbox,
+            # preserving the full ToolReturn metadata on the return part.
+            return_metadata: Any = None
+            if isinstance(result, ToolReturn):
+                return_metadata = result.metadata
+                result = result.return_value
+
+            nested_returns[tool_call_id] = ToolReturnPart(
+                tool_name=original_name,
+                content=result,
+                tool_call_id=tool_call_id,
+                metadata=return_metadata,
+            )
+
+            # Serialize to JSON-compatible form so Monty receives only plain data.
+            return _TOOL_RETURN_CONTENT_TA.dump_python(result)
+
+        # Static type checking on fresh REPL sessions (first call or after
+        # restart). Skipped on subsequent calls because accumulated REPL state
+        # (variables from prior snippets) is invisible to the stateless checker.
+        # Runs before REPL creation so that if this raises ModelRetry, the REPL
+        # stays None and the next retry still gets type-checked.
+        if fresh_repl and callable_defs:
+            self._type_check(code, callable_defs=callable_defs)
+
+        # Create the REPL after type checking passes.
+        if fresh_repl:
+            self._repl = MontyRepl()
+        assert self._repl is not None
+
+        capture = _PrintCapture()
+
+        try:
+            monty_state = self._repl.feed_start(code, print_callback=capture)
+            completed = await _execution_loop(
+                monty_state,
+                dispatch=dispatch_tool_call,
+                callable_defs=callable_defs,
+                sanitized_to_original=sanitized_to_original,
+                sequential_tools=sequential_tools,
+                global_sequential=global_sequential,
+            )
+        except MontySyntaxError as e:
+            raise ModelRetry(f'Syntax error in code:\n{_prepend_prints(e.display(), capture)}') from e
+        except MontyTypingError as e:  # pragma: no cover — MontyRepl.feed_start doesn't raise this
+            raise ModelRetry(f'Type error in code:\n{_prepend_prints(e.display(), capture)}') from e
+        except MontyRuntimeError as e:
+            # Exceptions raised inside dispatch_tool_call (e.g. UserError from
+            # ApprovalRequired, or ModelRetry from a wrapped tool) are passed
+            # back into Monty via ExternalException. Monty re-raises them at the
+            # await site; if the sandbox code doesn't catch them, they bubble up
+            # as MontyRuntimeError. The original exception message is preserved
+            # in the display string, so the model sees a useful error. This means
+            # ModelRetry from a wrapped tool gets double-wrapped
+            # (ModelRetry → MontyRuntimeError → ModelRetry), but the retry
+            # semantics are the same — the model gets another chance.
+            raise ModelRetry(f'Runtime error:\n{_prepend_prints(e.display(), capture)}') from e
+
+        result = completed.output
+        printed = capture.joined
+
+        # Validate result to reconstruct multimodal types (e.g. BinaryContent from
+        # serialized dicts) so they flow through to the model natively.
+        if result is not None:
+            result = _TOOL_RETURN_CONTENT_TA.validate_python(result)
+
+        # Build return value:
+        # - No print → return result directly (multimodal content stays top-level
+        #   so _split_content can extract it for native model delivery)
+        # - Print + multimodal result → list format so _split_content can extract files
+        # - Print + plain result → dict with output/result keys
+        if not printed:
+            return_value: Any = result if result is not None else {}
+        elif result is None:
+            return_value = {'output': printed}
+        elif _contains_multimodal(result):
+            # Flatten lists so _split_content can find each multimodal item at top level.
+            return_value = [printed, *result] if isinstance(result, list) else [printed, result]
+        else:
+            return_value = {'output': printed, 'result': result}
+
+        return ToolReturn(
+            return_value=return_value,
+            metadata={'code_mode': True, 'tool_calls': nested_calls, 'tool_returns': nested_returns},
+        )
+
+    def _partition_callable_tools(
+        self, wrapped_tools: dict[str, ToolsetTool[AgentDepsT]]
+    ) -> tuple[dict[str, ToolDefinition], dict[str, str], set[str]]:
+        """Return tool definitions that can be called from inside the sandbox.
+
+        Tool names that are not valid Python identifiers (e.g. MCP tools with
+        hyphens or dots like `get-weather`, `api.call`) are sanitized to
+        underscored forms and mapped back to their original names for dispatch.
+
+        Tools requiring deferred execution (kind `external`/`unapproved`) or
+        deferred loading (`defer_loading=True`) cannot run in the sandbox and
+        are excluded from `callable_defs`. Their names are returned in the
+        third element so the caller can promote them back to native tools.
+
+        Returns:
+            A tuple of `(callable_defs, sanitized_to_original, native_fallbacks)`
+            where `native_fallbacks` contains original tool names that should
+            be exposed as native tools instead of being sandboxed.
+        """
+        callable_defs: dict[str, ToolDefinition] = {}
+        sanitized_to_original: dict[str, str] = {}
+        native_fallbacks: set[str] = set()
+        for name, tool in wrapped_tools.items():
+            td = tool.tool_def
+            if td.defer:
+                if name not in self._warned_deferred:
+                    self._warned_deferred.add(name)
+                    warnings.warn(
+                        f'CodeMode: tool {name!r} requires deferred execution '
+                        f'(kind={td.kind!r}) and cannot be called from inside the '
+                        f'sandbox; it will be exposed as a native tool instead.',
+                        UserWarning,
+                        stacklevel=2,
+                    )
+                native_fallbacks.add(name)
+                continue
+            if td.defer_loading:
+                if name not in self._warned_deferred:
+                    self._warned_deferred.add(name)
+                    warnings.warn(
+                        f'CodeMode: tool {name!r} uses deferred loading (tool search) '
+                        f'and cannot be pre-registered in the sandbox; it will be '
+                        f'exposed as a native tool instead.',
+                        UserWarning,
+                        stacklevel=2,
+                    )
+                native_fallbacks.add(name)
+                continue
+
+            safe_name = _sanitize_tool_name(name)
+            if safe_name == _RUN_CODE_TOOL_NAME:
+                raise UserError(
+                    f"Tool name '{name}' (sanitized to '{safe_name}') conflicts with the code mode "
+                    f'meta-tool. Rename your tool to avoid conflicts.'
+                )
+            if safe_name in callable_defs:
+                existing = sanitized_to_original.get(safe_name, safe_name)
+                warnings.warn(
+                    f'CodeMode: tool {name!r} (sanitized to {safe_name!r}) collides '
+                    f'with {existing!r}; {name!r} will be hidden from the sandbox.',
+                    UserWarning,
+                    stacklevel=2,
+                )
+                continue
+            # Warn when a sandboxed tool has no return schema — the generated
+            # signature will show `-> Any`, giving the model no type information
+            # about the return shape, which limits code mode effectiveness.
+            if td.return_schema is None and name not in self._warned_deferred:
+                self._warned_deferred.add(name)
+                warnings.warn(
+                    f'CodeMode: tool {name!r} has no return schema; '
+                    f'its signature will show `-> Any`, which may reduce code mode effectiveness.',
+                    UserWarning,
+                    stacklevel=2,
+                )
+
+            if safe_name != name:
+                sanitized_to_original[safe_name] = name
+                td = replace(td, name=safe_name)
+
+            callable_defs[safe_name] = td
+        return callable_defs, sanitized_to_original, native_fallbacks
+
+    @staticmethod
+    def _build_description(callable_defs: dict[str, ToolDefinition]) -> str:
+        """Render the `run_code` description: base prose + TypedDicts + function signatures."""
+        if not callable_defs:
+            return _RUN_CODE_BASE_DESCRIPTION
+
+        sigs, conflicting = _get_sigs_and_conflicting(callable_defs)
+        type_blocks = FunctionSignature.render_type_definitions(sigs, conflicting)
+        function_blocks = [
+            td.render_signature('...', is_async=not td.sequential, conflicting_type_names=conflicting)
+            for td in callable_defs.values()
+        ]
+
+        has_sync = any(td.sequential for td in callable_defs.values())
+        has_async = any(not td.sequential for td in callable_defs.values())
+        header = _functions_header(has_sync=has_sync, has_async=has_async)
+
+        sections = [_RUN_CODE_BASE_DESCRIPTION, header]
+        if type_blocks:
+            sections.append('```python\n' + '\n\n'.join(type_blocks) + '\n```')
+        sections.append('```python\n' + '\n\n'.join(function_blocks) + '\n```')
+        return '\n\n'.join(sections)
+
+    @staticmethod
+    def _build_type_check_stubs(callable_defs: dict[str, ToolDefinition]) -> str:
+        """Build Python stubs for Monty's static type checker."""
+        sigs, conflicting = _get_sigs_and_conflicting(callable_defs)
+        parts = ['import asyncio\nfrom typing import Any, TypedDict, NotRequired, Literal']
+        type_blocks = FunctionSignature.render_type_definitions(sigs, conflicting)
+        parts.extend(type_blocks)
+        parts.extend(
+            td.render_signature(
+                'raise NotImplementedError()', is_async=not td.sequential, conflicting_type_names=conflicting
+            )
+            for td in callable_defs.values()
+        )
+        return '\n\n'.join(parts)
+
+    @staticmethod
+    def _type_check(code: str, *, callable_defs: dict[str, ToolDefinition]) -> None:
+        """Type-check a code snippet against tool signatures before execution.
+
+        Uses Monty's stateless type checker with function stubs. Only sound
+        when the REPL has no accumulated state (first call or after restart).
+
+        Raises:
+            ModelRetry: If the code has type errors or syntax errors.
+        """
+        stubs = CodeModeToolset._build_type_check_stubs(callable_defs)
+        try:
+            Monty(code, type_check=True, type_check_stubs=stubs)
+        except MontyTypingError as e:
+            raise ModelRetry(f'Type error in code:\n{e.display()}') from e
+        except MontySyntaxError as e:
+            raise ModelRetry(f'Syntax error in code:\n{e.display()}') from e
+
+
+def _get_sigs_and_conflicting(
+    callable_defs: dict[str, ToolDefinition],
+) -> tuple[list[FunctionSignature], frozenset[str]]:
+    """Extract FunctionSignatures and conflicting type names from tool definitions."""
+    sigs: list[FunctionSignature] = []
+    for td in callable_defs.values():
+        assert td.function_signature is not None, f'function_signature missing for tool {td.name!r}'
+        sigs.append(td.function_signature)
+    return sigs, FunctionSignature.get_conflicting_type_names(sigs)
+
+
+async def _execution_loop(
+    monty_state: FunctionSnapshot | FutureSnapshot | NameLookupSnapshot | MontyComplete,
+    *,
+    dispatch: _DispatchFn,
+    callable_defs: dict[str, ToolDefinition],
+    sanitized_to_original: dict[str, str],
+    sequential_tools: set[str],
+    global_sequential: bool,
+) -> MontyComplete:
+    """Drive the Monty REPL via the synchronous snapshot API until completion.
+
+    Uses Monty's `feed_start`/`resume` snapshot API instead of `feed_run_async`
+    to avoid background threads and `call_soon_threadsafe`. This makes it safe
+    to run inside restricted event loops like Temporal's workflow sandbox.
+
+    Tool calls are handled based on their execution mode:
+
+    - **Parallel tools** (``async def``): deferred via ``resume(future=...)``
+      and eagerly scheduled as ``asyncio.Task``s for concurrent execution.
+      Resolved at ``FutureSnapshot`` via ``asyncio.gather``.
+    - **Sequential tools** (``def``): resolved inline at ``FunctionSnapshot``
+      via ``resume(return_value=...)`` or ``resume(exception=...)``. Before
+      dispatching, any pending parallel tasks are awaited to maintain ordering.
+    - **Global sequential mode** (DBOS/Temporal): all tools are deferred via
+      ``resume(future=...)`` but stored as bare coroutines and awaited
+      one-at-a-time at ``FutureSnapshot`` to prevent interleaving.
+    """
+    pending: dict[int, asyncio.Task[Any] | Coroutine[Any, Any, Any]] = {}
+    # Results from parallel tasks that were awaited early (at a sequential-tool
+    # barrier) but whose FutureSnapshot hasn't been reached yet.
+    pre_resolved: dict[int, ExternalResult] = {}
+    try:
+        while not isinstance(monty_state, MontyComplete):
+            if isinstance(monty_state, NameLookupSnapshot):
+                monty_state = monty_state.resume()
+            elif isinstance(monty_state, FunctionSnapshot):
+                monty_state = await _handle_function_snapshot(
+                    monty_state,
+                    dispatch,
+                    callable_defs,
+                    sanitized_to_original,
+                    sequential_tools=sequential_tools,
+                    global_sequential=global_sequential,
+                    pending=pending,
+                    pre_resolved=pre_resolved,
+                )
+            else:
+                monty_state = await _resolve_future_snapshot(
+                    monty_state,
+                    pending=pending,
+                    pre_resolved=pre_resolved,
+                    global_sequential=global_sequential,
+                )
+    finally:
+        for item in pending.values():  # pragma: no cover
+            if isinstance(item, asyncio.Task):
+                item.cancel()
+            else:
+                item.close()
+
+    return monty_state
+
+
+async def _handle_function_snapshot(
+    snapshot: FunctionSnapshot,
+    dispatch: _DispatchFn,
+    callable_defs: dict[str, ToolDefinition],
+    sanitized_to_original: dict[str, str],
+    *,
+    sequential_tools: set[str],
+    global_sequential: bool,
+    pending: dict[int, asyncio.Task[Any] | Coroutine[Any, Any, Any]],
+    pre_resolved: dict[int, ExternalResult],
+) -> FunctionSnapshot | FutureSnapshot | NameLookupSnapshot | MontyComplete:
+    """Handle a single FunctionSnapshot from the Monty execution loop."""
+    fn_name = snapshot.function_name
+
+    if fn_name not in callable_defs:
+        return snapshot.resume(exception=NameError(f'Unknown function: {fn_name}'))
+
+    if snapshot.args:
+        return snapshot.resume(
+            exception=TypeError(f'{fn_name}() does not accept positional arguments; use keyword arguments')
+        )
+
+    original_name = sanitized_to_original.get(fn_name, fn_name)
+
+    if fn_name in sequential_tools:
+        # Per-tool sequential: rendered as `def` (sync), so must resolve inline —
+        # the sandbox code doesn't `await` the result. Await pending parallel
+        # tasks first (barrier) to maintain ordering.
+        for cid in list(pending):
+            pre_resolved[cid] = await _resolve_coro(pending.pop(cid))
+        outcome = await _resolve_coro(dispatch(original_name, snapshot.kwargs))
+        if 'return_value' in outcome:
+            return snapshot.resume(return_value=outcome['return_value'])
+        return snapshot.resume(exception=outcome['exception'])
+
+    # Deferred execution — store for later resolution at FutureSnapshot.
+    if global_sequential:
+        # Bare coroutine — don't schedule on the event loop yet.
+        pending[snapshot.call_id] = dispatch(original_name, snapshot.kwargs)
+    else:
+        # Eagerly schedule as a Task for concurrent execution.
+        pending[snapshot.call_id] = asyncio.ensure_future(dispatch(original_name, snapshot.kwargs))
+    return snapshot.resume(future=...)
+
+
+async def _resolve_future_snapshot(
+    snapshot: FutureSnapshot,
+    *,
+    pending: dict[int, asyncio.Task[Any] | Coroutine[Any, Any, Any]],
+    pre_resolved: dict[int, ExternalResult],
+    global_sequential: bool,
+) -> FunctionSnapshot | FutureSnapshot | NameLookupSnapshot | MontyComplete:
+    """Resolve pending tool calls at a FutureSnapshot."""
+    pending_ids = snapshot.pending_call_ids
+    if not pending_ids:  # pragma: no cover
+        return snapshot.resume(results={})
+
+    results: dict[int, ExternalResult] = {}
+    for cid in pending_ids:
+        if cid in pre_resolved:
+            results[cid] = pre_resolved.pop(cid)
+        elif global_sequential:
+            results[cid] = await _resolve_coro(pending.pop(cid))
+
+    # Gather remaining parallel tasks.
+    gather_ids = [cid for cid in pending_ids if cid not in results]
+    if gather_ids:
+        tasks = [pending[cid] for cid in gather_ids]
+        settled = await asyncio.gather(*tasks, return_exceptions=True)
+        for cid in gather_ids:
+            del pending[cid]
+        for cid, outcome in zip(gather_ids, settled):
+            results[cid] = _settle_outcome(outcome)
+
+    return snapshot.resume(results=results)
+
+
+async def _resolve_coro(
+    coro: Coroutine[Any, Any, Any] | asyncio.Task[Any],
+) -> ExternalReturnValue | ExternalException:
+    """Await a single coroutine/task and wrap the result for Monty."""
+    try:
+        result = await coro
+    except Exception as exc:
+        return ExternalException(exception=exc)
+    else:
+        return ExternalReturnValue(return_value=result)
+
+
+def _settle_outcome(outcome: Any) -> ExternalReturnValue | ExternalException:
+    """Wrap an `asyncio.gather(return_exceptions=True)` outcome for Monty."""
+    if isinstance(outcome, Exception):
+        return ExternalException(exception=outcome)
+    if isinstance(outcome, BaseException):  # pragma: no cover
+        raise outcome
+    return ExternalReturnValue(return_value=outcome)
+
+
+def _prepend_prints(error_message: str, capture: _PrintCapture) -> str:
+    """Prepend any captured print output to an error message.
+
+    When sandbox code prints debug output before crashing, this preserves
+    that output in the error so the model can use it for debugging.
+    """
+    printed = capture.joined.rstrip('\n')
+    if not printed:
+        return error_message
+    return f'[stdout before error]\n{printed}\n[/stdout before error]\n{error_message}'
+
+
+def _contains_multimodal(value: Any) -> bool:
+    """Check if a value is or directly contains multimodal content (images, audio, etc.)."""
+    if is_multi_modal_content(value):
+        return True
+    if isinstance(value, list):
+        return any(is_multi_modal_content(item) for item in value)  # pyright: ignore[reportUnknownVariableType]
+    return False
+
+
+class _PrintCapture:
+    """Accumulates print-callback chunks from the Monty REPL.
+
+    Pulled out to module scope (rather than a closure inside `call_tool`) so
+    the callback path is testable in isolation and visible to coverage.py.
+    """
+
+    def __init__(self) -> None:
+        self._chunks: list[str] = []
+
+    def __call__(self, _stream: str, text: str) -> None:
+        self._chunks.append(text)
+
+    @property
+    def joined(self) -> str:
+        return ''.join(self._chunks)

pyproject.toml

@@ -26,7 +26,20 @@ classifiers = [
     'Topic :: Software Development :: Libraries',
     'Typing :: Typed',
 ]
-dependencies = ['pydantic-ai-slim>=1.76.0']
+dependencies = [
+    'pydantic-ai-slim>=1.78.0',
+]
+
+[project.optional-dependencies]
+code-mode = [
+    'pydantic-monty>=0.0.10',
+]
+temporal = [
+    'pydantic-ai-slim[temporal]',
+]
+dbos = [
+    'pydantic-ai-slim[dbos]',
+]
 
 [project.urls]
 Homepage = 'https://github.com/pydantic/pydantic-harness'
@@ -35,16 +48,21 @@ Issues = 'https://github.com/pydantic/pydantic-harness/issues'
 
 [dependency-groups]
 dev = [
+    'pydantic-harness[code-mode]',
     'pytest',
     'anyio[trio]',
     'pytest-anyio',
     'coverage',
+    'logfire[httpx]>=4.31.0',
 ]
 lint = [
     'ruff>=0.14',
     'pyright>=1.1.408',
 ]
 
+[tool.uv.sources]
+pydantic-ai-slim = { git = 'https://github.com/pydantic/pydantic-ai.git', branch = 'main', subdirectory = 'pydantic_ai_slim' }
+
 [tool.hatch.version]
 source = 'uv-dynamic-versioning'
 
@@ -53,8 +71,9 @@ vcs = 'git'
 style = 'pep440'
 bump = true
 
+
 [tool.hatch.build.targets.wheel]
-packages = ['src/pydantic_harness']
+packages = ['pydantic_harness']
 
 [tool.ruff]
 line-length = 120
@@ -90,18 +109,26 @@ executionEnvironments = [
 [tool.pytest.ini_options]
 testpaths = ['tests']
 xfail_strict = true
-filterwarnings = ['error']
+filterwarnings = [
+    'error',
+    # DBOS's run_sync triggers this on Python 3.12+ — not our code.
+    'ignore:There is no current event loop:DeprecationWarning',
+]
 anyio_mode = 'auto'
 
 [tool.coverage.run]
 branch = true
 source = ['pydantic_harness', 'tests']
 
+[tool.coverage.paths]
+source = ['.', '/home/runner/work/pydantic-harness/pydantic-harness']
+
 [tool.coverage.report]
 fail_under = 100
 show_missing = true
 exclude_lines = [
     'pragma: no cover',
+    'pragma: lax no cover',
     'assert_never',
     'if TYPE_CHECKING:',
 ]

src/pydantic_harness/__init__.py

@@ -1,10 +0,0 @@
-"""Agent harness for composable, reusable AI agent capabilities, built on PydanticAI.
-
-Usage:
-    from pydantic_harness import Memory, Skills, Guardrails, ...
-"""
-
-# Each capability module is imported and re-exported here.
-# Capabilities are listed alphabetically.
-
-__all__: list[str] = []

tests/_code_mode/test_dbos.py

@@ -0,0 +1,149 @@
+"""DBOS integration tests for CodeMode.
+
+Verifies that the snapshot-based execution loop works inside a DBOS
+durable workflow. DBOS uses SQLite locally — no external services needed.
+
+DBOS defaults to `parallel_ordered_events` execution mode, which triggers
+the sequential FutureSnapshot resolution path in the execution loop.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Generator
+from typing import Any
+
+import pytest
+
+try:
+    from dbos import DBOS, DBOSConfig
+    from pydantic_ai.durable_exec.dbos import DBOSAgent
+except ImportError:  # pragma: lax no cover
+    pytest.skip('dbos not installed', allow_module_level=True)
+
+from pydantic_ai import Agent
+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_harness import CodeMode
+
+# ---------------------------------------------------------------------------
+# Fixtures
+# ---------------------------------------------------------------------------
+
+
+@pytest.fixture(scope='module')
+def dbos_instance(tmp_path_factory: pytest.TempPathFactory) -> Generator[DBOS, Any, None]:
+    dbos_sqlite_file = tmp_path_factory.mktemp('dbos') / 'dbostest.sqlite'
+    dbos_config: DBOSConfig = {
+        'name': 'pydantic_harness_dbos_tests',
+        'system_database_url': f'sqlite:///{dbos_sqlite_file}',
+        'run_admin_server': False,
+        'enable_otlp': False,
+    }
+    dbos = DBOS(config=dbos_config)
+    DBOS.launch()
+    try:
+        yield dbos
+    finally:
+        DBOS.destroy()
+
+
+# ---------------------------------------------------------------------------
+# Tools and agents (module-level)
+# ---------------------------------------------------------------------------
+
+
+def add(a: int, b: int) -> int:
+    """Add two numbers."""
+    return a + b
+
+
+def _code_mode_model(messages: list[ModelRequest | ModelResponse], info: AgentInfo) -> ModelResponse:
+    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}')])
+
+    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_dbos_agent',
+    toolsets=[FunctionToolset(tools=[add], id='math')],
+    capabilities=[CodeMode()],
+)
+
+dbos_code_mode_agent = DBOSAgent(code_mode_agent)
+
+
+# ---------------------------------------------------------------------------
+# Tests
+# ---------------------------------------------------------------------------
+
+
+def test_code_mode_runs_in_dbos_workflow(dbos_instance: DBOS) -> None:
+    """CodeMode's snapshot-based execution loop works inside a DBOS durable
+    workflow. DBOS defaults to `parallel_ordered_events` mode, which triggers
+    the sequential FutureSnapshot resolution path."""
+    result = dbos_code_mode_agent.run_sync('Calculate 3 + 4')
+    assert result.output == 'done: 7'
+
+    messages = result.all_messages()
+    assert len(messages) == 4
+
+    # 1. User prompt
+    assert isinstance(messages[0], ModelRequest)
+    user_part = messages[0].parts[0]
+    assert user_part.part_kind == 'user-prompt'
+    assert user_part.content == 'Calculate 3 + 4'  # pyright: ignore[reportUnknownMemberType]
+
+    # 2. Model response — run_code tool call
+    assert isinstance(messages[1], ModelResponse)
+    tc = messages[1].parts[0]
+    assert isinstance(tc, ToolCallPart)
+    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 isinstance(messages[2], ModelRequest)
+    tr = messages[2].parts[0]
+    assert isinstance(tr, ToolReturnPart)
+    assert tr.tool_name == 'run_code'
+    assert tr.content == 7
+    assert tr.tool_call_id == 'test_tc_1'
+
+    # Verify nested tool call/return metadata
+    assert tr.metadata is not None
+    assert tr.metadata['code_mode'] is True
+    nested_calls: dict[str, ToolCallPart] = tr.metadata['tool_calls']
+    nested_returns: dict[str, ToolReturnPart] = tr.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 isinstance(messages[3], ModelResponse)
+    final = messages[3].parts[0]
+    assert isinstance(final, TextPart)
+    assert final.content == 'done: 7'

tests/_code_mode/test_temporal.py

@@ -0,0 +1,208 @@
+"""Temporal integration tests for CodeMode.
+
+Verifies that the snapshot-based execution loop (`feed_start`/`resume`)
+works inside a Temporal workflow sandbox, which forbids threads and
+`call_soon_threadsafe`.
+
+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
+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 Worker
+    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
+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_harness import CodeMode
+
+pytestmark = pytest.mark.anyio
+
+TEMPORAL_PORT = 7244  # avoid conflict with other test suites
+TASK_QUEUE = 'pydantic-harness-code-mode-queue'
+BASE_ACTIVITY_CONFIG = ActivityConfig(
+    start_to_close_timeout=timedelta(seconds=60),
+    retry_policy=RetryPolicy(maximum_attempts=1),
+)
+
+
+# ---------------------------------------------------------------------------
+# Fixtures
+# ---------------------------------------------------------------------------
+
+
+@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
+
+
+# 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."""
+    # 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(),
+        }
+
+
+# ---------------------------------------------------------------------------
+# 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.
+    """
+    async with Worker(
+        client,
+        task_queue=TASK_QUEUE,
+        workflows=[CodeModeWorkflow],
+        plugins=[AgentPlugin(temporal_code_mode_agent)],
+    ):
+        result = await client.execute_workflow(
+            CodeModeWorkflow.run,
+            args=['Calculate 3 + 4'],
+            id='test_code_mode_temporal_1',
+            task_queue=TASK_QUEUE,
+        )
+
+    assert result['output'] == 'done: 7'
+
+    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'