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

Keep ModalSandbox compatible with current main

Browse Source
This commit is contained in:
Aditya Vardhan
2026-07-13 14:51:56 +05:30
parent 24569b3589 310ffadf09
commit a030d95d9f
160 changed files with 12274 additions and 755 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.github/workflows/main.yml
+41
View File
@@ -263,3 +263,44 @@ jobs:
- run: uv build
- uses: pypa/gh-action-pypi-publish@ed0c53931b1dc9bd32cbe73a98c7f6766f8a527e # v1.13.0
send-tweet:
name: Send tweet
needs: [release]
if: needs.release.result == 'success'
runs-on: ubuntu-latest
environment: release
steps:
- uses: actions/setup-python@ece7cb06caefa5fff74198d8649806c4678c61a1 # v6.3.0
with:
python-version: '3.12'
- name: Install dependencies
run: pip install tweepy==4.14.0
- name: Send tweet
shell: python
run: |
import os
import tweepy
client = tweepy.Client(
access_token=os.getenv("TWITTER_ACCESS_TOKEN"),
access_token_secret=os.getenv("TWITTER_ACCESS_TOKEN_SECRET"),
consumer_key=os.getenv("TWITTER_CONSUMER_KEY"),
consumer_secret=os.getenv("TWITTER_CONSUMER_SECRET"),
)
version = os.getenv("VERSION")
tweet = os.getenv("TWEET").format(version=version)
client.create_tweet(text=tweet)
env:
# The tag ref (e.g. "v1.2.3") already carries the leading v, so it is
# used verbatim in both the message and the release URL.
VERSION: ${{ github.ref_name }}
TWEET: |
Pydantic AI Harness {version} is out! 🎉
https://github.com/pydantic/pydantic-ai-harness/releases/tag/{version}
TWITTER_CONSUMER_KEY: ${{ secrets.TWITTER_CONSUMER_KEY }}
TWITTER_CONSUMER_SECRET: ${{ secrets.TWITTER_CONSUMER_SECRET }}
TWITTER_ACCESS_TOKEN: ${{ secrets.TWITTER_ACCESS_TOKEN }}
TWITTER_ACCESS_TOKEN_SECRET: ${{ secrets.TWITTER_ACCESS_TOKEN_SECRET }}
AGENTS.md
+6 -3
View File
@@ -44,9 +44,12 @@ Before implementing or reviewing a capability change:
signatures when needed. Do not assume a contributor's local checkout layout.
5. Use `pydantic_ai_harness.code_mode` as the exemplar for capability shape,
docs, tests, and public exports until another capability becomes a better
example. `code_mode` is a released top-level capability; new capabilities
start under `pydantic_ai_harness.experimental` (see
`agent_docs/capability-authoring.md`, "Experimental Vs Released Exports").
example. Capabilities live in their own top-level submodule
`pydantic_ai_harness/<name>/` (module name = capability name; one module per
capability or strategy) and are not re-exported from the root `__init__.py`,
so each keeps its own optional dependencies. The `experimental` tier is
retired; ACP is the sole remaining experimental capability (see
`agent_docs/capability-authoring.md`, "Capability Submodules And Exports").
## Capabilities API reference
README.md
+65 -15
View File
@@ -15,7 +15,7 @@ Pydantic AI's [capabilities](https://ai.pydantic.dev/capabilities/) and [hooks](
The [capability matrix](#capability-matrix) tracks where we are. [Tell us what to prioritize.](#help-us-prioritize)
**Contents:** [Installation](#installation) · [Quick start](#quick-start) · [Capability matrix](#capability-matrix) · [An ecosystem agent](#an-ecosystem-agent) · [Help us prioritize](#help-us-prioritize) · [Build your own](#build-your-own) · [Contributing](#contributing) · [Version policy](#version-policy) · [Pydantic AI references](#pydantic-ai-references) · [License](#license)
**Contents:** [Installation](#installation) · [Quick start](#quick-start) · [DynamicWorkflow](#orchestrating-sub-agents-dynamicworkflow) · [Capability matrix](#capability-matrix) · [An ecosystem agent](#an-ecosystem-agent) · [Help us prioritize](#help-us-prioritize) · [Build your own](#build-your-own) · [Contributing](#contributing) · [Version policy](#version-policy) · [Pydantic AI references](#pydantic-ai-references) · [License](#license)
## Installation
@@ -26,13 +26,16 @@ uv add pydantic-ai-harness
Extras for specific capabilities:
```bash
uv add "pydantic-ai-harness[codemode]" # CodeMode (adds the Monty sandbox)
uv add "pydantic-ai-harness[logfire]" # ManagedPrompt (Logfire-managed prompts)
uv add "pydantic-ai-harness[codemode]" # CodeMode (adds the Monty sandbox)
uv add "pydantic-ai-harness[dynamic-workflow]" # DynamicWorkflow (adds the Monty sandbox)
uv add "pydantic-ai-harness[modal]" # ModalSandboxCapability (adds the Modal SDK)
uv add "pydantic-ai-harness[logfire]" # ManagedPrompt (Logfire-managed prompts)
uv add "pydantic-ai-harness[acp]" # ACP (serve an agent to editors over the Agent Client Protocol)
```
The `code-mode` extra is also supported as an alias.
Requires Python 3.10+ and `pydantic-ai-slim>=1.95.1`.
Requires Python 3.10+ and `pydantic-ai-slim>=2.1.0`.
## Quick start
@@ -94,6 +97,46 @@ practices and warning of a "normalization of deviance" as engineers stop reviewi
**[See this run as a public Logfire trace →](https://logfire-us.pydantic.dev/public-trace/84bcf123-2106-49da-9f6f-5c26395339bb?spanId=7650806a0785b946)** Each `run_code` span fans out into the tool calls the model issued from inside the sandbox -- it's the easiest way to understand what code mode actually did.
## Orchestrating sub-agents: DynamicWorkflow
`CodeMode` gives the model one script for its *tools*. `DynamicWorkflow` does the same for *sub-agents*. Without it, an orchestrator delegates one tool call at a time: call a sub-agent, wait, read the result into context, think, call the next one. Ten delegations cost ten model round-trips, and every intermediate result flows through the orchestrator's context whether it needed to see it or not.
With it, the model writes one Python script in which each sub-agent is an async function, and the whole tree runs in a single tool call:
```python
from pydantic_ai import Agent
from pydantic_ai_harness.dynamic_workflow import DynamicWorkflow
reviewer = Agent('anthropic:claude-sonnet-4-6', name='reviewer', description='Reviews code for bugs.')
summarizer = Agent('anthropic:claude-sonnet-4-6', name='summarizer', description='Summarizes findings.')
orchestrator = Agent(
'anthropic:claude-opus-4-7',
capabilities=[DynamicWorkflow(agents=[reviewer, summarizer])],
)
```
The script the model writes looks like this -- fan out, chain, and only the last line's value returns to its context:
```python
import asyncio
reports = await asyncio.gather(
reviewer(task="Review auth.py for bugs:\n<file contents>"),
reviewer(task="Review parser.py for bugs:\n<file contents>"),
)
await summarizer(task="Summarize these findings:\n" + "\n\n".join(reports))
```
It composes with the rest of the harness:
- **Budgets**: `max_agent_calls` is an exact, host-enforced ceiling on sub-agent runs (it holds even under concurrent fan-out), and by default the whole tree's token spend lands on the parent run's `usage`.
- **On-demand**: `defer_loading=True` keeps the catalog out of the prompt until the model loads the capability, and `reveal()` adds a sub-agent mid-run without disturbing the prompt cache.
`DynamicWorkflow`'s API is subject to change while planned extensions (structured sub-agent inputs, durable workflows) settle the call contract. Breaking changes ship deprecation warnings where practical.
[Full tutorial →](pydantic_ai_harness/dynamic_workflow/)
## Capability matrix
We studied leading coding agents, agent frameworks, and Claw-style assistants to map every capability area that matters for production agents. Each one is tracked as an [issue](https://github.com/pydantic/pydantic-ai-harness/issues) in this repo.
@@ -106,23 +149,30 @@ We studied leading coding agents, agent frameworks, and Claw-style assistants to
| | **Tool search** | Progressive tool discovery for large tool sets | :white_check_mark: [Pydantic&nbsp;AI](https://pydantic.dev/docs/ai/tools-toolsets/toolsets/#deferred-loading) | |
| | **File system** | Read, write, edit, search files with path traversal prevention | :white_check_mark: [Docs](pydantic_ai_harness/filesystem/) | [pydantic-ai-backend](https://github.com/vstorm-co/pydantic-ai-backend) (vstorm&#8209;co) |
| | **Shell** | Execute commands with allowlists, denylists, and timeouts | :white_check_mark: [Docs](pydantic_ai_harness/shell/) | [pydantic-ai-backend](https://github.com/vstorm-co/pydantic-ai-backend) (vstorm&#8209;co) |
| | **Repo context injection** | Auto-load CLAUDE.md/AGENTS.md and repo structure | :construction: [PR&nbsp;#175](https://github.com/pydantic/pydantic-ai-harness/pull/175) | [pydantic-deep](https://github.com/vstorm-co/pydantic-deepagents) (vstorm&#8209;co) |
| | **Modal sandbox** | Run commands and manage files in an isolated [Modal](https://modal.com) cloud sandbox | :white_check_mark: [Docs](pydantic_ai_harness/modal_sandbox/) | |
| | **Repo context injection** | Auto-load CLAUDE.md/AGENTS.md and repo structure | :white_check_mark: [Docs](pydantic_ai_harness/context/) | [pydantic-deep](https://github.com/vstorm-co/pydantic-deepagents) (vstorm&#8209;co) |
| | **Docs lookup** | On-demand `read_pyai_docs` tool for Pydantic AI docs | :white_check_mark: [Docs](pydantic_ai_harness/docs/) | |
| | **Verification loop** | Run tests after edits, auto-fix failures | :construction: [PR&nbsp;#169](https://github.com/pydantic/pydantic-ai-harness/pull/169) | |
| **Context management** | **Sliding window** | Trim conversation history to stay within token limits | :construction: [PR&nbsp;#191](https://github.com/pydantic/pydantic-ai-harness/pull/191) | [summarization-pydantic-ai](https://github.com/vstorm-co/summarization-pydantic-ai) (vstorm&#8209;co) |
| | **Context compaction** | LLM-powered summarization of older messages | :construction: [PR&nbsp;#191](https://github.com/pydantic/pydantic-ai-harness/pull/191) | [summarization-pydantic-ai](https://github.com/vstorm-co/summarization-pydantic-ai) (vstorm&#8209;co) |
| | **Limit warnings** | Warn agent before hitting context/iteration limits | :construction: [PR&nbsp;#191](https://github.com/pydantic/pydantic-ai-harness/pull/191) | [summarization-pydantic-ai](https://github.com/vstorm-co/summarization-pydantic-ai) (vstorm&#8209;co) |
| | **Tool output management** | Truncate, summarize, or spill large tool outputs | :construction: [PR&nbsp;#185](https://github.com/pydantic/pydantic-ai-harness/pull/185) | |
| **Editor integration** | **ACP** | Serve an agent to editors (Zed, etc.) over the [Agent Client Protocol](https://agentclientprotocol.com) -- streamed text, diff-rendered edits, tool approval | :white_check_mark: [Docs](pydantic_ai_harness/experimental/acp/) (experimental) | |
| **Prompt management** | **Managed prompt** | Back an agent's instructions with a [Logfire](https://pydantic.dev/logfire)-managed prompt, editable without shipping code | :white_check_mark: [Docs](pydantic_ai_harness/logfire/) | |
| **Context management** | **Sliding window** | Trim conversation history to stay within token limits | :white_check_mark: [Docs](pydantic_ai_harness/compaction/) | [summarization-pydantic-ai](https://github.com/vstorm-co/summarization-pydantic-ai) (vstorm&#8209;co) |
| | **Context compaction** | LLM-powered summarization of older messages | :white_check_mark: [Docs](pydantic_ai_harness/compaction/) | [summarization-pydantic-ai](https://github.com/vstorm-co/summarization-pydantic-ai) (vstorm&#8209;co) |
| | **Limit warnings** | Warn agent before hitting context/iteration limits | :white_check_mark: [Docs](pydantic_ai_harness/compaction/) | [summarization-pydantic-ai](https://github.com/vstorm-co/summarization-pydantic-ai) (vstorm&#8209;co) |
| | **Tool output management** | Truncate, summarize, or spill large tool outputs | :white_check_mark: [Docs](pydantic_ai_harness/overflowing_tool_output/) | |
| | **System reminders** | Inject periodic reminders to counteract instruction drift | :construction: [PR&nbsp;#181](https://github.com/pydantic/pydantic-ai-harness/pull/181) | |
| **Memory &&nbsp;persistence** | **Memory** | Persistent key-value memory across sessions | :construction: [PR&nbsp;#179](https://github.com/pydantic/pydantic-ai-harness/pull/179) | [pydantic-deep](https://github.com/vstorm-co/pydantic-deepagents) (vstorm&#8209;co) |
| | **Session persistence** | Save and restore full conversation state | :construction: [PR&nbsp;#176](https://github.com/pydantic/pydantic-ai-harness/pull/176) | |
| | **Checkpointing** | Save, rewind, and fork conversation state | :memo: [#196](https://github.com/pydantic/pydantic-ai-harness/issues/196) | [pydantic-deep](https://github.com/vstorm-co/pydantic-deepagents) (vstorm&#8209;co) |
| **Agent orchestration** | **Sub-agents** | Delegate subtasks to specialized child agents | :construction: [PR&nbsp;#178](https://github.com/pydantic/pydantic-ai-harness/pull/178) | [subagents-pydantic-ai](https://github.com/vstorm-co/subagents-pydantic-ai) (vstorm&#8209;co) |
| | **Session persistence** | Save and restore full conversation state | :white_check_mark: [Docs](pydantic_ai_harness/step_persistence/) | |
| | **Checkpointing** | Snapshot, resume (`continue_run`), and fork (`fork_run`) a run | :white_check_mark: [Docs](pydantic_ai_harness/step_persistence/) | [pydantic-deep](https://github.com/vstorm-co/pydantic-deepagents) (vstorm&#8209;co) |
| | **Media externalization** | Offload large `BinaryContent` to content-addressed stores (building blocks) | :white_check_mark: [Docs](pydantic_ai_harness/media/) | |
| **Agent orchestration** | **Sub-agents** | Delegate subtasks to specialized child agents | :white_check_mark: [Docs](pydantic_ai_harness/subagents/) | [subagents-pydantic-ai](https://github.com/vstorm-co/subagents-pydantic-ai) (vstorm&#8209;co) |
| | **Dynamic workflow** | Orchestrate sub-agents from a model-written Python script -- fan-out, chaining, voting in one tool call | :white_check_mark: [Docs](pydantic_ai_harness/dynamic_workflow/) | |
| | **Skills** | Progressive tool loading -- search, activate, deactivate | :construction: [PR&nbsp;#183](https://github.com/pydantic/pydantic-ai-harness/pull/183) | [pydantic-ai-skills](https://github.com/DougTrajano/pydantic-ai-skills) (DougTrajano), [pydantic-deep](https://github.com/vstorm-co/pydantic-deepagents) (vstorm&#8209;co) |
| | **Planning** | Break complex tasks into structured plans before execution | :construction: [PR&nbsp;#180](https://github.com/pydantic/pydantic-ai-harness/pull/180) | |
| | **Planning** | Break complex tasks into structured plans before execution | :white_check_mark: [Docs](pydantic_ai_harness/planning/) | |
| | **Runtime authoring** | Let an agent author, validate, and load real capabilities at runtime | :white_check_mark: [Docs](pydantic_ai_harness/runtime_authoring/) | |
| | **Task tracking** | Track tasks, subtasks, and dependencies | :memo: [#65](https://github.com/pydantic/pydantic-ai-harness/issues/65) | [pydantic-ai-todo](https://github.com/vstorm-co/pydantic-ai-todo) (vstorm&#8209;co) |
| | **Teams** | Multi-agent teams with shared state and message bus | :memo: [#195](https://github.com/pydantic/pydantic-ai-harness/issues/195) | [pydantic-deep](https://github.com/vstorm-co/pydantic-deepagents) (vstorm&#8209;co) |
| **Safety &&nbsp;guardrails** | **Input guardrails** | Validate user input before the agent run starts | :construction: [PR&nbsp;#182](https://github.com/pydantic/pydantic-ai-harness/pull/182) | [pydantic-ai-shields](https://github.com/vstorm-co/pydantic-ai-shields) (vstorm&#8209;co) |
| | **Output guardrails** | Validate model output after the run completes | :construction: [PR&nbsp;#182](https://github.com/pydantic/pydantic-ai-harness/pull/182) | [pydantic-ai-shields](https://github.com/vstorm-co/pydantic-ai-shields) (vstorm&#8209;co) |
| **Safety &&nbsp;guardrails** | **Input guardrails** | Validate user input before the agent run starts | :white_check_mark: [Docs](pydantic_ai_harness/guardrails/) | [pydantic-ai-shields](https://github.com/vstorm-co/pydantic-ai-shields) (vstorm&#8209;co) |
| | **Output guardrails** | Validate model output after the run completes | :white_check_mark: [Docs](pydantic_ai_harness/guardrails/) | [pydantic-ai-shields](https://github.com/vstorm-co/pydantic-ai-shields) (vstorm&#8209;co) |
| | **Cost/token budgets** | Enforce token and cost limits per run | :construction: [PR&nbsp;#182](https://github.com/pydantic/pydantic-ai-harness/pull/182) | [pydantic-ai-shields](https://github.com/vstorm-co/pydantic-ai-shields) (vstorm&#8209;co) |
| | **Tool access control** | Block tools or require approval before execution | :construction: [PR&nbsp;#182](https://github.com/pydantic/pydantic-ai-harness/pull/182) | [pydantic-ai-shields](https://github.com/vstorm-co/pydantic-ai-shields) (vstorm&#8209;co) |
| | **Async guardrails** | Run validation concurrently with model requests | :construction: [PR&nbsp;#182](https://github.com/pydantic/pydantic-ai-harness/pull/182) | [pydantic-ai-shields](https://github.com/vstorm-co/pydantic-ai-shields) (vstorm&#8209;co) |
agent_docs/capability-authoring.md
+26 -12
View File
@@ -31,21 +31,35 @@ Each capability package should normally have:
The root `pydantic_ai_harness/__init__.py` should re-export stable public
capabilities. Keep implementation helpers private unless users need them.
### Experimental Vs Released Exports
### Capability Submodules And Exports
New capabilities land under `pydantic_ai_harness/experimental/`. Each
experimental package calls `warn_experimental('<name>')` in its `__init__.py` so
importing it emits `HarnessExperimentalWarning`. Anything under `experimental`
may change or be removed in any release, without a deprecation period.
The `experimental` tier is retired. ACP is the sole remaining experimental
capability (`pydantic_ai_harness/experimental/acp/`); do not add new capabilities
there.
Promote a capability to a top-level package and a top-level re-export in
`pydantic_ai_harness/__init__.py` only when its API is stable.
New capabilities land as a top-level submodule `pydantic_ai_harness/<name>/`.
They are not re-exported from the root `pydantic_ai_harness/__init__.py`: each
capability keeps its own optional dependencies, so importing the root package
must not pull in a capability's extras. Users import a capability from its
submodule (`from pydantic_ai_harness.<name> import ...`).
Top-level exports in `pydantic_ai_harness/__init__.py` are the intended public
surface. Once an export has shipped in a published release it is a
backward-compatibility commitment: do not move, rename, or break it. (`CodeMode`
is a shipped, released example.) Do not relocate an already-top-level capability
into `experimental`.
Naming: the module name is the capability name, one module per capability or
strategy. Prefer a longer descriptive name over a terse one (e.g.
`overflowing_tool_output`, not `overflow`). A known term is fine as-is (e.g.
`compaction`). If you are unsure what to name a capability, ask the user (via the
ask-user tool) rather than guessing -- a name is a public commitment once shipped.
When a capability's module path changes, keep the old path working as a
`DeprecationWarning` shim so existing imports do not break.
Top-level re-exports in `pydantic_ai_harness/__init__.py` (`CodeMode`,
`FileSystem`, `Shell`, `ManagedPrompt`) are the exception, not the rule. Once an
export has shipped in a published release it is a backward-compatibility
commitment: do not move, rename, or break it. Do not add new top-level
re-exports.
APIs are subject to change between releases; breaking changes ship deprecation
warnings where practical.
## API Design
pydantic_ai_harness/.agents/skills/pydantic-ai-harness/SKILL.md
+35 -9
View File
@@ -1,8 +1,8 @@
---
name: pydantic-ai-harness
description: Extend Pydantic AI agents with batteries-included capabilities from pydantic-ai-harness -- currently Code Mode, which collapses many tool calls into one sandboxed Python execution. Use when the user mentions pydantic-ai-harness, CodeMode, Monty, code mode, or tool sandboxing, when they want an agent to run agent-written Python, or when a Pydantic AI agent would benefit from orchestrating multiple tool calls in a single sandboxed script.
description: Extend Pydantic AI agents with batteries-included capabilities from pydantic-ai-harness -- Code Mode (collapse many tool calls into one sandboxed Python execution), a filesystem and shell, sub-agents, planning, context compaction, and more. Use when the user mentions pydantic-ai-harness, CodeMode, Monty, code mode, or tool sandboxing, when they want first-party filesystem/shell/sub-agent/planning/compaction capabilities for a Pydantic AI agent, when they want an agent to run agent-written Python, or when a Pydantic AI agent would benefit from orchestrating multiple tool calls in a single sandboxed script.
license: MIT
compatibility: Requires Python 3.10+ and pydantic-ai-slim>=1.95.1
compatibility: Requires Python 3.10+ and pydantic-ai-slim>=2.1.0
metadata:
version: "0.1.0"
author: pydantic
@@ -33,13 +33,39 @@ Do **not** use this skill for:
## Supported Capabilities
| Capability | Description | Reference |
|---|---|---|
| `CodeMode` | Wraps eligible tools into a single sandboxed `run_code` tool so the model orchestrates them in Python | [Code Mode](./references/CODE-MODE.md) |
`CodeMode` has a full reference below; it is the flagship capability and the one this skill goes deep on.
The rest ship today and each has its own README with API and examples.
More capability areas are tracked in the
[capability matrix](https://github.com/pydantic/pydantic-ai-harness#capability-matrix); as they stabilize,
this skill grows to cover them.
Each capability lives in its own submodule and is imported from there
(`from pydantic_ai_harness.<module> import ...`). Capabilities are not importable from the top-level
`pydantic_ai_harness` package by design, so each one keeps its own optional dependencies isolated.
`CodeMode`, `FileSystem`, `Shell`, and `ManagedPrompt` also have top-level re-exports (importable directly
from `pydantic_ai_harness`).
APIs are subject to change between releases; breaking changes ship deprecation warnings where practical.
| Capability | Module | Description |
|---|---|---|
| `CodeMode` | `pydantic_ai_harness.code_mode` (also top-level) | Wraps eligible tools into a single sandboxed `run_code` tool so the model orchestrates them in Python -- see [Code Mode](./references/CODE-MODE.md) |
| `FileSystem` | `pydantic_ai_harness.filesystem` (also top-level) | Read, write, edit, and search files under a root directory, with traversal prevention |
| `Shell` | `pydantic_ai_harness.shell` (also top-level) | Run commands in a subprocess with allowlists, a default denylist, timeouts, and env masking |
| `ManagedPrompt` | `pydantic_ai_harness.logfire` (also top-level) | Back an agent's instructions with a Logfire-managed prompt |
| `SubAgents` | `pydantic_ai_harness.subagents` | Delegate subtasks to specialized child agents |
| `DynamicWorkflow` | `pydantic_ai_harness.dynamic_workflow` | Orchestrate sub-agents from a model-written Python script |
| `Planning` | `pydantic_ai_harness.planning` | Break complex tasks into structured plans before execution |
| compaction family (`SlidingWindow`, `SummarizingCompaction`, ...) | `pydantic_ai_harness.compaction` | Trim or summarize conversation history to stay within token limits |
| `OverflowingToolOutput` | `pydantic_ai_harness.overflowing_tool_output` | Truncate, summarize, or spill large tool outputs |
| `RepoContext` | `pydantic_ai_harness.context` | Auto-load CLAUDE.md/AGENTS.md and repo structure |
| `StepPersistence` | `pydantic_ai_harness.step_persistence` | Save, restore, resume, and fork run state |
| `PyaiDocs` | `pydantic_ai_harness.docs` | On-demand `read_pyai_docs` tool for Pydantic AI docs |
| `RuntimeAuthoring` | `pydantic_ai_harness.runtime_authoring` | Let an agent author, validate, and load real capabilities at runtime |
| media externalization | `pydantic_ai_harness.media` | Offload large `BinaryContent` to content-addressed stores |
Still experimental: an ACP server adapter, imported from `pydantic_ai_harness.experimental.acp`. Importing it
emits a `HarnessExperimentalWarning`.
The full, current list with links and status is in the
[capability matrix](https://github.com/pydantic/pydantic-ai-harness#capability-matrix).
## Install
@@ -53,7 +79,7 @@ Each capability declares its own extra. Code Mode needs the Monty sandbox:
uv add "pydantic-ai-harness[codemode]" # `code-mode` is also accepted as an alias
```
Requires Python 3.10+ and `pydantic-ai-slim>=1.95.1`.
Requires Python 3.10+ and `pydantic-ai-slim>=2.1.0`.
## Quick Start
pydantic_ai_harness/__init__.py
+33
View File
@@ -5,23 +5,56 @@ from typing import TYPE_CHECKING
if TYPE_CHECKING:
from .code_mode import CodeMode
from .filesystem import FileSystem
from .guardrails import (
GuardrailError,
GuardResult,
InputBlocked,
InputGuard,
InputGuardFunc,
OutputBlocked,
OutputGuard,
OutputGuardFunc,
)
from .logfire import ManagedPrompt
from .shell import LLM_API_KEY_ENV_PATTERNS, Shell
__all__ = [
'CodeMode',
'FileSystem',
'GuardResult',
'GuardrailError',
'InputBlocked',
'InputGuard',
'InputGuardFunc',
'LLM_API_KEY_ENV_PATTERNS',
'ManagedPrompt',
'OutputBlocked',
'OutputGuard',
'OutputGuardFunc',
'Shell',
]
_GUARDRAIL_EXPORTS = {
'GuardResult',
'GuardrailError',
'InputBlocked',
'InputGuard',
'InputGuardFunc',
'OutputBlocked',
'OutputGuard',
'OutputGuardFunc',
}
def __getattr__(name: str) -> object:
if name == 'CodeMode':
from .code_mode import CodeMode
return CodeMode
if name in _GUARDRAIL_EXPORTS:
from . import guardrails
return getattr(guardrails, name)
if name == 'FileSystem':
from .filesystem import FileSystem
pydantic_ai_harness/_monty_exec.py
+224
View File
@@ -0,0 +1,224 @@
"""Shared Monty execution loop for code-execution capabilities.
Drives a Monty REPL via the synchronous snapshot API (`feed_start`/`resume`),
dispatching external function calls back to a host-supplied async callback.
Two capabilities build on this:
- `code_mode`: the dispatch callback runs the agent's own tools.
- `dynamic_workflow`: the dispatch callback runs sub-agents.
The synchronous snapshot API (rather than `feed_run_async`) is used deliberately:
it avoids background threads and `call_soon_threadsafe`, so the loop is safe inside
restricted event loops such as Temporal's workflow sandbox.
"""
from __future__ import annotations
import asyncio
from collections.abc import Callable, Container, Coroutine
from dataclasses import dataclass, field
from typing import Any
try:
from pydantic_monty import (
AbstractOS,
ExternalException,
ExternalResult,
ExternalReturnValue,
FunctionSnapshot,
FutureSnapshot,
MontyComplete,
MountDir,
NameLookupSnapshot,
OsFunction,
)
except ImportError as _import_error: # pragma: no cover
raise ImportError(
'pydantic-monty is required for code-execution capabilities. Install it with: '
'pip install "pydantic-ai-harness[code-mode]" or "pydantic-ai-harness[dynamic-workflow]"'
) from _import_error
# Dispatch callback: given the sandbox function name and keyword arguments,
# perform the host-side work (tool call or sub-agent run) and return the result.
DispatchFn = Callable[[str, dict[str, Any]], Coroutine[Any, Any, Any]]
MontyState = FunctionSnapshot | FutureSnapshot | NameLookupSnapshot | MontyComplete
# OS handler and host-directory mounts that route environment, clock, and filesystem
# calls from inside the sandbox. A raw callback or a ready-made `AbstractOS` is accepted.
MontyOSCallback = Callable[[OsFunction, tuple[object, ...], dict[str, object]], object]
MontyOS = AbstractOS | MontyOSCallback
MontyMount = MountDir | list[MountDir]
# A coroutine not yet scheduled on the event loop, or its running Task.
PendingCall = asyncio.Task[Any] | Coroutine[Any, Any, Any]
def is_sandbox_panic(exc: BaseException) -> bool:
"""Whether `exc` is a Rust-side sandbox panic surfacing through pyo3.
pyo3 raises `pyo3_runtime.PanicException`, a `BaseException` (not `Exception`) subclass
from a module that cannot be imported, so it is matched by name. The model can provoke a
panic from inside the sandbox (e.g. awaiting the same external call twice in one
`asyncio.gather`), so callers should convert it to a retry rather than let it tear down
the whole agent run.
"""
return type(exc).__name__ == 'PanicException'
class PrintCapture:
"""Accumulates print-callback chunks from a Monty REPL."""
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)
def prepend_to(self, error_message: str) -> str:
"""Prefix captured stdout to an error message, so the model sees what printed before the error."""
printed = self.joined.rstrip('\n')
if not printed:
return error_message
return f'[stdout before error]\n{printed}\n[/stdout before error]\n{error_message}'
@dataclass
class MontyExecutor:
"""Drives a Monty REPL to completion, dispatching external calls to a host callback.
Single-use: it accumulates per-run state in `_pending`/`_pre_resolved`, so construct a
fresh executor for each `run` rather than reusing or sharing one across concurrent runs.
External calls are handled by execution mode:
- **Parallel** (`async def`): deferred via `resume({'future': ...})` and eagerly
scheduled as `asyncio.Task`s. Resolved at `FutureSnapshot` via `asyncio.gather`.
- **Per-call sequential** (`def`, name in `sequential_names`): resolved inline at
`FunctionSnapshot`. Any pending parallel tasks are awaited first (barrier).
- **Global sequential** (DBOS/Temporal): all calls deferred but stored as bare
coroutines and awaited one-at-a-time to prevent interleaving.
"""
dispatch: DispatchFn
valid_names: Container[str]
sequential_names: set[str] = field(default_factory=set[str])
global_sequential: bool = False
# OS handler and mounts. Monty auto-dispatches OS calls only while every `resume`
# carries them, so they are threaded through each resume below, not just `feed_start`.
os_access: MontyOS | None = None
mount: MontyMount | None = None
# Parallel calls deferred but not yet resolved, keyed by Monty call id.
_pending: dict[int, PendingCall] = field(default_factory=dict[int, PendingCall], init=False)
# Parallel results awaited early at a sequential barrier, before their FutureSnapshot is reached.
_pre_resolved: dict[int, ExternalResult] = field(default_factory=dict[int, ExternalResult], init=False)
async def run(self, state: MontyState) -> MontyComplete:
"""Drive the REPL from `state` until it completes."""
try:
while not isinstance(state, MontyComplete):
if isinstance(state, NameLookupSnapshot):
state = state.resume(os=self.os_access, mount=self.mount)
elif isinstance(state, FunctionSnapshot):
state = await self._handle_function(state)
else:
state = await self._resolve_futures(state)
finally:
cancelled: list[asyncio.Task[Any]] = []
for call in self._pending.values():
if isinstance(call, asyncio.Task):
call.cancel()
cancelled.append(call)
else:
call.close()
if cancelled:
# `cancel()` only schedules a `CancelledError` at each task's next suspension
# point; await them so dispatched work (e.g. sub-agent runs mutating shared
# usage) has fully unwound before this returns. `return_exceptions=True` keeps
# one task's teardown error from masking the original exception, and the
# results are deliberately discarded.
await asyncio.gather(*cancelled, return_exceptions=True)
return state
async def _handle_function(self, snapshot: FunctionSnapshot) -> MontyState:
"""Dispatch (or defer) a single external function call."""
name = snapshot.function_name
if name not in self.valid_names:
return snapshot.resume(
{'exception': NameError(f'Unknown function: {name}')}, os=self.os_access, mount=self.mount
)
if snapshot.args:
return snapshot.resume(
{'exception': TypeError(f'{name}() does not accept positional arguments; use keyword arguments')},
os=self.os_access,
mount=self.mount,
)
if name in self.sequential_names:
# Rendered as `def` (sync), so the sandbox code doesn't `await` the result --
# resolve inline. Await pending parallel tasks first (barrier) for ordering.
# The dispatch coroutine is created only after the barrier: it is not in
# `_pending`, so if it existed while the barrier awaits and we were cancelled
# there, `run`'s cleanup would never close it.
for cid in list(self._pending):
self._pre_resolved[cid] = await _await_external(self._pending.pop(cid))
# The wrapped outcome (`{'return_value': ...}` / `{'exception': ...}`) is already
# exactly the payload `resume` expects.
return snapshot.resume(
await _await_external(self.dispatch(name, snapshot.kwargs)), os=self.os_access, mount=self.mount
)
# Deferred execution -- resolved later at FutureSnapshot.
call = self.dispatch(name, snapshot.kwargs)
if self.global_sequential:
# Keep the bare coroutine unscheduled; it's awaited one-at-a-time to avoid interleaving.
self._pending[snapshot.call_id] = call
else:
# Schedule now as a Task so concurrently-deferred calls actually run in parallel.
self._pending[snapshot.call_id] = asyncio.ensure_future(call)
return snapshot.resume({'future': ...}, os=self.os_access, mount=self.mount)
async def _resolve_futures(self, snapshot: FutureSnapshot) -> MontyState:
"""Resolve the deferred calls a `FutureSnapshot` is waiting on."""
pending_ids = snapshot.pending_call_ids
results: dict[int, ExternalResult] = {}
for cid in pending_ids:
if cid in self._pre_resolved:
results[cid] = self._pre_resolved.pop(cid)
elif self.global_sequential:
results[cid] = await _await_external(self._pending.pop(cid))
# Gather any remaining parallel tasks concurrently. They stay in `_pending` until
# gather returns, so the cleanup in `run` can still cancel them if this is cancelled.
gather_ids = [cid for cid in pending_ids if cid not in results]
if gather_ids:
settled = await asyncio.gather(*(self._pending[cid] for cid in gather_ids), return_exceptions=True)
for cid, outcome in zip(gather_ids, settled):
del self._pending[cid]
results[cid] = _wrap_gathered(outcome)
return snapshot.resume(results=results, os=self.os_access, mount=self.mount)
async def _await_external(call: PendingCall) -> ExternalReturnValue | ExternalException:
"""Await a single deferred call and wrap its outcome for Monty."""
try:
result = await call
except Exception as exc:
return ExternalException(exception=exc)
return ExternalReturnValue(return_value=result)
def _wrap_gathered(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)
pydantic_ai_harness/code_mode/_toolset.py
+41 -228
View File
@@ -2,12 +2,11 @@
from __future__ import annotations
import asyncio
import inspect
import keyword
import re
import warnings
from collections.abc import Callable, Coroutine, Sequence
from collections.abc import Callable, Sequence
from dataclasses import dataclass, field, replace
from typing import Annotated, Any
@@ -35,19 +34,12 @@ except ImportError: # pragma: no cover
try:
from pydantic_monty import (
AbstractOS,
ExternalException,
ExternalResult,
ExternalReturnValue,
FunctionSnapshot,
FutureSnapshot,
Monty,
MontyComplete,
MontyRepl,
MontyRuntimeError,
MontySyntaxError,
MontyTypingError,
MountDir,
NameLookupSnapshot,
OsFunction,
)
except ImportError as _import_error: # pragma: no cover
@@ -56,8 +48,7 @@ except ImportError as _import_error: # pragma: no cover
) 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]]
from pydantic_ai_harness._monty_exec import MontyExecutor, PrintCapture, is_sandbox_panic
# A raw OS callback. Return `pydantic_monty.NOT_HANDLED` to defer the call to the
# sandbox's default, which leaves it unavailable.
@@ -183,6 +174,18 @@ _TOOL_SEARCH_ADDENDUM = (
_INVALID_IDENT_CHARS = re.compile(r'[^a-zA-Z0-9_]')
def _is_code_execution_tool(tool_def: ToolDefinition) -> bool:
"""Whether a tool is itself a code-execution sandbox that takes a code string.
Such tools (this `run_code`, or DynamicWorkflow's `run_workflow`) carry `code_arg_name`
metadata -- the same marker instrumentation reads to render the argument as code. They must
not be folded into `run_code`: nesting one code sandbox inside another would make the model
write a script that passes a second script as a string literal. They stay native so the two
code surfaces sit side by side.
"""
return bool(tool_def.metadata and 'code_arg_name' in tool_def.metadata)
def _sanitize_tool_name(name: str) -> str:
"""Turn a tool name into a valid Python identifier.
@@ -352,6 +355,10 @@ class CodeModeToolset(WrapperToolset[AgentDepsT]):
# Keep the local fallback native so `Model.prepare_request` can drop it
# when the provider supports the native tool.
native_tools[name] = tool
elif _is_code_execution_tool(tool.tool_def):
# A tool that is itself a code-execution sandbox (e.g. DynamicWorkflow's
# `run_workflow`) is a peer of `run_code`, not something to fold inside it.
native_tools[name] = tool
elif await matches_tool_selector(self.tool_selector, ctx, tool.tool_def):
sandboxed_tools[name] = tool
else:
@@ -459,7 +466,7 @@ class CodeModeToolset(WrapperToolset[AgentDepsT]):
nested_returns: dict[str, ToolReturnPart] = {}
call_counter = 0
async def dispatch_tool_call(original_name: str, kwargs: dict[str, Any]) -> Any:
async def dispatch_tool_call(sandbox_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
@@ -468,6 +475,7 @@ class CodeModeToolset(WrapperToolset[AgentDepsT]):
`await` site.
"""
nonlocal call_counter
original_name = sanitized_to_original.get(sandbox_name, sandbox_name)
call_counter += 1
parent_id = ctx.tool_call_id or 'pyd_ai_code_mode'
tool_call_id = f'{parent_id}__{call_counter}'
@@ -531,24 +539,22 @@ class CodeModeToolset(WrapperToolset[AgentDepsT]):
self._repl = MontyRepl()
assert self._repl is not None
capture = _PrintCapture()
capture = PrintCapture()
try:
monty_state = self._repl.feed_start(code, print_callback=capture, os=self.os_access, mount=self.mount)
completed = await _execution_loop(
monty_state,
completed = await MontyExecutor(
dispatch=dispatch_tool_call,
callable_defs=callable_defs,
sanitized_to_original=sanitized_to_original,
sequential_tools=sequential_tools,
valid_names=callable_defs,
sequential_names=sequential_tools,
global_sequential=global_sequential,
os_access=self.os_access,
mount=self.mount,
)
).run(monty_state)
except MontySyntaxError as e:
raise ModelRetry(f'Syntax error in code:\n{_prepend_prints(e.display(), capture)}') from e
raise ModelRetry(f'Syntax error in code:\n{capture.prepend_to(e.display())}') 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
raise ModelRetry(f'Type error in code:\n{capture.prepend_to(e.display())}') 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
@@ -559,7 +565,20 @@ class CodeModeToolset(WrapperToolset[AgentDepsT]):
# 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
raise ModelRetry(f'Runtime error:\n{capture.prepend_to(e.display())}') from e
except BaseException as e:
# Convert a model-provokable sandbox panic to a retry (see `is_sandbox_panic`);
# anything else (CancelledError, ...) re-raises unchanged.
if not is_sandbox_panic(e):
raise
# The panic aborts the VM mid-execution, so the REPL's accumulated state cannot
# be trusted; drop it so the retry starts from a fresh, type-checked session.
self._repl = None
raise ModelRetry(
'The code aborted inside the sandbox and the session was reset. This can happen '
'when the same tool call is awaited more than once in one asyncio.gather -- give '
'each gathered call its own invocation. Revise the code and try again.'
) from e
result = completed.output
printed = capture.joined
@@ -721,194 +740,6 @@ def _get_sigs_and_conflicting(
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,
os_access: CodeModeOS | None,
mount: CodeModeMount | None,
) -> 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.
`os`/`mount` must be passed to every `resume` call (not just `feed_start`):
Monty's auto-dispatch of OS calls stops the moment a resume omits them.
"""
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(os=os_access, mount=mount)
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,
os_access=os_access,
mount=mount,
)
else:
monty_state = await _resolve_future_snapshot(
monty_state,
pending=pending,
pre_resolved=pre_resolved,
global_sequential=global_sequential,
os_access=os_access,
mount=mount,
)
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],
os_access: CodeModeOS | None,
mount: CodeModeMount | None,
) -> 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}')}, os=os_access, mount=mount)
if snapshot.args:
return snapshot.resume(
{'exception': TypeError(f'{fn_name}() does not accept positional arguments; use keyword arguments')},
os=os_access,
mount=mount,
)
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']}, os=os_access, mount=mount)
return snapshot.resume({'exception': outcome['exception']}, os=os_access, mount=mount)
# 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': ...}, os=os_access, mount=mount)
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,
os_access: CodeModeOS | None,
mount: CodeModeMount | None,
) -> 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={}, os=os_access, mount=mount)
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, os=os_access, mount=mount)
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):
@@ -916,21 +747,3 @@ def _contains_multimodal(value: Any) -> bool:
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)
pydantic_ai_harness/experimental/compaction/README.md → pydantic_ai_harness/compaction/README.md
+7 -18
View File
@@ -1,23 +1,13 @@
# Compaction capabilities
> [!WARNING]
> **Experimental.** These capabilities live under `pydantic_ai_harness.experimental` and may
> change or be removed in any release, without a deprecation period. Import them from the
> experimental path -- there is no top-level export:
> [!NOTE]
> Import these capabilities from their submodule -- there is no top-level `pydantic_ai_harness` re-export:
>
> ```python
> from pydantic_ai_harness.experimental.compaction import TieredCompaction
> from pydantic_ai_harness.compaction import TieredCompaction
> ```
>
> Importing any experimental capability emits a `HarnessExperimentalWarning`. Silence **all**
> harness experimental warnings with a single filter (no per-capability lines needed):
>
> ```python
> import warnings
> from pydantic_ai_harness.experimental import HarnessExperimentalWarning
>
> warnings.filterwarnings('ignore', category=HarnessExperimentalWarning)
> ```
> The API may change between releases. Where practical, breaking changes ship with a deprecation warning.
A menu of strategies for keeping an agent's conversation history within a model's context
window. Each is a Pydantic AI `Capability` that runs in the `before_model_request` hook; edits
@@ -62,7 +52,7 @@ low-entropy repetition, so a head/tail slice loses little.
```python
from pydantic_ai import Agent
from pydantic_ai_harness.experimental.compaction import ClampOversizedMessages
from pydantic_ai_harness.compaction import ClampOversizedMessages
agent = Agent(
'openai:gpt-4o',
@@ -89,7 +79,7 @@ user input should not be silently rewritten, and oversized tool returns are the
Use it as the first tier of `TieredCompaction`, before `ClearToolResults`:
```python
from pydantic_ai_harness.experimental.compaction import (
from pydantic_ai_harness.compaction import (
ClampOversizedMessages,
ClearToolResults,
TieredCompaction,
@@ -113,7 +103,7 @@ that is not enough -- which is exactly what `TieredCompaction` encodes:
```python
from pydantic_ai import Agent
from pydantic_ai_harness.experimental.compaction import (
from pydantic_ai_harness.compaction import (
ClearToolResults,
DeduplicateFileReads,
SummarizingCompaction,
@@ -168,7 +158,6 @@ call is not a file read:
```python
from pydantic_ai.messages import ToolCallPart
def my_file_key(call: ToolCallPart) -> str | None:
if call.tool_name != 'read_file':
return None
pydantic_ai_harness/compaction/__init__.py
+27
View File
@@ -0,0 +1,27 @@
"""Compaction capabilities: keep an agent's conversation history within the context window.
Each capability lives in its own module; shared utilities (token estimation, the
`CompactionStrategy` protocol, tool-pair-safe cutoffs, in-place clearing) live in `_shared`.
"""
from pydantic_ai_harness.compaction._clamp_oversized_messages import ClampOversizedMessages
from pydantic_ai_harness.compaction._clear_tool_results import ClearToolResults
from pydantic_ai_harness.compaction._deduplicate_file_reads import DeduplicateFileReads
from pydantic_ai_harness.compaction._limit_warner import LimitWarner, WarningKind
from pydantic_ai_harness.compaction._shared import CompactionStrategy, estimate_token_count
from pydantic_ai_harness.compaction._sliding_window import SlidingWindow
from pydantic_ai_harness.compaction._summarizing_compaction import SummarizingCompaction
from pydantic_ai_harness.compaction._tiered_compaction import TieredCompaction
__all__ = [
'ClampOversizedMessages',
'ClearToolResults',
'CompactionStrategy',
'DeduplicateFileReads',
'LimitWarner',
'SlidingWindow',
'SummarizingCompaction',
'TieredCompaction',
'WarningKind',
'estimate_token_count',
]
pydantic_ai_harness/experimental/compaction/_clamp_oversized_messages.py → pydantic_ai_harness/compaction/_clamp_oversized_messages.py
+2 -3
View File
@@ -17,12 +17,11 @@ from pydantic_ai.messages import (
)
from pydantic_ai.tools import RunContext
from pydantic_ai_harness.experimental.compaction._shared import compact_with_span, estimate_text_tokens
from pydantic_ai_harness.compaction._shared import compact_with_span, estimate_text_tokens
if TYPE_CHECKING:
from pydantic_ai.models import ModelRequestContext
_CLAMP_MARKER = '\n[clamped: removed {removed} of {original} characters]\n'
"""Inserted between the head and tail slices of a clamped part. ``{removed}`` and ``{original}``
are filled with character counts."""
@@ -71,7 +70,7 @@ class ClampOversizedMessages(AbstractCapability[AgentDepsT]):
Example:
```python
from pydantic_ai import Agent
from pydantic_ai_harness.experimental.compaction import ClampOversizedMessages
from pydantic_ai_harness.compaction import ClampOversizedMessages
agent = Agent(
'openai:gpt-4o',
pydantic_ai_harness/experimental/compaction/_clear_tool_results.py → pydantic_ai_harness/compaction/_clear_tool_results.py
+2 -2
View File
@@ -11,7 +11,7 @@ from pydantic_ai.capabilities import AbstractCapability
from pydantic_ai.messages import ModelMessage
from pydantic_ai.tools import RunContext
from pydantic_ai_harness.experimental.compaction._shared import (
from pydantic_ai_harness.compaction._shared import (
compact_with_span,
estimate_token_count,
exceeds,
@@ -43,7 +43,7 @@ class ClearToolResults(AbstractCapability[AgentDepsT]):
Example:
```python
from pydantic_ai import Agent
from pydantic_ai_harness.experimental.compaction import ClearToolResults
from pydantic_ai_harness.compaction import ClearToolResults
agent = Agent(
'openai:gpt-4o',
pydantic_ai_harness/experimental/compaction/_deduplicate_file_reads.py → pydantic_ai_harness/compaction/_deduplicate_file_reads.py
+2 -4
View File
@@ -11,7 +11,7 @@ from pydantic_ai.capabilities import AbstractCapability
from pydantic_ai.messages import ModelMessage, ToolCallPart
from pydantic_ai.tools import RunContext
from pydantic_ai_harness.experimental.compaction._shared import (
from pydantic_ai_harness.compaction._shared import (
compact_with_span,
exceeds,
iter_tool_pairs,
@@ -39,8 +39,7 @@ class DeduplicateFileReads(AbstractCapability[AgentDepsT]):
```python
from pydantic_ai import Agent
from pydantic_ai.messages import ToolCallPart
from pydantic_ai_harness.experimental.compaction import DeduplicateFileReads
from pydantic_ai_harness.compaction import DeduplicateFileReads
def file_key(call: ToolCallPart) -> str | None:
if call.tool_name != 'read_file':
@@ -48,7 +47,6 @@ class DeduplicateFileReads(AbstractCapability[AgentDepsT]):
args = call.args_as_dict()
return args.get('path')
agent = Agent('openai:gpt-4o', capabilities=[DeduplicateFileReads(file_key=file_key)])
```
"""
pydantic_ai_harness/experimental/compaction/_limit_warner.py → pydantic_ai_harness/compaction/_limit_warner.py
+2 -2
View File
@@ -10,7 +10,7 @@ from pydantic_ai.capabilities import AbstractCapability
from pydantic_ai.messages import ModelMessage, ModelRequest, SystemPromptPart, UserPromptPart
from pydantic_ai.tools import RunContext
from pydantic_ai_harness.experimental.compaction._shared import estimate_token_count
from pydantic_ai_harness.compaction._shared import estimate_token_count
if TYPE_CHECKING:
from pydantic_ai.models import ModelRequestContext
@@ -43,7 +43,7 @@ class LimitWarner(AbstractCapability[AgentDepsT]):
Example:
```python
from pydantic_ai import Agent
from pydantic_ai_harness.experimental.compaction import LimitWarner
from pydantic_ai_harness.compaction import LimitWarner
agent = Agent(
'openai:gpt-4o',
pydantic_ai_harness/experimental/compaction/_shared.py → pydantic_ai_harness/compaction/_shared.py
-1
View File
@@ -331,7 +331,6 @@ def prepend_first_user_message(
# Tool-pair inspection and in-place clearing
# ---------------------------------------------------------------------------
_CLEARED_TOOL_ARGS = '{}'
"""Replacement for cleared tool-call arguments.
pydantic_ai_harness/experimental/compaction/_sliding_window.py → pydantic_ai_harness/compaction/_sliding_window.py
+2 -2
View File
@@ -11,7 +11,7 @@ from pydantic_ai.capabilities import AbstractCapability
from pydantic_ai.messages import ModelMessage
from pydantic_ai.tools import RunContext
from pydantic_ai_harness.experimental.compaction._shared import (
from pydantic_ai_harness.compaction._shared import (
compact_with_span,
exceeds,
find_safe_cutoff,
@@ -37,7 +37,7 @@ class SlidingWindow(AbstractCapability[AgentDepsT]):
Example:
```python
from pydantic_ai import Agent
from pydantic_ai_harness.experimental.compaction import SlidingWindow
from pydantic_ai_harness.compaction import SlidingWindow
agent = Agent(
'openai:gpt-4o',
pydantic_ai_harness/experimental/compaction/_summarizing_compaction.py → pydantic_ai_harness/compaction/_summarizing_compaction.py
+2 -2
View File
@@ -20,7 +20,7 @@ from pydantic_ai.messages import (
)
from pydantic_ai.tools import RunContext
from pydantic_ai_harness.experimental.compaction._shared import (
from pydantic_ai_harness.compaction._shared import (
compact_with_span,
exceeds,
find_first_user_message,
@@ -151,7 +151,7 @@ class SummarizingCompaction(AbstractCapability[AgentDepsT]):
Example:
```python
from pydantic_ai import Agent
from pydantic_ai_harness.experimental.compaction import SummarizingCompaction
from pydantic_ai_harness.compaction import SummarizingCompaction
agent = Agent(
'openai:gpt-4o',
pydantic_ai_harness/experimental/compaction/_tiered_compaction.py → pydantic_ai_harness/compaction/_tiered_compaction.py
+2 -2
View File
@@ -11,7 +11,7 @@ from pydantic_ai.capabilities import AbstractCapability
from pydantic_ai.messages import ModelMessage
from pydantic_ai.tools import RunContext
from pydantic_ai_harness.experimental.compaction._shared import (
from pydantic_ai_harness.compaction._shared import (
CompactionStrategy,
compact_with_span,
estimate_token_count,
@@ -36,7 +36,7 @@ class TieredCompaction(AbstractCapability[AgentDepsT]):
Example:
```python
from pydantic_ai import Agent
from pydantic_ai_harness.experimental.compaction import (
from pydantic_ai_harness.compaction import (
ClearToolResults,
SummarizingCompaction,
TieredCompaction,
pydantic_ai_harness/experimental/context/README.md → pydantic_ai_harness/context/README.md
+5 -15
View File
@@ -1,23 +1,13 @@
# RepoContext
> [!WARNING]
> **Experimental.** This capability lives under `pydantic_ai_harness.experimental` and may
> change or be removed in any release, without a deprecation period. Import it from the
> experimental path -- there is no top-level export:
> [!NOTE]
> Import this capability from its submodule -- there is no top-level `pydantic_ai_harness` re-export:
>
> ```python
> from pydantic_ai_harness.experimental.context import RepoContext
> from pydantic_ai_harness.context import RepoContext
> ```
>
> Importing any experimental capability emits a `HarnessExperimentalWarning`. Silence **all**
> harness experimental warnings with a single filter (no per-capability lines needed):
>
> ```python
> import warnings
> from pydantic_ai_harness.experimental import HarnessExperimentalWarning
>
> warnings.filterwarnings('ignore', category=HarnessExperimentalWarning)
> ```
> The API may change between releases. Where practical, breaking changes ship with a deprecation warning.
Discover and load a repo's accumulated coding-assistant context engineering (CE).
@@ -37,7 +27,7 @@ idea the rest of the setup exists, so it can neither honor it nor translate it.
from pathlib import Path
from pydantic_ai import Agent
from pydantic_ai_harness.experimental.context import RepoContext
from pydantic_ai_harness.context import RepoContext
agent = Agent(
'anthropic:claude-sonnet-4-6',
pydantic_ai_harness/context/__init__.py
+8
View File
@@ -0,0 +1,8 @@
"""Context capability: discover and load a repo's accumulated context engineering."""
from pydantic_ai_harness.context._capability import RepoContext
from pydantic_ai_harness.context._inventory import AgentContextInventory, AssetRoot
from pydantic_ai_harness.context._loader import ContextFile
from pydantic_ai_harness.context._toolset import RepoContextToolset
__all__ = ['AgentContextInventory', 'AssetRoot', 'ContextFile', 'RepoContext', 'RepoContextToolset']
pydantic_ai_harness/experimental/context/_capability.py → pydantic_ai_harness/context/_capability.py
+3 -4
View File
@@ -12,19 +12,18 @@ from pydantic_ai.messages import ToolCallPart
from pydantic_ai.tools import AgentDepsT, RunContext, ToolDefinition
from pydantic_ai.toolsets import AgentToolset
from pydantic_ai_harness.experimental.context._loader import (
from pydantic_ai_harness.context._loader import (
ContextFile,
discover_instruction_files,
find_dir_context_file,
render_context_file,
render_context_files,
)
from pydantic_ai_harness.experimental.context._toolset import RepoContextToolset
from pydantic_ai_harness.context._toolset import RepoContextToolset
if TYPE_CHECKING:
from pydantic_ai._instructions import AgentInstructions
_INVENTORY_HINT = (
'Call `{tool_name}` to map where this repo keeps its coding-assistant setup '
'(instruction dirs, skills, sub-agents, and hooks) so you can read and translate it.'
@@ -63,7 +62,7 @@ class RepoContext(AbstractCapability[AgentDepsT]):
from pathlib import Path
from pydantic_ai import Agent
from pydantic_ai_harness.experimental.context import RepoContext
from pydantic_ai_harness.context import RepoContext
agent = Agent(
'anthropic:claude-sonnet-4-6',
pydantic_ai_harness/experimental/context/_inventory.py → pydantic_ai_harness/context/_inventory.py
View File
pydantic_ai_harness/experimental/context/_loader.py → pydantic_ai_harness/context/_loader.py
View File
pydantic_ai_harness/experimental/context/_toolset.py → pydantic_ai_harness/context/_toolset.py
+1 -1
View File
@@ -8,7 +8,7 @@ from pathlib import Path
from pydantic_ai.tools import AgentDepsT
from pydantic_ai.toolsets import FunctionToolset
from pydantic_ai_harness.experimental.context._inventory import AgentContextInventory, scan_assets
from pydantic_ai_harness.context._inventory import AgentContextInventory, scan_assets
class RepoContextToolset(FunctionToolset[AgentDepsT]):
pydantic_ai_harness/experimental/docs/README.md → pydantic_ai_harness/docs/README.md
+5 -15
View File
@@ -1,23 +1,13 @@
# PyaiDocs
> [!WARNING]
> **Experimental.** This capability lives under `pydantic_ai_harness.experimental` and may
> change or be removed in any release, without a deprecation period. Import it from the
> experimental path -- there is no top-level export:
> [!NOTE]
> Import this capability from its submodule -- there is no top-level `pydantic_ai_harness` re-export:
>
> ```python
> from pydantic_ai_harness.experimental.docs import PyaiDocs
> from pydantic_ai_harness.docs import PyaiDocs
> ```
>
> Importing any experimental capability emits a `HarnessExperimentalWarning`. Silence **all**
> harness experimental warnings with a single filter (no per-capability lines needed):
>
> ```python
> import warnings
> from pydantic_ai_harness.experimental import HarnessExperimentalWarning
>
> warnings.filterwarnings('ignore', category=HarnessExperimentalWarning)
> ```
> The API may change between releases. Where practical, breaking changes ship with a deprecation warning.
Give an agent a tool that locates and returns Pydantic AI documentation on demand.
@@ -40,7 +30,7 @@ Topics: `capabilities`, `hooks`, `tools`, `tools-advanced`, `toolsets`, `agent`.
from pathlib import Path
from pydantic_ai import Agent
from pydantic_ai_harness.experimental.docs import PyaiDocs
from pydantic_ai_harness.docs import PyaiDocs
agent = Agent(
'anthropic:claude-sonnet-4-6',
pydantic_ai_harness/docs/__init__.py
+6
View File
@@ -0,0 +1,6 @@
"""Docs capability: an on-demand tool that locates Pydantic AI documentation."""
from pydantic_ai_harness.docs._capability import PyaiDocs
from pydantic_ai_harness.docs._toolset import PyaiDocsToolset, PyaiDocsTopic
__all__ = ['PyaiDocs', 'PyaiDocsToolset', 'PyaiDocsTopic']
pydantic_ai_harness/experimental/docs/_capability.py → pydantic_ai_harness/docs/_capability.py
+2 -2
View File
@@ -11,7 +11,7 @@ from pydantic_ai.capabilities import AbstractCapability
from pydantic_ai.tools import AgentDepsT
from pydantic_ai.toolsets import AgentToolset
from pydantic_ai_harness.experimental.docs._toolset import PyaiDocsToolset, PyaiDocsTopic
from pydantic_ai_harness.docs._toolset import PyaiDocsToolset, PyaiDocsTopic
if TYPE_CHECKING:
from pydantic_ai._instructions import AgentInstructions
@@ -46,7 +46,7 @@ class PyaiDocs(AbstractCapability[AgentDepsT]):
from pathlib import Path
from pydantic_ai import Agent
from pydantic_ai_harness.experimental.docs import PyaiDocs
from pydantic_ai_harness.docs import PyaiDocs
agent = Agent(
'anthropic:claude-sonnet-4-6',
pydantic_ai_harness/experimental/docs/_toolset.py → pydantic_ai_harness/docs/_toolset.py
View File
pydantic_ai_harness/dynamic_workflow/README.md
+535
View File
@@ -0,0 +1,535 @@
# Dynamic Workflow
Let one agent coordinate a whole team of sub-agents by writing a small Python script.
> [!NOTE]
> Import this capability from its submodule -- there is no top-level `pydantic_ai_harness` re-export:
>
> ```python
> from pydantic_ai_harness.dynamic_workflow import DynamicWorkflow
> ```
>
> The API may change between releases. Where practical, breaking changes ship with a deprecation warning.
> The extensions planned in [What is coming](#what-is-coming), structured sub-agent inputs and durable
> workflows, touch the sub-agent call contract itself.
## The idea
Say you have a few specialist agents. One reviews code. One summarizes findings. One writes the
final note. Each one is easy to call on its own. The hard part is the choreography between them. You
want to review three files at once, keep only the reports that found something, summarize those, and
hand the summary to the writer.
The usual way to do this is one tool call per step. The agent calls the reviewer and waits. It reads
the result, calls the reviewer again, and waits again. And so on. Every intermediate result travels
back into the agent's context, and every step that depends on the previous one is a separate model
turn.
`DynamicWorkflow` takes a different route. You hand it a catalog of named sub-agents, and it gives
the model a single tool, `run_workflow`. Inside that tool the model writes ordinary Python. Each of
your sub-agents is an `async` function it can call, loop over, and combine. The script runs to
completion in one tool call, and only its final value comes back to the model.
The choreography moves out of the conversation and into code.
Claude Code ships this same idea as a feature, also called dynamic workflows. Jarred Sumner used it
to port Bun from Zig to Rust: around 750,000 lines of Rust, 99.8% of the existing test suite passing,
and eleven days from first commit to merge. One workflow mapped the right Rust lifetime for every
struct field. The next wrote every file as a behavior-identical port, with hundreds of agents in
parallel and two reviewers on each file. A fix loop then drove the build and tests until both ran
clean. Claude Code runs that at the scale of a whole session. This capability brings the same idea
into your own Pydantic AI agents, inside a single `run_workflow` tool call. You can
[read the Bun story here](https://bun.com/blog/bun-in-rust).
> **Tip**
>
> If you have met [Code Mode](../../code_mode/README.md), this will feel familiar. It is the same
> sandbox and the same idea: write a script instead of many tool calls. The difference is what the
> script gets to call. In Code Mode it calls the agent's own tools. Here it calls whole sub-agents.
## How this relates to SubAgents
The harness has two delegation capabilities. They trade in the same currency, named and isolated
sub-agent runs, but they work at different altitudes:
- [`SubAgents`](../subagents/README.md) exposes one `delegate_task(agent_name, task)` tool. Each
delegation is its own tool call and its own model turn. The parent calls, waits, reads the result
into context, then decides the next step. It is simple to reason about. It is the right fit when
delegations are occasional, or when each result needs the parent's judgment before the next one.
- `DynamicWorkflow` moves the choreography into a script. Fan-out, chaining, voting, and retry loops
all run inside one tool call, and intermediate results never enter the parent's context. It is the
right fit when the coordination between sub-agents is the actual work.
Start with `SubAgents` if you are not sure. A `delegate_task` orchestrator converts to a workflow
catalog without changing the sub-agents themselves.
## Install
The script runs inside the [Monty](https://github.com/pydantic/monty) sandbox, so install the extra:
```bash
uv add "pydantic-ai-harness[dynamic-workflow]"
```
## Your first workflow
Let's build the smallest thing that works. Two sub-agents, one orchestrator.
```python
from pydantic_ai import Agent
from pydantic_ai_harness.dynamic_workflow import DynamicWorkflow
reviewer = Agent('openai:gpt-5', name='reviewer', description='Reviews code for bugs.')
summarizer = Agent('openai:gpt-5', name='summarizer', description='Summarizes findings.')
orchestrator = Agent(
'openai:gpt-5',
capabilities=[DynamicWorkflow(agents=[reviewer, summarizer])],
)
```
That is the whole setup. Here is what each piece does:
1. `reviewer` and `summarizer` are plain agents. There is nothing special about them. They are the
same `Agent` you already know.
2. Their `name` becomes the function name the model calls in the script, so pick names that are
valid Python identifiers.
3. Their `description` tells the model what each one is for. The model reads these to decide how to
wire them together, so write them the way you would document a function.
4. `DynamicWorkflow(agents=[...])` bundles them into one capability and hands the orchestrator a
single `run_workflow` tool.
## What the model does with it
When the orchestrator decides to use the tool, it does not call your sub-agents one at a time. It
writes a script. For a "review these two files and summarize" task, the script it writes looks like
this:
```python
import asyncio
reports = await asyncio.gather(
reviewer(task="Review auth.py for bugs:\n<file contents>"),
reviewer(task="Review parser.py for bugs:\n<file contents>"),
)
await summarizer(task="Summarize these review findings:\n" + "\n\n".join(reports))
```
Here are the parts that matter most, because they are the core of how you use this capability:
- Each sub-agent is an `async` function. You call it with `await`.
- You pass the work as a single keyword argument, `task`. Always by keyword. Write
`reviewer(task="...")`, not `reviewer("...")`.
- `asyncio.gather(...)` runs the two reviews at the same time instead of one after the other.
- The last line's value becomes the result the model sees. The intermediate `reports` list never
leaves the sandbox.
> **Info: what "call a sub-agent" actually means**
>
> Each call is a full `Agent.run`. It has its own model loop, its own message history, its own
> tools, and its own typed output. It is not a lightweight function. It is a real agent doing real
> work. Two things follow from that, and both matter when you write or debug workflows:
>
> - **Calls are isolated.** A sub-agent remembers nothing from an earlier call. Put everything it
> needs into `task`.
> - **Calls cost tokens and take time.** That is why this capability gives you budgets, which we
> get to below.
Because the coordination is ordinary Python, it scales past one-shot fan-out. Take a rule like
"re-dispatch only the files that failed review, with the reviewer's issues attached, for up to two
more rounds." That is a `for` loop over `asyncio.gather`, with the retry task text rebuilt from each
failed review. Without this, the model would run that control flow turn by turn, one round-trip per
step, with every intermediate draft flowing through its context.
## Sub-agents can return structured data
A sub-agent returns whatever its `output_type` produces. By default that is a string. But give a
sub-agent a Pydantic model, and the script receives a `dict`:
```python
from pydantic import BaseModel
class Score(BaseModel):
value: int
reason: str
critic = Agent('openai:gpt-5', name='critic', description='Scores an answer 0-10.', output_type=Score)
```
Inside the script, the model reads the fields by subscript, the way you read a JSON object:
```python
result = await critic(task="Score this answer: ...")
result["value"] # not result.value
```
The catalog the model sees renders each output type as a `TypedDict`, so it knows the fields and
reads them by subscript on its own.
## A complete, runnable example
Now let's put it together into something you can actually run. This orchestrator runs a small
tournament. It drafts three candidate answers in parallel, scores each one, picks the winner, and
refines it. All of that happens in a single tool call.
You need an Anthropic key and the `anthropic` package:
```bash
export ANTHROPIC_API_KEY=sk-...
uv run --with 'pydantic-ai-harness[dynamic-workflow]' --with anthropic --with logfire python wf.py
```
```python
# wf.py
import asyncio
import logfire
from pydantic import BaseModel
from pydantic_ai import Agent
from pydantic_ai.usage import UsageLimits
from pydantic_ai_harness.dynamic_workflow import DynamicWorkflow
# With Logfire configured, the trace shows the orchestrator turn, the run_workflow call (including
# the exact script the model wrote), and every sub-agent run nested underneath it.
logfire.configure(send_to_logfire='if-token-present', service_name='dynamic-workflow')
logfire.instrument_pydantic_ai()
MODEL = 'anthropic:claude-sonnet-4-6' # or 'anthropic:claude-opus-4-8'
class Score(BaseModel):
value: int # 0-10
reason: str
drafter = Agent(
MODEL,
name='drafter',
description='Writes one candidate answer to a task.',
instructions='Write one concise candidate answer to the task.',
)
critic = Agent(
MODEL,
name='critic',
description='Scores a candidate answer 0-10, returns {value, reason}.',
output_type=Score,
instructions='Score the candidate 0-10 with a one-line reason.',
)
editor = Agent(
MODEL,
name='editor',
description='Improves an answer given a critique.',
instructions='Improve the given answer using the critique. Return only the answer.',
)
orchestrator = Agent(
MODEL,
instructions=(
'Use run_workflow to: draft 3 candidate answers in parallel, score each with the critic, '
'pick the highest-scoring one, then have the editor refine it using that critique. '
'Return the refined answer.'
),
capabilities=[DynamicWorkflow(agents=[drafter, critic, editor])],
)
async def main() -> None:
result = await orchestrator.run(
'Explain, for a new hire, why our service uses idempotency keys on payment requests.',
usage_limits=UsageLimits(request_limit=20),
)
logfire.info('done', answer=result.output, requests=result.usage.requests)
asyncio.run(main())
```
Given just those three sub-agents and the instructions, the model writes and runs a script along
these lines:
```python
import asyncio
# 1. Draft three candidates at the same time.
drafts = await asyncio.gather(
drafter(task="explain idempotency keys on payments"),
drafter(task="explain idempotency keys on payments"),
drafter(task="explain idempotency keys on payments"),
)
# 2. Score each one. Structured output arrives as {"value": int, "reason": str}.
scores = await asyncio.gather(*[critic(task="Score this answer:\n" + d) for d in drafts])
# 3. Pick the winner and refine it, all in plain Python, no extra model turns.
best = max(range(len(drafts)), key=lambda i: scores[i]["value"])
await editor(task="Answer:\n" + drafts[best] + "\n\nCritique:\n" + scores[best]["reason"])
```
Read that script and notice what did not happen. The three drafts, the three scores, and the
selection logic never traveled back through the orchestrator's context. The model issued one tool
call and got back one answer. The comparison, the `max(...)`, and the string assembly are ordinary
Python running in the sandbox.
> **Tip**
>
> The [Logfire](https://pydantic.dev/logfire) trace is the best way to see what a workflow did.
> Each sub-agent run appears nested under the `run_workflow` span, and the span carries the exact
> `code` argument the model wrote, so you can read the script it actually ran.
## How results come back
The value of the script's last expression becomes the tool result. The model does not `print()` it.
For the common cases, that is all you need to know. If you want the exact rules, including what
happens when the script also prints for debugging, here they are:
> **Info: the precise return shape**
>
> | Sub-agent `output_type` | Value inside the script |
> | --- | --- |
> | `str` (the default) | the string |
> | a Pydantic model | a `dict`, read as `r['field']` |
> | list or scalar | the list or scalar |
>
> And here is how the final tool result is shaped:
>
> | The script... | The model receives |
> | --- | --- |
> | ends in a value, no print | that value directly (or `{}` if it is `None`) |
> | prints and ends in a value | `{"output": "<printed text>", "result": <value>}` |
> | prints and ends in `None` | `{"output": "<printed text>"}` |
>
> `print()` is for debug logging. It stringifies, so it is the wrong tool for returning structured
> data. Let the last expression carry the real result.
## Choosing sub-agent models
By default, each sub-agent uses the model it was constructed with. Set `inherit_model=True` when the
host passes a per-run model override to the parent agent, for example from a `/model` command, and
every sub-agent dispatch should follow that resolved parent model. Leave it `False` when a sub-agent
is deliberately pinned to a different model.
## Keeping it safe: budgets
A sub-agent is non-deterministic and costs tokens, and it can fan out into more sub-agents. So a
workflow needs two kinds of ceiling. It needs a cap on *how many* sub-agent runs happen, and a cap
on *how much* they spend. `DynamicWorkflow` gives you both, plus a guard against runaway sandbox
scripts.
### `max_agent_calls`: an exact count
```python
DynamicWorkflow(agents=[...], max_agent_calls=50) # 50 is the default
```
This is a hard, host-enforced ceiling on the number of sub-agent runs in one parent run. It is one
budget shared across every `run_workflow` call in that run, not a per-script allowance. It holds
exactly, even when the script fans out with `asyncio.gather`. When the budget runs out, the workflow
stops calling sub-agents and returns a terminal result that tells the model to conclude with what it
has. That result includes the sub-agent results that did complete, so nothing you already paid for
is wasted.
> **Note**
>
> `max_agent_calls` is the only knob that bounds the number of runs exactly. Reach for it when you
> need a guarantee. The token-based limits below are budgets, not guarantees.
### `sub_agent_usage_limits` and `forward_usage`: bounding cost
`sub_agent_usage_limits` is a `UsageLimits` applied to each sub-agent run. How tight a ceiling it
gives depends on `forward_usage`, which controls whether the whole tree shares one usage counter:
| `forward_usage` | Counter | What the limit means |
| --- | --- | --- |
| `True` (default) | the parent's `usage` is shared across the tree | a tree-wide cap, checked against the shared counter. Under concurrent fan-out it is best-effort: several sub-agents can pass the check before any of them adds to the count. |
| `False` | each sub-agent run counts on its own | per-run limits. A per-run `total_tokens_limit` of `T` with `max_agent_calls` of `N` bounds the tree to roughly `N * T` tokens. |
> **Warning**
>
> The `usage_limits` you pass to the parent `run()` is not forwarded into sub-agents. Core does not
> expose that limit value to the capability, so it is re-checked only at the parent's own request
> boundaries. If you want to bound sub-agents, set `sub_agent_usage_limits`. If you want an exact
> ceiling on the number of runs, use `max_agent_calls`.
### `resource_limits`: guarding the script itself
These limits guard the orchestration script's own memory and allocations, not the sub-agents it
calls. The default backstop is 256 MB and 50 million allocations, with no time limit.
```python
DynamicWorkflow(agents=[...], resource_limits={'max_duration_secs': 30})
```
There is no default duration cap. To see why, it helps to know what the timer actually measures.
> **Info: what `max_duration_secs` measures**
>
> Monty checks this limit once per bytecode step. So it measures the time your script spends
> running sandbox code. It does not measure wall-clock time.
>
> This is the part that matters for sub-agents. While your script waits on a sub-agent, it is
> suspended on the host. It is not running. That time does not count. The same is true whether you
> await one sub-agent at a time or fan several out with `asyncio.gather`. A normal workflow spends
> most of its time waiting, so the cap will not fire on it, no matter how long the sub-agents take.
>
> So what is the cap for? One thing: a pure-CPU runaway. Picture a `while True:` loop that never
> awaits. It burns a whole core and blocks the event loop, and none of the sub-agent budgets
> (`max_agent_calls`, `sub_agent_usage_limits`) can stop it, because it never calls a sub-agent. If
> you want that guard, set `max_duration_secs` yourself.
>
> Two more knobs. Pass `'unlimited'` to remove every limit. Pass a partial dict like
> `{'max_memory': ...}` and it merges onto the backstop, so you override only the caps you name and
> the rest keep their defaults.
### Workflows do not nest
A sub-agent cannot start its own workflow. If one tries, the nested `run_workflow` call returns a
terminal error instead of running.
> **Tip**
>
> Here is the practical rule: do not give the sub-agents in your catalog the `DynamicWorkflow`
> capability. They are the leaves of the orchestration, not orchestrators themselves.
## Renaming a sub-agent: `WorkflowAgent`
By default, a sub-agent shows up in the script under its own `name` and `description`. Sometimes you
want a different name or a different description for one particular workflow, without editing the
agent itself. Wrap it in a `WorkflowAgent`:
```python
from pydantic_ai_harness.dynamic_workflow import WorkflowAgent
DynamicWorkflow(
agents=[
WorkflowAgent(
reviewer,
name='check',
description='Checks one code change and returns actionable review findings.',
),
],
)
```
Now the model calls `check(task=...)` instead of `reviewer(task=...)`. Passing a bare agent is just
shorthand for wrapping it in a `WorkflowAgent` with no overrides.
## Adding sub-agents while a run is going: `reveal()`
The catalog is fixed when a run starts, which keeps it in the prompt-cache prefix across turns. But
sometimes you learn during a run that a new sub-agent should be available, say once a fixer agent has
been provisioned. Keep a reference to the `DynamicWorkflow` instance and call `reveal()`:
```python
workflow = DynamicWorkflow(agents=[reviewer])
orchestrator = Agent('openai:gpt-5', deps_type=MyDeps, capabilities=[workflow])
# later, from the host or from another tool:
workflow.reveal(fixer)
```
The revealed sub-agent becomes callable on the next step. The model learns about it through a short
announcement message that carries the new function's signature. The `run_workflow` description itself
stays frozen at the agents present when the run started, so even a runtime reveal never moves the
prompt-cache prefix.
> **Note**
>
> `reveal()` is append-only. Once a sub-agent appears it stays for the rest of the run, and there
> is no way to remove or hide it again. Plan the catalog as something that only grows.
>
> It validates right away. A missing name, an invalid identifier, a reserved keyword, or a name
> collision raises `UserError` at the call site. And if you share one `DynamicWorkflow` instance
> across concurrent runs, `reveal()` reaches all in-flight runs and joins the baseline for runs that
> start afterward.
## Loading it only when needed: `defer_loading`
`DynamicWorkflow` carries a fair amount of instruction text, and most turns do not need it. You can
keep it collapsed to a one-line entry until the model actually loads it. That pays close to zero
tokens on turns that never orchestrate:
```python
DynamicWorkflow(
agents=[reviewer, summarizer],
id='workflow',
defer_loading=True,
)
```
`defer_loading=True` needs a stable `id`. See
[on-demand capabilities](https://pydantic.dev/docs/ai/core-concepts/capabilities/#on-demand-capabilities)
for the full picture.
## What runs in the sandbox
The script runs in Monty, a subset of Python. The subset is what makes the sandbox safe, so it is
worth knowing where the edges are:
- No class definitions, and no third-party libraries.
- Useful standard-library modules: `asyncio`, `math`, `json`, `re`, `typing`. Import what you use.
Other modules are unavailable or stubbed.
- No wall-clock or timing primitives. There is no `asyncio.sleep`, no `datetime.now()`, and no
`time` module.
- `asyncio.gather(...)` runs sub-agents concurrently, but it does not support
`return_exceptions=True`.
Before a script runs, it is statically type-checked against the sub-agent signatures. A misspelled
function, a positional `task`, or a wrong-typed argument costs one retry, but no sub-agent budget and
no sandbox execution.
> **Warning: errors abort the whole script**
>
> A sub-agent that raises cannot be caught inside the script. One failure aborts the whole script,
> and the model retries it. So write scripts where sub-agents do not depend on catching each other's
> errors. If a script does fail after some sub-agents already finished, the retry prompt lists those
> completed results, so the model can reuse them as plain values instead of paying for the same
> calls again.
## What is coming
A suspended Monty program is a small serializable value you can dump, reload, and fork. That points
at two patterns that do not ship yet. The first is forking one expensive shared prefix into N
best-of-N branches. The second is durable workflows that resume from a persisted snapshot after a
crash or a redeploy. Two smaller extensions are also planned: structured sub-agent inputs (a
`parameters` schema per `WorkflowAgent`, instead of only `task: str`) and first-class progress
streaming. Until then, set `event_stream_handler` on each sub-agent `Agent`, or use Logfire, to
watch sub-agent runs inside the one tool call.
## API
```python
DynamicWorkflow( # all parameters are keyword-only
agents=[...], # Sequence[AbstractAgent | WorkflowAgent], required
tool_name='run_workflow',
max_agent_calls=50,
max_retries=3,
forward_usage=True,
inherit_model=False, # True -> sub-agents run with the parent run's resolved model
sub_agent_usage_limits=None, # UsageLimits per sub-agent run; None -> pydantic-ai default
resource_limits=None, # None -> backstop (256 MB, 50M allocs, no time cap);
# 'unlimited' -> off; a dict is merged onto the backstop
id=None, # required when defer_loading=True
description=None, # one-line catalog entry shown while deferred
defer_loading=False,
)
workflow.reveal(agent) # AbstractAgent | WorkflowAgent; validates before appending
WorkflowAgent(
agent, # Agent, required, positional
name=None, # sandbox function name; falls back to agent.name
description=None, # function docstring; falls back to agent.description
)
```
## Further reading
- [Code Mode](../../code_mode/README.md), the same sandbox, calling the agent's own tools instead of
sub-agents.
- [SubAgents](../subagents/README.md), one-delegation-per-tool-call sub-agents, without the
scripted choreography.
- [Tool use via code](https://www.anthropic.com/engineering/code-execution-with-mcp) (Anthropic),
the mechanism this applies to sub-agents.
- [Building Effective Agents](https://www.anthropic.com/engineering/building-effective-agents)
(Anthropic), the orchestration patterns a script can express.
- [Rewriting Bun in Rust](https://bun.com/blog/bun-in-rust) (Bun), the same pattern at scale: Jarred
Sumner's port of Bun from Zig to Rust with Claude Code dynamic workflows.
- [Capabilities](https://pydantic.dev/docs/ai/core-concepts/capabilities/) and
[on-demand capabilities](https://pydantic.dev/docs/ai/core-concepts/capabilities/#on-demand-capabilities).
- [Monty](https://github.com/pydantic/monty), the sandbox.
pydantic_ai_harness/dynamic_workflow/__init__.py
+10
View File
@@ -0,0 +1,10 @@
"""Dynamic workflow capability: orchestrate sub-agents from a sandboxed Python script."""
from pydantic_ai_harness.dynamic_workflow._capability import DynamicWorkflow
from pydantic_ai_harness.dynamic_workflow._toolset import (
DynamicWorkflowToolset,
WorkflowAgent,
WorkflowResourceLimits,
)
__all__ = ['DynamicWorkflow', 'DynamicWorkflowToolset', 'WorkflowAgent', 'WorkflowResourceLimits']
pydantic_ai_harness/dynamic_workflow/_capability.py
+162
View File
@@ -0,0 +1,162 @@
"""Dynamic workflow capability: orchestrate sub-agents from a sandboxed Python script."""
from __future__ import annotations
from collections.abc import Sequence
from dataclasses import dataclass, field
from typing import Literal
from pydantic_ai.agent.abstract import AbstractAgent
from pydantic_ai.capabilities import AbstractCapability
from pydantic_ai.tools import AgentDepsT
from pydantic_ai.usage import UsageLimits
from pydantic_ai_harness.dynamic_workflow._toolset import (
DynamicWorkflowToolset,
WorkflowAgent,
WorkflowResourceLimits,
index_workflow_agents,
validate_workflow_agent,
)
@dataclass(kw_only=True)
class DynamicWorkflow(AbstractCapability[AgentDepsT]):
"""Capability that lets the model orchestrate named sub-agents from a Python script.
Instead of one sub-agent per tool call, the model writes a single Python script (run in a
Monty sandbox) that calls each sub-agent as an async function and composes the results: fan
out with `asyncio.gather`, chain one agent's output into the next, vote, or loop until done.
```python
from pydantic_ai import Agent
from pydantic_ai_harness.dynamic_workflow import DynamicWorkflow
reviewer = Agent('openai:gpt-5', name='reviewer', description='Review code for bugs.')
summarizer = Agent('openai:gpt-5', name='summarizer', description='Summarize findings.')
orchestrator = Agent(
'openai:gpt-5',
capabilities=[DynamicWorkflow(agents=[reviewer, summarizer])],
)
```
Each sub-agent runs isolated (its own message history) with the parent's `deps` forwarded;
by default the parent's `usage` accumulator is shared so the whole tree's spend is tallied in
one place. Use `max_agent_calls` for a hard, host-enforced ceiling on sub-agent runs. Workflows
do not nest. Set `defer_loading=True` (with a stable `id`) to keep the tool out of the prompt
until the model loads the capability.
"""
agents: Sequence[AbstractAgent[AgentDepsT, object] | WorkflowAgent[AgentDepsT]]
"""Sub-agents the orchestration script can call as async functions.
Read at construction only; later mutation of the passed sequence is ignored. A raw agent is
shorthand for `WorkflowAgent(agent)`, using the agent's own `name` and `description`; use a
`WorkflowAgent` entry for a per-use-site override. Use `reveal()` to add a sub-agent after
construction.
"""
_catalog: list[WorkflowAgent[AgentDepsT]] = field(init=False, repr=False)
"""Normalized catalog passed by reference to toolsets."""
tool_name: str = 'run_workflow'
"""Name of the orchestration tool exposed to the model."""
max_agent_calls: int = 50
"""Maximum total sub-agent runs per agent run: an exact, host-enforced ceiling that holds even
under concurrent fan-out (unlike a parent `usage_limits`)."""
max_retries: int = 3
"""Maximum retries for the orchestration tool (syntax/runtime errors count as retries)."""
forward_usage: bool = True
"""Share the parent run's `usage` accumulator with sub-agents, tallying the whole tree's
token and request spend in one place.
This does **not** forward the parent's `usage_limits` into sub-agent runs (`RunContext` does
not expose the limit value): set `sub_agent_usage_limits` to bound sub-agents, or
`max_agent_calls` for an exact ceiling on the number of runs.
"""
inherit_model: bool = False
"""Run every sub-agent with the parent run's resolved model instead of its constructed model.
Use this when the host can switch models per run, such as a `/model` command that passes a
run-level model override to the parent agent. Without it, that per-run choice silently leaves
catalog sub-agents on the model they were bound to when constructed; `inherit_model=True` makes
the workflow crew follow the parent run's resolved model. Keep `False` to pin sub-agents to
their own configured models.
"""
sub_agent_usage_limits: UsageLimits | None = None
"""`UsageLimits` applied to every sub-agent run, replacing pydantic-ai's default.
With `forward_usage=False`, a per-run `total_tokens_limit` of `T` plus `max_agent_calls` of
`N` bounds the tree to roughly `N * T` tokens (each run can overshoot by its final response,
since core checks token limits after a response arrives). With `forward_usage=True` the limit
is checked against the shared counter -- a tree-wide cap, best-effort under concurrent fan-out.
`None` keeps the default (`request_limit=50`, no token limit).
"""
resource_limits: WorkflowResourceLimits | Literal['unlimited'] | None = None
"""Sandbox limits guarding the orchestration script's own memory/allocations.
`None` applies a safe backstop (256 MB, 50M allocations, no execution-time cap); `'unlimited'`
removes all limits; a `WorkflowResourceLimits` mapping is merged onto the backstop, overriding
only the caps it names. There is no default `max_duration_secs`: Monty's timer bounds in-sandbox
execution time and time awaiting sub-agents does not count against it, so set one only to guard
a pure-CPU `while True` loop.
"""
@classmethod
def get_serialization_name(cls) -> str | None:
# Not spec-serializable: `agents` holds live Agent objects, not YAML-expressible config.
return None
def __post_init__(self) -> None:
catalog = [self._normalize_workflow_agent(entry) for entry in self.agents]
index_workflow_agents(catalog)
self._catalog = catalog
def _normalize_workflow_agent(
self, entry: AbstractAgent[AgentDepsT, object] | WorkflowAgent[AgentDepsT]
) -> WorkflowAgent[AgentDepsT]:
"""Normalize a public catalog entry to the internal wrapper form."""
if isinstance(entry, WorkflowAgent):
return entry
return WorkflowAgent(agent=entry)
def reveal(self, agent: AbstractAgent[AgentDepsT, object] | WorkflowAgent[AgentDepsT]) -> None:
"""Reveal a sub-agent on the next model step (the supported runtime API for doing so).
The sub-agent is announced to the model on the next step and becomes callable then; the
`run_workflow` description stays frozen at the agents present when the run started. Reveal
is append-only: a revealed sub-agent cannot be removed for the rest of the run. Its resolved
name must be a valid, unique sandbox function name, or this raises `UserError` at the call
site. If one `DynamicWorkflow` instance is shared across concurrent runs, `reveal()` reaches
all in-flight runs and joins the baseline catalog for runs that start afterwards.
"""
entry = self._normalize_workflow_agent(agent)
existing_names: set[str] = set()
for catalog_entry in self._catalog:
existing_names.add(validate_workflow_agent(catalog_entry, existing_names))
validate_workflow_agent(entry, existing_names)
self._catalog.append(entry)
def get_toolset(self) -> DynamicWorkflowToolset[AgentDepsT]:
"""Provide the orchestration toolset to the agent."""
return DynamicWorkflowToolset(
# Toolsets keep this same list object; `reveal()` appends to it so in-flight
# toolsets can fold in the new sub-agent on the next step.
agents=self._catalog,
tool_name=self.tool_name,
max_agent_calls=self.max_agent_calls,
max_retries=self.max_retries,
forward_usage=self.forward_usage,
inherit_model=self.inherit_model,
sub_agent_usage_limits=self.sub_agent_usage_limits,
resource_limits=self.resource_limits,
toolset_id=self.id,
owning_capability=self,
)
pydantic_ai_harness/dynamic_workflow/_toolset.py
+723
View File
@@ -0,0 +1,723 @@
"""Toolset for the `DynamicWorkflow` capability.
Exposes a single `run_workflow` tool: the model writes a Python orchestration
script (run in a Monty sandbox) that calls named sub-agents as async functions
and composes their results -- fan-out, chaining, voting, loops -- in one step.
"""
from __future__ import annotations
import contextvars
import copy
import json
import keyword
import warnings
from collections.abc import Mapping, Sequence
from dataclasses import dataclass, field
from typing import Annotated, Any, Generic, Literal
from pydantic import Field, TypeAdapter
from pydantic_ai import AbstractToolset, RunContext, ToolDefinition
from pydantic_ai.agent.abstract import AbstractAgent
from pydantic_ai.capabilities import AbstractCapability, WrapperCapability
from pydantic_ai.exceptions import ModelRetry, UsageLimitExceeded, UserError
from pydantic_ai.function_signature import FunctionSignature
from pydantic_ai.tools import AgentDepsT
from pydantic_ai.toolsets.abstract import SchemaValidatorProt, ToolsetTool
from pydantic_ai.usage import UsageLimits
from pydantic_core import to_jsonable_python
from typing_extensions import Self, TypedDict
try:
from pydantic_monty import Monty, MontyRepl, MontyRuntimeError, MontySyntaxError, MontyTypingError, ResourceLimits
except ImportError as _import_error: # pragma: no cover
raise ImportError(
'pydantic-monty is required for DynamicWorkflow. '
'Install it with: uv add "pydantic-ai-harness[dynamic-workflow]"'
) from _import_error
from pydantic_ai_harness._monty_exec import MontyExecutor, PrintCapture, is_sandbox_panic
# Set while a workflow script is executing, so a sub-agent that itself tries to run a workflow can
# be refused -- workflows do not nest. asyncio copies the context into each task `asyncio.gather`
# schedules, so concurrently-dispatched sub-agents inherit this flag (the capability is asyncio-only).
_in_workflow: contextvars.ContextVar[bool] = contextvars.ContextVar('pydantic_ai_harness_in_workflow', default=False)
class WorkflowResourceLimits(TypedDict, total=False):
"""Caps on the orchestration script's own sandbox resources (not sub-agent latency).
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 script spends executing sandbox code. Monty checks it per bytecode
step, so time spent awaiting sub-agents does not count against it -- neither a single `await`
nor a concurrent `asyncio.gather` batch, because during that wait the script 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 -- the one
runaway the sub-agent budgets do not catch."""
max_memory: int
"""Maximum sandbox memory, in bytes."""
max_allocations: int
"""Maximum number of sandbox allocations."""
def _default_resource_limits() -> ResourceLimits:
"""Backstop sandbox limits; no duration cap -- see `WorkflowResourceLimits.max_duration_secs`."""
return {
'max_memory': 256 * 1024 * 1024,
'max_allocations': 50_000_000,
}
# The keys `WorkflowResourceLimits` 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 -- quietly disabling the only guard against a pure-CPU `while True`. We reject unknowns.
_RESOURCE_LIMIT_KEYS = frozenset(WorkflowResourceLimits.__annotations__)
_MODEL_SAFE_EXCEPTION_MESSAGE_TYPES = (UsageLimitExceeded,)
_MAX_COMPLETED_DISPATCHES = 20
_MAX_TASK_PREVIEW_CHARS = 120
_MAX_RESULT_PREVIEW_CHARS = 300
_TRUNCATED_MARKER = ' ... [truncated]'
def _resolve_resource_limits(limits: WorkflowResourceLimits | 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 silently drops the allocations backstop. Full semantics: `DynamicWorkflow.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 _WorkflowArguments(TypedDict):
code: Annotated[str, Field(description='The Python orchestration script to execute in the sandbox.')]
_WORKFLOW_ARGS_ADAPTER = TypeAdapter(_WorkflowArguments)
_WORKFLOW_ARGS_JSON_SCHEMA = _WORKFLOW_ARGS_ADAPTER.json_schema()
_WORKFLOW_ARGS_VALIDATOR: SchemaValidatorProt = _WORKFLOW_ARGS_ADAPTER.validator # pyright: ignore[reportAssignmentType]
class _BudgetExhausted(RuntimeError):
"""The run's `max_agent_calls` budget is spent; no further sub-agent runs are allowed."""
def __init__(self, max_agent_calls: int) -> None:
super().__init__(f'sub-agent call budget ({max_agent_calls}) exhausted')
_WORKFLOW_BASE_DESCRIPTION = """\
Write and run a Python orchestration script in a sandbox to coordinate multiple sub-agents.
Use this to break a task across specialized sub-agents and combine their results in a single step --
fan work out in parallel, chain one agent's output into the next, vote across several, or loop until
done -- instead of delegating to one sub-agent at a time.
The sandbox uses Monty, a subset of Python. Key restrictions:
- **No classes** and **no third-party libraries**.
- **Useful standard-library modules**: `asyncio`, `math`, `json`, `re`, `typing`. Import what you use
at the top of the script. Other modules are unavailable or stubbed -- don't rely on them.
- **No wall-clock or timing primitives** (`asyncio.sleep`, `datetime.now()`, the `time` module).
Each sub-agent below is an async function. Call it with the `task` keyword argument -- write
`reviewer(task="...")`, not `reviewer("...")`; all parameters are keyword-only. A sub-agent returns
that agent's output: a string by default, or -- if it has a structured `output_type` -- a dict, whose
fields you read by subscript (`r["field"]`), not attribute (`r.field`). Each sub-agent call is an
independent run with no memory of earlier calls; include all needed context in `task`. Run several
at once with `asyncio.gather` rather than awaiting each sequentially:
```python
import asyncio
reviews = await asyncio.gather(reviewer(task="check auth"), reviewer(task="check parsing"))
```
`asyncio.gather` does **not** support `return_exceptions=True`, and a sub-agent that raises cannot be
caught inside the script: one failure aborts the whole script and you retry it. Design the script so
sub-agents don't depend on catching each other's errors.
The last expression's value is captured as the result -- you do **not** need to `print()` it, and
printing produces a string representation, not structured data. Use `print()` only for debug logging.
Return shapes: no print returns the last expression value (or `{}` if it is `None`); print plus a
non-`None` value returns `{"output": "<printed text>", "result": <last expression>}`; print plus
`None` returns `{"output": "<printed text>"}`. If a script fails after some sub-agent calls complete,
those completed results are reported back so a retry can reuse them.\
"""
def _is_valid_sandbox_name(name: str) -> bool:
"""Whether `name` can be exposed as a sandbox function: a non-keyword Python identifier.
`str.isidentifier()` alone is not enough -- Python keywords (`for`, `class`, `async`, ...) are
valid identifiers but cannot be used as function names, so the model could never call them.
Callers guard the empty/`None` case before this is reached.
"""
return name.isidentifier() and not keyword.iskeyword(name)
# Every sub-agent is exposed with the same fixed parameters -- `(*, task: str)` -- and a per-agent
# return schema. Render each catalog entry through core's `FunctionSignature` (the renderer
# code_mode and Pydantic AI already use). This keeps the catalog format consistent across
# capabilities, forces keyword-only `task` to match `dispatch` (which reads `kwargs['task']`), and
# renders docstrings safely -- a hand-rolled f-string breaks on a newline or a quote inside a
# description.
_SUB_AGENT_PARAMS_SCHEMA: dict[str, Any] = {
'type': 'object',
'properties': {'task': {'type': 'string'}},
'required': ['task'],
}
_NO_CONFLICTING_TYPE_NAMES: frozenset[str] = frozenset()
def _agent_return_schema(agent: AbstractAgent[AgentDepsT, object]) -> dict[str, Any] | None:
"""The sub-agent's output JSON schema, or `None` when it cannot be derived."""
try:
return agent.output_json_schema()
except Exception:
return None
def _agent_signature(name: str, agent: AbstractAgent[AgentDepsT, object]) -> FunctionSignature:
"""Build the sandbox function signature for one sub-agent."""
return FunctionSignature.from_schema(
name=name,
parameters_schema=_SUB_AGENT_PARAMS_SCHEMA,
return_schema=_agent_return_schema(agent),
)
def _render_agent_block(
signature: FunctionSignature,
description: str | None,
*,
conflicting_type_names: frozenset[str] = _NO_CONFLICTING_TYPE_NAMES,
) -> str:
"""Render one sub-agent as the async function signature shown to the model."""
return signature.render(
'...',
description=description,
is_async=True,
conflicting_type_names=conflicting_type_names,
)
def _render_catalog(catalog: Mapping[str, WorkflowAgent[AgentDepsT]], *, max_agent_calls: int) -> str:
"""Render the available sub-agents as async function signatures for the tool description."""
signatures = {name: _agent_signature(name, entry.agent) for name, entry in catalog.items()}
signature_list = list(signatures.values())
conflicting = FunctionSignature.get_conflicting_type_names(signature_list)
type_blocks = FunctionSignature.render_type_definitions(signature_list, conflicting)
function_blocks = [
_render_agent_block(signatures[name], entry.resolved_description, conflicting_type_names=conflicting)
for name, entry in catalog.items()
]
listing = '```python\n' + '\n\n'.join([*type_blocks, *function_blocks]) + '\n```'
budget = (
f'This run can make at most {max_agent_calls} sub-agent calls in total -- one budget shared across '
'every `run_workflow` call in the run, not per script; plan fan-out width accordingly.'
)
return f'{_WORKFLOW_BASE_DESCRIPTION}\n\n{budget}\n\nAvailable sub-agents:\n\n{listing}'
def _render_reveal(
name: str,
catalog: Mapping[str, WorkflowAgent[AgentDepsT]],
tool_name: str,
) -> str:
"""Announcement enqueued when a sub-agent is revealed mid-run.
Delivered as a conversation message (not folded into the cached tool description), so the
prompt-cache prefix stays stable while the model still learns the agent is now callable.
"""
signatures = {agent_name: _agent_signature(agent_name, entry.agent) for agent_name, entry in catalog.items()}
signature = signatures[name]
conflicting = FunctionSignature.get_conflicting_type_names(list(signatures.values()))
type_blocks = FunctionSignature.render_type_definitions([signature], conflicting)
function_block = _render_agent_block(
signature, catalog[name].resolved_description, conflicting_type_names=conflicting
)
block = '\n\n'.join([*type_blocks, function_block])
return f'A new sub-agent is now available to call from inside the `{tool_name}` script:\n\n```python\n{block}\n```'
def _workflow_result(result: object, printed: str) -> object:
"""Shape the tool return: the script's result, its captured `print()` output, or both."""
if not printed:
return result if result is not None else {}
if result is None:
return {'output': printed}
return {'output': printed, 'result': result}
@dataclass(frozen=True)
class _CompletedDispatch:
"""One sub-agent result completed by the failed workflow script."""
agent_name: str
task: str
result: object
def _truncate_preview(value: str, max_chars: int) -> str:
"""Trim a preview with an explicit marker so the model can see it is incomplete."""
if len(value) <= max_chars:
return value
return value[: max_chars - len(_TRUNCATED_MARKER)] + _TRUNCATED_MARKER
def _json_preview(value: object, max_chars: int) -> str:
"""Render a compact JSON preview of a completed sub-agent result."""
rendered = json.dumps(value, ensure_ascii=True, separators=(',', ':'))
return _truncate_preview(rendered, max_chars)
def _completed_dispatch_lines(completed: list[_CompletedDispatch]) -> list[str]:
"""Format the most recent completed sub-agent results for model-facing salvage."""
shown = completed[-_MAX_COMPLETED_DISPATCHES:]
lines = [
f'{entry.agent_name}(task={json.dumps(_truncate_preview(entry.task, _MAX_TASK_PREVIEW_CHARS), ensure_ascii=True)})'
f' -> {_json_preview(entry.result, _MAX_RESULT_PREVIEW_CHARS)}'
for entry in shown
]
omitted = len(completed) - len(shown)
if omitted:
lines.insert(0, f'... {omitted} earlier completed result(s) omitted ...')
return lines
def _completed_retry_section(completed: list[_CompletedDispatch]) -> str:
"""Build the optional retry-message section listing salvageable completed results."""
lines = _completed_dispatch_lines(completed)
if not lines:
return ''
listing = '\n'.join(f'- {line}' for line in lines)
return (
'\n\nCompleted sub-agent results from the failed script '
'(reuse these values instead of re-calling them; their budget was already spent):\n'
f'{listing}'
)
def _budget_terminal_result(
*,
max_agent_calls: int,
last_error: str,
completed_dispatches: list[_CompletedDispatch],
) -> dict[str, object]:
"""Build the terminal result returned after the exact sub-agent-call budget is exhausted."""
return {
'error': (
f'This run exhausted its sub-agent call budget ({max_agent_calls}). '
'Conclude using the results already gathered; further sub-agent calls in '
'this run will be refused.'
),
'last_error': last_error,
'completed': _completed_dispatch_lines(completed_dispatches),
}
@dataclass(frozen=True)
class WorkflowAgent(Generic[AgentDepsT]):
"""One sub-agent exposed to the orchestration script as an async function.
`WorkflowAgent` is the per-use-site override for when the agent's own `name`
or `description` is not what this workflow should show, such as renaming the
sandbox function or re-describing the agent for this catalog. Passing a bare
agent to `DynamicWorkflow(agents=[...])` is equivalent to `WorkflowAgent(agent)`.
"""
agent: AbstractAgent[AgentDepsT, object]
"""The sub-agent to run when the script calls this function."""
name: str | None = None
"""Sandbox function name; must be a valid Python identifier and unique across the
workflow. Falls back to the agent's `name`."""
description: str | None = None
"""Description shown to the model in the sub-agent catalog, rendered as the sandbox
function's docstring. An explicit value overrides the agent's own `description`.
When neither is set, the model sees only the bare signature."""
@property
def resolved_name(self) -> str | None:
"""The sandbox function name: the explicit `name`, else the agent's `name`."""
return self.name or self.agent.name
@property
def resolved_description(self) -> str | None:
"""The catalog description: the explicit `description`, else the agent's own `description`."""
return self.description or self.agent.description
def validate_workflow_agent(entry: WorkflowAgent[AgentDepsT], existing_names: set[str]) -> str:
"""Validate one sub-agent entry against the names already taken."""
name = entry.resolved_name
if not name:
raise UserError(
'DynamicWorkflow sub-agent has no `name` and its agent has no `name`; '
'set `WorkflowAgent(name=...)` so it can be exposed as a sandbox function.'
)
if not _is_valid_sandbox_name(name):
raise UserError(
f'DynamicWorkflow sub-agent name {name!r} cannot be exposed as a sandbox function: '
'it must be a Python identifier that is not a reserved keyword. Rename it.'
)
if name in existing_names:
raise UserError(f'DynamicWorkflow has two sub-agents named {name!r}; names must be unique.')
return name
def index_workflow_agents(
agents: Sequence[WorkflowAgent[AgentDepsT]],
) -> dict[str, WorkflowAgent[AgentDepsT]]:
"""Index validated sub-agent entries by resolved sandbox name."""
if not agents:
raise UserError('DynamicWorkflow requires at least one sub-agent in `agents`.')
by_name: dict[str, WorkflowAgent[AgentDepsT]] = {}
existing_names: set[str] = set()
for entry in agents:
name = validate_workflow_agent(entry, existing_names)
existing_names.add(name)
by_name[name] = entry
return by_name
@dataclass(kw_only=True)
class DynamicWorkflowToolset(AbstractToolset[AgentDepsT]):
"""Single-tool toolset that runs sub-agent orchestration scripts in a Monty sandbox."""
agents: list[WorkflowAgent[AgentDepsT]]
"""Sub-agents callable from the orchestration script, each as an async function.
`DynamicWorkflow.reveal()` is the only supported way to add a sub-agent mid-run.
The toolset observes appends to this list as the reveal channel. Any entry that
arrives invalid is a contract violation that raises."""
tool_name: str = 'run_workflow'
"""Name of the tool exposed to the model."""
max_agent_calls: int = 50
"""Maximum total sub-agent runs per agent run (an exact, host-enforced ceiling)."""
max_retries: int = 3
"""Maximum retries for the `run_workflow` tool (syntax/runtime errors count as retries)."""
forward_usage: bool = True
"""Share the parent run's `usage` accumulator with sub-agents. See
`DynamicWorkflow.forward_usage` for what is and is not forwarded."""
inherit_model: bool = False
"""Run every sub-agent with the parent run's resolved model instead of its constructed model.
See `DynamicWorkflow.inherit_model` for when to use this."""
sub_agent_usage_limits: UsageLimits | None = None
"""`UsageLimits` applied to every sub-agent run, replacing pydantic-ai's default.
See `DynamicWorkflow.sub_agent_usage_limits` for the budgeting semantics."""
resource_limits: WorkflowResourceLimits | Literal['unlimited'] | None = None
"""Sandbox limits guarding the orchestration script's own memory/allocations (not sub-agents).
See `DynamicWorkflow.resource_limits` for the `None`/`'unlimited'`/partial-dict semantics."""
toolset_id: str | None = None
"""Stable toolset id; defaults to the tool name."""
owning_capability: AbstractCapability[AgentDepsT] | None = None
"""The `DynamicWorkflow` capability this toolset belongs to, when capability-provided.
Used to resolve whether the capability is visible to the model on the current step (deferred
and unloaded means hidden), by identity against the run's capability registry -- ids cannot be
used for this, because an id-less capability is registered under a run-generated key and a
wrapper capability is registered in place of what it wraps. `None` (a toolset used directly,
outside any capability) is always visible."""
# Per-run count of sub-agent calls; reset on `for_run`.
_call_count: int = field(default=0, init=False, repr=False)
# Sub-agents indexed by resolved sandbox name; seeded from `agents` in `__post_init__` and
# extended in place as runtime appends to `agents` are revealed (`_reveal_pending`).
_by_name: dict[str, WorkflowAgent[AgentDepsT]] = field(init=False, repr=False)
# Tool description, frozen at run start. Rendered from the agents present when the run began
# and never re-rendered after a reveal, so the description -- and thus the prompt-cache
# prefix -- never changes mid-run.
_description: str = field(init=False, repr=False)
def __post_init__(self) -> None:
if self.max_agent_calls < 1:
raise UserError('DynamicWorkflow `max_agent_calls` must be at least 1.')
_resolve_resource_limits(self.resource_limits) # validate keys now, not at the first tool call
self._rebuild()
def _rebuild(self) -> None:
"""Rebuild the name index and the frozen tool description from the current `agents`.
Unusable entries raise `UserError`. `DynamicWorkflow.__post_init__` and
`DynamicWorkflow.reveal()` validate entries eagerly, so this fails only when the
lower-level toolset list was mutated outside those APIs.
"""
by_name = index_workflow_agents(self.agents)
self._by_name = by_name
self._description = _render_catalog(by_name, max_agent_calls=self.max_agent_calls)
@property
def id(self) -> str | None:
return self.toolset_id or self.tool_name
async def for_run(self, ctx: RunContext[AgentDepsT]) -> Self:
"""Fresh instance per run so the sub-agent-call budget is per-run.
Clone shallowly, keeping `agents` shared so reveals stay visible to this run,
reset the per-run call budget, and rebuild the per-run index. Entries are
validated eagerly by `DynamicWorkflow.__post_init__` and `DynamicWorkflow.reveal()`,
so a strict rebuild here cannot fail unless this list was mutated outside those
APIs. That is unsupported and fails fast.
"""
clone = copy.copy(self)
clone._call_count = 0
clone._rebuild()
return clone
def _fold_reveals(self, ctx: RunContext[AgentDepsT]) -> None:
"""Fold sub-agents appended to `agents` since the run started into the name index.
Diffs the live `agents` list against the names already known (`_by_name`, which holds the
baseline plus anything revealed so far) and folds each newcomer in -- so `dispatch` resolves
it -- enqueuing an announcement for the model. The frozen `_description` is untouched, so the
cached prompt prefix is unaffected. Re-seeing an already-known entry (a baseline agent, or
one revealed on an earlier step) is a no-op, identified by object identity so it is never
mistaken for a name collision.
A newcomer whose name is missing, invalid, or already taken by a different
sub-agent is a contract violation and raises `UserError`.
Synchronous and await-free between snapshotting `agents` and mutating `_by_name`, so a
concurrently-running `dispatch` never observes a half-revealed agent -- the same
await-free-critical-section reasoning that keeps `max_agent_calls` exact under fan-out.
"""
for entry in tuple(self.agents):
name = entry.resolved_name
existing = self._by_name.get(name) if name else None
if existing is entry:
continue
name = validate_workflow_agent(entry, set(self._by_name))
self._by_name[name] = entry
try:
ctx.enqueue(_render_reveal(name, self._by_name, self.tool_name))
except UserError as exc:
warnings.warn(
f'DynamicWorkflow revealed sub-agent {name!r}, but could not enqueue its announcement: '
f'{exc}. It is callable in this run, but the model will not see the reveal message.',
stacklevel=2,
)
def _visible_to_model(self, ctx: RunContext[AgentDepsT]) -> bool:
"""Whether this toolset's capability is visible to the model on this step.
A deferred capability's toolset still gets `get_tools` calls while unloaded (its tool
is indexed for `load_capability`/`search_tools`), but announcing a reveal then would
leak the sub-agent signature into the conversation while the catalog is still hidden.
The owning capability is resolved against the run's registry by identity, unwrapping
`WrapperCapability` chains -- a wrapper over a leaf capability is registered in place of
the leaf, and its `defer_loading`/registry id are what core keys loading on. Matching by
id instead would misfire both ways: an id-less capability is registered under a
run-generated key (so a lookup by tool name can hit an unrelated capability), and a
deferred wrapper hides a non-deferred inner capability. No registry match (a toolset used
directly, outside any capability) means always visible.
"""
owner = self.owning_capability
if owner is None:
return True
for capability_id, registered in ctx.capabilities.items():
candidate: AbstractCapability[AgentDepsT] | None = registered
while candidate is not None:
if candidate is owner:
if registered.defer_loading is not True:
return True
return capability_id in ctx.loaded_capability_ids
candidate = candidate.wrapped if isinstance(candidate, WrapperCapability) else None
return True
async def get_tools(self, ctx: RunContext[AgentDepsT]) -> dict[str, ToolsetTool[AgentDepsT]]:
# Skip reveal processing while the capability is deferred and unloaded: the newcomers
# stay pending in `agents` and are folded in and announced on the first step after the
# model loads the capability.
if self._visible_to_model(ctx):
self._fold_reveals(ctx)
return {
self.tool_name: ToolsetTool(
toolset=self,
tool_def=ToolDefinition(
name=self.tool_name,
description=self._description,
parameters_json_schema=_WORKFLOW_ARGS_JSON_SCHEMA,
metadata={'code_arg_name': 'code', 'code_arg_language': 'python'},
sequential=True,
),
max_retries=self.max_retries,
args_validator=_WORKFLOW_ARGS_VALIDATOR,
)
}
async def _run_one(self, agent_name: str, task: str, ctx: RunContext[AgentDepsT]) -> Any:
"""Run one sub-agent against the shared per-run budget.
The budget check + increment must stay suspension-free: there must be no `await`
between them. asyncio only switches tasks at suspension points, so an await-free
check-then-increment is atomic across the concurrently-gathered dispatches, which
is what makes `max_agent_calls` an exact ceiling under fan-out. Insert an `await`
here (e.g. an async permission check) and the count can race past the limit; you
would then need an explicit reservation instead.
This exists precisely because `usage_limits` cannot give an exact ceiling here:
core's own limit check is split from its increment by the model-request `await`
(a TOCTOU race -- N gathered sub-agents all pass the check before any increments;
measured ~20x overshoot), and `RunContext` exposes `usage` but not `usage_limits`,
so the parent's configured limit can't be forwarded to sub-agents at all.
TODO: file upstream on pydantic-ai -- (a) expose `usage_limits` on `RunContext`,
(b) atomic reserve-then-request in the run loop. Until then, tree-wide token caps
stay best-effort (see `forward_usage` / `sub_agent_usage_limits` docstrings).
"""
if self._call_count >= self.max_agent_calls:
raise _BudgetExhausted(self.max_agent_calls)
self._call_count += 1
try:
result = await self._by_name[agent_name].agent.run(
task,
deps=ctx.deps,
model=ctx.model if self.inherit_model else None,
usage=ctx.usage if self.forward_usage else None,
usage_limits=self.sub_agent_usage_limits,
)
return to_jsonable_python(result.output)
except Exception as exc:
# Don't leak host internals (file paths, deps/agent reprs) to the model;
# surface the failing agent and error type by default.
message = f'sub-agent {agent_name!r} raised {type(exc).__name__}'
if isinstance(exc, _MODEL_SAFE_EXCEPTION_MESSAGE_TYPES):
message = f'{message}: {exc}'
raise RuntimeError(message) from exc
def _type_check(self, code: str, capture: PrintCapture) -> None:
"""Statically check the script against the sub-agent signatures before running it.
Each `run_workflow` call is a fresh sandbox with no accumulated state, so the check is
always sound. Catching a positional `task`, a misspelled function, or a wrong-typed
argument here costs a retry but no sub-agent budget.
"""
signatures = [_agent_signature(name, entry.agent) for name, entry in self._by_name.items()]
conflicting = FunctionSignature.get_conflicting_type_names(signatures)
parts = ['import asyncio\nfrom typing import Any, TypedDict, NotRequired, Literal']
parts.extend(FunctionSignature.render_type_definitions(signatures, conflicting))
parts.extend(
signature.render('raise NotImplementedError()', is_async=True, conflicting_type_names=conflicting)
for signature in signatures
)
try:
Monty(code, type_check=True, type_check_stubs='\n\n'.join(parts))
except MontyTypingError as e:
raise ModelRetry(f'Type error in workflow:\n{capture.prepend_to(e.display())}') from e
except MontySyntaxError as e:
raise ModelRetry(f'Syntax error in workflow:\n{capture.prepend_to(e.display())}') from e
async def call_tool(
self, name: str, tool_args: dict[str, Any], ctx: RunContext[AgentDepsT], tool: ToolsetTool[AgentDepsT]
) -> Any:
if _in_workflow.get():
return {
'error': (
'Workflows do not nest: this sub-agent was invoked from a workflow and cannot start '
'its own. Return your result to the orchestrating workflow instead.'
)
}
code = tool_args['code']
budget_exhausted = False
completed_dispatches: list[_CompletedDispatch] = []
async def dispatch(agent_name: str, kwargs: dict[str, Any]) -> Any:
nonlocal budget_exhausted
# The sandbox signature is `(*, task: str)`, but Monty does not validate kwargs against
# it at runtime, and the static check can be evaded through `Any` (e.g. `json.loads`
# results) -- so check here: a dropped extra kwarg or a non-string `task` would
# otherwise run the sub-agent on silently-wrong input. Each raises before the budget
# is touched.
if 'task' not in kwargs:
raise TypeError(f'{agent_name}() missing required keyword argument: task')
extra = sorted(set(kwargs) - {'task'})
if extra:
raise TypeError(f'{agent_name}() got unexpected keyword argument(s): {", ".join(extra)}; only task')
task = kwargs['task']
if not isinstance(task, str):
raise TypeError(f'{agent_name}() task must be a string, got {type(task).__name__}')
try:
output = await self._run_one(agent_name, task, ctx)
except _BudgetExhausted:
budget_exhausted = True
raise
completed_dispatches.append(_CompletedDispatch(agent_name=agent_name, task=task, result=output))
return output
limits = _resolve_resource_limits(self.resource_limits)
capture = PrintCapture()
self._type_check(code, capture)
in_workflow_token = _in_workflow.set(True)
try:
repl = MontyRepl(limits=limits)
monty_state = repl.feed_start(code, print_callback=capture)
# `_by_name` is not mutated while a script executes (reveals land in `get_tools`,
# which does not interleave with `call_tool`), so it is a stable name registry for
# the whole script. Sub-agents always run concurrently (the executor's defaults);
# durable ordering (global_sequential) lands with durability.
completed = await MontyExecutor(dispatch=dispatch, valid_names=self._by_name).run(monty_state)
except MontySyntaxError as e: # pragma: no cover -- backstop; `_type_check` rejects syntax errors first
raise ModelRetry(f'Syntax error in workflow:\n{capture.prepend_to(e.display())}') from e
except MontyRuntimeError as e:
if budget_exhausted:
# On this capability's deferred-future path, host-raised exceptions cannot be
# caught inside the sandbox (Monty's inline resume path can catch them). The
# flag proves this script hit the budget; under gather, the displayed error may
# be another independently surfaced failure from the same batch.
return _budget_terminal_result(
max_agent_calls=self.max_agent_calls,
last_error=capture.prepend_to(e.display()),
completed_dispatches=completed_dispatches,
)
raise ModelRetry(
f'Runtime error in workflow:\n{capture.prepend_to(e.display())}'
f'{_completed_retry_section(completed_dispatches)}'
) from e
except BaseException as e:
# Convert a model-provokable sandbox panic to a retry (see `is_sandbox_panic`);
# anything else (CancelledError, ...) re-raises unchanged.
if not is_sandbox_panic(e):
raise
if budget_exhausted:
return _budget_terminal_result(
max_agent_calls=self.max_agent_calls,
last_error='The workflow script aborted inside the sandbox after exhausting the sub-agent budget.',
completed_dispatches=completed_dispatches,
)
raise ModelRetry(
'The workflow script aborted inside the sandbox. This can happen when the same '
'sub-agent call is awaited more than once in one asyncio.gather -- give each gathered '
'call its own invocation. Revise the script and try again.'
f'{_completed_retry_section(completed_dispatches)}'
) from e
finally:
_in_workflow.reset(in_workflow_token)
return _workflow_result(completed.output, capture.joined)
pydantic_ai_harness/experimental/_warn.py
+15
View File
@@ -38,3 +38,18 @@ def warn_experimental(feature: str) -> None:
category=HarnessExperimentalWarning,
stacklevel=2,
)
def warn_moved(old: str, new: str) -> None:
"""Emit a `DeprecationWarning` that `experimental.old` now lives at top-level `new`.
Left behind at each old `pydantic_ai_harness.experimental.<name>` path when a capability
graduates out of `experimental`, so existing imports keep working with a clear pointer to
the new location.
"""
warnings.warn(
f'`pydantic_ai_harness.experimental.{old}` has moved to `pydantic_ai_harness.{new}`. '
f'Update your imports; this compatibility shim will be removed in a future release.',
category=DeprecationWarning,
stacklevel=2,
)
pydantic_ai_harness/experimental/acp/README.md
+254
View File
@@ -0,0 +1,254 @@
# ACP (Agent Client Protocol)
> [!WARNING]
> **Experimental.** This capability lives under `pydantic_ai_harness.experimental` and may
> change or be removed in any release, without a deprecation period. Import it from the
> experimental path -- there is no top-level export:
>
> ```python
> from pydantic_ai_harness.experimental.acp import run_acp_stdio_sync
> ```
>
> Importing any experimental capability emits a `HarnessExperimentalWarning`. Silence **all**
> harness experimental warnings with a single filter (no per-capability lines needed):
>
> ```python
> import warnings
> from pydantic_ai_harness.experimental import HarnessExperimentalWarning
>
> warnings.filterwarnings('ignore', category=HarnessExperimentalWarning)
> ```
Expose a Pydantic AI agent to editors and terminal UIs over the [Agent Client Protocol](https://agentclientprotocol.com).
## The problem
Editors like [Zed](https://zed.dev/docs/ai/external-agents) speak ACP: a stdio JSON-RPC protocol that lets a TUI or editor drive an external coding agent -- streaming its text, rendering its file edits as diffs, and prompting the user to approve sensitive tool calls. To plug a Pydantic AI agent into one of these editors you would otherwise have to implement the ACP server side yourself.
## The solution
`run_acp_stdio` serves any Pydantic AI `Agent` as an ACP agent over stdin/stdout. The editor launches your script as a subprocess and talks to it; the adapter translates between ACP and the agent's run loop:
| ACP needs | The adapter provides |
|---|---|
| Streamed assistant text and reasoning | Agent text/thinking deltas, chunked under the wire limit |
| Rich tool calls (`kind`, file `locations`, diffs) | A presenter that recognizes `FileSystem`/`Shell` tool calls |
| Human-in-the-loop tool approval | Maps ACP permission requests to Pydantic AI's deferred-approval tools |
| Per-workspace sessions | A `session_config` hook to root tools at the client's working directory |
| Cancellation, multi-turn history, session close | Handled per session |
## Installation
```bash
uv add "pydantic-ai-harness[acp]"
```
This pulls in the [`agent-client-protocol`](https://pypi.org/project/agent-client-protocol/) SDK. The rest of the harness does not depend on it -- only `pydantic_ai_harness.experimental.acp` does.
## Quick start
Write a script that builds your agent and serves it:
```python
# my_acp_agent.py
from pydantic_ai import Agent
from pydantic_ai_harness.experimental.acp import run_acp_stdio_sync
def build_agent() -> Agent[None, str]:
return Agent('anthropic:claude-sonnet-4-6', instructions='You are a coding assistant.')
if __name__ == '__main__':
run_acp_stdio_sync(build_agent())
```
`run_acp_stdio_sync` blocks for the lifetime of the connection -- it is the `main()` of an agent the editor launches. Inside an existing event loop, use the async `run_acp_stdio` instead.
## Connecting from an editor
ACP clients launch the agent as a subprocess. In Zed, register it as an [external agent](https://zed.dev/docs/ai/external-agents) in `settings.json`:
```json
{
"agent_servers": {
"My Pydantic AI Agent": {
"type": "custom",
"command": "python",
"args": ["/absolute/path/to/my_acp_agent.py"],
"env": { "ANTHROPIC_API_KEY": "..." }
}
}
}
```
Any ACP-compatible client works the same way -- point it at `python my_acp_agent.py`. Refer to your editor's external-agent documentation for the exact config location.
The provider environment must be available to the launched subprocess. GUI editors and SDK-based
test wrappers may not source your interactive shell startup files, and the ACP Python SDK's
`spawn_agent_process` helper starts from a trimmed default environment unless you pass `env`
explicitly. If a real-model agent exits before initialize or fails provider auth, first verify that
the command's process can see variables such as `ANTHROPIC_API_KEY`.
## Rooting tools at the workspace
A coding agent should read and write files in the workspace the editor opened, not wherever the subprocess started. ACP gives each session a working directory (`cwd`); a `session_config` factory turns that into per-session tools:
```python
from pydantic_ai import Agent
from pydantic_ai_harness.experimental.acp import AcpSession, AcpSessionConfig, run_acp_stdio_sync
from pydantic_ai_harness.filesystem import FileSystem
from pydantic_ai_harness.shell import Shell
agent = Agent('anthropic:claude-sonnet-4-6')
def session_config(session: AcpSession) -> AcpSessionConfig[None]:
# Root file and shell tools at the workspace the client opened.
return AcpSessionConfig(
deps=None,
toolsets=[
FileSystem[None](root_dir=session.cwd).get_toolset(),
Shell[None](cwd=session.cwd).get_toolset(),
],
)
if __name__ == '__main__':
run_acp_stdio_sync(agent, session_config=session_config)
```
The factory runs once per session with the client's [`AcpSession`][pydantic_ai_harness.experimental.acp.AcpSession] setup (its `cwd`, `mcp_servers`, and capabilities) and returns an [`AcpSessionConfig`][pydantic_ai_harness.experimental.acp.AcpSessionConfig] whose `deps` and `toolsets` apply to every run in that session. This is correct across multiple concurrent sessions in one process, where a single static `FileSystem` could not be.
## Editor-native filesystem and shell (optional)
The local `FileSystem` and `Shell` above operate on the agent process's own disk and subprocesses. An editor's source of truth is different: it has unsaved buffers, the file layout it considers the workspace, and -- for a remote or containerized editor -- the machine the code actually lives on. When the client advertises support, [`acp_filesystem`][pydantic_ai_harness.experimental.acp.acp_filesystem] and [`acp_terminal`][pydantic_ai_harness.experimental.acp.acp_terminal] give the agent `read_file`/`write_file`/`run_command` tools that route through the client, so it acts where the user is:
```python
from pydantic_ai_harness.experimental.acp import AcpSession, AcpSessionConfig, acp_filesystem, acp_terminal
from pydantic_ai_harness.filesystem import FileSystem
from pydantic_ai_harness.shell import Shell
def session_config(session: AcpSession) -> AcpSessionConfig[None]:
# Use the editor's filesystem/terminal when offered; otherwise fall back to local.
fs = acp_filesystem(session) or FileSystem[None](root_dir=session.cwd).get_toolset()
shell = acp_terminal(session) or Shell[None](cwd=session.cwd).get_toolset()
return AcpSessionConfig(deps=None, toolsets=[fs, shell])
```
Each helper returns `None` when the client did not advertise the capability, so the `or` falls back to local and the agent works either way. The tool names match the local `FileSystem`/`Shell`, so rich rendering (next section) is identical. `acp_terminal` runs the command in the editor's environment and returns its captured output (see [Limitations](#cancellation-and-limitations)).
If a client advertises filesystem *reads* but not *writes*, `acp_filesystem` keeps editor-native reads and sends writes to the local `FileSystem` rooted at `session.cwd` -- coherent only when the agent shares the workspace disk with the editor (same machine, or an agent inside the editor's container); for a remote editor those writes land on the agent's disk, not the editor's.
## Tool approval
Mark a tool to require approval and ACP relays the decision to the client, which shows the user an approve/reject prompt:
```python
@agent.tool_plain(requires_approval=True)
def delete_file(path: str) -> str:
...
```
The lifecycle the client sees is `pending` (awaiting approval) → `in_progress` (granted, running) → `completed`/`failed`, so an unapproved action is never shown as already running. "Always allow"/"always reject" decisions are remembered for the session, scoped by default to the exact call (tool name plus arguments) so approving one call never silently approves a different one. Pass `permission_policy` to widen or narrow that scope.
## Rich tool rendering
By default the adapter recognizes the harness `FileSystem` and `Shell` tool calls by name and annotates them with an ACP `kind` (`read`/`edit`/`search`/`execute`), the file `locations` they touch, and an inline diff for edits -- so the editor renders click-to-file links and diff views instead of opaque JSON. Pass `tool_presenter` to add rendering for your own tools (optionally with `chain_presenters` ahead of the default), or `lambda _call: None` to disable it.
## MCP servers
An ACP client may offer MCP servers during session setup. This adapter does not connect them itself; a `session_config` is the place to turn `session.mcp_servers` into Pydantic AI toolsets. If a client sends MCP servers and no `session_config` is installed to consume them, the session request is rejected (rather than silently ignoring them) so the mismatch is visible. The spec expects every agent to accept stdio MCP servers, so an agent meant for arbitrary editors should install a `session_config` that connects them (for example with `pydantic_ai.mcp.MCPServerStdio`).
A spec-following client only sends HTTP/SSE MCP servers when the agent advertises support for those transports during `initialize` (stdio servers are not capability-gated). When your `session_config` connects them, say so:
```python
PydanticAIACPAgent(
agent,
session_config=connect_mcp_servers,
mcp_capabilities=schema.McpCapabilities(http=True, sse=True),
)
```
## Prompt content types
The agent advertises which prompt content it accepts. The default is **text only**, so a client is not invited to send blocks a text model cannot handle. Enable the kinds your model supports:
```python
from acp import schema
run_acp_stdio_sync(agent, prompt_capabilities=schema.PromptCapabilities(image=True, embedded_context=True))
```
## Session persistence
Pass a `session_store` to let a client reopen a past conversation with `session/load`. Each committed turn is persisted as two parts -- the model's message history and the client-visible transcript (the user's messages plus everything streamed back) -- and reopening restores the history into the agent and replays the transcript to the client, so its UI is rebuilt as the user last saw it. Without a store, `session/load` is advertised as unsupported.
```python
from pydantic_ai_harness.experimental.acp import InMemorySessionStore
run_acp_stdio_sync(agent, session_store=InMemorySessionStore())
```
`InMemorySessionStore` keeps sessions for the lifetime of the process. Implement the `SessionStore` protocol (`save`/`load` a `StoredSession`) over a file or database to make them survive a restart -- the stored values are Pydantic models, so they serialize with Pydantic.
Session persistence is for *reopening a conversation*; it is orthogonal to per-run durability. To also make individual turns crash-resilient (or resume a long sub-agent run), add a step-durability capability to the agent -- each ACP turn is one agent run, so the two layers compose with no glue.
## Model selection
Pass `models` to advertise a stable ACP session config option named `model` (using Pydantic AI [model names](https://ai.pydantic.dev/models/)). The first is each session's default. A selection is applied as a per-run override -- the shared agent is never mutated -- and is persisted with the session when a `session_store` is set.
```python
run_acp_stdio_sync(agent, models=['anthropic:claude-sonnet-4-6', 'anthropic:claude-opus-4-8', 'openai:gpt-4o'])
```
A model id is any string a Pydantic AI model accepts, so newer models not yet in `KnownModelName` work too. Pass `models='all'` to offer every model Pydantic AI knows (its default is then the first known model, so curate the list when you want a specific default). Without `models`, no model config option is advertised.
To advertise ids Pydantic AI's `infer_model` does not understand (for example OAuth or subscription models), pass `model_resolver` to map the selected id to a prebuilt `Model`; returning the id unchanged falls back to `infer_model`.
## Token usage
Each completed turn reports its token counts (input/output/total, plus cached tokens) on the ACP `PromptResponse`, summed across any approval pauses. This is an UNSTABLE ACP field, so clients that don't support it simply ignore it.
## Cancellation and limitations
- **Cancellation.** `session/cancel` and `session/close` cancel the in-flight turn; close waits for it to unwind before returning. Cooperative async tools stop promptly. A synchronous tool already running in a worker thread cannot be force-stopped, so its side effects may complete after the turn reports `cancelled` -- prefer async tools for cancellation-sensitive work.
- **Approval detection.** Tools that require approval are recognized when they live in a `FunctionToolset` (which the harness `FileSystem`/`Shell` and `@agent.tool` both use). A tool whose approval requirement is decided dynamically per call (by raising `ApprovalRequired` from its body) starts as `in_progress`, and any side effects it ran *before* raising have already happened by the time the client is asked -- use an `ApprovalRequiredToolset` (which gates before the tool body runs) for actions that must not partially execute before approval.
- **Overwrite diffs.** `write_file` renders an overwrite as if creating a new file (no prior contents), so the diff understates what it replaced.
- **Live terminal panes.** `acp_terminal` returns a command's captured output; it does not embed a *live* terminal pane in the tool call, which would need the terminal id at call-start, before the command runs.
- **Images.** Prompt image blocks are off by default and must be enabled via `prompt_capabilities` with a model that accepts them (see [Prompt content types](#prompt-content-types)). The harness `FileSystem.read_file` is text-only, so the agent cannot open image files from the workspace itself.
- **Slash commands.** The adapter does not yet advertise any commands (`available_commands`), so no slash commands appear in the client. Planned.
- **MCP servers.** Client-offered MCP servers are surfaced to your `session_config` to turn into toolsets (advertise the transports with `mcp_capabilities`); the adapter does not auto-connect them. Resource metadata is not yet wired end-to-end.
## API
```python
run_acp_stdio( # async; serve until the client disconnects
agent,
*,
deps=None,
name=None, # advertised name; defaults to the agent's name
version='0.1.0',
session_config=None, # per-session deps/toolsets from the client's setup
permission_policy=None, # scope of remembered "always" approval decisions
prompt_capabilities=None, # defaults to text-only
mcp_capabilities=None, # MCP transports to advertise; needs a session_config to connect them
tool_presenter=None, # defaults to the FileSystem/Shell presenter
session_store=None, # enables session/load by persisting each session
models=None, # models offered as the `model` config option ('all' for every known model)
model_resolver=None, # maps an advertised model id to the Model used for the run
usage_limits=None, # per-run request/token ceilings
)
run_acp_stdio_sync(...) # synchronous wrapper, same arguments
PydanticAIACPAgent(agent, *, ...) # the ACP agent object, to embed in a custom server
```
## Further reading
- [Agent Client Protocol](https://agentclientprotocol.com) -- protocol specification
- [Zed external agents](https://zed.dev/docs/ai/external-agents) -- editor-side configuration
- [Human-in-the-loop tool approval](https://pydantic.dev/docs/ai/tools-toolsets/deferred-tools/#human-in-the-loop-tool-approval) (Pydantic AI)
- [Pydantic AI capabilities](https://ai.pydantic.dev/capabilities/)
pydantic_ai_harness/experimental/acp/__init__.py
+48
View File
@@ -0,0 +1,48 @@
"""Expose a Pydantic AI agent over the Agent Client Protocol (ACP).
ACP lets terminal UIs and editors (such as Zed and Toad) drive a coding agent over stdio
JSON-RPC. [`PydanticAIACPAgent`][pydantic_ai_harness.experimental.acp.PydanticAIACPAgent] adapts a
Pydantic AI [`Agent`][pydantic_ai.Agent] to that interface, and
[`run_acp_stdio`][pydantic_ai_harness.experimental.acp.run_acp_stdio] serves it over stdio.
"""
from pydantic_ai_harness.experimental._warn import warn_experimental
from pydantic_ai_harness.experimental.acp._adapter import PydanticAIACPAgent
from pydantic_ai_harness.experimental.acp._client_toolsets import (
AcpFileSystemToolset,
AcpTerminalToolset,
acp_filesystem,
acp_terminal,
)
from pydantic_ai_harness.experimental.acp._permission import ToolCallPermission, default_permission_scope
from pydantic_ai_harness.experimental.acp._presentation import (
ToolCallPresentation,
chain_presenters,
default_coding_presenter,
)
from pydantic_ai_harness.experimental.acp._server import run_acp_stdio, run_acp_stdio_sync
from pydantic_ai_harness.experimental.acp._session import AcpSession, AcpSessionConfig, McpServer
from pydantic_ai_harness.experimental.acp._store import InMemorySessionStore, SessionStore, StoredSession
warn_experimental('acp')
__all__ = [
'AcpFileSystemToolset',
'AcpSession',
'AcpSessionConfig',
'AcpTerminalToolset',
'InMemorySessionStore',
'McpServer',
'PydanticAIACPAgent',
'SessionStore',
'StoredSession',
'ToolCallPermission',
'ToolCallPresentation',
'acp_filesystem',
'acp_terminal',
'chain_presenters',
'default_coding_presenter',
'default_permission_scope',
'run_acp_stdio',
'run_acp_stdio_sync',
]
pydantic_ai_harness/experimental/acp/_adapter.py
+951
View File
@@ -0,0 +1,951 @@
"""Adapt a Pydantic AI agent to the Agent Client Protocol (ACP) agent interface.
ACP (https://agentclientprotocol.com) lets a code editor or terminal UI (the *client*)
drive a coding agent (the *server*) over stdio JSON-RPC. This module exposes a Pydantic AI
[`Agent`][pydantic_ai.Agent] as such a server: it streams the agent's output to the client
as `session/update` notifications and bridges ACP permission requests to Pydantic AI's
human-in-the-loop tool approval.
"""
from __future__ import annotations
import asyncio
import contextlib
import json
import logging
from collections.abc import Callable, Hashable, Sequence
from dataclasses import dataclass, field
from inspect import isawaitable
from typing import Generic, Literal
from uuid import uuid4
import acp
import anyio
from acp import schema
from acp.interfaces import Client
from pydantic_ai import DeferredToolRequests, DeferredToolResults, ToolDenied, UsageLimitExceeded
from pydantic_ai.agent import AbstractAgent
from pydantic_ai.messages import (
AgentStreamEvent,
FinishReason,
FunctionToolCallEvent,
FunctionToolResultEvent,
ModelMessage,
PartDeltaEvent,
PartStartEvent,
RetryPromptPart,
TextPart,
TextPartDelta,
ThinkingPart,
ThinkingPartDelta,
ToolCallPart,
UserContent,
)
from pydantic_ai.models import KnownModelName, Model, known_model_names
from pydantic_ai.output import OutputDataT
from pydantic_ai.run import AgentRunResultEvent
from pydantic_ai.tools import AgentDepsT
from pydantic_ai.toolsets import AbstractToolset, FunctionToolset
from pydantic_ai.usage import RunUsage, UsageLimits
from pydantic_ai_harness.experimental.acp._content import PromptContentBlock, prompt_blocks_to_user_content
from pydantic_ai_harness.experimental.acp._permission import (
PermissionPolicy,
ToolCallPermission,
default_permission_scope,
)
from pydantic_ai_harness.experimental.acp._presentation import (
ToolCallContent,
ToolCallPresentation,
ToolCallPresenter,
absolutize,
default_coding_presenter,
)
from pydantic_ai_harness.experimental.acp._serialize import bounded_jsonable, chunk_text, jsonable
from pydantic_ai_harness.experimental.acp._session import (
AcpSession,
AcpSessionConfig,
McpServers,
SessionConfigFunc,
SessionState,
SessionUpdate,
)
from pydantic_ai_harness.experimental.acp._store import SessionStore, StoredSession
# Version advertised to the client when the caller does not supply one.
DEFAULT_VERSION = '0.1.0'
_logger = logging.getLogger(__name__)
_MODEL_CONFIG_ID = 'model'
def _all_known_model_names() -> tuple[str, ...]:
"""Every model name Pydantic AI exposes through the public `known_model_names()` API."""
return known_model_names()
def _finish_reason_to_stop_reason(finish_reason: FinishReason | None) -> schema.StopReason:
"""Map the model's terminal `finish_reason` to the ACP stop reason for a completed turn.
Only `length` and `content_filter` carry a distinct ACP meaning; every other reason (a normal
stop, a final tool call, an unknown/missing reason) is reported as a plain `end_turn`.
Cancellation is handled by the caller and never reaches here.
"""
if finish_reason == 'length':
return 'max_tokens'
if finish_reason == 'content_filter':
return 'refusal'
return 'end_turn'
def _usage_limit_stop_reason(exc: UsageLimitExceeded) -> schema.StopReason:
"""Map a run's exceeded usage limit to the ACP stop reason ending the turn.
Token-based limits report `max_tokens`; the request/tool-call count limits report
`max_turn_requests`. The exception carries no structured detail, so this reads its message; an
unrecognized wording falls back to `max_turn_requests` (the limit pydantic-ai's default
configuration can hit).
"""
return 'max_tokens' if 'tokens_limit' in str(exc) else 'max_turn_requests'
def _to_acp_usage(usage: RunUsage) -> schema.Usage:
"""Map a Pydantic AI `RunUsage` to ACP's per-turn `Usage` (an UNSTABLE field on `PromptResponse`)."""
return schema.Usage(
input_tokens=usage.input_tokens,
output_tokens=usage.output_tokens,
total_tokens=usage.total_tokens,
cached_read_tokens=usage.cache_read_tokens,
cached_write_tokens=usage.cache_write_tokens,
)
class _TurnCancelled(Exception):
"""Raised internally to unwind a turn the client cancelled mid-flight (e.g. via a permission dialog)."""
@dataclass(kw_only=True)
class _TurnState:
"""Per-turn state shared by the stream handlers, created once per turn in `_run_turn`."""
conn: Client
session_id: str
cwd: str
approval_names: frozenset[str]
# Tool-call ids accumulated across the turn's approval-resume passes: `started` so a call
# paused for approval is announced only once, `denied` so a rejected call's result is failed,
# `resulted` so a cancelled turn can fail only the calls still left without a terminal status.
started: set[str] = field(default_factory=set[str])
denied: set[str] = field(default_factory=set[str])
resulted: set[str] = field(default_factory=set[str])
# The `session/update`s sent this turn; appended to the session transcript only on commit.
updates: list[SessionUpdate] = field(default_factory=list[SessionUpdate])
@dataclass(frozen=True, kw_only=True)
class _ToolCallFields:
"""The ACP tool-call fields derived from the presenter, shared by the start and permission paths."""
title: str
kind: schema.ToolKind | None
content: list[ToolCallContent] | None
locations: list[schema.ToolCallLocation] | None
class PydanticAIACPAgent(acp.Agent, Generic[AgentDepsT, OutputDataT]):
"""An ACP agent backed by a Pydantic AI [`Agent`][pydantic_ai.Agent].
Each ACP session is an independent conversation, keyed by session ID. Pass the instance to
[`run_acp_stdio`][pydantic_ai_harness.experimental.acp.run_acp_stdio] (or the lower-level `acp.run_agent`)
to serve it over stdio. When calling `acp.run_agent` yourself, pass `use_unstable_protocol=True`
as `run_acp_stdio` does, or the SDK router rejects `session/close` (still UNSTABLE in the SDK)
with `method_not_found`.
A tool that [requires approval](https://pydantic.dev/docs/ai/tools-toolsets/deferred-tools/)
pauses the run and asks the client via `session/request_permission`; "always allow"/"always
reject" decisions are remembered for the rest of the session. Per-session workspace setup
(the client's `cwd` and MCP servers) is surfaced through the optional `session_config` factory.
`session/close` cancels any in-flight turn and discards the session. `session/load` (with a
`session_store`) and model selection via `session/set_config_option` (with `models`) are
supported when configured; fork, resume, and modes are advertised as unsupported. As with any
`output_type` override, this adapter is incompatible with agents that define output validators.
"""
def __init__(
self,
agent: AbstractAgent[AgentDepsT, OutputDataT],
*,
deps: AgentDepsT = None,
name: str | None = None,
version: str = DEFAULT_VERSION,
session_config: SessionConfigFunc[AgentDepsT] | None = None,
permission_policy: PermissionPolicy | None = None,
prompt_capabilities: schema.PromptCapabilities | None = None,
mcp_capabilities: schema.McpCapabilities | None = None,
tool_presenter: ToolCallPresenter | None = None,
session_store: SessionStore | None = None,
models: Sequence[KnownModelName | str] | Literal['all'] | None = None,
model_resolver: Callable[[str], Model | str] | None = None,
usage_limits: UsageLimits | None = None,
) -> None:
"""Build the adapter.
Args:
agent: The Pydantic AI agent to expose. Its model, tools, and instructions are used as-is.
deps: Dependencies passed to every agent run, equivalent to `Agent.run(..., deps=deps)`.
Used for sessions that `session_config` does not configure (or when it is not given).
name: Name advertised to the client. Defaults to the agent's name, then `'pydantic-ai-agent'`.
version: Version advertised to the client.
session_config: Optional factory called once per session with the client's
[`AcpSession`][pydantic_ai_harness.experimental.acp.AcpSession] setup (its `cwd`, MCP servers,
and capabilities). It returns an
[`AcpSessionConfig`][pydantic_ai_harness.experimental.acp.AcpSessionConfig] whose `deps` and
`toolsets` are applied to every run in that session. May be sync or async.
permission_policy: Optional function deciding the scope under which an "always
allow"/"always reject" decision is remembered. Defaults to the exact call (tool
name plus arguments).
prompt_capabilities: Prompt content the agent advertises support for (image, audio,
embedded context). Defaults to text-only; enable what the backing model handles,
e.g. `schema.PromptCapabilities(image=True, embedded_context=True)`.
mcp_capabilities: MCP transports the agent advertises support for, e.g.
`schema.McpCapabilities(http=True, sse=True)`. A spec-following client only sends
HTTP/SSE MCP servers when these are advertised (stdio servers are not gated), so
set this when the `session_config` connects them. Defaults to neither; requires a
`session_config`, since without one MCP servers are rejected at `session/new`.
tool_presenter: Maps a tool call to its rich ACP presentation (`kind`, file
`locations`, and diff `content`). Defaults to
[`default_coding_presenter`][pydantic_ai_harness.experimental.acp.default_coding_presenter],
which recognizes the `FileSystem`/`Shell` tools by name (a custom tool sharing a
name is also matched). Pass your own presenter (optionally chained ahead of the
default with [`chain_presenters`][pydantic_ai_harness.experimental.acp.chain_presenters]), or
`lambda _call: None` to disable rich rendering.
session_store: Optional [`SessionStore`][pydantic_ai_harness.experimental.acp.SessionStore] enabling
`session/load`: each committed turn is persisted (model history plus client
transcript) and a stored session can be reopened. Defaults to `None`, advertising
`session/load` as unsupported.
models: Optional models the client may switch between with the stable ACP session
config option `model`. The first is each session's default. A selection applies as
a per-run override (the shared agent is never mutated) and is persisted with the
session. Pass `'all'` to offer every model Pydantic AI knows (its default is then
the first known model, so the user should pick one). Defaults to `None`,
advertising no model config option.
model_resolver: Optional hook mapping an advertised model id to the `Model` (or model
string) used for that run, applied just before each run's per-run override. Lets a
host advertise ids Pydantic AI's `infer_model` does not understand (e.g. OAuth/
subscription models) and supply a pre-built `Model` for them. Returning the id
unchanged falls back to `infer_model`. Defaults to identity.
usage_limits: Optional per-run ceilings (`request_limit`, `tool_calls_limit`, token
limits) applied to every agent run, equivalent to `Agent.run(..., usage_limits=...)`.
A turn that resumes after a tool approval starts a fresh run, so the limits bound
each approval-to-approval segment, not the whole turn. An exceeded limit ends the
turn with a `max_tokens`/`max_turn_requests` stop reason, never a request error.
Defaults to Pydantic AI's own defaults (50 requests per run).
"""
self._agent = agent
self._deps = deps
self._name = name or agent.name or 'pydantic-ai-agent'
self._version = version
self._session_config = session_config
self._permission_policy = permission_policy or default_permission_scope
# Text-only by default, so a client is not invited to send image/audio/resource blocks
# the backing model may not accept.
self._prompt_capabilities = prompt_capabilities or schema.PromptCapabilities()
if mcp_capabilities is not None and session_config is None:
# Advertising HTTP/SSE MCP support would invite server definitions that
# `_build_config` then rejects; fail at construction instead of per session.
raise ValueError('mcp_capabilities requires a session_config that connects the MCP servers')
self._mcp_capabilities = mcp_capabilities or schema.McpCapabilities()
self._tool_presenter = tool_presenter or default_coding_presenter
self._session_store = session_store
self._models: tuple[str, ...] = _all_known_model_names() if models == 'all' else tuple(models) if models else ()
self._model_resolver = model_resolver
self._usage_limits = usage_limits
self._client_capabilities: schema.ClientCapabilities | None = None
self._conn: Client | None = None
# All live sessions, keyed by session id. Each owns its own turn lock (the dispatcher
# delivers requests concurrently), so turns within a session are serialized.
self._sessions: dict[str, SessionState[AgentDepsT]] = {}
def on_connect(self, conn: Client) -> None:
"""Capture the outbound connection used to stream updates and request permission."""
self._conn = conn
async def initialize(
self,
protocol_version: int,
client_capabilities: schema.ClientCapabilities | None = None,
client_info: schema.Implementation | None = None,
**kwargs: object,
) -> schema.InitializeResponse:
"""Negotiate the protocol version and advertise the agent's capabilities."""
self._client_capabilities = client_capabilities
# Echo a supported client version; otherwise answer with the latest version we speak so
# an out-of-range client does not negotiate an unsupported one.
if 1 <= protocol_version <= acp.PROTOCOL_VERSION:
negotiated_version = protocol_version
else:
negotiated_version = acp.PROTOCOL_VERSION
return schema.InitializeResponse(
protocol_version=negotiated_version,
agent_capabilities=schema.AgentCapabilities(
load_session=self._session_store is not None,
prompt_capabilities=self._prompt_capabilities,
mcp_capabilities=self._mcp_capabilities,
session_capabilities=schema.SessionCapabilities(close=schema.SessionCloseCapabilities()),
),
agent_info=schema.Implementation(name=self._name, version=self._version),
)
async def new_session(
self,
cwd: str,
additional_directories: list[str] | None = None,
mcp_servers: McpServers = None,
**kwargs: object,
) -> schema.NewSessionResponse:
"""Start a new session with an empty conversation history.
The session's `cwd` is the workspace the client opened: root the agent's tools at it via
`session_config` (see the package README) so file edits and tool-call locations line up
with that workspace. A request with MCP servers but no `session_config` is rejected (see
`_build_config`).
"""
# The spec requires `cwd` to be absolute, a MUST on the client; neither the SDK router
# nor this adapter re-validates it. `additional_directories` is
# part of the `acp.Agent` interface but not consumed: the capability is not advertised, so
# a conformant client never sends extra roots, and they are not surfaced to `session_config`.
session_id = uuid4().hex
config = await self._build_config(session_id, cwd, mcp_servers)
# The first configured model is the session default; `None` (no models) uses the agent's own.
default_model = self._models[0] if self._models else None
state = SessionState(session_id=session_id, config=config, cwd=cwd, model=default_model)
self._sessions[session_id] = state
# Persist the empty session so the client can reopen it even before its first turn.
await self._persist(state)
return schema.NewSessionResponse(session_id=session_id, config_options=self._model_config_options(state.model))
async def load_session(
self,
cwd: str,
session_id: str,
mcp_servers: McpServers = None,
additional_directories: list[str] | None = None,
**kwargs: object,
) -> schema.LoadSessionResponse | None:
"""Reopen a stored session: restore its model history and replay its transcript to the client.
Only called when a `session_store` was given (`session/load` is otherwise advertised as
unsupported). The run configuration is rebuilt from `cwd`/`mcp_servers` via
`session_config`, as in `new_session`.
Raises:
acp.RequestError: if no session with `session_id` is stored, or the stored session
cannot be read.
"""
if self._session_store is None:
# Advertised off, but the SDK router routes `session/load` regardless; a client
# calling it anyway gets the rejection the advertisement implies.
raise acp.RequestError.method_not_found('session/load')
if self._conn is None: # pragma: no cover - on_connect always runs before load_session
raise RuntimeError('load_session called before on_connect()')
try:
stored = await self._session_store.load(session_id)
except Exception as exc:
# A read or deserialization failure is a server-side durability problem, not a bad
# client request: surface a clear, purpose-built error rather than leaking the store's
# low-level exception (a corrupt payload would otherwise reach the client as raw
# pydantic validation detail).
raise acp.RequestError.internal_error(
{'session_id': session_id, 'reason': 'stored session could not be read'}
) from exc
if stored is None:
raise acp.RequestError.invalid_params(
{'session_id': session_id, 'reason': 'no stored session with this id'}
)
# A client may load a session id that is still open (a reconnecting editor, or a double
# load). Tear down any live turn first so the orphaned turn cannot later persist its stale
# state over the transcript and history we are about to restore.
prior = self._sessions.pop(session_id, None)
if prior is not None:
await self._cancel_active_turn(prior)
config = await self._build_config(session_id, cwd, mcp_servers)
state = SessionState(
session_id=session_id,
config=config,
cwd=cwd,
history=list(stored.messages),
transcript=list(stored.updates),
model=stored.model,
)
self._sessions[session_id] = state
for update in stored.updates:
await self._conn.session_update(session_id=session_id, update=update)
return schema.LoadSessionResponse(config_options=self._model_config_options(state.model))
async def _persist(self, state: SessionState[AgentDepsT]) -> None:
"""Save a session's committed state (history, transcript, selected model), if a store is configured.
A store failure is logged and swallowed rather than failing the operation that triggered the
save: the turn (or session) has already streamed and committed in memory, so a durable-write
error must not surface as a failure for work the user already saw succeed. The session stays
usable and the next successful save catches the store back up.
"""
if self._session_store is None:
return
try:
await self._session_store.save(
state.session_id,
StoredSession(messages=list(state.history), updates=list(state.transcript), model=state.model),
)
except Exception:
_logger.exception('failed to persist ACP session %s; durable state is now behind', state.session_id)
def _model_option(self, current_model_id: str | None) -> schema.SessionConfigOptionSelect | None:
"""The model config option advertised to the client, or `None` when none is configured."""
if current_model_id is None or not self._models:
return None
return schema.SessionConfigOptionSelect(
id=_MODEL_CONFIG_ID,
name='Model',
type='select',
current_value=current_model_id,
options=[schema.SessionConfigSelectOption(value=model, name=model) for model in self._models],
)
def _model_config_options(
self, current_model_id: str | None
) -> list[schema.SessionConfigOptionSelect | schema.SessionConfigOptionBoolean] | None:
"""The full session config option list, or `None` when model switching is not configured."""
option = self._model_option(current_model_id)
return None if option is None else [option]
def session_history(self, session_id: str) -> list[ModelMessage] | None:
"""A snapshot of a resident session's committed model history (what the next turn will send).
Returns a shallow copy: the list is the caller's to keep, but the `ModelMessage` objects are
shared with the live session -- treat them as read-only. `None` when the session is not
resident in this adapter.
"""
state = self._sessions.get(session_id)
return list(state.history) if state is not None else None
def session_model_option(self, session_id: str) -> schema.SessionConfigOptionSelect | None:
"""The model config option for a resident session -- its option set + current selection.
The public counterpart to the `config_options` entry `new_session`/`load_session` return,
for callers that re-surface model selection on a later interaction (e.g. an attach/resume
reply) without re-creating or reloading the session. `None` when the session is unknown to
this adapter or no switch-set was configured.
"""
state = self._sessions.get(session_id)
return self._model_option(state.model) if state is not None else None
async def _build_config(self, session_id: str, cwd: str, mcp_servers: McpServers) -> AcpSessionConfig[AgentDepsT]:
"""Derive a session's run configuration, calling the `session_config` factory if present.
Raises:
acp.RequestError: if the client provides MCP servers but no `session_config` is
installed to connect them.
"""
if self._session_config is None:
if mcp_servers:
raise acp.RequestError.invalid_params(
{
'reason': 'this agent does not connect MCP servers; provide a session_config '
'that turns mcp_servers into toolsets to use them',
'mcp_server_count': len(mcp_servers),
}
)
return AcpSessionConfig(deps=self._deps)
if self._conn is None: # pragma: no cover - on_connect always runs before session setup
raise RuntimeError('_build_config called before on_connect()')
session = AcpSession(
cwd=cwd,
mcp_servers=mcp_servers or [],
client_capabilities=self._client_capabilities,
client=self._conn,
session_id=session_id,
)
result = self._session_config(session)
return await result if isawaitable(result) else result
async def prompt(
self,
session_id: str,
prompt: list[PromptContentBlock],
**kwargs: object,
) -> schema.PromptResponse:
"""Run one user turn, streaming the agent's output to the client.
Turns for a session are serialized. The turn runs in a cancellable task so a concurrent
`session/cancel` (or a cancelled permission dialog) stops it promptly with a `cancelled`
stop reason; a cancelled turn's prompt and partial output are not committed to the
session history.
"""
if self._conn is None: # pragma: no cover - on_connect always runs first in practice
raise RuntimeError('prompt() called before on_connect()')
state = self._sessions.get(session_id)
if state is None:
raise acp.RequestError.invalid_params({'session_id': session_id})
user_content = prompt_blocks_to_user_content(prompt)
async with state.lock:
if self._sessions.get(session_id) is not state:
# The session was closed or reloaded while this prompt waited its turn on the
# lock. Running now would resurrect the discarded state: `cancel` could no longer
# reach the turn, and its commit would persist the orphaned history over the
# live session.
raise acp.RequestError.invalid_params(
{'session_id': session_id, 'reason': 'session was closed while the prompt was queued'}
)
state.cancel_requested = False
turn = asyncio.ensure_future(self._run_turn(state, prompt, user_content))
state.active_turn = turn
try:
# Shielded so this coroutine's *own* cancellation (connection teardown) surfaces
# here instead of propagating into the turn: with a bare `await turn` the two are
# indistinguishable, and answering a request whose handler the connection already
# cancelled would deadlock its shutdown (the response future is never resolved).
return await asyncio.shield(turn)
except _TurnCancelled:
return schema.PromptResponse(stop_reason='cancelled')
except asyncio.CancelledError:
# `turn.done()` cannot fully prove the cancellation came from the turn: in the one
# tick between the turn completing and this coroutine resuming, a teardown cancel
# aimed at *this* handler is indistinguishable on 3.10 (3.11's `Task.cancelling()`
# would settle it). Accepted residual; every wider window is covered below.
if turn.done() and state.cancel_requested:
return schema.PromptResponse(stop_reason='cancelled')
if state.cancel_requested:
# `cancel()` already delivered cancellation to the turn. Do not inject a second
# one while Pydantic AI's stream context is closing, or its internal event stream
# can be interrupted before its receive side closes.
with contextlib.suppress(asyncio.CancelledError, _TurnCancelled):
await asyncio.shield(turn)
raise
# The prompt coroutine itself was cancelled (e.g. connection teardown): stop the
# child turn before propagating, so it is not left running on a closing connection.
turn.cancel()
with contextlib.suppress(asyncio.CancelledError, _TurnCancelled):
await turn
raise
finally:
state.active_turn = None
state.cancel_requested = False
async def cancel(self, session_id: str, **kwargs: object) -> None:
"""Cancel the in-flight turn for a session, if any."""
state = self._sessions.get(session_id)
if state is None:
return
state.cancel_requested = True
if state.active_turn is not None and not state.active_turn.done():
state.active_turn.cancel()
async def close_session(self, session_id: str, **kwargs: object) -> schema.CloseSessionResponse | None:
"""Cancel any in-flight turn and discard all state for the session.
Closing an unknown (or already-closed) session is a no-op.
"""
state = self._sessions.pop(session_id, None)
if state is not None:
await self._cancel_active_turn(state)
return schema.CloseSessionResponse()
async def _cancel_active_turn(self, state: SessionState[AgentDepsT]) -> None:
"""End a session's in-flight turn (if any) with a `cancelled` stop reason and await its unwind.
Awaiting the turn keeps the caller from racing the turn's teardown -- the prompt has fully
returned `cancelled` and released the session's resources before this returns.
"""
turn = state.active_turn
if turn is not None and not turn.done():
state.cancel_requested = True
turn.cancel()
try:
# Shielded for the same reason as in `prompt`: only the turn's own unwind may be
# swallowed here, never this coroutine's cancellation (see the comment there).
await asyncio.shield(turn)
except (asyncio.CancelledError, _TurnCancelled):
# As in `prompt`, `turn.done()` is a heuristic with a one-tick residual: a
# teardown cancel landing between the turn completing and this resuming is
# swallowed here on 3.10. Accepted.
if not turn.done():
raise
async def _run_turn(
self,
state: SessionState[AgentDepsT],
prompt: list[PromptContentBlock],
user_content: list[UserContent],
) -> schema.PromptResponse:
"""Drive the agent (resuming across tool-approval pauses), stream updates, and build the response.
The response carries the ACP stop reason mapped from the final model response, plus the
turn's total token usage (summed across every run pass, one per approval pause). A turn
ended by a usage limit rolls back like a cancellation, so its response omits usage while
still answering with the limit's `max_tokens`/`max_turn_requests` stop reason.
"""
conn = self._conn
assert conn is not None
history = state.history
config = state.config
usage = RunUsage()
stop_reason: schema.StopReason = 'end_turn'
deferred_results: DeferredToolResults | None = None
run_input: list[UserContent] | None = user_content
output_type = [self._agent.output_type, DeferredToolRequests]
turn = _TurnState(
conn=conn, session_id=state.session_id, cwd=state.cwd, approval_names=self._approval_tool_names(config)
)
# Recorded for replay only, never sent live: the client renders its own prompt during
# the turn, but `session/load` must rebuild the user side of the conversation too.
# Going through `turn.updates` keeps the cancel semantics: a rolled-back turn does not
# persist its user message either.
turn.updates.extend(acp.update_user_message(block) for block in prompt)
try:
while True:
result = None
async with self._agent.run_stream_events(
run_input,
message_history=history,
deferred_tool_results=deferred_results,
output_type=output_type,
deps=config.deps,
toolsets=config.toolsets,
# Per-run override for the client's model config choice; `None` uses the
# agent's own model, never mutating the shared agent. A `model_resolver` (if
# given) maps the advertised id to a pre-built `Model` for ids `infer_model`
# can't parse.
model=self._resolve_run_model(state.model),
usage_limits=self._usage_limits,
) as stream:
event: AgentStreamEvent | AgentRunResultEvent[object]
async for event in stream:
if isinstance(event, AgentRunResultEvent):
result = event.result
else:
await self._emit_event(turn, event)
assert result is not None, 'run_stream_events always yields a final result event'
history = result.all_messages()
usage += result.usage # pydantic-ai 2.0: `usage` is a property, not a method
output = result.output
if isinstance(output, DeferredToolRequests):
if not output.approvals and not output.calls: # pragma: no cover
break # defensive: the run never yields an empty DeferredToolRequests
deferred_results = await self._resolve_approvals(turn, state, output)
run_input = None
continue
if output is not None and not isinstance(output, str):
# Structured output isn't streamed as text parts, so deliver it as a final message.
# (`str` output already streamed via text deltas; `None` has nothing to show.)
await self._emit_text(turn, json.dumps(jsonable(output)), thought=False)
stop_reason = _finish_reason_to_stop_reason(result.response.finish_reason)
break
except (asyncio.CancelledError, _TurnCancelled):
# The turn is ending without finishing its tool calls; close out any still shown as
# pending/in_progress so the client stops rendering them as running. Shielded because an
# asyncio cancellation would otherwise abort the sends, and live-only: a cancelled turn
# never commits its transcript.
with anyio.CancelScope(shield=True):
await self._fail_outstanding_tool_calls(turn)
raise
except UsageLimitExceeded as exc:
# ACP models a turn ending at a limit as a normal stop reason, not a request error.
# The raising run's partial messages are not retrievable, so nothing is committed --
# the turn rolls back to the prior state, like a cancellation, and the response must
# not report uncommitted usage.
await self._fail_outstanding_tool_calls(turn)
return schema.PromptResponse(stop_reason=_usage_limit_stop_reason(exc))
except Exception:
# The turn is failing with the error the client receives as the prompt's response;
# close out its announced tool calls so they are not left rendering as running.
await self._fail_outstanding_tool_calls(turn)
raise
state.history = history
# Commit and persist this turn's updates alongside the history; a cancelled turn never
# reaches here, so only committed turns are persisted.
if self._session_store is not None:
state.transcript.extend(turn.updates)
try:
await self._persist(state)
except asyncio.CancelledError:
# The turn is already committed: a cancel landing inside the store's save came
# too late to roll anything back. The spec still requires the prompt to answer
# `cancelled`, so that stop reason is reported alongside committed usage. The
# interrupted save is the same benign failure `_persist` swallows; the next
# successful save catches the store up.
_logger.warning(
'persisting ACP session %s was cancelled; durable state is now behind', state.session_id
)
stop_reason = 'cancelled'
return schema.PromptResponse(stop_reason=stop_reason, usage=_to_acp_usage(usage))
async def _fail_outstanding_tool_calls(self, turn: _TurnState) -> None:
"""Drive every announced-but-unfinished tool call to a terminal `failed` status.
Called when a turn is cancelled, so a client does not keep rendering a `pending`/`in_progress`
tool call as running. Sent directly (not recorded for the transcript) and best-effort: the
connection may already be going away.
"""
for tool_call_id in turn.started - turn.resulted:
with contextlib.suppress(Exception):
await turn.conn.session_update(
session_id=turn.session_id,
update=acp.update_tool_call(tool_call_id=tool_call_id, status='failed'),
)
async def _send_update(self, turn: _TurnState, update: SessionUpdate) -> None:
"""Send one `session/update` to the client, recording it for the transcript."""
turn.updates.append(update)
await turn.conn.session_update(session_id=turn.session_id, update=update)
async def _emit_text(self, turn: _TurnState, text: str, *, thought: bool) -> None:
"""Stream text to the client as one or more chunked `session/update` notifications."""
for chunk in chunk_text(text):
update = acp.update_agent_thought_text(chunk) if thought else acp.update_agent_message_text(chunk)
await self._send_update(turn, update)
def _tool_call_fields(
self, call: ToolCallPart, cwd: str, *, default_kind: schema.ToolKind | None
) -> _ToolCallFields:
"""Derive the ACP tool-call fields shared by the start and permission paths via the presenter.
Falls back to the tool name for the title, and to `default_kind` when the presenter
offers no `kind`.
"""
presentation = self._tool_presenter(call)
presentation = absolutize(presentation, cwd) if presentation is not None else ToolCallPresentation()
return _ToolCallFields(
title=presentation.title or call.tool_name,
kind=presentation.kind or default_kind,
content=list(presentation.content) or None,
locations=list(presentation.locations) or None,
)
async def _emit_event(self, turn: _TurnState, event: AgentStreamEvent) -> None:
"""Translate a single Pydantic AI stream event into an ACP `session/update`."""
if isinstance(event, PartStartEvent):
part = event.part
if isinstance(part, TextPart) and part.content:
await self._emit_text(turn, part.content, thought=False)
elif isinstance(part, ThinkingPart) and part.content:
await self._emit_text(turn, part.content, thought=True)
elif isinstance(event, PartDeltaEvent):
delta = event.delta
if isinstance(delta, TextPartDelta) and delta.content_delta:
await self._emit_text(turn, delta.content_delta, thought=False)
elif isinstance(delta, ThinkingPartDelta) and delta.content_delta:
await self._emit_text(turn, delta.content_delta, thought=True)
elif isinstance(event, FunctionToolCallEvent):
call = event.part
if call.tool_call_id in turn.started:
return # the model re-emits the call event when resuming after approval
turn.started.add(call.tool_call_id)
fields = self._tool_call_fields(call, turn.cwd, default_kind=None)
# A call awaiting approval is not running yet: it starts `pending` and
# `_resolve_approvals` promotes it once approved. Any other call is already executing.
status: schema.ToolCallStatus = 'pending' if call.tool_name in turn.approval_names else 'in_progress'
await self._send_update(
turn,
acp.start_tool_call(
tool_call_id=call.tool_call_id,
title=fields.title,
kind=fields.kind,
status=status,
content=fields.content,
locations=fields.locations,
raw_input=bounded_jsonable(call.args),
),
)
elif isinstance(event, FunctionToolResultEvent):
part = event.part
if isinstance(part, RetryPromptPart):
status, raw_output = 'failed', part.model_response()
elif part.tool_call_id in turn.denied:
status, raw_output = 'failed', part.content
else:
status, raw_output = 'completed', part.content
turn.resulted.add(part.tool_call_id)
await self._send_update(
turn,
acp.update_tool_call(
tool_call_id=part.tool_call_id,
status=status,
raw_output=bounded_jsonable(raw_output),
),
)
async def _resolve_approvals(
self, turn: _TurnState, state: SessionState[AgentDepsT], requests: DeferredToolRequests
) -> DeferredToolResults:
"""Ask the client to approve each pending tool call and collect the decisions.
Records rejected call IDs in `turn.denied` so the eventual result event is reported as failed.
Raises:
_TurnCancelled: if the client cancels a permission request.
acp.RequestError: if the agent requests external tool execution, which is unsupported.
"""
if requests.calls:
names = sorted({call.tool_name for call in requests.calls})
raise acp.RequestError.internal_error(
{'reason': 'external tool execution is not supported by the ACP adapter', 'tools': names}
)
results = DeferredToolResults()
for call in requests.approvals:
# `args_as_dict()` canonicalizes the call's arguments to a dict: a model may deliver them
# as a JSON string (the OpenAI default, and how streamed calls accumulate), and a raw
# string would make the scope key sensitive to key order, defeating a remembered "always"
# decision for what is the same logical call.
scope = self._permission_policy(
ToolCallPermission(tool_name=call.tool_name, tool_call_id=call.tool_call_id, args=call.args_as_dict())
)
if scope in state.always_allow:
results.approvals[call.tool_call_id] = True
await self._mark_running(turn, call.tool_call_id)
continue
if scope in state.always_reject:
results.approvals[call.tool_call_id] = ToolDenied('Rejected by the client.')
turn.denied.add(call.tool_call_id)
continue
# Only an approved call is promoted to `in_progress`; a rejected one stays `pending`
# until its `failed` result, never shown as running.
decision = await self._request_permission(turn, state, call, scope)
results.approvals[call.tool_call_id] = decision
if isinstance(decision, ToolDenied):
turn.denied.add(call.tool_call_id)
else:
await self._mark_running(turn, call.tool_call_id)
return results
async def _mark_running(self, turn: _TurnState, tool_call_id: str) -> None:
"""Promote an approved tool call from `pending` to `in_progress`, just before it executes."""
await self._send_update(turn, acp.update_tool_call(tool_call_id=tool_call_id, status='in_progress'))
def _approval_tool_names(self, config: AcpSessionConfig[AgentDepsT]) -> frozenset[str]:
"""Names of the tools that pause for the client's approval, announced `pending` in `_emit_event`.
Only `FunctionToolset`-held tools expose `requires_approval` without a live run context;
tools from other toolset types are treated as not requiring approval.
"""
toolsets: list[AbstractToolset[AgentDepsT]] = [*self._agent.toolsets, *(config.toolsets or [])]
names: set[str] = set()
for toolset in toolsets:
if isinstance(toolset, FunctionToolset):
names.update(name for name, tool in toolset.tools.items() if tool.requires_approval)
return frozenset(names)
async def _request_permission(
self, turn: _TurnState, state: SessionState[AgentDepsT], call: ToolCallPart, scope: Hashable
) -> bool | ToolDenied:
"""Ask the client to approve a single tool call, remembering "always" decisions by `scope`."""
fields = self._tool_call_fields(call, turn.cwd, default_kind='execute')
response = await turn.conn.request_permission(
session_id=turn.session_id,
tool_call=schema.ToolCallUpdate(
tool_call_id=call.tool_call_id,
title=fields.title,
kind=fields.kind,
status='pending',
content=fields.content,
locations=fields.locations,
raw_input=bounded_jsonable(call.args),
),
options=[
schema.PermissionOption(kind='allow_once', name='Allow', option_id='allow_once'),
schema.PermissionOption(kind='allow_always', name='Always allow', option_id='allow_always'),
schema.PermissionOption(kind='reject_once', name='Reject', option_id='reject_once'),
schema.PermissionOption(kind='reject_always', name='Always reject', option_id='reject_always'),
],
)
outcome = response.outcome
if isinstance(outcome, schema.DeniedOutcome):
# ACP signals a cancelled turn via a "cancelled" permission outcome.
raise _TurnCancelled
if outcome.option_id == 'allow_always':
state.always_allow.add(scope)
elif outcome.option_id == 'reject_always':
state.always_reject.add(scope)
if outcome.option_id in ('allow_once', 'allow_always'):
return True
return ToolDenied('Rejected by the client.')
# --- Optional ACP methods not supported by this adapter ------------------------------
# The capabilities for these are advertised as off in `initialize`, but the SDK router still
# routes the methods here (`acp.Agent` is a Protocol whose inherited stub bodies would answer
# success-`null`), so these raises are what a client calling them anyway actually receives.
async def authenticate(self, method_id: str, **kwargs: object) -> schema.AuthenticateResponse | None:
"""No authentication is required, so this is a no-op."""
return None
async def list_sessions(
self,
cwd: str | None = None,
cursor: str | None = None,
**kwargs: object,
) -> schema.ListSessionsResponse:
raise acp.RequestError.method_not_found('session/list')
def _resolve_run_model(self, model_id: str | None) -> Model | str | None:
"""Map an advertised model id to the per-run override, applying `model_resolver` if set.
`None` (no advertised models) is passed through so the run uses the agent's own model.
"""
if model_id is None or self._model_resolver is None:
return model_id
return self._model_resolver(model_id)
async def set_session_mode(
self, session_id: str, mode_id: str, **kwargs: object
) -> schema.SetSessionModeResponse | None:
raise acp.RequestError.method_not_found('session/set_mode')
async def set_config_option(
self, config_id: str, session_id: str, value: str | bool, **kwargs: object
) -> schema.SetSessionConfigOptionResponse | None:
"""Switch the session's model through ACP's stable session config option surface."""
state = self._sessions.get(session_id)
if state is None:
raise acp.RequestError.invalid_params({'session_id': session_id})
if config_id != _MODEL_CONFIG_ID:
raise acp.RequestError.invalid_params({'config_id': config_id, 'reason': 'unknown config option'})
if not isinstance(value, str) or value not in self._models:
raise acp.RequestError.invalid_params({'model_id': value, 'reason': 'not an advertised model'})
state.model = value
await self._persist(state)
options = self._model_config_options(state.model)
assert options is not None
return schema.SetSessionConfigOptionResponse(config_options=options)
async def fork_session(
self,
session_id: str,
cwd: str,
additional_directories: list[str] | None = None,
mcp_servers: McpServers = None,
**kwargs: object,
) -> schema.ForkSessionResponse:
raise acp.RequestError.method_not_found('session/fork')
async def resume_session(
self,
session_id: str,
cwd: str,
additional_directories: list[str] | None = None,
mcp_servers: McpServers = None,
**kwargs: object,
) -> schema.ResumeSessionResponse:
raise acp.RequestError.method_not_found('session/resume')
async def ext_method(self, method: str, params: dict[str, object]) -> dict[str, object]:
raise acp.RequestError.method_not_found(method)
async def ext_notification(self, method: str, params: dict[str, object]) -> None:
return None
pydantic_ai_harness/experimental/acp/_client_toolsets.py
+233
View File
@@ -0,0 +1,233 @@
"""Filesystem and shell tools that route through the ACP client instead of local disk/processes.
The local [`FileSystem`][pydantic_ai_harness.FileSystem] and [`Shell`][pydantic_ai_harness.Shell]
capabilities operate on the agent process's own disk and subprocesses. In an editor, that misses
the source of truth: unsaved buffers, the file layout the editor (not the launching shell)
considers the workspace, and -- for a remote or containerized editor -- the machine the code
actually lives on. ACP lets the agent ask the *client* to do the I/O: `fs/read_text_file` /
`fs/write_text_file` for files, and the terminal lifecycle (`terminal/create`, `terminal/output`,
`terminal/wait_for_exit`, `terminal/release`) for commands.
[`AcpFileSystemToolset`][pydantic_ai_harness.experimental.acp.AcpFileSystemToolset] and
[`AcpTerminalToolset`][pydantic_ai_harness.experimental.acp.AcpTerminalToolset] are those editor-native
counterparts. Build them per session with
[`acp_filesystem`][pydantic_ai_harness.experimental.acp.acp_filesystem] /
[`acp_terminal`][pydantic_ai_harness.experimental.acp.acp_terminal], which return the toolset only when the
client advertised the matching capability and otherwise return `None` so the caller can fall back
to the local capability.
"""
# A read-only ACP client gets editor-native reads with writes delegated to the local `FileSystem`
# capability; otherwise these toolsets are self-contained. A future shared I/O-backend seam is
# expected to subsume them.
from __future__ import annotations
import asyncio
import contextlib
import os.path
from collections.abc import Awaitable
from typing import Protocol
import anyio
from acp import Client, schema
from pydantic_ai.tools import AgentDepsT
from pydantic_ai.toolsets import FunctionToolset
from pydantic_ai_harness.experimental.acp._session import AcpSession
from pydantic_ai_harness.filesystem import FileSystem
class _LocalFileWriter(Protocol):
"""Something that can write a file on the local disk -- structurally satisfied by `FileSystemToolset`."""
def write_file(self, path: str, content: str) -> Awaitable[str]: ... # pragma: no cover - structural protocol
class AcpFileSystemToolset(FunctionToolset[AgentDepsT]):
"""`read_file`/`write_file` tools backed by the ACP client connection.
Each call invokes the client's `fs/read_text_file` / `fs/write_text_file`, so the agent sees
the editor's live view of the workspace (including unsaved buffers). The tool names match the
local `FileSystem` capability, so the default presenter renders them the same way.
ACP requires absolute paths, but models routinely produce workspace-relative ones (the local
`FileSystem` tools take them, and the same agent may run against either backend): when `cwd`
is set, relative paths are resolved against it before reaching the wire. The client still
resolves and authorizes every path itself (this toolset adds no sandboxing of its own).
If `local_writer` is set (a read-only client -- see [`acp_filesystem`][pydantic_ai_harness.experimental.acp.acp_filesystem]),
`write_file` goes there instead of to the client, while reads still route through the editor.
"""
def __init__(
self, *, client: Client, session_id: str, cwd: str | None = None, local_writer: _LocalFileWriter | None = None
) -> None:
super().__init__()
self._client = client
self._session_id = session_id
self._cwd = cwd
self._local_writer = local_writer
self.add_function(self.read_file, name='read_file')
self.add_function(self.write_file, name='write_file')
def _absolute(self, path: str) -> str:
if self._cwd is None or os.path.isabs(path):
return path
return os.path.normpath(os.path.join(self._cwd, path))
async def read_file(self, path: str) -> str:
"""Read a text file's contents through the editor.
Args:
path: Path to the file; resolved against the session workspace when relative.
"""
response = await self._client.read_text_file(path=self._absolute(path), session_id=self._session_id)
return response.content
async def write_file(self, path: str, content: str) -> str:
"""Write a text file's full contents through the editor.
Args:
path: Path to the file; resolved against the session workspace when relative.
content: The complete new contents of the file.
"""
path = self._absolute(path)
if self._local_writer is not None:
return await self._local_writer.write_file(path, content)
await self._client.write_text_file(content=content, path=path, session_id=self._session_id)
return f'Wrote {path} ({len(content)} characters).'
def acp_filesystem(session: AcpSession) -> AcpFileSystemToolset[None] | None:
"""Build an ACP-client-backed filesystem toolset for `session`, or `None` if unsupported.
Returns an [`AcpFileSystemToolset`][pydantic_ai_harness.experimental.acp.AcpFileSystemToolset] whenever the
client advertised `fs/read_text_file` during `initialize`:
- read + write advertised: reads and writes both route through the editor.
- read only (no `fs/write_text_file`): reads route through the editor, while writes go to the
local [`FileSystem`][pydantic_ai_harness.FileSystem] rooted at `session.cwd`. This is coherent
only when the agent shares the workspace disk with the editor (same machine, or an agent
running inside the editor's container) -- for a *remote* editor the writes land on the agent's
disk, not the editor's.
Returns `None` only when the client advertised no readable filesystem, so the caller can fall
back to a fully local toolset:
```python
def session_config(session: AcpSession) -> AcpSessionConfig[None]:
fs = acp_filesystem(session) or FileSystem(root_dir=session.cwd).get_toolset()
return AcpSessionConfig(deps=None, toolsets=[fs])
```
For an agent with non-`None` deps, construct `AcpFileSystemToolset[YourDeps](...)` directly
(the toolset ignores deps); this helper covers the common no-deps case.
"""
capabilities = session.client_capabilities
fs = capabilities.fs if capabilities is not None else None
if fs is None or not fs.read_text_file:
return None
local_writer = None if fs.write_text_file else FileSystem(root_dir=session.cwd).get_toolset()
return AcpFileSystemToolset[None](
client=session.client, session_id=session.session_id, cwd=session.cwd, local_writer=local_writer
)
def _format_terminal_output(result: schema.TerminalOutputResponse) -> str:
"""Render a terminal's captured output plus a trailing note for truncation or non-success exit."""
parts = [result.output]
if result.truncated:
parts.append('[output truncated]')
status = result.exit_status
if status is not None and status.exit_code not in (None, 0):
parts.append(f'[exited with code {status.exit_code}]')
elif status is not None and status.signal is not None:
parts.append(f'[terminated by signal {status.signal}]')
return '\n'.join(parts)
class AcpTerminalToolset(FunctionToolset[AgentDepsT]):
"""A `run_command` tool backed by the ACP client's terminal lifecycle.
Each call asks the client to create a terminal, waits for it to exit, reads its output, and
releases it, so the command runs in the editor's environment rather than as a local
subprocess of the agent. The tool name matches the local `Shell` capability so the default
presenter renders it as an `execute` call. If the call is cancelled while the command is
running, the terminal is killed and released.
"""
def __init__(self, *, client: Client, session_id: str, cwd: str | None = None) -> None:
super().__init__()
self._client = client
self._session_id = session_id
self._cwd = cwd
self.add_function(self.run_command, name='run_command')
# Embedding a live terminal pane in the tool call would require the terminal id at
# call-start, before the command runs; the captured output is returned instead.
async def run_command(self, command: str) -> str:
"""Run a shell command in the editor's terminal and return its captured output.
Args:
command: The shell command line to execute.
"""
# The create runs as its own task so a cancellation landing mid-flight cannot abandon the
# request: it may already be on the wire (the client then starts the command regardless),
# so its response must still be read to learn the id and clean up. A raw `task.cancel()`
# (how the adapter and pydantic-ai deliver cancellation) pierces anyio shields, so the
# create must live outside this task; `asyncio.wait` rather than `asyncio.shield` because
# shield on 3.12+ reports a late create failure to the loop exception handler even when
# the cleanup below retrieves it.
create = asyncio.ensure_future(
self._client.create_terminal(command=command, session_id=self._session_id, cwd=self._cwd)
)
terminal_id: str | None = None
try:
await asyncio.wait([create])
terminal_id = create.result().terminal_id
await self._client.wait_for_terminal_exit(session_id=self._session_id, terminal_id=terminal_id)
result = await self._client.terminal_output(session_id=self._session_id, terminal_id=terminal_id)
return _format_terminal_output(result)
except asyncio.CancelledError:
# Kill the still-running terminal before unwinding, shielded so the cleanup completes
# even though this task is being cancelled (when the cancellation landed during the
# create, the id is learned from the still-running create first). Suppress failures: a
# client error here must not replace the `CancelledError` the caller needs to see
# (the spec requires the turn to end with a `cancelled` stop reason).
with anyio.CancelScope(shield=True):
if terminal_id is None:
with contextlib.suppress(Exception):
terminal_id = (await create).terminal_id
if terminal_id is not None:
with contextlib.suppress(Exception):
await self._client.kill_terminal(session_id=self._session_id, terminal_id=terminal_id)
raise
finally:
# Always release the terminal (if one came into existence); suppress failures so a
# release error never masks the exception (or successful return) already in flight.
if terminal_id is not None:
with anyio.CancelScope(shield=True), contextlib.suppress(Exception):
await self._client.release_terminal(session_id=self._session_id, terminal_id=terminal_id)
def acp_terminal(session: AcpSession) -> AcpTerminalToolset[None] | None:
"""Build an ACP-client-backed terminal toolset for `session`, or `None` if unsupported.
Returns an [`AcpTerminalToolset`][pydantic_ai_harness.experimental.acp.AcpTerminalToolset] only when the
client advertised terminal support during `initialize`; otherwise returns `None` so the caller
can fall back to a local toolset:
```python
def session_config(session: AcpSession) -> AcpSessionConfig[None]:
shell = acp_terminal(session) or Shell(cwd=session.cwd).get_toolset()
return AcpSessionConfig(deps=None, toolsets=[shell])
```
For an agent with non-`None` deps, construct `AcpTerminalToolset[YourDeps](...)` directly (the
toolset ignores deps); this helper covers the common no-deps case.
"""
capabilities = session.client_capabilities
if capabilities is None or not capabilities.terminal:
return None
return AcpTerminalToolset[None](client=session.client, session_id=session.session_id, cwd=session.cwd)
pydantic_ai_harness/experimental/acp/_content.py
+60
View File
@@ -0,0 +1,60 @@
"""Translation between ACP content blocks and Pydantic AI user content."""
from __future__ import annotations
import base64
from collections.abc import Sequence
from acp import schema
from pydantic_ai.messages import BinaryContent, ImageUrl, UserContent
# The content block variants ACP may send in a `session/prompt` request.
PromptContentBlock = (
schema.TextContentBlock
| schema.ImageContentBlock
| schema.AudioContentBlock
| schema.ResourceContentBlock
| schema.EmbeddedResourceContentBlock
)
# Used when an embedded binary resource arrives without a declared media type; an empty
# media type would later raise when the content is formatted for a model request.
_DEFAULT_BINARY_MEDIA_TYPE = 'application/octet-stream'
def prompt_blocks_to_user_content(blocks: Sequence[PromptContentBlock]) -> list[UserContent]:
"""Convert ACP prompt content blocks into Pydantic AI user content.
Text and embedded text resources become plain strings; images and audio become
[`BinaryContent`][pydantic_ai.messages.BinaryContent] (or [`ImageUrl`][pydantic_ai.messages.ImageUrl]
when the image is referenced by URL); a resource link contributes its URI as text, prefixed with
its title or name when the client provided one (for example `My File (file:///a.txt)`).
"""
content: list[UserContent] = []
for block in blocks:
if isinstance(block, schema.TextContentBlock):
content.append(block.text)
elif isinstance(block, schema.ImageContentBlock):
# ACP requires inline `data`; `uri` is only an optional reference to the image's source.
# Prefer the bytes the client actually sent, falling back to the URL only when no inline
# data is present (a client sending both must not have its image silently replaced by a
# link the model may be unable to fetch).
if not block.data and block.uri is not None:
content.append(ImageUrl(url=block.uri))
else:
content.append(BinaryContent(data=base64.b64decode(block.data), media_type=block.mime_type))
elif isinstance(block, schema.AudioContentBlock):
content.append(BinaryContent(data=base64.b64decode(block.data), media_type=block.mime_type))
elif isinstance(block, schema.EmbeddedResourceContentBlock):
resource = block.resource
if isinstance(resource, schema.TextResourceContents):
content.append(resource.text)
else:
media_type = resource.mime_type or _DEFAULT_BINARY_MEDIA_TYPE
content.append(BinaryContent(data=base64.b64decode(resource.blob), media_type=media_type))
else:
# The only remaining variant is a resource link, which carries no inline content;
# pass its URI through as text, prefixed with a human-readable label when present.
label = block.title or block.name
content.append(f'{label} ({block.uri})' if label else block.uri)
return content
pydantic_ai_harness/experimental/acp/_permission.py
+37
View File
@@ -0,0 +1,37 @@
"""Tool-approval permission types: the call under review and the policy that scopes decisions."""
from __future__ import annotations
import json
from collections.abc import Callable, Hashable, Mapping
from dataclasses import dataclass
from pydantic_ai_harness.experimental.acp._serialize import jsonable
@dataclass(frozen=True, kw_only=True)
class ToolCallPermission:
"""A tool call awaiting the client's approval, passed to a `permission_policy`.
`args` is the tool's arguments, canonicalized to a mapping (so a policy can scope by a
specific argument without narrowing first).
"""
tool_name: str
tool_call_id: str
args: Mapping[str, object]
# Maps a tool call to the scope key under which an "always allow"/"always reject" decision is
# remembered for the session. Two calls with the same key share a remembered decision.
PermissionPolicy = Callable[[ToolCallPermission], Hashable]
def default_permission_scope(call: ToolCallPermission) -> Hashable:
"""Scope an "always" decision to the exact call: the same tool with the same arguments.
This keeps "always allow `delete_file(path='tmp')`" from also approving
`delete_file(path='.env')`. Pass `permission_policy` to widen the scope (for example to the
tool name alone) or narrow it.
"""
return (call.tool_name, json.dumps(jsonable(call.args), sort_keys=True))
pydantic_ai_harness/experimental/acp/_presentation.py
+218
View File
@@ -0,0 +1,218 @@
"""Map Pydantic AI tool calls onto ACP's rich tool-call presentation fields.
ACP lets an agent annotate a tool call with a `kind` (read/edit/execute/...), the file
`locations` it touches, and `content` such as an inline diff. A TUI renders those as
click-to-file links and diff views instead of opaque JSON. This module turns a recognized
`FileSystem`/`Shell` tool call into that presentation; unrecognized calls (or ones whose
arguments do not match the expected shape) return `None` so the adapter falls back to its
generic rendering.
A presenter sees a tool call's (workspace-relative) arguments; ACP requires absolute paths,
so the adapter resolves locations and diffs against the session's working directory via
`absolutize` before sending.
"""
from __future__ import annotations
import os.path
from collections.abc import Callable, Mapping
from dataclasses import dataclass, replace
import acp
from acp import schema
from pydantic_ai.messages import ToolCallPart
# The tool-call content variants ACP accepts on a `session/update` (inline content, a file
# diff, or a reference to a live terminal).
ToolCallContent = schema.ContentToolCallContent | schema.FileEditToolCallContent | schema.TerminalToolCallContent
@dataclass(frozen=True, kw_only=True)
class ToolCallPresentation:
"""How a tool call should be shown in an ACP client.
Mirrors the rich fields of ACP's tool-call updates. Every field is optional: a presenter
fills in only what it knows, and the adapter defaults the title to the tool name. Build
`content` with ACP's helpers, e.g. `acp.tool_diff_content(path, new_text, old_text)` for a
file edit. Paths may be workspace-relative; the adapter makes them absolute.
"""
kind: schema.ToolKind | None = None
title: str | None = None
locations: tuple[schema.ToolCallLocation, ...] = ()
content: tuple[ToolCallContent, ...] = ()
# A function deciding how to present a tool call, or `None` to leave it to generic rendering.
# `ToolCallPresentation` is defined above, so the reference is eager (not a string forward ref):
# the alias stays fully resolved when it appears in another module's annotations under
# `typing.get_type_hints`.
ToolCallPresenter = Callable[[ToolCallPart], ToolCallPresentation | None]
def _nonempty_str(args: Mapping[str, object], key: str) -> str | None:
"""Return `args[key]` when it is a non-empty string, else `None`."""
value = args.get(key)
return value if isinstance(value, str) and value else None
def _any_str(args: Mapping[str, object], key: str) -> str | None:
"""Return `args[key]` when it is a string -- empty is legitimate for content text -- else `None`."""
value = args.get(key)
return value if isinstance(value, str) else None
def _location(path: str) -> schema.ToolCallLocation:
return schema.ToolCallLocation(path=path)
def _present_read(args: Mapping[str, object]) -> ToolCallPresentation | None:
path = _nonempty_str(args, 'path')
if path is None:
return None
return ToolCallPresentation(kind='read', locations=(_location(path),))
def _present_edit(args: Mapping[str, object]) -> ToolCallPresentation | None:
path = _nonempty_str(args, 'path')
old_text = _nonempty_str(args, 'old_text')
new_text = _any_str(args, 'new_text')
if path is None or old_text is None or new_text is None:
return None
diff = acp.tool_diff_content(path=path, new_text=new_text, old_text=old_text)
return ToolCallPresentation(kind='edit', locations=(_location(path),), content=(diff,))
def _present_write(args: Mapping[str, object]) -> ToolCallPresentation | None:
path = _nonempty_str(args, 'path')
content = _any_str(args, 'content')
if path is None or content is None:
return None
# `write_file` is usually a create, where omitting `old_text` (the ACP convention for a new
# file) is correct. On an overwrite the prior contents are unknown from the args alone, so
# the diff understates what is replaced -- an accepted limitation until reads route through
# the client.
diff = acp.tool_diff_content(path=path, new_text=content)
return ToolCallPresentation(kind='edit', locations=(_location(path),), content=(diff,))
def _present_create_directory(args: Mapping[str, object]) -> ToolCallPresentation | None:
path = _nonempty_str(args, 'path')
if path is None:
return None
return ToolCallPresentation(kind='other', locations=(_location(path),))
def _present_list(args: Mapping[str, object]) -> ToolCallPresentation | None:
# `list_directory` takes only an optional `path`, so it has no required argument to validate
# against -- it is matched by name alone (a same-named custom tool would also render here).
path = _nonempty_str(args, 'path')
locations = (_location(path),) if path is not None else ()
return ToolCallPresentation(kind='search', locations=locations)
def _present_grep(args: Mapping[str, object]) -> ToolCallPresentation | None:
if _nonempty_str(args, 'pattern') is None:
return None
path = _nonempty_str(args, 'path')
locations = (_location(path),) if path is not None else ()
return ToolCallPresentation(kind='search', locations=locations)
def _present_run(args: Mapping[str, object]) -> ToolCallPresentation | None:
if _nonempty_str(args, 'command') is None:
return None
return ToolCallPresentation(kind='execute')
def _present_command_ref(args: Mapping[str, object]) -> ToolCallPresentation | None:
if _nonempty_str(args, 'command_id') is None:
return None
return ToolCallPresentation(kind='execute')
# Recognized `FileSystem`/`Shell` tool names mapped to their presenters. Coupling is by tool
# name (a rename in those capabilities silently falls back to generic rendering).
_HANDLERS: dict[str, Callable[[Mapping[str, object]], ToolCallPresentation | None]] = {
'read_file': _present_read,
'file_info': _present_read,
'edit_file': _present_edit,
'write_file': _present_write,
'create_directory': _present_create_directory,
'list_directory': _present_list,
'search_files': _present_grep,
'find_files': _present_grep,
'run_command': _present_run,
'start_command': _present_run,
'check_command': _present_command_ref,
'stop_command': _present_command_ref,
}
def default_coding_presenter(call: ToolCallPart) -> ToolCallPresentation | None:
"""Present a recognized `FileSystem`/`Shell` tool call, else `None`.
Matches the tool by name and validates its argument shape; a mismatch returns `None` so the
call falls back to generic rendering. The exception is `list_directory`, whose only argument
is optional, so it is matched by name alone.
"""
handler = _HANDLERS.get(call.tool_name)
if handler is None:
return None
# Malformed arguments surface from `args_as_dict` as a sentinel dict no handler matches.
return handler(call.args_as_dict())
def chain_presenters(*presenters: ToolCallPresenter) -> ToolCallPresenter:
"""Combine presenters, using the first one that returns a presentation.
Each presenter returns `None` to mean "I do not handle this call"; the chain tries them in
order and returns the first non-`None` result (or `None` if none match). Use it to add
rendering for your own tools while keeping the built-in one:
`chain_presenters(my_presenter, default_coding_presenter)`.
"""
def presenter(call: ToolCallPart) -> ToolCallPresentation | None:
for candidate in presenters:
result = candidate(call)
if result is not None:
return result
return None
return presenter
def _resolve_within(path: str, cwd: str) -> str | None:
"""Resolve a relative `path` against `cwd`, or `None` if it escapes it; absolute paths pass through."""
if os.path.isabs(path):
# May name an advertised additional directory, which the presenter cannot see to validate.
return path
resolved = os.path.normpath(os.path.join(cwd, path))
relative = os.path.relpath(resolved, os.path.normpath(cwd))
if relative == os.pardir or relative.startswith(os.pardir + os.sep):
return None
return resolved
def _within_content(item: ToolCallContent, cwd: str) -> ToolCallContent | None:
if isinstance(item, schema.FileEditToolCallContent):
resolved = _resolve_within(item.path, cwd)
return None if resolved is None else item.model_copy(update={'path': resolved})
return item
def absolutize(presentation: ToolCallPresentation, cwd: str) -> ToolCallPresentation:
"""Resolve a presentation's workspace-relative paths against the session `cwd`.
ACP requires tool-call locations and file-edit diffs to carry absolute paths. Absolute paths
are left unchanged; a relative path that escapes `cwd` via `..` is dropped rather than
absolutized, so the editor is never pointed outside the workspace. Assumes the agent's
filesystem tools are rooted at `cwd`.
"""
locations = tuple(
loc.model_copy(update={'path': resolved})
for loc in presentation.locations
if (resolved := _resolve_within(loc.path, cwd)) is not None
)
content = tuple(c for c in (_within_content(item, cwd) for item in presentation.content) if c is not None)
return replace(presentation, locations=locations, content=content)
pydantic_ai_harness/experimental/acp/_serialize.py
+73
View File
@@ -0,0 +1,73 @@
"""Coerce tool input/output into JSON-able payloads sized for ACP `session/update` notifications."""
from __future__ import annotations
import json
from collections.abc import Iterator
from pydantic_core import to_jsonable_python
# ACP stdio is newline-delimited JSON, and clients read it with a bounded buffer (asyncio's
# default `StreamReader` limit is 64 KiB). A single oversized notification overruns that buffer
# and drops the connection, so long text is split across several updates.
#
# The SDK serializes outbound text with `json.dumps(..., ensure_ascii=True)`, so a non-ASCII code
# point expands *inside* the JSON string: a BMP char to `\uXXXX` (6 bytes) and an astral char to a
# surrogate pair `\uXXXX\uXXXX` (12 bytes). A character-count cap therefore can't bound the wire
# size -- 8K emoji serialize to ~96 KiB and drop the connection. We chunk by escaped byte length
# instead (see `chunk_text`), leaving headroom below 64 KiB for the notification envelope.
MAX_TEXT_UPDATE_BYTES = 48 * 1024
# Tool input/output is sent whole in one notification (it cannot be chunked across updates like
# streamed text), so an oversized payload is truncated to keep the notification under the buffer.
MAX_RAW_FIELD_CHARS = 16 * 1024
def _escaped_len(char: str) -> int:
"""Bytes `char` occupies inside a `json.dumps(..., ensure_ascii=True)` string."""
if char in '"\\\b\f\n\r\t':
return 2 # short escapes: `\"`, `\\`, `\n`, ...
code = ord(char)
if code < 0x20:
return 6 # other control chars -> `\u00XX`
if code < 0x80:
return 1 # printable ASCII
if code < 0x10000:
return 6 # BMP non-ASCII -> `\uXXXX`
return 12 # astral plane -> surrogate pair `\uXXXX\uXXXX`
def chunk_text(text: str, budget: int = MAX_TEXT_UPDATE_BYTES) -> Iterator[str]:
"""Split `text` so each chunk's JSON-escaped byte length stays within `budget`.
Bounds the serialized size of each `session/update` regardless of how the text escapes, so a
single notification can't overrun the client's read buffer (see `MAX_TEXT_UPDATE_BYTES`).
"""
chunk: list[str] = []
size = 0
for char in text:
char_size = _escaped_len(char)
if chunk and size + char_size > budget:
yield ''.join(chunk)
chunk = []
size = 0
chunk.append(char)
size += char_size
if chunk:
yield ''.join(chunk)
def jsonable(value: object) -> object:
"""Coerce arbitrary tool input/output into something an ACP `session/update` can serialize."""
# `bytes_mode='base64'` keeps raw (non-UTF-8) bytes from raising; `fallback=str` covers types
# pydantic cannot otherwise serialize.
return to_jsonable_python(value, fallback=str, bytes_mode='base64')
def bounded_jsonable(value: object) -> object:
"""`jsonable`, but replace an oversized payload with a truncated marker (see `MAX_RAW_FIELD_CHARS`)."""
payload = jsonable(value)
serialized = json.dumps(payload)
if len(serialized) <= MAX_RAW_FIELD_CHARS:
return payload
return {'truncated': True, 'original_length': len(serialized), 'preview': serialized[:MAX_RAW_FIELD_CHARS]}
pydantic_ai_harness/experimental/acp/_server.py
+126
View File
@@ -0,0 +1,126 @@
"""Entry points for serving a Pydantic AI agent over ACP."""
from __future__ import annotations
import asyncio
from collections.abc import Callable, Sequence
from typing import Literal
import acp
from acp import schema
from pydantic_ai.agent import AbstractAgent
from pydantic_ai.models import KnownModelName, Model
from pydantic_ai.output import OutputDataT
from pydantic_ai.tools import AgentDepsT
from pydantic_ai.usage import UsageLimits
from pydantic_ai_harness.experimental.acp._adapter import DEFAULT_VERSION, PydanticAIACPAgent
from pydantic_ai_harness.experimental.acp._permission import PermissionPolicy
from pydantic_ai_harness.experimental.acp._presentation import ToolCallPresenter
from pydantic_ai_harness.experimental.acp._session import SessionConfigFunc
from pydantic_ai_harness.experimental.acp._store import SessionStore
async def run_acp_stdio(
agent: AbstractAgent[AgentDepsT, OutputDataT],
*,
deps: AgentDepsT = None,
name: str | None = None,
version: str = DEFAULT_VERSION,
session_config: SessionConfigFunc[AgentDepsT] | None = None,
permission_policy: PermissionPolicy | None = None,
prompt_capabilities: schema.PromptCapabilities | None = None,
mcp_capabilities: schema.McpCapabilities | None = None,
tool_presenter: ToolCallPresenter | None = None,
session_store: SessionStore | None = None,
models: Sequence[KnownModelName | str] | Literal['all'] | None = None,
model_resolver: Callable[[str], Model | str] | None = None,
usage_limits: UsageLimits | None = None,
) -> None:
"""Serve `agent` as an ACP agent over stdin/stdout until the client disconnects.
This is the entry point an ACP client (such as an editor or terminal UI) launches as a
subprocess. It blocks for the lifetime of the connection.
Args:
agent: The Pydantic AI agent to expose over ACP.
deps: Dependencies passed to every agent run.
name: Name advertised to the client. Defaults to the agent's name, then `'pydantic-ai-agent'`.
version: Version advertised to the client.
session_config: Per-session factory deriving deps/toolsets from the client's workspace setup.
See [`PydanticAIACPAgent`][pydantic_ai_harness.experimental.acp.PydanticAIACPAgent].
permission_policy: Scopes how "always allow"/"always reject" decisions are remembered.
See [`PydanticAIACPAgent`][pydantic_ai_harness.experimental.acp.PydanticAIACPAgent].
prompt_capabilities: Prompt content types the agent advertises support for.
See [`PydanticAIACPAgent`][pydantic_ai_harness.experimental.acp.PydanticAIACPAgent].
mcp_capabilities: MCP transports the agent advertises support for. Requires a
`session_config` that connects them. See
[`PydanticAIACPAgent`][pydantic_ai_harness.experimental.acp.PydanticAIACPAgent].
tool_presenter: Maps tool calls to rich ACP presentation (kind, file locations, diffs).
See [`PydanticAIACPAgent`][pydantic_ai_harness.experimental.acp.PydanticAIACPAgent].
session_store: Enables `session/load` by persisting each session. See
[`PydanticAIACPAgent`][pydantic_ai_harness.experimental.acp.PydanticAIACPAgent].
models: Models the client may switch between with the `model` session config option. See
[`PydanticAIACPAgent`][pydantic_ai_harness.experimental.acp.PydanticAIACPAgent].
model_resolver: Maps an advertised model id to the `Model` (or model string) used for the
run. See [`PydanticAIACPAgent`][pydantic_ai_harness.experimental.acp.PydanticAIACPAgent].
usage_limits: Per-run request/token ceilings applied to every agent run. See
[`PydanticAIACPAgent`][pydantic_ai_harness.experimental.acp.PydanticAIACPAgent].
"""
adapter = PydanticAIACPAgent(
agent,
deps=deps,
name=name,
version=version,
session_config=session_config,
permission_policy=permission_policy,
prompt_capabilities=prompt_capabilities,
mcp_capabilities=mcp_capabilities,
tool_presenter=tool_presenter,
session_store=session_store,
models=models,
model_resolver=model_resolver,
usage_limits=usage_limits,
)
# `session/close` is still UNSTABLE in the ACP SDK, and the SDK's router rejects unstable
# methods with `method_not_found` unless this flag is set. Keep enabled until close stabilizes.
await acp.run_agent(adapter, use_unstable_protocol=True)
def run_acp_stdio_sync(
agent: AbstractAgent[AgentDepsT, OutputDataT],
*,
deps: AgentDepsT = None,
name: str | None = None,
version: str = DEFAULT_VERSION,
session_config: SessionConfigFunc[AgentDepsT] | None = None,
permission_policy: PermissionPolicy | None = None,
prompt_capabilities: schema.PromptCapabilities | None = None,
mcp_capabilities: schema.McpCapabilities | None = None,
tool_presenter: ToolCallPresenter | None = None,
session_store: SessionStore | None = None,
models: Sequence[KnownModelName | str] | Literal['all'] | None = None,
model_resolver: Callable[[str], Model | str] | None = None,
usage_limits: UsageLimits | None = None,
) -> None:
"""Synchronous wrapper around [`run_acp_stdio`][pydantic_ai_harness.experimental.acp.run_acp_stdio].
Convenient as the `main()` of an ACP agent script, which clients launch as a subprocess.
"""
asyncio.run(
run_acp_stdio(
agent,
deps=deps,
name=name,
version=version,
session_config=session_config,
permission_policy=permission_policy,
prompt_capabilities=prompt_capabilities,
mcp_capabilities=mcp_capabilities,
tool_presenter=tool_presenter,
session_store=session_store,
models=models,
model_resolver=model_resolver,
usage_limits=usage_limits,
)
)
pydantic_ai_harness/experimental/acp/_session.py
+108
View File
@@ -0,0 +1,108 @@
"""Per-session value types: the client's session setup, its run configuration, and live state."""
from __future__ import annotations
import asyncio
from collections.abc import Awaitable, Hashable, Sequence
from dataclasses import dataclass, field
from typing import Generic, Protocol, TypeAlias
from acp import Client, schema
from pydantic_ai.messages import ModelMessage
from pydantic_ai.tools import AgentDepsT
from pydantic_ai.toolsets import AbstractToolset
# A single MCP server configuration, in any of the transports ACP carries.
McpServer: TypeAlias = schema.HttpMcpServer | schema.SseMcpServer | schema.AcpMcpServer | schema.McpServerStdio
# MCP server configuration list, as carried by ACP session methods (matches `acp.Agent`).
McpServers: TypeAlias = list[McpServer] | None
# A single `session/update` payload, as accepted by the client's `session_update`. The adapter
# records these per session (the client-visible transcript) so `session/load` can replay them.
SessionUpdate: TypeAlias = (
schema.UserMessageChunk
| schema.AgentMessageChunk
| schema.AgentThoughtChunk
| schema.ToolCallStart
| schema.ToolCallProgress
| schema.AgentPlanUpdate
| schema.AvailableCommandsUpdate
| schema.CurrentModeUpdate
| schema.ConfigOptionUpdate
| schema.SessionInfoUpdate
| schema.UsageUpdate
)
@dataclass(frozen=True, kw_only=True)
class AcpSession:
"""The setup an ACP client provides when it starts a session.
Passed to the adapter's `session_config` factory so an embedder can derive per-session run
configuration from the workspace the client opened -- for example rooting the agent's
dependencies at `cwd` or turning the client-provided `mcp_servers` into toolsets.
`client_capabilities` is whatever the client advertised during `initialize` (for example
filesystem or terminal support).
`client` and `session_id` are the live connection handle for this session: pass them to a
client-backed toolset (see [`acp_filesystem`][pydantic_ai_harness.experimental.acp.acp_filesystem]) to
route the agent's I/O through the editor rather than local disk.
"""
cwd: str
mcp_servers: list[McpServer]
client_capabilities: schema.ClientCapabilities | None
client: Client
session_id: str
@dataclass(frozen=True, kw_only=True)
class AcpSessionConfig(Generic[AgentDepsT]):
"""Per-session run configuration returned by a `session_config` factory.
`deps` and `toolsets` are applied to every agent run in that session, mirroring
`Agent.run(..., deps=..., toolsets=...)`. `toolsets` is added to the agent's own toolsets
rather than replacing them. `deps` is required; pass `deps=None` for an agent with no
dependencies.
"""
deps: AgentDepsT
toolsets: Sequence[AbstractToolset[AgentDepsT]] | None = None
# A `Protocol` rather than a `Callable` alias so it stays subscriptable (`SessionConfigFunc[MyDeps]`)
# and evaluable by `typing.get_type_hints` in the public constructor's annotations.
class SessionConfigFunc(Protocol[AgentDepsT]):
"""Factory mapping an ACP session's setup to its run configuration; may be sync or async.
Any callable taking the client's `AcpSession` and returning an `AcpSessionConfig` (or an
awaitable of one) satisfies it.
"""
# `session` is positional-only (`/`) so any callable taking one `AcpSession` satisfies the
# protocol regardless of how it names the parameter -- matching a plain `Callable` alias.
def __call__(
self, session: AcpSession, /
) -> (
AcpSessionConfig[AgentDepsT] | Awaitable[AcpSessionConfig[AgentDepsT]]
): ... # pragma: no cover - structural protocol; the body never runs
@dataclass(kw_only=True)
class SessionState(Generic[AgentDepsT]):
"""All per-session adapter state; `new_session`/`load_session` create one, `close_session` removes it."""
session_id: str
config: AcpSessionConfig[AgentDepsT]
cwd: str
history: list[ModelMessage] = field(default_factory=list[ModelMessage])
transcript: list[SessionUpdate] = field(default_factory=list[SessionUpdate])
# The model the client selected for this session via the `model` config option, or `None` to
# use the agent's own model. Applied as a per-run override so the shared agent is never mutated.
model: str | None = None
lock: asyncio.Lock = field(default_factory=asyncio.Lock)
active_turn: asyncio.Task[schema.PromptResponse] | None = None
cancel_requested: bool = False
always_allow: set[Hashable] = field(default_factory=set[Hashable])
always_reject: set[Hashable] = field(default_factory=set[Hashable])
pydantic_ai_harness/experimental/acp/_store.py
+68
View File
@@ -0,0 +1,68 @@
"""Pluggable persistence for ACP sessions, so `session/load` can reopen a past conversation.
Persistence is opt-in: pass a `SessionStore` to the adapter to advertise and support
`session/load`. [`InMemorySessionStore`][pydantic_ai_harness.experimental.acp.InMemorySessionStore] keeps
sessions for the process's lifetime; implement the protocol over a file or database for
durability.
"""
from __future__ import annotations
from dataclasses import dataclass, field
from typing import Protocol
from pydantic_ai.messages import ModelMessage
from pydantic_ai_harness.experimental.acp._session import SessionUpdate
# The two views are stored separately because neither can be derived from the other: the messages
# lack a tool call's rendered title/diff, and the transcript lacks what the model saw.
@dataclass(frozen=True, kw_only=True)
class StoredSession:
"""A session's persisted state: the model's message history and the client-visible transcript.
On `session/load`, `messages` is restored into the agent and `updates` is replayed verbatim
to the client. Both hold Pydantic models, so a durable store can serialize them with Pydantic.
"""
messages: list[ModelMessage] = field(default_factory=list[ModelMessage])
updates: list[SessionUpdate] = field(default_factory=list[SessionUpdate])
# The model selected via the `model` config option, restored so a reopened session keeps it.
model: str | None = None
class SessionStore(Protocol):
"""Where the adapter saves and restores sessions so `session/load` can reopen them.
`save` is called after each committed turn (and once when the session is created); `load`
returns a previously saved session or `None` if the id is unknown.
Failure handling: a `save` that raises is logged and swallowed, never failing the turn or
session operation that triggered it -- that work already streamed and committed in memory, so a
durable-write error must not surface as a failure for what the user saw succeed; the next
successful save catches the store up. A `load` that raises (a read error or a corrupt, unparsable
payload) fails `session/load` with an `internal_error`, since a session that cannot be read cannot
be reopened. Implementations therefore do not need to translate their own errors into ACP errors.
"""
async def save(self, session_id: str, session: StoredSession) -> None: ... # pragma: no cover - protocol stub
async def load(self, session_id: str) -> StoredSession | None: ... # pragma: no cover - protocol stub
class InMemorySessionStore:
"""A [`SessionStore`][pydantic_ai_harness.experimental.acp.SessionStore] holding sessions in a dict.
Sessions can be reopened within one process but do not survive a restart; back the store
with a file or database for that.
"""
def __init__(self) -> None:
self._sessions: dict[str, StoredSession] = {}
async def save(self, session_id: str, session: StoredSession) -> None:
self._sessions[session_id] = session
async def load(self, session_id: str) -> StoredSession | None:
return self._sessions.get(session_id)
pydantic_ai_harness/experimental/authoring/__init__.py
+12 -7
View File
@@ -1,16 +1,21 @@
"""Runtime capability authoring: let an agent write, validate, and register real capabilities."""
"""Deprecated import location for `pydantic_ai_harness.runtime_authoring`.
from pydantic_ai_harness.experimental._warn import warn_experimental
from pydantic_ai_harness.experimental.authoring._capability import RuntimeAuthoring
from pydantic_ai_harness.experimental.authoring._store import AuthoredCapability, CapabilityStore
from pydantic_ai_harness.experimental.authoring._toolset import AuthoringToolset
from pydantic_ai_harness.experimental.authoring._validate import (
This capability graduated out of `experimental`; importing from here still works but
emits a `DeprecationWarning`. Import from `pydantic_ai_harness.runtime_authoring` instead.
"""
from pydantic_ai_harness.experimental._warn import warn_moved
from pydantic_ai_harness.runtime_authoring import (
AuthoredCapability,
AuthoringToolset,
CapabilityStore,
CapabilityValidationError,
RuntimeAuthoring,
load_capability_instance,
validate_capability_file,
)
warn_experimental('authoring')
warn_moved('authoring', 'runtime_authoring')
__all__ = [
'AuthoredCapability',
pydantic_ai_harness/experimental/compaction/__init__.py
+17 -13
View File
@@ -1,20 +1,24 @@
"""Compaction capabilities: keep an agent's conversation history within the context window.
"""Deprecated import location for `pydantic_ai_harness.compaction`.
Each capability lives in its own module; shared utilities (token estimation, the
`CompactionStrategy` protocol, tool-pair-safe cutoffs, in-place clearing) live in `_shared`.
This capability graduated out of `experimental`; importing from here still works but
emits a `DeprecationWarning`. Import from `pydantic_ai_harness.compaction` instead.
"""
from pydantic_ai_harness.experimental._warn import warn_experimental
from pydantic_ai_harness.experimental.compaction._clamp_oversized_messages import ClampOversizedMessages
from pydantic_ai_harness.experimental.compaction._clear_tool_results import ClearToolResults
from pydantic_ai_harness.experimental.compaction._deduplicate_file_reads import DeduplicateFileReads
from pydantic_ai_harness.experimental.compaction._limit_warner import LimitWarner, WarningKind
from pydantic_ai_harness.experimental.compaction._shared import CompactionStrategy, estimate_token_count
from pydantic_ai_harness.experimental.compaction._sliding_window import SlidingWindow
from pydantic_ai_harness.experimental.compaction._summarizing_compaction import SummarizingCompaction
from pydantic_ai_harness.experimental.compaction._tiered_compaction import TieredCompaction
from pydantic_ai_harness.compaction import (
ClampOversizedMessages,
ClearToolResults,
CompactionStrategy,
DeduplicateFileReads,
LimitWarner,
SlidingWindow,
SummarizingCompaction,
TieredCompaction,
WarningKind,
estimate_token_count,
)
from pydantic_ai_harness.experimental._warn import warn_moved
warn_experimental('compaction')
warn_moved('compaction', 'compaction')
__all__ = [
'ClampOversizedMessages',
pydantic_ai_harness/experimental/context/__init__.py
+21 -8
View File
@@ -1,11 +1,24 @@
"""Context capability: discover and load a repo's accumulated context engineering."""
"""Deprecated import location for `pydantic_ai_harness.context`.
from pydantic_ai_harness.experimental._warn import warn_experimental
from pydantic_ai_harness.experimental.context._capability import RepoContext
from pydantic_ai_harness.experimental.context._inventory import AgentContextInventory, AssetRoot
from pydantic_ai_harness.experimental.context._loader import ContextFile
from pydantic_ai_harness.experimental.context._toolset import RepoContextToolset
This capability graduated out of `experimental`; importing from here still works but
emits a `DeprecationWarning`. Import from `pydantic_ai_harness.context` instead.
"""
warn_experimental('context')
from pydantic_ai_harness.context import (
AgentContextInventory,
AssetRoot,
ContextFile,
RepoContext,
RepoContextToolset,
)
from pydantic_ai_harness.experimental._warn import warn_moved
__all__ = ['AgentContextInventory', 'AssetRoot', 'ContextFile', 'RepoContext', 'RepoContextToolset']
warn_moved('context', 'context')
__all__ = [
'AgentContextInventory',
'AssetRoot',
'ContextFile',
'RepoContext',
'RepoContextToolset',
]
pydantic_ai_harness/experimental/docs/__init__.py
+17 -6
View File
@@ -1,9 +1,20 @@
"""Docs capability: an on-demand tool that locates Pydantic AI documentation."""
"""Deprecated import location for `pydantic_ai_harness.docs`.
from pydantic_ai_harness.experimental._warn import warn_experimental
from pydantic_ai_harness.experimental.docs._capability import PyaiDocs
from pydantic_ai_harness.experimental.docs._toolset import PyaiDocsToolset, PyaiDocsTopic
This capability graduated out of `experimental`; importing from here still works but
emits a `DeprecationWarning`. Import from `pydantic_ai_harness.docs` instead.
"""
warn_experimental('docs')
from pydantic_ai_harness.docs import (
PyaiDocs,
PyaiDocsToolset,
PyaiDocsTopic,
)
from pydantic_ai_harness.experimental._warn import warn_moved
__all__ = ['PyaiDocs', 'PyaiDocsToolset', 'PyaiDocsTopic']
warn_moved('docs', 'docs')
__all__ = [
'PyaiDocs',
'PyaiDocsToolset',
'PyaiDocsTopic',
]
pydantic_ai_harness/experimental/dynamic_workflow/__init__.py
+22
View File
@@ -0,0 +1,22 @@
"""Deprecated import location for `pydantic_ai_harness.dynamic_workflow`.
This capability graduated out of `experimental`; importing from here still works but
emits a `DeprecationWarning`. Import from `pydantic_ai_harness.dynamic_workflow` instead.
"""
from pydantic_ai_harness.dynamic_workflow import (
DynamicWorkflow,
DynamicWorkflowToolset,
WorkflowAgent,
WorkflowResourceLimits,
)
from pydantic_ai_harness.experimental._warn import warn_moved
warn_moved('dynamic_workflow', 'dynamic_workflow')
__all__ = [
'DynamicWorkflow',
'DynamicWorkflowToolset',
'WorkflowAgent',
'WorkflowResourceLimits',
]
pydantic_ai_harness/experimental/media/__init__.py
+9 -10
View File
@@ -1,28 +1,27 @@
"""Content-addressed media stores for offloading large binary parts.
"""Deprecated import location for `pydantic_ai_harness.media`.
Used by `pydantic_ai_harness.experimental.step_persistence` to keep snapshots small when
messages carry `BinaryContent` payloads. A forthcoming `MediaExternalizer`
capability will reuse these stores for in-flight wire-payload reduction
(rewriting `BinaryContent` to URL parts before the model sees them).
This capability graduated out of `experimental`; importing from here still works but
emits a `DeprecationWarning`. Import from `pydantic_ai_harness.media` instead.
"""
from pydantic_ai_harness.experimental._warn import warn_experimental
from pydantic_ai_harness.experimental.media._s3 import S3MediaStore
from pydantic_ai_harness.experimental.media._store import (
from pydantic_ai_harness.experimental._warn import warn_moved
from pydantic_ai_harness.media import (
DiskMediaStore,
KeyStrategy,
MediaContext,
MediaStore,
PublicUrlResolver,
S3MediaStore,
SqliteMediaStore,
default_key_strategy,
externalize_media,
make_static_public_url,
media_uri_for,
parse_media_uri,
restore_media,
)
from pydantic_ai_harness.experimental.media._walker import externalize_media, restore_media
warn_experimental('media')
warn_moved('media', 'media')
__all__ = [
'DiskMediaStore',
pydantic_ai_harness/experimental/overflow/__init__.py
+11 -15
View File
@@ -1,30 +1,26 @@
"""Overflow capability: reduce oversized tool returns at production time.
"""Deprecated import location for `pydantic_ai_harness.overflowing_tool_output`.
`OverflowingToolOutput` intercepts a tool return when it is produced and reduces it --
truncating, spilling to a queryable file, or summarizing -- so an oversized payload does
not persist in history and get re-sent on every later model request. Combine the three
modes through an ordered list of size `bands`.
Spilled payloads are read back on demand through the registered `read_tool_result` tool;
the `OverflowStore` protocol is the seam for a durable backend (the local-file default
ships for single-process runs).
This capability graduated out of `experimental`; importing from here still works but
emits a `DeprecationWarning`. Import from `pydantic_ai_harness.overflowing_tool_output` instead.
"""
from pydantic_ai_harness.experimental._warn import warn_experimental
from pydantic_ai_harness.experimental.overflow._bands import (
from pydantic_ai_harness.experimental._warn import warn_moved
from pydantic_ai_harness.overflowing_tool_output import (
READ_TOOL_NAME,
Action,
Band,
LocalFileStore,
OverflowingToolOutput,
OverflowStore,
Passthrough,
Spill,
Summarize,
SummarizeFunc,
Truncate,
TruncationStrategy,
)
from pydantic_ai_harness.experimental.overflow._capability import READ_TOOL_NAME, OverflowingToolOutput
from pydantic_ai_harness.experimental.overflow._payload import TruncationStrategy
from pydantic_ai_harness.experimental.overflow._store import LocalFileStore, OverflowStore
warn_experimental('overflow')
warn_moved('overflow', 'overflowing_tool_output')
__all__ = [
'READ_TOOL_NAME',
pydantic_ai_harness/experimental/planning/__init__.py
+19 -6
View File
@@ -1,9 +1,22 @@
"""Planning capability: model-owned, cache-friendly task planning for agents."""
"""Deprecated import location for `pydantic_ai_harness.planning`.
from pydantic_ai_harness.experimental._warn import warn_experimental
from pydantic_ai_harness.experimental.planning._capability import Planning
from pydantic_ai_harness.experimental.planning._toolset import PlanItem, PlanningToolset, TaskStatus
This capability graduated out of `experimental`; importing from here still works but
emits a `DeprecationWarning`. Import from `pydantic_ai_harness.planning` instead.
"""
warn_experimental('planning')
from pydantic_ai_harness.experimental._warn import warn_moved
from pydantic_ai_harness.planning import (
PlanItem,
Planning,
PlanningToolset,
TaskStatus,
)
__all__ = ['PlanItem', 'Planning', 'PlanningToolset', 'TaskStatus']
warn_moved('planning', 'planning')
__all__ = [
'PlanItem',
'Planning',
'PlanningToolset',
'TaskStatus',
]
pydantic_ai_harness/experimental/step_persistence/__init__.py
+19 -19
View File
@@ -1,29 +1,29 @@
"""Step-event persistence: append-only event log, continuable snapshots, tool-effect ledger."""
"""Deprecated import location for `pydantic_ai_harness.step_persistence`.
from pydantic_ai_harness.experimental._warn import warn_experimental
from pydantic_ai_harness.experimental.step_persistence._capability import StepPersistence
from pydantic_ai_harness.experimental.step_persistence._helpers import (
This capability graduated out of `experimental`; importing from here still works but
emits a `DeprecationWarning`. Import from `pydantic_ai_harness.step_persistence` instead.
"""
from pydantic_ai_harness.experimental._warn import warn_moved
from pydantic_ai_harness.step_persistence import (
ContinuableSnapshot,
EventKind,
FileStepStore,
InMemoryStepStore,
RunRecord,
SqliteStepStore,
StepEvent,
StepPersistence,
StepStore,
ToolEffectRecord,
ToolEffectStatus,
annotate_tool_effect,
continue_run,
fork_run,
is_provider_valid,
)
from pydantic_ai_harness.experimental.step_persistence._store import (
FileStepStore,
InMemoryStepStore,
SqliteStepStore,
StepStore,
)
from pydantic_ai_harness.experimental.step_persistence._types import (
ContinuableSnapshot,
EventKind,
RunRecord,
StepEvent,
ToolEffectRecord,
ToolEffectStatus,
)
warn_experimental('step_persistence')
warn_moved('step_persistence', 'step_persistence')
__all__ = [
'ContinuableSnapshot',
pydantic_ai_harness/experimental/subagents/__init__.py
+16 -7
View File
@@ -1,12 +1,21 @@
"""Sub-agent capability: delegate self-contained tasks to named child agents."""
"""Deprecated import location for `pydantic_ai_harness.subagents`.
from pydantic_ai_harness.experimental._warn import warn_experimental
from pydantic_ai_harness.experimental.subagents._capability import SubAgents, ToolResolver
from pydantic_ai_harness.experimental.subagents._disk import AgentOverride
from pydantic_ai_harness.experimental.subagents._effort import MINIMUM_EFFORT_FLOOR, clamp_effort
from pydantic_ai_harness.experimental.subagents._toolset import SubAgent, SubAgentToolset
This capability graduated out of `experimental`; importing from here still works but
emits a `DeprecationWarning`. Import from `pydantic_ai_harness.subagents` instead.
"""
warn_experimental('subagents')
from pydantic_ai_harness.experimental._warn import warn_moved
from pydantic_ai_harness.subagents import (
MINIMUM_EFFORT_FLOOR,
AgentOverride,
SubAgent,
SubAgents,
SubAgentToolset,
ToolResolver,
clamp_effort,
)
warn_moved('subagents', 'subagents')
__all__ = [
'MINIMUM_EFFORT_FLOOR',
pydantic_ai_harness/filesystem/_capability.py
+1 -2
View File
@@ -9,7 +9,6 @@ from typing import Any
from pydantic_ai.capabilities import AbstractCapability
from pydantic_ai.tools import AgentDepsT
from pydantic_ai.toolsets import AgentToolset
from pydantic_ai_harness.filesystem._toolset import FileSystemToolset
@@ -68,7 +67,7 @@ class FileSystem(AbstractCapability[AgentDepsT]):
if not isinstance(value, int) or value <= 0:
raise ValueError(f'{name} must be a positive integer, got {value!r}')
def get_toolset(self) -> AgentToolset[AgentDepsT]:
def get_toolset(self) -> FileSystemToolset[AgentDepsT]:
"""Build and return the filesystem toolset."""
return FileSystemToolset[AgentDepsT](
root_dir=Path(self.root_dir),
pydantic_ai_harness/guardrails/README.md
+181
View File
@@ -0,0 +1,181 @@
# Guardrails
Validate the user prompt before it reaches the model, and the model output before it reaches the caller.
## The problem
Agents take unstructured input from users and return unstructured output to callers. Without a validation layer, a prompt injection attempt, PII-laden message, or off-topic question goes to the model as-is, and any output the model produces is returned verbatim. The framework does not reason about "this is unsafe to send" or "this is unsafe to show".
## The solution
Two capabilities — `InputGuard` and `OutputGuard` — each backed by a `guard` callable you supply. The guard inspects a value (the prompt, or the output) and returns one of four outcomes:
| Outcome | `InputGuard` | `OutputGuard` |
|---|---|---|
| **allow** | send the prompt to the model | return the output to the caller |
| **block** | skip the model call; a refusal message becomes the response (`SkipModelRequest`) | raise `OutputBlocked` |
| **replace** | rewrite the prompt sent to the model (redaction) | substitute a sanitized output |
| **retry** | — (not valid for input) | send the output back to the model to try again (`ModelRetry`) |
A guard that raises an exception instead propagates it as a hard failure. The asymmetry between input `block` and output `block` is intentional: blocking the input spends no tokens, so a graceful refusal is almost always right; blocking the output means the model already produced something you do not want exposed, so raising forces the caller to decide what to do next.
## Usage
A guard returns a bare `bool` (`True` = allow, `False` = block) for the simple case, or a `GuardResult` for the richer outcomes.
```python
from pydantic_ai import Agent
from pydantic_ai_harness import GuardResult, InputGuard, OutputGuard
def no_secrets(prompt: str) -> bool:
return 'api_key' not in prompt.lower()
def no_pii(output: object) -> GuardResult:
if 'SSN' in str(output):
return GuardResult.block('The response contained personal data.')
return GuardResult.allow()
agent = Agent(
'openai:gpt-5.4',
capabilities=[
InputGuard(guard=no_secrets),
OutputGuard(guard=no_pii),
],
)
```
`OutputGuard` receives the output unchanged — no automatic stringification. For a string output the guard reads it directly; for a typed (Pydantic model) output the guard gets the model instance, so pick the serialization that fits the check (read a field, or call `output.model_dump_json()` for JSON text). This avoids the trap of `str(MyModel(...))` producing a `MyModel(field=...)` repr that hides field contents from regex-based checks.
Guards may also be async — return an awaitable `bool`/`GuardResult`, e.g. to call a moderation API.
## `GuardResult`
Construct a `GuardResult` with its classmethods, not the raw fields:
```python
from pydantic_ai_harness import GuardResult
GuardResult.allow() # let the value through
GuardResult.block('reason') # refuse; `reason` is optional (a default is used otherwise)
GuardResult.replace(cleaned_value) # substitute a sanitized value and continue
GuardResult.retry('instruction') # OutputGuard only: ask the model to redo the output
```
The block/retry message is produced at the moment the guard decides, so it can carry the guard's own reasoning rather than a string frozen at construction time.
## Redaction (`replace`)
Return `GuardResult.replace(value)` to sanitize rather than refuse. `InputGuard` rewrites the prompt sent to the model; `OutputGuard` substitutes the output returned to the caller.
```python
def scrub_emails(text: str) -> GuardResult:
cleaned = EMAIL_RE.sub('[email]', text)
return GuardResult.replace(cleaned) if cleaned != text else GuardResult.allow()
agent = Agent(
'openai:gpt-5.4',
capabilities=[
InputGuard(guard=scrub_emails), # strip PII before it reaches the model
OutputGuard(guard=scrub_emails), # strip PII before it reaches the caller
],
)
```
Input redaction requires sequential mode — it is incompatible with `parallel=True`, since a parallel guard runs alongside a model call that has already started with the original prompt.
## Retry (`retry`)
`OutputGuard` can send a bad output back to the model instead of blocking it. Return `GuardResult.retry(instruction)` — the instruction is the retry prompt the model sees. This reuses pydantic-ai's normal retry machinery and counts against the run's output-retry budget.
```python
def must_cite_sources(output: object) -> GuardResult:
if not has_citations(output):
return GuardResult.retry('Include at least one source citation.')
return GuardResult.allow()
OutputGuard(guard=must_cite_sources)
```
## Streaming
`OutputGuard` inspects the **final** output only — during `run_stream()` partial chunks reach the caller before the guard runs, so a `block` or `replace` verdict cannot un-send content already streamed. Use `run()` / `run_sync()` when the output must be screened before any of it is exposed. `GuardResult.retry()` is **not** supported under `run_stream()` — pydantic-ai does not retry output during streaming, and a `retry` verdict there surfaces as `UnexpectedModelBehavior`. `InputGuard` (including `parallel=True`) works the same in streamed and non-streamed runs.
## Tracing
`replace` and `block` are recorded as spans on the active OpenTelemetry tracer, so a redaction or refusal shows up in Logfire traces (`guardrail redacted input`, `guardrail blocked output`, etc.) with `guardrail.*` attributes. Content attributes — the original/replacement values for a redaction and the refusal `message` for a block — are attached **only** when `RunContext.trace_include_content` is enabled, since these can quote the very content the guard exists to keep out of traces. `retry` needs no special tracing: the retried model request appears in the trace on its own.
`OutputGuard` declares `position='outermost', wrapped_by=[Instrumentation]` so its block/redact spans are always captured by an enclosing `Instrumentation` span regardless of how the user orders capabilities. `InputGuard` declares `position='innermost'` so any capability that morphs messages (a prompt rewriter, a context manager) runs first and the guard sees the final prompt the model will receive.
## Parallel input guards
A slow guard (an LLM classifier, a network call) run sequentially adds its latency to every turn. Set `parallel=True` to run the guard concurrently with the model call instead, overlapping the two so the guard adds no latency on the pass path. The model call is cancelled the moment the guard reports a violation.
```python
InputGuard(guard=slow_async_classifier, parallel=True)
```
Parallel mode trades tokens for latency: sequential mode never calls the model when the guard blocks, but parallel mode has already started the model call — if the guard trips only after the model has responded, those tokens were spent. For fast local checks (regex, keyword lookup) sequential is the better default. `replace` is not available under `parallel=True` (see [Redaction](#redaction-replace)).
## Accessing run context
A guard may take a `RunContext` as its first parameter when it needs run state — `deps` for tenant- or role-aware policy, message history for conversation-aware checks. The parameter is detected from the signature, so prompt-only guards need not declare it:
```python
from pydantic_ai import RunContext
from pydantic_ai_harness import InputGuard
def tenant_policy(ctx: RunContext[MyDeps], prompt: str) -> bool:
return ctx.deps.tier == 'pro' or 'advanced-feature' not in prompt
InputGuard(guard=tenant_policy)
```
## Hard-fail path
`block` is the graceful path. To make the caller see an exception instead, raise from the guard:
```python
from pydantic_ai_harness import InputBlocked
def strict_guard(prompt: str) -> bool:
if contains_credentials(prompt):
raise InputBlocked('credentials detected')
return True
```
Any exception raised by the guard propagates as-is — use `InputBlocked` / `OutputBlocked` from this module, or your own exception types.
## API
```python
@dataclass
class GuardResult:
action: Literal['allow', 'block', 'replace', 'retry']
message: str | None = None
replacement: object | None = None
# classmethods: allow(), block(message=None), replace(value), retry(message)
InputGuard(
guard: Callable[..., bool | GuardResult | Awaitable[bool | GuardResult]],
parallel: bool = False,
)
OutputGuard(
guard: Callable[..., bool | GuardResult | Awaitable[bool | GuardResult]],
)
```
The guard callable takes the inspected value — the prompt for `InputGuard`, the output for `OutputGuard` — optionally preceded by a `RunContext`.
## Relationship to `pydantic-ai-shields`
`pydantic-ai-shields` provides opinionated implementations on top of these primitives (prompt-injection detectors, PII scrubbers, keyword blocklists, etc.). Use the guardrails here when you want to plug in your own validation logic; reach for shields when you need a batteries-included detector.
pydantic_ai_harness/guardrails/__init__.py
+25
View File
@@ -0,0 +1,25 @@
"""Input and output guardrails for Pydantic AI agents."""
from pydantic_ai_harness.guardrails._capability import (
GuardResult,
InputGuard,
InputGuardFunc,
OutputGuard,
OutputGuardFunc,
)
from pydantic_ai_harness.guardrails._exceptions import (
GuardrailError,
InputBlocked,
OutputBlocked,
)
__all__ = [
'GuardResult',
'GuardrailError',
'InputBlocked',
'InputGuard',
'InputGuardFunc',
'OutputBlocked',
'OutputGuard',
'OutputGuardFunc',
]
pydantic_ai_harness/guardrails/_capability.py
+456
View File
@@ -0,0 +1,456 @@
"""Input and output guardrail capabilities.
`InputGuard` intercepts the first model request and lets a user-supplied
callable decide what to do with the user prompt. `OutputGuard` runs as the
model output is processed and decides what to do with the agent output.
A guard returns a bare `bool` (`True` = allow) or a
[`GuardResult`][pydantic_ai_harness.guardrails.GuardResult] — one of four
outcomes:
- `allow` — let the value through unchanged.
- `block` — refuse: `InputGuard` short-circuits the model call with a refusal
message via [`SkipModelRequest`][pydantic_ai.exceptions.SkipModelRequest];
`OutputGuard` raises [`OutputBlocked`][pydantic_ai_harness.guardrails.OutputBlocked].
- `replace` — substitute a sanitized value (redaction) and continue.
- `retry` — send the output back to the model to try again (`OutputGuard` only).
A guard that raises propagates the exception so the caller sees a hard
failure. Guards may be sync or async and may optionally take a
[`RunContext`][pydantic_ai.tools.RunContext] as their first argument.
`replace` and `block` are recorded as spans on the active OpenTelemetry
tracer, so a redaction or refusal is visible in Logfire traces. Content
attributes (the original/replacement values and the refusal message) are
attached only when `RunContext.trace_include_content` is set.
"""
from __future__ import annotations
import asyncio
import inspect
from collections.abc import Awaitable, Callable, Sequence
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any, Literal
from pydantic_ai.capabilities import AbstractCapability, CapabilityOrdering, Instrumentation, WrapModelRequestHandler
from pydantic_ai.exceptions import ModelRetry, SkipModelRequest, UserError
from pydantic_ai.messages import ModelMessage, ModelResponse, TextPart, UserPromptPart
from pydantic_ai.tools import AgentDepsT, RunContext
from typing_extensions import assert_never
from pydantic_ai_harness.guardrails._exceptions import OutputBlocked
if TYPE_CHECKING:
from pydantic_ai.models import ModelRequestContext
from pydantic_ai.output import OutputContext
_DEFAULT_INPUT_BLOCK_MESSAGE = 'Request blocked by input guardrail.'
_DEFAULT_OUTPUT_BLOCK_MESSAGE = 'Output blocked by output guardrail.'
_DEFAULT_OUTPUT_RETRY_MESSAGE = 'Output rejected by output guardrail.'
@dataclass(frozen=True, kw_only=True)
class GuardResult:
"""The outcome a guard reports for the value it inspected.
Construct one with the classmethods — `GuardResult.allow()`,
`GuardResult.block()`, `GuardResult.replace()`, `GuardResult.retry()` —
rather than the raw fields. A guard may also return a bare `bool`: `True`
is `allow()`, `False` is `block()`.
"""
action: Literal['allow', 'block', 'replace', 'retry']
"""What the capability should do with the inspected value."""
message: str | None = None
"""For `block`, the refusal text. For `retry`, the instruction sent back to the model."""
replacement: object | None = None
"""For `replace`, the value substituted for the inspected one."""
def __post_init__(self) -> None:
"""Reject field combinations the four-outcome contract does not allow."""
match self.action:
case 'allow':
if self.message is not None or self.replacement is not None:
raise UserError("GuardResult(action='allow') must not set `message` or `replacement`.")
case 'replace':
if self.replacement is None:
raise UserError("GuardResult(action='replace') requires a `replacement` value.")
case 'retry':
if self.message is None:
raise UserError("GuardResult(action='retry') requires a `message`.")
case 'block':
# `message=None` is valid: a default is supplied at the use site.
pass
case _: # pragma: no cover - assert_never exhaustiveness guard
assert_never(self.action)
@classmethod
def allow(cls) -> GuardResult:
"""Let the value through unchanged."""
return cls(action='allow')
@classmethod
def block(cls, message: str | None = None) -> GuardResult:
"""Refuse the value. `message` is the refusal text; `None` uses a default."""
return cls(action='block', message=message)
@classmethod
def replace(cls, value: object) -> GuardResult:
"""Substitute `value` for the inspected one and continue.
For `InputGuard`, `value` is the replacement prompt text sent to the
model. For `OutputGuard`, it is the agent output returned to the caller.
"""
return cls(action='replace', replacement=value)
@classmethod
def retry(cls, message: str) -> GuardResult:
"""Send the output back to the model to try again — `OutputGuard` only.
`message` is the instruction the model sees on the retry.
"""
return cls(action='retry', message=message)
GuardOutcome = bool | GuardResult
"""What a guard callable returns: a bare `bool` (`True` = allow), or a `GuardResult`."""
InputGuardFunc = (
Callable[[str], GuardOutcome | Awaitable[GuardOutcome]]
| Callable[[RunContext[AgentDepsT], str], GuardOutcome | Awaitable[GuardOutcome]]
)
"""Signature of the callable passed to `InputGuard`.
The callable receives the user prompt and returns `True` / `GuardResult`. It
may optionally take a [`RunContext`][pydantic_ai.tools.RunContext] as a first
argument — for `deps`, message history, or other run state — and may be sync
or async. Raising an exception is treated as a hard failure and propagates up
to the caller.
"""
OutputGuardFunc = (
Callable[[object], GuardOutcome | Awaitable[GuardOutcome]]
| Callable[[RunContext[AgentDepsT], object], GuardOutcome | Awaitable[GuardOutcome]]
)
"""Signature of the callable passed to `OutputGuard`.
The callable receives the agent output unchanged — for typed outputs this is
the Pydantic model — and returns `True` / `GuardResult`. It may optionally take
a [`RunContext`][pydantic_ai.tools.RunContext] first, and may be sync or async.
"""
def _takes_ctx(func: Callable[..., object]) -> bool:
"""Return `True` when `func` declares a leading `RunContext` parameter.
Detected by parameter count, not annotation: a guard always takes the
guarded value, so a second parameter means it also wants the run context.
This matches pydantic-ai's own optional-`ctx` convention for output
validators. A callable whose signature cannot be introspected is treated
as taking the value only.
"""
try:
parameters = inspect.signature(func).parameters
except ValueError: # pragma: no cover - callable without an introspectable signature
return False
return len(parameters) > 1
async def _evaluate(
guard: Callable[..., GuardOutcome | Awaitable[GuardOutcome]],
ctx: RunContext[AgentDepsT],
value: object,
) -> GuardResult:
"""Call `guard` (passing `ctx` when declared), await it, and normalize to `GuardResult`."""
outcome = guard(ctx, value) if _takes_ctx(guard) else guard(value)
if inspect.isawaitable(outcome):
outcome = await outcome
if isinstance(outcome, GuardResult):
return outcome
return GuardResult.allow() if outcome else GuardResult.block()
def _extract_prompt(ctx: RunContext[AgentDepsT], messages: Sequence[ModelMessage]) -> str | None:
"""Return the text of the most recent user prompt, or `None` if absent.
Prefers `ctx.prompt` (set at run start) and falls back to scanning the
message history for the last [`UserPromptPart`][pydantic_ai.messages.UserPromptPart]
so that sub-agent calls or resumed runs without a fresh prompt still work.
"""
if ctx.prompt is not None:
return ctx.prompt if isinstance(ctx.prompt, str) else str(ctx.prompt)
for message in reversed(messages):
for part in reversed(message.parts):
if isinstance(part, UserPromptPart):
return part.content if isinstance(part.content, str) else str(part.content)
return None
def _replace_prompt(messages: Sequence[ModelMessage], new_content: str) -> bool:
"""Rewrite the most recent user prompt to `new_content`. Returns whether one was found."""
for message in reversed(messages):
for part in reversed(message.parts):
if isinstance(part, UserPromptPart):
part.content = new_content
return True
return False
def _trace_block(ctx: RunContext[AgentDepsT], *, direction: str, message: str) -> None:
"""Record a zero-duration span marking a guardrail refusal.
The refusal message is attached only when `ctx.trace_include_content` is
set — it can quote sensitive content from the guarded value, and ops
audiences are broader than the user who sees the refusal text.
"""
attributes: dict[str, str] = {'guardrail.direction': direction, 'guardrail.action': 'block'}
if ctx.trace_include_content:
attributes['guardrail.message'] = message
ctx.tracer.start_span(f'guardrail blocked {direction}', attributes=attributes).end()
def _trace_redaction(ctx: RunContext[AgentDepsT], *, direction: str, original: object, replacement: object) -> None:
"""Record a zero-duration span marking a guardrail redaction.
The original and replacement values are attached only when
`ctx.trace_include_content` is set, since a redacted value is often the
sensitive content the guard exists to keep out of traces.
"""
attributes: dict[str, str] = {'guardrail.direction': direction, 'guardrail.action': 'replace'}
if ctx.trace_include_content:
attributes['guardrail.original'] = str(original)
attributes['guardrail.replacement'] = str(replacement)
ctx.tracer.start_span(f'guardrail redacted {direction}', attributes=attributes).end()
@dataclass
class InputGuard(AbstractCapability[AgentDepsT]):
"""Validate the user prompt before it reaches the model.
The `guard` callable receives the prompt text and returns one of the four
outcomes (see the module docstring). `replace` rewrites the prompt sent to
the model and also overwrites the original in the run's message history,
so a redacted secret is not retained; a `str` replacement overwrites a
multimodal prompt's other parts. `retry` is not valid for an input guard.
```python
from pydantic_ai import Agent
from pydantic_ai_harness import GuardResult, InputGuard
def no_secrets(prompt: str) -> GuardResult:
if 'api_key' in prompt.lower():
return GuardResult.block('Your message looks like it contains an API key.')
return GuardResult.allow()
agent = Agent('openai:gpt-5.4', capabilities=[InputGuard(guard=no_secrets)])
```
The guard may take a [`RunContext`][pydantic_ai.tools.RunContext] as a
first parameter when it needs run state — `deps` for tenant/role-aware
policy, `messages` for conversation-aware checks. The parameter is detected
from the signature, so prompt-only guards stay as-is.
Set `parallel=True` to run the guard concurrently with the model call
rather than before it, overlapping a slow guard (an LLM classifier, a
network call) with the model round-trip so it adds no latency on the pass
path. The model call is cancelled the moment the guard reports a
violation. Trade-off: sequential mode never calls the model on a blocked
prompt, whereas parallel mode has already started it — if the guard trips
only after the model has responded, those tokens were still spent. Prefer
sequential for fast local guards. `replace` (redaction) is incompatible
with `parallel=True`, since the model call has already started with the
original prompt.
Scope: the guard runs exactly once per run — on the first model request —
and evaluates the original user prompt. Subsequent model requests in the
same run (e.g. after tool calls) are not re-checked, since the user input
has not changed.
Ordering: declares `position='innermost'` so any capability that morphs
the messages (a prompt rewriter, a context manager) runs first and the
guard sees the final prompt that will reach the model.
"""
guard: InputGuardFunc[AgentDepsT]
"""Callable that decides what to do with the prompt before it reaches the model."""
parallel: bool = False
"""Run the guard concurrently with the model request and cancel the model call on failure."""
def get_ordering(self) -> CapabilityOrdering:
"""Sit innermost so message-morphing capabilities run first and the guard sees the final prompt."""
return CapabilityOrdering(position='innermost')
async def _run_guard(
self,
ctx: RunContext[AgentDepsT],
request_context: ModelRequestContext,
prompt: str,
) -> None:
"""Evaluate the guard and act on its verdict.
`allow` returns; `block` raises `SkipModelRequest`; `replace` rewrites
the prompt in `request_context`; `retry` and `replace` under
`parallel=True` raise `UserError`.
"""
verdict = await _evaluate(self.guard, ctx, prompt)
match verdict.action:
case 'allow':
return
case 'retry':
raise UserError(
'An InputGuard guard cannot return GuardResult.retry() — retry applies to model output only.'
)
case 'block':
message = verdict.message or _DEFAULT_INPUT_BLOCK_MESSAGE
_trace_block(ctx, direction='input', message=message)
raise SkipModelRequest(ModelResponse(parts=[TextPart(content=message)]))
case 'replace':
if self.parallel:
raise UserError(
'InputGuard(parallel=True) is incompatible with GuardResult.replace(): the model call has '
'already started with the original prompt. Use sequential mode for prompt redaction.'
)
replacement = verdict.replacement
if not isinstance(replacement, str):
raise UserError(
'GuardResult.replace() for an input guard must provide replacement prompt text (str).'
)
if not _replace_prompt(request_context.messages, replacement):
raise UserError('InputGuard could not find a user prompt to redact in the request.')
_trace_redaction(ctx, direction='input', original=prompt, replacement=replacement)
case _: # pragma: no cover - assert_never exhaustiveness guard
assert_never(verdict.action)
async def wrap_model_request(
self,
ctx: RunContext[AgentDepsT],
*,
request_context: ModelRequestContext,
handler: WrapModelRequestHandler,
) -> ModelResponse:
"""Check the prompt before the first model call.
Sequential mode runs the guard then the model. `parallel=True` races
the guard against the model call and cancels it on a violation.
"""
if ctx.run_step > 1:
return await handler(request_context)
prompt = _extract_prompt(ctx, request_context.messages)
if prompt is None:
return await handler(request_context)
if not self.parallel:
await self._run_guard(ctx, request_context, prompt)
return await handler(request_context)
async def run_handler() -> ModelResponse:
return await handler(request_context)
guard_task: asyncio.Task[None] = asyncio.create_task(self._run_guard(ctx, request_context, prompt))
handler_task: asyncio.Task[ModelResponse] = asyncio.create_task(run_handler())
try:
done, _ = await asyncio.wait(
{guard_task, handler_task},
return_when=asyncio.FIRST_COMPLETED,
)
if guard_task in done:
await guard_task
return await handler_task
response = await handler_task
await guard_task
return response
finally:
for task in (guard_task, handler_task):
if not task.done():
task.cancel()
await asyncio.gather(guard_task, handler_task, return_exceptions=True)
@dataclass
class OutputGuard(AbstractCapability[AgentDepsT]):
"""Validate the agent output as it is produced.
The `guard` callable receives the output — no automatic stringification, so
a typed output arrives as the Pydantic model instance — and returns one of
the four outcomes (see the module docstring): `allow` exposes the output,
`block` raises [`OutputBlocked`][pydantic_ai_harness.guardrails.OutputBlocked],
`replace` substitutes a sanitized output (redaction), and `retry` sends the
output back to the model to try again.
```python
from pydantic_ai import Agent
from pydantic_ai_harness import GuardResult, OutputGuard
def no_pii(output: object) -> GuardResult:
if 'SSN' in str(output):
return GuardResult.retry('Do not include personal identifiers.')
return GuardResult.allow()
agent = Agent('openai:gpt-5.4', capabilities=[OutputGuard(guard=no_pii)])
```
The guard runs as the output is processed, so `retry` reuses pydantic-ai's
normal retry machinery and counts against the run's output-retry budget.
Like `InputGuard`, the guard may take a
[`RunContext`][pydantic_ai.tools.RunContext] as a first parameter; it is
detected from the signature.
Ordering: declares `position='outermost'` so the guard sees the final
output after every inner capability has processed it, and `wrapped_by=
[Instrumentation]` so an enclosing `Instrumentation` span always captures
the guard's block/redact spans regardless of user list order.
Streaming caveats. `retry` is supported with
[`run()`][pydantic_ai.Agent.run] / `run_sync()` only — pydantic-ai does not
support output retries during [`run_stream()`][pydantic_ai.Agent.run_stream],
where a `retry` verdict surfaces as
[`UnexpectedModelBehavior`][pydantic_ai.exceptions.UnexpectedModelBehavior].
The guard inspects only the final output, not partial chunks, so during a
streamed run the caller may already have received partial output before a
`block` or `replace` verdict is reached — use `run()` when the output must
be screened before any of it is exposed.
"""
guard: OutputGuardFunc[AgentDepsT]
"""Callable that decides what to do with the agent output."""
def get_ordering(self) -> CapabilityOrdering:
"""Sit outermost (inside `Instrumentation`) so the guard sees the final processed output."""
return CapabilityOrdering(position='outermost', wrapped_by=[Instrumentation])
async def after_output_process(
self,
ctx: RunContext[AgentDepsT],
*,
output_context: OutputContext,
output: Any,
) -> Any:
"""Evaluate the guard against the processed output and act on its verdict."""
if ctx.partial_output:
return output
verdict = await _evaluate(self.guard, ctx, output)
match verdict.action:
case 'allow':
return output
case 'block':
message = verdict.message or _DEFAULT_OUTPUT_BLOCK_MESSAGE
_trace_block(ctx, direction='output', message=message)
raise OutputBlocked(message)
case 'retry':
raise ModelRetry(verdict.message or _DEFAULT_OUTPUT_RETRY_MESSAGE)
case 'replace':
_trace_redaction(ctx, direction='output', original=output, replacement=verdict.replacement)
return verdict.replacement
case _: # pragma: no cover - assert_never exhaustiveness guard
assert_never(verdict.action)
pydantic_ai_harness/guardrails/_exceptions.py
+20
View File
@@ -0,0 +1,20 @@
"""Exceptions raised by the guardrail capabilities."""
from __future__ import annotations
class GuardrailError(Exception):
"""Base exception for guardrail violations."""
class InputBlocked(GuardrailError):
"""Raised by a user-supplied input guard to hard-fail a run.
Prefer returning `False` from the guard callable to trigger a graceful
refusal via [`SkipModelRequest`][pydantic_ai.exceptions.SkipModelRequest].
Raise this explicitly when the caller should have to handle the failure.
"""
class OutputBlocked(GuardrailError):
"""Raised by [`OutputGuard`][pydantic_ai_harness.OutputGuard] when the final output fails validation."""
pydantic_ai_harness/media/README.md
+64
View File
@@ -0,0 +1,64 @@
# Media
> [!NOTE]
> Import these helpers from their submodule -- there is no top-level `pydantic_ai_harness` re-export:
>
> ```python
> from pydantic_ai_harness.media import (
> DiskMediaStore,
> S3MediaStore,
> SqliteMediaStore,
> externalize_media,
> restore_media,
> )
> ```
>
> The API may change between releases. Where practical, breaking changes ship with a deprecation warning.
Content-addressed stores and walker helpers that move large binary payloads out of message history and put them back on demand.
These are building blocks, not a capability. There is no class you add to `Agent(capabilities=[...])` yet. [`StepPersistence`](../step_persistence/) uses them to keep snapshots small when messages carry `BinaryContent`. A forthcoming `MediaExternalizer` capability will reuse the same stores to rewrite `BinaryContent` into URL parts before the model sees them.
## Why
A conversation that carries images, audio, or other `BinaryContent` inlines those bytes into every message. Persist that history and each snapshot re-serializes the payloads. Content-addressed storage writes each payload once, keyed by its own hash, and leaves a short `media://` URI in its place. The same bytes are stored once no matter how many messages or snapshots reference them.
## Stores
Every store implements the `MediaStore` protocol: `put`, `get`, `exists`, `public_url`, and `get_metadata`, all async and content-addressed (the URI is derived from the payload hash, so identical bytes deduplicate).
| Store | Backed by | Use when |
|---|---|---|
| `DiskMediaStore(directory=...)` | A directory on disk | Local runs and tests |
| `SqliteMediaStore(...)` | A SQLite database | A single-file store that travels with the data |
| `S3MediaStore(...)` | S3 or an S3-compatible bucket | Shared or production storage |
A `KeyStrategy` controls the on-store layout, and a `PublicUrlResolver` (or `make_static_public_url`) turns a stored URI into a public URL when the store is served over HTTP.
## Walker helpers
`externalize_media` and `restore_media` walk a message node and swap payloads for URIs and back.
```python
store = DiskMediaStore(directory='./media')
# Replace BinaryContent larger than the threshold with media:// URIs.
lean = await externalize_media(message, media_store=store, threshold_bytes=32_000)
# Later, rehydrate the URIs back into BinaryContent.
full = await restore_media(lean, media_store=store)
```
`externalize_media` only externalizes payloads over `threshold_bytes`; smaller ones stay inline. `media_uri_for` and `parse_media_uri` give you the raw URI round-trip if you need to key media yourself.
## API
| Symbol | Purpose |
|---|---|
| `MediaStore` | Async content-addressed store protocol (`put` / `get` / `exists` / `public_url` / `get_metadata`) |
| `DiskMediaStore`, `SqliteMediaStore`, `S3MediaStore` | Concrete stores |
| `MediaContext` | Per-call context (e.g. tenant) threaded through store operations |
| `KeyStrategy`, `default_key_strategy` | On-store key layout |
| `PublicUrlResolver`, `make_static_public_url` | Resolve a stored URI to a public URL |
| `externalize_media`, `restore_media` | Walk a message node to externalize / rehydrate payloads |
| `media_uri_for`, `parse_media_uri` | Compute and parse a `media://` URI |
pydantic_ai_harness/media/__init__.py
+38
View File
@@ -0,0 +1,38 @@
"""Content-addressed media stores for offloading large binary parts.
Used by `pydantic_ai_harness.step_persistence` to keep snapshots small when
messages carry `BinaryContent` payloads. A forthcoming `MediaExternalizer`
capability will reuse these stores for in-flight wire-payload reduction
(rewriting `BinaryContent` to URL parts before the model sees them).
"""
from pydantic_ai_harness.media._s3 import S3MediaStore
from pydantic_ai_harness.media._store import (
DiskMediaStore,
KeyStrategy,
MediaContext,
MediaStore,
PublicUrlResolver,
SqliteMediaStore,
default_key_strategy,
make_static_public_url,
media_uri_for,
parse_media_uri,
)
from pydantic_ai_harness.media._walker import externalize_media, restore_media
__all__ = [
'DiskMediaStore',
'KeyStrategy',
'MediaContext',
'MediaStore',
'PublicUrlResolver',
'S3MediaStore',
'SqliteMediaStore',
'default_key_strategy',
'externalize_media',
'make_static_public_url',
'media_uri_for',
'parse_media_uri',
'restore_media',
]
pydantic_ai_harness/experimental/media/_s3.py → pydantic_ai_harness/media/_s3.py
+1 -2
View File
@@ -22,7 +22,7 @@ from datetime import datetime, timezone
from typing import TYPE_CHECKING
from urllib.parse import quote
from pydantic_ai_harness.experimental.media._store import (
from pydantic_ai_harness.media._store import (
_EMPTY_CONTEXT, # pyright: ignore[reportPrivateUsage]
KeyStrategy,
MediaContext,
@@ -36,7 +36,6 @@ from pydantic_ai_harness.experimental.media._store import (
if TYPE_CHECKING:
import httpx
_ALGORITHM = 'AWS4-HMAC-SHA256'
_SERVICE = 's3'
# Conservative subset for x-amz-meta-* keys: ASCII letters/digits/dash; lowercase.
pydantic_ai_harness/experimental/media/_store.py → pydantic_ai_harness/media/_store.py
-1
View File
@@ -61,7 +61,6 @@ class MediaContext:
_EMPTY_CONTEXT = MediaContext()
PublicUrlResolver = Callable[[str, MediaContext], 'str | None | Awaitable[str | None]']
"""User-supplied callable that turns a `media+sha256://<hex>` URI into a fetchable URL.
pydantic_ai_harness/experimental/media/_walker.py → pydantic_ai_harness/media/_walker.py
+1 -1
View File
@@ -16,7 +16,7 @@ from __future__ import annotations
import base64
from typing import TypeGuard
from pydantic_ai_harness.experimental.media._store import MediaContext, MediaStore
from pydantic_ai_harness.media._store import MediaContext, MediaStore
_EXTERNAL_MARKER = '__harness_external_media__'
pydantic_ai_harness/experimental/modal_sandbox/README.md → pydantic_ai_harness/modal_sandbox/README.md
+4 -4
View File
@@ -20,7 +20,7 @@ container is the isolation boundary.
```python
from pydantic_ai import Agent
from pydantic_ai_harness.experimental.modal_sandbox import ModalSandboxCapability
from pydantic_ai_harness.modal_sandbox import ModalSandboxCapability
agent = Agent('anthropic:claude-sonnet-4-6', capabilities=[ModalSandboxCapability()])
@@ -100,7 +100,7 @@ terminates it, so the owner decides when the sandbox goes away, and can read its
```python
from pydantic_ai import Agent
from pydantic_ai_harness.experimental.modal_sandbox import ModalSandboxCapability, ModalSandboxSession
from pydantic_ai_harness.modal_sandbox import ModalSandboxCapability, ModalSandboxSession
async with ModalSandboxSession(image='python:3.12-slim') as session:
print(session.sandbox_id) # the running sandbox id
@@ -140,7 +140,7 @@ internals can change without affecting the tools. The session is also usable on
its own as a lower-level async context manager:
```python
from pydantic_ai_harness.experimental.modal_sandbox import ModalSandboxSession
from pydantic_ai_harness.modal_sandbox import ModalSandboxSession
async with ModalSandboxSession(image='python:3.12-slim') as session:
result = await session.exec(['echo', 'hello'])
@@ -223,7 +223,7 @@ capabilities:
```python
from pydantic_ai import Agent
from pydantic_ai_harness.experimental.modal_sandbox import ModalSandboxCapability
from pydantic_ai_harness.modal_sandbox import ModalSandboxCapability
agent = Agent.from_file('agent.yaml', custom_capability_types=[ModalSandboxCapability])
```
pydantic_ai_harness/experimental/modal_sandbox/__init__.py → pydantic_ai_harness/modal_sandbox/__init__.py
+3 -6
View File
@@ -9,17 +9,14 @@ can change without disturbing the capability or its tool surface. Treat the
lower-level pieces as more likely to change than `ModalSandboxCapability` itself.
"""
from pydantic_ai_harness.experimental._warn import warn_experimental
from pydantic_ai_harness.experimental.modal_sandbox._capability import ModalSandboxCapability
from pydantic_ai_harness.experimental.modal_sandbox._session import (
from pydantic_ai_harness.modal_sandbox._capability import ModalSandboxCapability
from pydantic_ai_harness.modal_sandbox._session import (
ModalSandboxError,
ModalSandboxSession,
ModalSandboxTerminalError,
ModalSandboxUnavailableError,
)
from pydantic_ai_harness.experimental.modal_sandbox._toolset import ModalSandboxToolset
warn_experimental('modal_sandbox')
from pydantic_ai_harness.modal_sandbox._toolset import ModalSandboxToolset
__all__ = [
'ModalSandboxCapability',
pydantic_ai_harness/experimental/modal_sandbox/_capability.py → pydantic_ai_harness/modal_sandbox/_capability.py
+3 -3
View File
@@ -9,8 +9,8 @@ from pydantic_ai.tools import AgentDepsT
from pydantic_ai.toolsets import AgentToolset
from pydantic_ai_harness._tool_output import DEFAULT_MAX_LINES
from pydantic_ai_harness.experimental.modal_sandbox._session import ModalSandboxSession
from pydantic_ai_harness.experimental.modal_sandbox._toolset import ModalSandboxToolset
from pydantic_ai_harness.modal_sandbox._session import ModalSandboxSession
from pydantic_ai_harness.modal_sandbox._toolset import ModalSandboxToolset
# Defaults shared by the field declarations and the validation below, so the two cannot
# drift: a setting is "left at its default" iff it equals the constant here.
@@ -53,7 +53,7 @@ class ModalSandboxCapability(AbstractCapability[AgentDepsT]):
```python
from pydantic_ai import Agent
from pydantic_ai_harness.experimental.modal_sandbox import ModalSandboxCapability
from pydantic_ai_harness.modal_sandbox import ModalSandboxCapability
agent = Agent('anthropic:claude-sonnet-4-6', capabilities=[ModalSandboxCapability()])
result = agent.run_sync('Write a Python script that prints the first 10 primes and run it.')
pydantic_ai_harness/experimental/modal_sandbox/_session.py → pydantic_ai_harness/modal_sandbox/_session.py
View File
pydantic_ai_harness/experimental/modal_sandbox/_toolset.py → pydantic_ai_harness/modal_sandbox/_toolset.py
+1 -1
View File
@@ -12,7 +12,7 @@ from pydantic_ai.toolsets import AbstractToolset, FunctionToolset
from typing_extensions import Self
from pydantic_ai_harness._tool_output import guard_read_size, render_file_window, truncate_output
from pydantic_ai_harness.experimental.modal_sandbox._session import (
from pydantic_ai_harness.modal_sandbox._session import (
ModalSandboxError,
ModalSandboxSession,
ModalSandboxTerminalError,
pydantic_ai_harness/experimental/overflow/README.md → pydantic_ai_harness/overflowing_tool_output/README.md
+5 -15
View File
@@ -1,23 +1,13 @@
# Overflow capability
> [!WARNING]
> **Experimental.** This capability lives under `pydantic_ai_harness.experimental` and may
> change or be removed in any release, without a deprecation period. Import it from the
> experimental path -- there is no top-level export:
> [!NOTE]
> Import this capability from its submodule -- there is no top-level `pydantic_ai_harness` re-export:
>
> ```python
> from pydantic_ai_harness.experimental.overflow import OverflowingToolOutput
> from pydantic_ai_harness.overflowing_tool_output import OverflowingToolOutput
> ```
>
> Importing any experimental capability emits a `HarnessExperimentalWarning`. Silence **all**
> harness experimental warnings with a single filter (no per-capability lines needed):
>
> ```python
> import warnings
> from pydantic_ai_harness.experimental import HarnessExperimentalWarning
>
> warnings.filterwarnings('ignore', category=HarnessExperimentalWarning)
> ```
> The API may change between releases. Where practical, breaking changes ship with a deprecation warning.
A tool can return a payload large enough to dominate the context window. Tool returns
persist in history as `ToolReturnPart`s, so an oversized one is re-sent on every later
@@ -60,7 +50,7 @@ that fits wins; anything below the smallest threshold passes through.
```python
from pydantic_ai import Agent
from pydantic_ai_harness.experimental.overflow import (
from pydantic_ai_harness.overflowing_tool_output import (
Band,
OverflowingToolOutput,
Spill,
pydantic_ai_harness/overflowing_tool_output/__init__.py
+39
View File
@@ -0,0 +1,39 @@
"""Overflow capability: reduce oversized tool returns at production time.
`OverflowingToolOutput` intercepts a tool return when it is produced and reduces it --
truncating, spilling to a queryable file, or summarizing -- so an oversized payload does
not persist in history and get re-sent on every later model request. Combine the three
modes through an ordered list of size `bands`.
Spilled payloads are read back on demand through the registered `read_tool_result` tool;
the `OverflowStore` protocol is the seam for a durable backend (the local-file default
ships for single-process runs).
"""
from pydantic_ai_harness.overflowing_tool_output._bands import (
Action,
Band,
Passthrough,
Spill,
Summarize,
SummarizeFunc,
Truncate,
)
from pydantic_ai_harness.overflowing_tool_output._capability import READ_TOOL_NAME, OverflowingToolOutput
from pydantic_ai_harness.overflowing_tool_output._payload import TruncationStrategy
from pydantic_ai_harness.overflowing_tool_output._store import LocalFileStore, OverflowStore
__all__ = [
'READ_TOOL_NAME',
'Action',
'Band',
'LocalFileStore',
'OverflowStore',
'OverflowingToolOutput',
'Passthrough',
'Spill',
'Summarize',
'SummarizeFunc',
'Truncate',
'TruncationStrategy',
]
pydantic_ai_harness/experimental/overflow/_bands.py → pydantic_ai_harness/overflowing_tool_output/_bands.py
+1 -1
View File
@@ -16,7 +16,7 @@ from collections.abc import Awaitable, Callable
from dataclasses import dataclass
from typing import TYPE_CHECKING
from pydantic_ai_harness.experimental.overflow._payload import TruncationStrategy
from pydantic_ai_harness.overflowing_tool_output._payload import TruncationStrategy
if TYPE_CHECKING:
from pydantic_ai.models import Model
pydantic_ai_harness/experimental/overflow/_capability.py → pydantic_ai_harness/overflowing_tool_output/_capability.py
+4 -4
View File
@@ -14,7 +14,7 @@ from pydantic_ai.messages import ToolCallPart, ToolReturn, ToolReturnContent, Us
from pydantic_ai.tools import AgentDepsT, RunContext, ToolDefinition, ToolSelector, matches_tool_selector
from pydantic_ai.toolsets import AgentToolset
from pydantic_ai_harness.experimental.overflow._bands import (
from pydantic_ai_harness.overflowing_tool_output._bands import (
Action,
Band,
Passthrough,
@@ -22,7 +22,7 @@ from pydantic_ai_harness.experimental.overflow._bands import (
Summarize,
Truncate,
)
from pydantic_ai_harness.experimental.overflow._payload import (
from pydantic_ai_harness.overflowing_tool_output._payload import (
is_binary,
json_sketch,
measure,
@@ -31,7 +31,7 @@ from pydantic_ai_harness.experimental.overflow._payload import (
to_text,
truncate_text,
)
from pydantic_ai_harness.experimental.overflow._store import LocalFileStore, OverflowStore
from pydantic_ai_harness.overflowing_tool_output._store import LocalFileStore, OverflowStore
READ_TOOL_NAME = 'read_tool_result'
"""Name of the registered read-back tool. Its own returns are exempt from reduction."""
@@ -96,7 +96,7 @@ class OverflowingToolOutput(AbstractCapability[AgentDepsT]):
Example:
```python
from pydantic_ai import Agent
from pydantic_ai_harness.experimental.overflow import (
from pydantic_ai_harness.overflowing_tool_output import (
Band,
OverflowingToolOutput,
Spill,
pydantic_ai_harness/experimental/overflow/_payload.py → pydantic_ai_harness/overflowing_tool_output/_payload.py
+1 -1
View File
@@ -15,7 +15,7 @@ from typing import TypeGuard
from pydantic_ai.messages import ModelMessage, ModelRequest, SystemPromptPart
from pydantic_core import to_json
from pydantic_ai_harness.experimental.compaction._shared import estimate_token_count
from pydantic_ai_harness.compaction._shared import estimate_token_count
class TruncationStrategy(str, Enum):
pydantic_ai_harness/experimental/overflow/_store.py → pydantic_ai_harness/overflowing_tool_output/_store.py
View File
pydantic_ai_harness/experimental/planning/README.md → pydantic_ai_harness/planning/README.md
+6 -16
View File
@@ -1,23 +1,13 @@
# Planning
> [!WARNING]
> **Experimental.** This capability lives under `pydantic_ai_harness.experimental` and may
> change or be removed in any release, without a deprecation period. Import it from the
> experimental path -- there is no top-level export:
> [!NOTE]
> Import this capability from its submodule -- there is no top-level `pydantic_ai_harness` re-export:
>
> ```python
> from pydantic_ai_harness.experimental.planning import Planning
> from pydantic_ai_harness.planning import Planning
> ```
>
> Importing any experimental capability emits a `HarnessExperimentalWarning`. Silence **all**
> harness experimental warnings with a single filter (no per-capability lines needed):
>
> ```python
> import warnings
> from pydantic_ai_harness.experimental import HarnessExperimentalWarning
>
> warnings.filterwarnings('ignore', category=HarnessExperimentalWarning)
> ```
> The API may change between releases. Where practical, breaking changes ship with a deprecation warning.
Give an agent a structured, self-updating task plan -- without ever invalidating the prompt cache.
@@ -36,7 +26,7 @@ So the plan stays current in the model's view while the cached prefix is never i
```python
from pydantic_ai import Agent
from pydantic_ai_harness.experimental.planning import Planning
from pydantic_ai_harness.planning import Planning
agent = Agent('anthropic:claude-sonnet-4-6', capabilities=[Planning()])
@@ -109,7 +99,7 @@ capabilities:
```python
from pydantic_ai import Agent
from pydantic_ai_harness.experimental.planning import Planning
from pydantic_ai_harness.planning import Planning
agent = Agent.from_file('agent.yaml', custom_capability_types=[Planning])
```
pydantic_ai_harness/planning/__init__.py
+6
View File
@@ -0,0 +1,6 @@
"""Planning capability: model-owned, cache-friendly task planning for agents."""
from pydantic_ai_harness.planning._capability import Planning
from pydantic_ai_harness.planning._toolset import PlanItem, PlanningToolset, TaskStatus
__all__ = ['PlanItem', 'Planning', 'PlanningToolset', 'TaskStatus']
pydantic_ai_harness/experimental/planning/_capability.py → pydantic_ai_harness/planning/_capability.py
+2 -3
View File
@@ -10,14 +10,13 @@ from pydantic_ai.messages import CachePoint, ModelRequest, ModelResponse, UserPr
from pydantic_ai.tools import AgentDepsT, RunContext
from pydantic_ai.toolsets import AgentToolset
from pydantic_ai_harness.experimental.planning._toolset import PlanningToolset, PlanState, render_plan
from pydantic_ai_harness.planning._toolset import PlanningToolset, PlanState, render_plan
if TYPE_CHECKING:
from pydantic_ai._instructions import AgentInstructions
from pydantic_ai.capabilities.abstract import WrapModelRequestHandler
from pydantic_ai.models import ModelRequestContext
_DEFAULT_GUIDANCE = (
'You have a planning tool, `write_plan`. For multi-step work, call it first to lay out the '
'steps, then call it again to update statuses as you start and finish each step. Pass the '
@@ -41,7 +40,7 @@ class Planning(AbstractCapability[AgentDepsT]):
```python
from pydantic_ai import Agent
from pydantic_ai_harness.experimental.planning import Planning
from pydantic_ai_harness.planning import Planning
agent = Agent('anthropic:claude-sonnet-4-6', capabilities=[Planning()])
```
pydantic_ai_harness/experimental/planning/_toolset.py → pydantic_ai_harness/planning/_toolset.py
View File
pydantic_ai_harness/experimental/authoring/README.md → pydantic_ai_harness/runtime_authoring/README.md
+5 -15
View File
@@ -1,23 +1,13 @@
# RuntimeAuthoring
> [!WARNING]
> **Experimental.** This capability lives under `pydantic_ai_harness.experimental` and may
> change or be removed in any release, without a deprecation period. Import it from the
> experimental path -- there is no top-level export:
> [!NOTE]
> Import this capability from its submodule -- there is no top-level `pydantic_ai_harness` re-export:
>
> ```python
> from pydantic_ai_harness.experimental.authoring import RuntimeAuthoring
> from pydantic_ai_harness.runtime_authoring import RuntimeAuthoring
> ```
>
> Importing any experimental capability emits a `HarnessExperimentalWarning`. Silence **all**
> harness experimental warnings with a single filter (no per-capability lines needed):
>
> ```python
> import warnings
> from pydantic_ai_harness.experimental import HarnessExperimentalWarning
>
> warnings.filterwarnings('ignore', category=HarnessExperimentalWarning)
> ```
> The API may change between releases. Where practical, breaking changes ship with a deprecation warning.
Let an agent author, validate, and persist real pydantic-ai capabilities at runtime.
@@ -50,7 +40,7 @@ authoring a hook means authoring a capability that overrides one lifecycle metho
from pathlib import Path
from pydantic_ai import Agent
from pydantic_ai_harness.experimental.authoring import RuntimeAuthoring
from pydantic_ai_harness.runtime_authoring import RuntimeAuthoring
authoring = RuntimeAuthoring(directory=Path('.authored'))
agent = Agent('anthropic:claude-sonnet-4-6', capabilities=[authoring])
pydantic_ai_harness/runtime_authoring/__init__.py
+20
View File
@@ -0,0 +1,20 @@
"""Runtime capability authoring: let an agent write, validate, and register real capabilities."""
from pydantic_ai_harness.runtime_authoring._capability import RuntimeAuthoring
from pydantic_ai_harness.runtime_authoring._store import AuthoredCapability, CapabilityStore
from pydantic_ai_harness.runtime_authoring._toolset import AuthoringToolset
from pydantic_ai_harness.runtime_authoring._validate import (
CapabilityValidationError,
load_capability_instance,
validate_capability_file,
)
__all__ = [
'AuthoredCapability',
'AuthoringToolset',
'CapabilityStore',
'CapabilityValidationError',
'RuntimeAuthoring',
'load_capability_instance',
'validate_capability_file',
]
pydantic_ai_harness/experimental/authoring/_capability.py → pydantic_ai_harness/runtime_authoring/_capability.py
+3 -4
View File
@@ -10,13 +10,12 @@ from pydantic_ai.capabilities import AbstractCapability
from pydantic_ai.tools import AgentDepsT
from pydantic_ai.toolsets import AgentToolset
from pydantic_ai_harness.experimental.authoring._store import CapabilityStore
from pydantic_ai_harness.experimental.authoring._toolset import AuthoringToolset
from pydantic_ai_harness.runtime_authoring._store import CapabilityStore
from pydantic_ai_harness.runtime_authoring._toolset import AuthoringToolset
if TYPE_CHECKING:
from pydantic_ai._instructions import AgentInstructions
_DEFAULT_GUIDANCE = (
'You can author new pydantic-ai capabilities at runtime with `author_capability(name, code)`. '
'A capability is a subclass of `pydantic_ai.capabilities.AbstractCapability` that constructs with '
@@ -48,7 +47,7 @@ class RuntimeAuthoring(AbstractCapability[AgentDepsT]):
from pathlib import Path
from pydantic_ai import Agent
from pydantic_ai_harness.experimental.authoring import RuntimeAuthoring
from pydantic_ai_harness.runtime_authoring import RuntimeAuthoring
authoring = RuntimeAuthoring(directory=Path('.authored'))
agent = Agent('anthropic:claude-sonnet-4-6', capabilities=[authoring])
pydantic_ai_harness/experimental/authoring/_store.py → pydantic_ai_harness/runtime_authoring/_store.py
+1 -1
View File
@@ -19,7 +19,7 @@ from typing import Literal
from pydantic import BaseModel, ValidationError
from pydantic_ai.capabilities import AbstractCapability
from pydantic_ai_harness.experimental.authoring._validate import (
from pydantic_ai_harness.runtime_authoring._validate import (
CapabilityValidationError,
load_capability_instance,
validate_capability_file,
pydantic_ai_harness/experimental/authoring/_toolset.py → pydantic_ai_harness/runtime_authoring/_toolset.py
+1 -1
View File
@@ -6,7 +6,7 @@ from pydantic_ai.exceptions import ModelRetry
from pydantic_ai.tools import AgentDepsT
from pydantic_ai.toolsets import FunctionToolset
from pydantic_ai_harness.experimental.authoring._store import CapabilityStore
from pydantic_ai_harness.runtime_authoring._store import CapabilityStore
class AuthoringToolset(FunctionToolset[AgentDepsT]):
pydantic_ai_harness/experimental/authoring/_validate.py → pydantic_ai_harness/runtime_authoring/_validate.py
View File
pydantic_ai_harness/shell/_capability.py
+1 -2
View File
@@ -8,7 +8,6 @@ from pathlib import Path
from pydantic_ai.capabilities import AbstractCapability
from pydantic_ai.tools import AgentDepsT
from pydantic_ai.toolsets import AgentToolset
from pydantic_ai_harness.shell._toolset import ShellToolset
@@ -100,7 +99,7 @@ class Shell(AbstractCapability[AgentDepsT]):
`LLM_API_KEY_ENV_PATTERNS` for a ready-made provider-credential denylist.
"""
def get_toolset(self) -> AgentToolset[AgentDepsT]:
def get_toolset(self) -> ShellToolset[AgentDepsT]:
"""Build and return the shell toolset."""
return ShellToolset[AgentDepsT](
cwd=Path(self.cwd),
pydantic_ai_harness/experimental/step_persistence/README.md → pydantic_ai_harness/step_persistence/README.md
+14 -23
View File
@@ -1,24 +1,15 @@
# StepPersistence
> [!WARNING]
> **Experimental.** `StepPersistence` and the `media` stores live under
> `pydantic_ai_harness.experimental` and may change or be removed in any release, without a
> deprecation period. Import them from the experimental path -- there is no top-level export:
> [!NOTE]
> Import `StepPersistence` and the `media` stores from their submodules -- there is no top-level
> `pydantic_ai_harness` re-export:
>
> ```python
> from pydantic_ai_harness.experimental.step_persistence import StepPersistence
> from pydantic_ai_harness.experimental.media import S3MediaStore
> from pydantic_ai_harness.step_persistence import StepPersistence
> from pydantic_ai_harness.media import S3MediaStore
> ```
>
> Importing either package emits a `HarnessExperimentalWarning`. Silence **all** harness
> experimental warnings with a single filter (no per-capability lines needed):
>
> ```python
> import warnings
> from pydantic_ai_harness.experimental import HarnessExperimentalWarning
>
> warnings.filterwarnings('ignore', category=HarnessExperimentalWarning)
> ```
> The API may change between releases. Where practical, breaking changes ship with a deprecation warning.
`StepPersistence` records what an agent did at each boundary, separate from
whether the run can be safely resumed. It is the persistence substrate for
@@ -52,7 +43,7 @@ snapshots, and graph-node resume are out of scope and tracked separately
```python
from pydantic_ai import Agent
from pydantic_ai_harness.experimental.step_persistence import StepPersistence, InMemoryStepStore
from pydantic_ai_harness.step_persistence import StepPersistence, InMemoryStepStore
store = InMemoryStepStore()
librarian = Agent(
@@ -128,7 +119,7 @@ context". `StepPersistence` does not introduce a parallel mechanism -- it
exposes one helper that loads the most recent provider-valid snapshot:
```python
from pydantic_ai_harness.experimental.step_persistence import continue_run
from pydantic_ai_harness.step_persistence import continue_run
# Earlier: tag the first turn with a conversation id so the follow-up can find it.
await librarian.run(
@@ -262,7 +253,7 @@ write external state should annotate their in-flight `ToolEffectRecord`
via `annotate_tool_effect`:
```python
from pydantic_ai_harness.experimental.step_persistence import annotate_tool_effect
from pydantic_ai_harness.step_persistence import annotate_tool_effect
@orchestrator.tool
async def set_label(ctx: RunContext[Deps], issue: int, label: str) -> str:
@@ -330,8 +321,8 @@ original bytes.
Override the destination by passing your own `MediaStore`:
```python
from pydantic_ai_harness.experimental.step_persistence import FileStepStore
from pydantic_ai_harness.experimental.media import S3MediaStore
from pydantic_ai_harness.step_persistence import FileStepStore
from pydantic_ai_harness.media import S3MediaStore
store = FileStepStore(
'runs',
@@ -379,7 +370,7 @@ without re-encoding bytes into the request body.
Static base URL (public R2 bucket, CDN):
```python
from pydantic_ai_harness.experimental.media import S3MediaStore, make_static_public_url
from pydantic_ai_harness.media import S3MediaStore, make_static_public_url
store = S3MediaStore(
bucket='my-bucket',
@@ -395,7 +386,7 @@ Presigned / rotating-signature URL -- pass any async callable that takes
`(uri, MediaContext)`:
```python
from pydantic_ai_harness.experimental.media import MediaContext, S3MediaStore
from pydantic_ai_harness.media import MediaContext, S3MediaStore
async def presign(uri: str, ctx: MediaContext) -> str:
key = 'media/' + uri.removeprefix('media+sha256://') + '.bin'
@@ -444,7 +435,7 @@ primary key is the digest, so a user-chosen key would either break
dedup or be a no-op):
```python
from pydantic_ai_harness.experimental.media import DiskMediaStore, MediaContext
from pydantic_ai_harness.media import DiskMediaStore, MediaContext
def by_media_type(uri: str, ctx: MediaContext) -> str:
digest = uri.removeprefix('media+sha256://')
pydantic_ai_harness/step_persistence/__init__.py
+41
View File
@@ -0,0 +1,41 @@
"""Step-event persistence: append-only event log, continuable snapshots, tool-effect ledger."""
from pydantic_ai_harness.step_persistence._capability import StepPersistence
from pydantic_ai_harness.step_persistence._helpers import (
annotate_tool_effect,
continue_run,
fork_run,
is_provider_valid,
)
from pydantic_ai_harness.step_persistence._store import (
FileStepStore,
InMemoryStepStore,
SqliteStepStore,
StepStore,
)
from pydantic_ai_harness.step_persistence._types import (
ContinuableSnapshot,
EventKind,
RunRecord,
StepEvent,
ToolEffectRecord,
ToolEffectStatus,
)
__all__ = [
'ContinuableSnapshot',
'EventKind',
'FileStepStore',
'InMemoryStepStore',
'RunRecord',
'SqliteStepStore',
'StepEvent',
'StepPersistence',
'StepStore',
'ToolEffectRecord',
'ToolEffectStatus',
'annotate_tool_effect',
'continue_run',
'fork_run',
'is_provider_valid',
]
pydantic_ai_harness/experimental/step_persistence/_capability.py → pydantic_ai_harness/step_persistence/_capability.py
+7 -7
View File
@@ -15,10 +15,10 @@ from pydantic_ai.models import ModelRequestContext
from pydantic_ai.run import AgentRunResult
from pydantic_ai.tools import AgentDepsT, RunContext, ToolDefinition
from pydantic_ai_harness.experimental.step_persistence._context import current_run_id, snapshot_saved
from pydantic_ai_harness.experimental.step_persistence._helpers import is_provider_valid
from pydantic_ai_harness.experimental.step_persistence._store import InMemoryStepStore, StepStore
from pydantic_ai_harness.experimental.step_persistence._types import (
from pydantic_ai_harness.step_persistence._context import current_run_id, snapshot_saved
from pydantic_ai_harness.step_persistence._helpers import is_provider_valid
from pydantic_ai_harness.step_persistence._store import InMemoryStepStore, StepStore
from pydantic_ai_harness.step_persistence._types import (
ContinuableSnapshot,
EventKind,
RunRecord,
@@ -49,7 +49,7 @@ class StepPersistence(AbstractCapability[AgentDepsT]):
```python
from pydantic_ai import Agent
from pydantic_ai_harness.experimental.step_persistence import StepPersistence, InMemoryStepStore
from pydantic_ai_harness.step_persistence import StepPersistence, InMemoryStepStore
store = InMemoryStepStore()
librarian = Agent(
@@ -124,12 +124,12 @@ class StepPersistence(AbstractCapability[AgentDepsT]):
if backend == 'memory':
return cls(store=InMemoryStepStore(), **kwargs)
if backend == 'file':
from pydantic_ai_harness.experimental.step_persistence._store import FileStepStore
from pydantic_ai_harness.step_persistence._store import FileStepStore
directory = kwargs.pop('directory', '.step-persistence')
return cls(store=FileStepStore(directory), **kwargs)
if backend == 'sqlite':
from pydantic_ai_harness.experimental.step_persistence._store import SqliteStepStore
from pydantic_ai_harness.step_persistence._store import SqliteStepStore
database = kwargs.pop('database', '.step-persistence.db')
return cls(store=SqliteStepStore(database=database), **kwargs)
pydantic_ai_harness/experimental/step_persistence/_context.py → pydantic_ai_harness/step_persistence/_context.py
+2 -2
View File
@@ -5,7 +5,7 @@ from __future__ import annotations
from contextvars import ContextVar
current_run_id: ContextVar[str | None] = ContextVar(
'pydantic_ai_harness.experimental.step_persistence.current_run_id',
'pydantic_ai_harness.step_persistence.current_run_id',
default=None,
)
"""Async-context-local pointer to the active `StepPersistence` `run_id`.
@@ -20,7 +20,7 @@ import.
"""
snapshot_saved: ContextVar[bool] = ContextVar(
'pydantic_ai_harness.experimental.step_persistence.snapshot_saved',
'pydantic_ai_harness.step_persistence.snapshot_saved',
default=False,
)
"""Async-context-local flag: did `after_node_run` already save a snapshot this run?
pydantic_ai_harness/experimental/step_persistence/_helpers.py → pydantic_ai_harness/step_persistence/_helpers.py
+3 -3
View File
@@ -13,9 +13,9 @@ from pydantic_ai.messages import (
)
from pydantic_ai.tools import RunContext
from pydantic_ai_harness.experimental.step_persistence._context import current_run_id
from pydantic_ai_harness.experimental.step_persistence._store import StepStore
from pydantic_ai_harness.experimental.step_persistence._types import ToolEffectRecord
from pydantic_ai_harness.step_persistence._context import current_run_id
from pydantic_ai_harness.step_persistence._store import StepStore
from pydantic_ai_harness.step_persistence._types import ToolEffectRecord
def is_provider_valid(messages: list[ModelMessage]) -> bool:
pydantic_ai_harness/experimental/step_persistence/_store.py → pydantic_ai_harness/step_persistence/_store.py
+2 -2
View File
@@ -14,14 +14,14 @@ import anyio.to_thread
from pydantic import TypeAdapter
from pydantic_ai.messages import ModelMessage, ModelMessagesTypeAdapter
from pydantic_ai_harness.experimental.media import (
from pydantic_ai_harness.media import (
DiskMediaStore,
MediaStore,
SqliteMediaStore,
externalize_media,
restore_media,
)
from pydantic_ai_harness.experimental.step_persistence._types import (
from pydantic_ai_harness.step_persistence._types import (
ContinuableSnapshot,
EventKind,
RunRecord,
pydantic_ai_harness/experimental/step_persistence/_types.py → pydantic_ai_harness/step_persistence/_types.py
View File
pydantic_ai_harness/experimental/subagents/README.md → pydantic_ai_harness/subagents/README.md
+13 -18
View File
@@ -1,23 +1,13 @@
# SubAgents
> [!WARNING]
> **Experimental.** This capability lives under `pydantic_ai_harness.experimental` and may
> change or be removed in any release, without a deprecation period. Import it from the
> experimental path -- there is no top-level export:
> [!NOTE]
> Import this capability from its submodule -- there is no top-level `pydantic_ai_harness` re-export:
>
> ```python
> from pydantic_ai_harness.experimental.subagents import SubAgent, SubAgents
> from pydantic_ai_harness.subagents import SubAgent, SubAgents
> ```
>
> Importing any experimental capability emits a `HarnessExperimentalWarning`. Silence **all**
> harness experimental warnings with a single filter (no per-capability lines needed):
>
> ```python
> import warnings
> from pydantic_ai_harness.experimental import HarnessExperimentalWarning
>
> warnings.filterwarnings('ignore', category=HarnessExperimentalWarning)
> ```
> The API may change between releases. Where practical, breaking changes ship with a deprecation warning.
Let an agent delegate self-contained tasks to named child agents.
@@ -31,7 +21,7 @@ A single agent that does everything accumulates a large tool set and a long cont
```python
from pydantic_ai import Agent
from pydantic_ai_harness.experimental.subagents import SubAgent, SubAgents
from pydantic_ai_harness.subagents import SubAgent, SubAgents
researcher = Agent('anthropic:claude-sonnet-4-6', name='researcher', description='Researches a topic and reports findings')
writer = Agent('anthropic:claude-sonnet-4-6', name='writer', description='Turns notes into polished prose')
@@ -71,7 +61,7 @@ Each `SubAgent` carries its own budgets, so one delegate's controls do not touch
```python
from pydantic_ai.usage import UsageLimits
from pydantic_ai_harness.experimental.subagents import SubAgent, SubAgents
from pydantic_ai_harness.subagents import SubAgent, SubAgents
# reproducer and librarian are Agent instances, as in the example above.
orchestrator = Agent(
@@ -93,6 +83,7 @@ orchestrator = Agent(
| `timeout_seconds` | A wall-clock budget for one delegation. When the child exceeds it, its run is cancelled and the parent gets a soft steering message instead of hanging on the child. The cancelled child's `event_stream_handler` (if any) stops receiving events without a terminal event. |
| `max_calls` | The maximum number of delegations to this sub-agent per parent run. Once reached, further delegations return a soft budget-exhausted message without running the child. Counts are scoped to one `Agent.run` (a `run_id`) and cleared when it ends, so each parent run and each level of a nested tree budgets independently. |
| `on_failure` | A steering message returned to the parent for any soft degradation of this delegate, in place of the built-in default. Setting it also makes child failures soft (see below). |
| `contain_errors` | Whether an unexpected crash in this delegate is caught and returned to the parent as a bounded `ModelRetry` instead of aborting the parent run (see below). Unset inherits the `SubAgents(contain_errors=...)` default (off). |
## Failure handling
@@ -102,6 +93,8 @@ A sub-agent run that fails with a *soft model error* (`ModelRetry`, `UnexpectedM
Hard errors propagate to stop the whole run. A `UsageLimitExceeded` from a child that has *no* per-delegate `usage_limits` (so it shares the parent's accounting) means the whole tree is out of budget and propagates; a child reaching its *own* `usage_limits` is soft, as above.
An *unexpected crash* -- any other exception the child raises, such as a provider `ModelAPIError`/`FallbackExceptionGroup` or a plain `ValueError` from a bad tool argument -- propagates by default and aborts the parent run. Set `contain_errors=True` (per delegate, or as the `SubAgents` default) to catch it and return it to the parent as a bounded `ModelRetry` instead, so one delegate crash cannot kill the whole run. Containment stays loud: the exception rides the retry message (`Sub-agent '<name>' crashed: ...`), it is logged via the standard `logging` module, and `tool_retries` still bounds consecutive crashes into an abort. This is orthogonal to `on_failure` -- a contained crash always raises the loud retry, never the soft `on_failure` return, so a genuine bug is never masked as success. Cancellation, a shared `UsageLimitExceeded`, pydantic-ai control-flow signals (`CallDeferred`, `ApprovalRequired`, the `Skip*` signals), and `UserError` always propagate regardless of `contain_errors`.
## Discovery
The sub-agents are listed in the system prompt via `get_instructions`, using each agent's `description` (or a `SubAgent(description=...)` override). A sub-agent with no description is listed by name alone.
@@ -112,7 +105,7 @@ A repo's markdown agent definitions become delegates without writing any `Agent`
```python
from pydantic_ai import Agent
from pydantic_ai_harness.experimental.subagents import SubAgents
from pydantic_ai_harness.subagents import SubAgents
orchestrator = Agent(
'anthropic:claude-opus-4-7',
@@ -152,7 +145,7 @@ Frontmatter is read by a small, dependency-free parser limited to those keys (`p
Disk agents inherit the parent run's model by default. Per agent, the caller can override the model and set a thinking/effort level via `agent_overrides`, keyed by the agent's name:
```python
from pydantic_ai_harness.experimental.subagents import AgentOverride, SubAgents
from pydantic_ai_harness.subagents import AgentOverride, SubAgents
SubAgents(
agent_folders='agents',
@@ -191,6 +184,7 @@ SubAgents(
event_stream_handler=None, # forwarded to each sub-agent run to stream its events
tool_name='delegate_task',
tool_retries=2, # extra delegate-tool attempts after a sub-agent error before aborting (None inherits the agent default)
contain_errors=False, # default for SubAgent.contain_errors: contain an unexpected crash as a bounded retry
)
```
@@ -203,6 +197,7 @@ SubAgent(
timeout_seconds=None, # per-delegation wall-clock budget
max_calls=None, # max delegations to this sub-agent per parent run
on_failure=None, # steering message for soft degradations of this delegate
contain_errors=None, # contain an unexpected crash as a bounded retry; None inherits the SubAgents default
)
```
pydantic_ai_harness/subagents/__init__.py
+16
View File
@@ -0,0 +1,16 @@
"""Sub-agent capability: delegate self-contained tasks to named child agents."""
from pydantic_ai_harness.subagents._capability import SubAgents, ToolResolver
from pydantic_ai_harness.subagents._disk import AgentOverride
from pydantic_ai_harness.subagents._effort import MINIMUM_EFFORT_FLOOR, clamp_effort
from pydantic_ai_harness.subagents._toolset import SubAgent, SubAgentToolset
__all__ = [
'MINIMUM_EFFORT_FLOOR',
'AgentOverride',
'SubAgent',
'SubAgentToolset',
'SubAgents',
'ToolResolver',
'clamp_effort',
]
pydantic_ai_harness/experimental/subagents/_capability.py → pydantic_ai_harness/subagents/_capability.py
+12 -4
View File
@@ -14,14 +14,14 @@ from pydantic_ai.settings import ModelSettings
from pydantic_ai.tools import AgentDepsT, RunContext
from pydantic_ai.toolsets import AgentToolset
from pydantic_ai_harness.experimental.subagents._disk import (
from pydantic_ai_harness.subagents._disk import (
AgentOverride,
ParsedAgent,
parse_agent_markdown,
resolve_folders,
)
from pydantic_ai_harness.experimental.subagents._effort import clamp_effort
from pydantic_ai_harness.experimental.subagents._toolset import SubAgent, SubAgentToolset
from pydantic_ai_harness.subagents._effort import clamp_effort
from pydantic_ai_harness.subagents._toolset import SubAgent, SubAgentToolset
if TYPE_CHECKING:
from pydantic_ai._instructions import AgentInstructions
@@ -66,7 +66,7 @@ class SubAgents(AbstractCapability[AgentDepsT]):
```python
from pydantic_ai import Agent
from pydantic_ai_harness.experimental.subagents import SubAgent, SubAgents
from pydantic_ai_harness.subagents import SubAgent, SubAgents
researcher = Agent('anthropic:claude-sonnet-4-6', name='researcher', description='Researches topics')
writer = Agent('anthropic:claude-sonnet-4-6', name='writer', description='Writes prose')
@@ -140,6 +140,13 @@ class SubAgents(AbstractCapability[AgentDepsT]):
repeated flaky sub-agent does not abort the parent run on its first repeat;
set `None` to inherit the parent agent's default tool retries instead."""
contain_errors: bool = False
"""Default for `SubAgent.contain_errors`: whether an unexpected sub-agent crash
is caught and returned to the parent as a bounded `ModelRetry` instead of
aborting the parent run. Off by default, so a crash propagates. Any `SubAgent`
can override this per delegate. See `SubAgent.contain_errors` for the
containment contract and what always propagates regardless."""
_by_name: dict[str, SubAgent[AgentDepsT]] = field(
default_factory=dict[str, 'SubAgent[AgentDepsT]'], init=False, repr=False, compare=False
)
@@ -273,6 +280,7 @@ class SubAgents(AbstractCapability[AgentDepsT]):
event_stream_handler=self.event_stream_handler,
tool_name=self.tool_name,
tool_retries=self.tool_retries,
contain_errors=self.contain_errors,
call_counts=self._call_counts,
)
pydantic_ai_harness/experimental/subagents/_disk.py → pydantic_ai_harness/subagents/_disk.py
View File
pydantic_ai_harness/experimental/subagents/_effort.py → pydantic_ai_harness/subagents/_effort.py
View File
pydantic_ai_harness/experimental/subagents/_toolset.py → pydantic_ai_harness/subagents/_toolset.py
+58 -1
View File
@@ -3,13 +3,24 @@
from __future__ import annotations
import asyncio
import logging
from collections.abc import Mapping, Sequence
from dataclasses import dataclass
from typing import Any, Generic
from pydantic_ai.agent import AbstractAgent, EventStreamHandler
from pydantic_ai.capabilities import AgentCapability
from pydantic_ai.exceptions import ModelRetry, UnexpectedModelBehavior, UsageLimitExceeded
from pydantic_ai.exceptions import (
ApprovalRequired,
CallDeferred,
ModelRetry,
SkipModelRequest,
SkipToolExecution,
SkipToolValidation,
UnexpectedModelBehavior,
UsageLimitExceeded,
UserError,
)
from pydantic_ai.tools import AgentDepsT, RunContext
from pydantic_ai.toolsets import AbstractToolset, FunctionToolset
@@ -18,6 +29,23 @@ from pydantic_ai.toolsets import AbstractToolset, FunctionToolset
from pydantic_ai.toolsets._capability_owned import CapabilityOwnedToolset
from pydantic_ai.usage import UsageLimits
logger = logging.getLogger(__name__)
# Signals that must always reach the parent run, even when a delegate has
# `contain_errors` on. Containing the first five would break the agent graph
# (deferred/approval/skip control-flow); a `UserError` is a setup bug that no
# retry can fix, so masking it into a retry only delays and obscures it.
# Cancellation (`asyncio.CancelledError`, a `BaseException`) is out of `except
# Exception`'s reach already, and a shared `UsageLimitExceeded` has its own clause.
_ALWAYS_PROPAGATE: tuple[type[Exception], ...] = (
CallDeferred,
ApprovalRequired,
SkipModelRequest,
SkipToolValidation,
SkipToolExecution,
UserError,
)
@dataclass(frozen=True)
class SubAgent(Generic[AgentDepsT]):
@@ -68,6 +96,19 @@ class SubAgent(Generic[AgentDepsT]):
failures soft: a child error returns this message as a normal tool result
instead of raising a parent `ModelRetry`."""
contain_errors: bool | None = None
"""Whether an unexpected sub-agent crash is contained instead of aborting the
parent run. When `True`, an exception the child raises that is not an expected
soft degradation (a provider `ModelAPIError`/`FallbackExceptionGroup`, a plain
`ValueError` from a bad tool argument, etc.) is caught and returned to the parent
as a bounded `ModelRetry`, so one delegate crash cannot kill the whole run. It
stays loud: the exception rides the retry message and is logged, and
`tool_retries` still bounds consecutive crashes into an abort. Cancellation, a
shared usage-limit, pydantic-ai control-flow signals, and `UserError` always
propagate regardless. Unset inherits `SubAgents.contain_errors` (default off).
Orthogonal to `on_failure`, which only sets the message for expected soft
degradations; a contained crash always raises the loud `ModelRetry`."""
@property
def resolved_name(self) -> str | None:
"""The delegate's name: `name` if set, else the agent's own `name`."""
@@ -108,6 +149,7 @@ class SubAgentToolset(FunctionToolset[AgentDepsT]):
event_stream_handler: EventStreamHandler[AgentDepsT] | None,
tool_name: str,
tool_retries: int | None,
contain_errors: bool,
call_counts: dict[str, dict[str, int]],
) -> None:
super().__init__()
@@ -117,6 +159,7 @@ class SubAgentToolset(FunctionToolset[AgentDepsT]):
self._shared_capabilities = list(shared_capabilities)
self._event_stream_handler = event_stream_handler
self._tool_name = tool_name
self._contain_errors = contain_errors
# Run-scoped delegation counts, keyed by run_id then sub-agent name.
# Shared with the capability, which clears each run's entry in wrap_run.
self._call_counts = call_counts
@@ -233,6 +276,20 @@ class SubAgentToolset(FunctionToolset[AgentDepsT]):
return sub_agent.on_failure
# Soft sub-agent failures come back to the parent as a retry it can react to.
raise ModelRetry(f'Sub-agent {agent_name!r} failed: {exc}') from exc
except _ALWAYS_PROPAGATE:
raise
except Exception as exc:
contain = sub_agent.contain_errors if sub_agent.contain_errors is not None else self._contain_errors
if not contain:
raise
# Contain the crash so it cannot abort the parent, but keep it loud: the
# exception rides the retry message and is logged, and `tool_retries`
# bounds consecutive crashes into an abort.
logger.warning('Contained crash from sub-agent %r', agent_name, exc_info=exc)
raise ModelRetry(
f'Sub-agent {agent_name!r} crashed: {type(exc).__name__}: {exc}. '
f'Treat this as a recoverable failure and decide from existing evidence.'
) from exc
return str(result.output)
@staticmethod
pyproject.toml
+7
View File
@@ -53,6 +53,13 @@ logfire = [
modal = [
"modal>=1.5.0",
]
dynamic-workflow = [
"pydantic-monty>=0.0.16",
]
acp = [
# The SDK breaks across minors, so cap per minor.
"agent-client-protocol>=0.11,<0.12",
]
[project.urls]
Homepage = 'https://github.com/pydantic/pydantic-ai-harness'

Some files were not shown because too many files have changed in this diff Show More