mirror of
https://github.com/pydantic/pydantic-ai-harness.git
synced 2026-07-21 02:45:34 +00:00
feat: serve a Pydantic AI agent to editors over ACP (experimental) (#274)
* feat: add ACP capability to serve agents to editors Editors and TUIs that speak the Agent Client Protocol (Zed and others) can drive an external coding agent, but plugging a Pydantic AI agent into one previously meant implementing the ACP server side by hand. run_acp_stdio serves any Agent over stdio: streamed text and thinking, file edits rendered as diffs, human-in-the-loop tool approval mapped to deferred-approval tools, per-workspace sessions via a session_config hook, model switching, cancellation, and optional session persistence. FileSystem.get_toolset/Shell.get_toolset now return their concrete toolset types so the ACP presenter can recognize their tool calls by name and annotate them with kinds, locations, and diffs. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * fix: skip ACP tests at collection when the acp extra is not installed The slim CI jobs sync without extras, so `agent-client-protocol` is absent and the ACP test modules failed at import during collection. Ignore them via conftest when `acp` can't be found; `test_packaging.py` stays collected since package metadata holds on base installs. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * fix(acp): enable unstable routing and bound streamed updates by byte size session/set_model and session/close are advertised at initialize, but the ACP SDK router rejects them as unstable unless run_agent is given use_unstable_protocol=True -- so the model picker and session-close affordance returned method_not_found over a real connection. Enable the flag. Streamed text was chunked by character count, but the SDK serializes with ensure_ascii=True, so a non-ASCII code point expands up to 12 bytes (a surrogate pair) inside the JSON string. A single agent_message_chunk of emoji/CJK could exceed the client's 64 KiB read buffer and drop the connection. Chunk by escaped byte length instead. Both gaps were invisible to the suite: the close/set_model tests called the adapter directly (bypassing the router) and the large-output test used only ASCII. Add a through-the-router stdio test and a non-ASCII chunking test as guards. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> * feat(acp): per-turn usage, read-resilient terminal cleanup, persistence + native-toolset tests Address review follow-ups on the ACP capability: - Report per-turn token usage on `PromptResponse.usage`, summed across approval passes (UNSTABLE ACP field; clients that don't support it ignore it). - Suppress client errors during the shielded terminal kill/release so a failing cleanup call can't mask the in-flight CancelledError (the spec requires the turn to end with a `cancelled` stop reason). - Cover the persistence guarantees that previously had no store-active test: a cancelled turn commits nothing, an approval-resume turn persists each update once (no duplicate tool-call start), and `StoredSession` round-trips through Pydantic across the full `SessionUpdate` union. - Add a through-the-wire stdio test that the editor-native fs/terminal toolsets route to the client when mounted per session. - Document the dynamic-`ApprovalRequired` partial-side-effect nuance, per-turn usage, and that session persistence composes with per-run step durability. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> * feat(acp): editor-native reads for read-only fs clients, writes delegated locally A client that advertises filesystem reads but not writes previously got no editor-native filesystem at all (acp_filesystem returned None, so callers fell back to a fully-local toolset and lost the editor's live view, e.g. unsaved buffers). It now returns a combined toolset: reads route through the editor, and writes are delegated to the local FileSystem capability rooted at the session cwd (reusing its path sandboxing rather than reimplementing writes). This is 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 the writes land on the agent's disk. Documented as such at the helper and README. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> * docs(acp): point _all_known_model_names at the now-merged known_model_names() pydantic-ai#5803 added `pydantic_ai.models.known_model_names()`, the public replacement for the `KnownModelName.__value__` introspection. It isn't in a released `pydantic-ai-slim` yet, so the swap (and the floor bump it requires) waits for that release; update the breadcrumb so it's actioned then. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> * fix(acp): replay user turns on session/load, advertise MCP capabilities Two spec-conformance gaps from PR review: - The transcript only ever recorded agent-direction updates, so a reopened session replayed a one-sided conversation. ACP requires session/load to replay the entire conversation, user turns included. The prompt's content blocks are now seeded into the turn's update list -- recorded for replay but never sent live (the client renders its own prompt; the prompt-turn spec sends no user echo) -- riding the existing commit path so a cancelled turn still rolls its user message back. - initialize left mcp_capabilities at the default (http/sse false), and the spec forbids clients from sending HTTP/SSE MCP servers that are not advertised -- making the documented session_config -> mcp_servers path unreachable for conforming clients over those transports (stdio is not gated). Advertising is now an explicit opt-in mirroring prompt_capabilities, since only the embedder knows which transports their session_config connects; setting it without a session_config fails at construction rather than inviting servers that session/new would reject. Also adds the guard test from review follow-up: _all_known_model_names fails loudly if pyai ever recomposes the KnownModelName alias. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * fix(acp): make the collection guard coverable in a single environment The acp-extra collection guard was an `if` statement whose true arm only runs on slim installs, so local `make testcov` (always all-extras) could never reach 100% -- only CI's combined slim+all-extras matrix could. A conditional expression has no statement arc for branch coverage to miss, so the documented local gate works again without excluding a line CI genuinely covers. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * docs(acp): record why _all_known_model_names avoids the unreleased public API The breadcrumb read as if known_model_names() were already usable; make the deliberate avoidance explicit. The public enumerator is merged upstream but not in any released pydantic-ai-slim, and the test-floor CI job runs against the floor release, so calling it would break there. Move the rationale from the docstring to a code comment (docstring stays contract-only) and spell out the swap-and-delete trigger: a release that ships it plus a floor bump. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * test(acp): add spec-conformance suite driven over a real in-memory wire The existing tests invoke adapter methods directly, below the SDK's JSON-RPC router and serialization -- a boundary that cannot see two bug classes that already bit this adapter: a method the router gates before the adapter runs (the use_unstable_protocol reachability bug) and a frame whose serialized bytes overrun the client buffer (the ensure_ascii chunking bug). `tests/acp/_wire.py` closes that gap without a subprocess: a real ClientSideConnection talks to acp.run_agent across a socket.socketpair in one event loop, so every request crosses the router and codec and every update arrives as bytes the client re-parses, with asyncio's default 64 KiB reader limit standing in for the stdio buffer. `tests/acp/test_conformance.py` is organized by spec clause, not by adapter method, each with an oracle built from the spec/input rather than the adapter's own output: - version negotiation echoes the requested version (not a literal) - capabilities are advertised iff supported (load_session, mcp, auth_methods, session list/fork/resume, modes, config options) -- read back off the wire - unstable methods route through with the flag and are method_not_found without it (the load-bearing reason run_acp_stdio enables it) - error codes distinguish method_not_found from invalid_params - session/load replays the entire conversation including user turns - large non-ASCII output reassembles intact within the client's read buffer CONFORMANCE.md records the clause-to-test matrix, the open gaps (none a known bug), the two verified deviations, and the N/A ledger, derived from a spec-page sweep with adversarial verification of every candidate finding. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * refactor(acp): stop surfacing unconsumed additionalDirectories The adapter accepted the client's additionalDirectories and forwarded them to the session_config AcpSession, but never advertised sessionCapabilities.additionalDirectories -- so a conformant client never sent them, and nothing downstream consumed them (the FileSystem capability is single-root). The ACP capability is also still UNSTABLE. Drop the dead surface: remove the field from AcpSession and stop forwarding it. The acp.Agent base signature still carries the parameter, so the session methods accept and ignore it. Supporting it properly (multi-root filesystem + advertisement) is recorded as a deferred feature. Also trims CONFORMANCE.md to a public supported-capability summary. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * fix(acp): keep inline image bytes and let run_acp_stdio advertise MCP An ACP image block always carries inline `data`; `uri` is only an optional source reference, so preferring the URL dropped the bytes the client actually sent in favour of a link the model may be unable to fetch. Prefer the inline data, matching the reference ACP agents. `run_acp_stdio`/`run_acp_stdio_sync` forwarded every adapter option except `mcp_capabilities`, so advertising MCP transports was only reachable by constructing `PydanticAIACPAgent` by hand. Forward it for parity. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * fix(acp): reliable approval scope, mapped stop reasons, safe session reload "Always allow"/"always reject" keyed the scope on the raw tool arguments, which a streaming model (and OpenAI by default) delivers as a JSON string; re-ordered keys then produced a different key and silently re-prompted for a call already decided. Canonicalize via `args_as_dict()` so the same logical call shares one remembered decision. Completed turns always reported `end_turn`, hiding `max_tokens`/`refusal` from the client. Map the model's finish reason to the ACP stop reason. `session/load` overwrote a still-open session without cancelling its in-flight turn, leaking a task that could later persist stale state over the restored transcript. Tear it down first, sharing the close_session teardown path. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * docs(acp): correct the API reference and capability notes The README API block omitted `session_store`/`models` (and now `mcp_capabilities`), the limitations list contradicted the documented MCP-via-`session_config` support, and the overwrite-diff note misdescribed what the code emits. Also drop a dangling internal-doc reference from a source comment. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * docs(acp): record the unhandled SessionStore failure contract A durable SessionStore can fail on save or return a corrupt payload on load, but the adapter currently propagates those exceptions rather than handling them. Document the gap on the Protocol where store implementers will see it, so the behaviour is a known boundary rather than a surprise -- without adding handling we have no consumer for yet. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * docs(acp): correct the SessionStore failure note to match observed behaviour The previous note assumed store exceptions surface "raw". Verified against the SDK (connection.py `_run_request`) and over a real in-memory wire with a raising stub store: they are converted to JSON-RPC errors before reaching the client -- a `pydantic.ValidationError` to `invalid_params` (-32602), anything else to a generic internal error (-32603). Record what the client actually receives rather than what was guessed. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * fix(acp): handle SessionStore failures instead of leaking them `session_store` is a public, documented extension point, so a durable store that fails on save or returns a corrupt payload on load is reachable today -- not hypothetical. A failed save previously errored the very turn the user had already watched stream to completion; a corrupt load leaked raw pydantic/IO detail to the client. Make save failures non-fatal: log and swallow them so the turn (or session) that already committed in memory still succeeds, and the next save catches the store up. Make load failures fail `session/load` with a clear internal error rather than a leaked exception, since a session that cannot be read cannot be reopened. Verified over the in-memory wire with a failing stub store. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * fix(acp): close out in-flight tool calls when a turn is cancelled A tool call announced as pending/in_progress was never driven to a terminal status when its turn was cancelled (via session/cancel or a dismissed permission dialog), so a client kept rendering it as running after the turn ended cancelled. On the cancel unwind, fail any tool call still lacking a terminal result -- sent live (a cancelled turn never commits its transcript) and shielded so the asyncio cancellation cannot abort the send. Verified over the in-memory wire. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * fix(acp): reject prompts queued behind a session that closed or reloaded A prompt waiting on the session's turn lock held a reference to the old SessionState; close_session/load_session only cancel the *active* turn, so the queued prompt would run against the discarded state - invisible to session/cancel - and persist its orphaned history over the closed (or just restored) session. Re-checking liveness after acquiring the lock turns that zombie turn into the invalid_params the client already gets for a prompt sent after close. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * fix(acp): never swallow a handler's own teardown cancellation prompt() and _cancel_active_turn awaited the turn task bare, so the CancelledError from connection teardown was indistinguishable from the turn's own cancellation and got converted into a response (or, via a dismissed permission dialog, replaced by the internal _TurnCancelled). The SDK cancels each handler task exactly once on shutdown and its sender is already closed by then, so a handler that survives that one cancel and tries to answer hangs Connection.close() forever. Awaiting through asyncio.shield keeps the two cancellation sources apart: the turn's end is read off the (done) task, while the handler's own cancellation propagates. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * fix(acp): don't report rollback when a cancel lands in the post-commit save The store save after a turn commits was the one suspension point left between commit and return: a session/cancel delivered there escaped the turn task and the prompt answered 'cancelled' (the rollback signal, no user_message_id) while the session's in-memory history and transcript already contained the whole turn - the next prompt would build on history the client believed was discarded. The InMemorySessionStore never suspends, which is why the existing cancel-persistence test could not see this; any real (file/db) store can. A cancel that arrives after commit has simply lost the race, so the interrupted save is treated like the write failures _persist already swallows. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * fix(acp): don't leak a terminal when run_command is cancelled mid-create The terminal/create call sat before the cleanup try-block, so a cancellation landing there unwound without a kill or release - but the request may already be on the wire (the SDK sender flushes queued payloads even when the awaiting future is cancelled), leaving the command running client-side with nobody holding its id. Running the create as a shielded task lets the cleanup await the response late, learn the id, and kill/release as usual. Folding the late-id case into one kill path and one release path also keeps the existing guarantee that a failing kill never skips the release. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * fix(acp): end limit-hit turns with max_turn_requests, close out tool calls on errors pydantic-ai's default UsageLimits(request_limit=50) means a long tool loop routinely raises UsageLimitExceeded - which escaped the turn as a JSON-RPC internal_error, the exact condition ACP defines the max_turn_requests stop reason for (token limits map to max_tokens). The raising run's partial messages are not retrievable, so such a turn rolls back like a cancellation and says so (no user_message_id, no usage). Turns failing with any other exception now also close out their announced tool calls before the error reaches the client, matching the cancellation path, so the client does not render them as running next to the turn's failure. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * fix(acp): resolve workspace-relative paths before they reach fs/read|write ACP requires absolute paths on fs/read_text_file and fs/write_text_file, but models routinely emit workspace-relative ones - the local FileSystem tools document relative paths and share these tools' names, and the presenter layer already absolutizes for exactly that reason. The client-backed toolset forwarded them raw, producing non-conformant requests the client may reject. acp_filesystem now hands the toolset the session cwd so relative paths resolve before the wire; absolute paths and directly-constructed cwd-less toolsets are untouched. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * docs(acp): true up the conformance claims, comments, and doc links The stop-reason note predated the finish-reason mapping and said the opposite of what the code (and its tests) do; the prompt-capabilities bullet read as inbound enforcement when the spec puts that restriction on the client; the stdio MCP gap (the spec's unconditional MUST) was undisclosed; and two code comments claimed validation/rejection the SDK router does not actually perform - the adapter's own raises are the load-bearing behavior. Doc links now use the canonical pydantic.dev paths both old URLs redirect to. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * refactor(acp): shield the terminal create in place instead of via a side task The ensure_future + asyncio.shield shape from the previous commit hit Python 3.12+'s shield behavior: once the outer await is cancelled, a late failure of the inner create is reported to the loop exception handler even when the cleanup retrieves it, which spams production logs and fails under pytest-anyio. An anyio shield around the create await itself is simpler and leans on the same anyio-mediated cancellation the kill/release cleanup already depends on: the cancellation defers to the next await, by which point the terminal id is known and the existing cleanup kills and releases as usual. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * refactor(acp): small type and surface cleanups from review - ToolCallPermission.args was typed object although the only construction site passes args_as_dict(); Mapping[str, object] lets a custom policy scope by an argument without isinstance gymnastics. - default_permission_scope is referenced by the permission_policy docs as the fallback to compose with, so it must be importable: exported. - McpServer names the per-server union once instead of spelling it in two places; McpServers gains the TypeAlias marker its sibling had. - _model_state folds the None case so both call sites drop a repeated conditional; the load_session store guard's no-cover pragma was wrong (the SDK router routes session/load regardless of the advertisement, and test_persistence already executes the branch). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * test(acp): pin claimed behavior the suite asserted too weakly or not at all - The unsupported-method test asserted only RequestError despite name and CONFORMANCE.md promising method_not_found; codes now pinned (-32601, and -32602 for the unknown-session prompt). - The model-override test was tautological: both models answered successfully either way. It now asserts the override's distinct canned output (mutation-checked: dropping the per-run override fails it). - Usage summation across approval passes had no test; an approval turn's output tokens must exceed the resume pass alone (mutation-checked against 'usage = result.usage()'). - New protocol-contract coverage: cancel racing an unanswered permission dialog (turn ends cancelled, pending call driven to failed), double cancel idempotency, image block arriving in model history as decoded BinaryContent, the denial message reaching the failed update's raw_output, chunk_text's exact-budget boundary, and a set_model save failure being swallowed like every other persist failure. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * fix(acp): survive raw cancellation during terminal create; answer cancelled after a post-commit cancel An adversarial re-review of this branch's own fixes caught a regression: the anyio shield around terminal/create only blocks anyio-mediated cancellation, but the adapter and pydantic-ai deliver raw task.cancel() (turn.cancel, cancel_and_drain), which pierces it - reintroducing the leaked-terminal window end to end. The create now runs as its own task awaited via asyncio.wait: a raw cancel hits the waiter, not the create, and unlike asyncio.shield (3.12+) a late create failure is not reported to the loop exception handler when the cleanup retrieves it. The leak scenario is pinned at both the toolset level and through the real adapter/agent path. Also from the re-review: a cancel that lands after the turn committed (inside the store save) now answers 'cancelled' as the spec requires - while keeping the committed signals (user_message_id, usage, history) - and the one-tick teardown-vs-turn ambiguity that 3.10 cannot resolve is documented at both await sites. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * docs(acp): record the post-commit-cancel semantics in CONFORMANCE Also drops an unbounded poll loop in the permission-dialog race test for a single deterministic tick (prompt sets active_turn before its first suspension point), which was the one partial branch left repo-wide. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * refactor(acp): let the turn build its own response and drop the commit-flag tuple _run_turn returned (usage, stop_reason, committed) so prompt() could reassemble a PromptResponse it had all the pieces for; building the response in the turn removes that protocol and the committed flag. Recording turn.updates is now unconditional (only the commit gate matters), and the presenter-fields/acp_filesystem branches collapse to single construction paths. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * docs(acp): appease codespell (unparseable -> unparsable) Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * Support Pydantic AI 2.x in ACP * feat(acp): forward usage limits into agent runs Let hosts bound each ACP run segment with Pydantic AI usage ceilings while preserving the default run behavior when no explicit limits are configured. * feat(acp): resolve selected session models Let embedders map advertised ACP model ids to concrete Pydantic AI models and re-surface a resident session's current model state. * feat(acp): expose read-only session history Hosts can gauge context with the same estimate_token_count path the compaction tiers use by reading committed resident-session history. The returned list is a shallow snapshot, and callers must treat the shared ModelMessage objects as read-only. * Track ACP 0.11 session config models * chore(acp): fix coverage, trim redundancy, clarify module names Get the ACP PR to green CI and simpler to review, without changing behavior. - Coverage: the only uncovered line was a dead `return None` branch in the `_model_option` test helper (source was already 100%). Assert the option is present instead, dropping the untested branch and its now-redundant guards. - Remove CONFORMANCE.md: its supported/not-supported matrix duplicated the package README's feature sections and "Cancellation and limitations". Point the two references at the README. - Dedup the test ACP clients: RecordingClient, FakeClient, and WireClient each re-spelled the same ~13-method "unused capability" stub block to satisfy the SDK Client interface. Extract it once into RecordingClientBase; each client now subclasses it and overrides only what it exercises. Interface drift is one edit. - Rename two opaque modules for clarity (files only; public symbols unchanged): _present.py -> _presentation.py (matches its ToolCallPresentation exports) and _native.py -> _client_toolsets.py (client-routed filesystem/shell toolsets), with tests/acp/test_native.py -> test_client_toolsets.py to mirror the source. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * refactor(acp): move under experimental and shrink the public surface New capabilities start under pydantic_ai_harness.experimental, so ACP follows: importing warns HarnessExperimentalWarning and the API may change without deprecation. __all__ drops the typing-only aliases (25 to 18 names); default_permission_scope and McpServer stay because policies compose with the default scope and AcpSession.mcp_servers is typed by McpServer. Tests move to tests/experimental/acp to mirror the source tree; the experimental warning test lives in the SDK-gated test_acp.py so slim (no-extras) installs stay green. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * feat(acp): expose model_resolver and usage_limits on the stdio entry points PydanticAIACPAgent already accepted both, so wanting a per-run token ceiling or host-defined model ids forced users off run_acp_stdio onto the class and a hand-rolled acp.run_agent call, where forgetting use_unstable_protocol=True silently loses session/close. Forward the two params so the one-liner covers every adapter option. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Co-authored-by: David SF <david.sanchez@pydantic.dev>
This commit is contained in:
Aditya Vardhan
co-authored by
Claude Opus 4.8
David SF
Claude Opus 4.8
David SF
parent
b5b93704c3
commit
b4365440b5
@@ -28,6 +28,7 @@ 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[acp]" # ACP (serve an agent to editors over the Agent Client Protocol)
|
||||
```
|
||||
|
||||
The `code-mode` extra is also supported as an alias.
|
||||
@@ -108,6 +109,7 @@ We studied leading coding agents, agent frameworks, and Claw-style assistants to
|
||||
| | **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‑co) |
|
||||
| | **Repo context injection** | Auto-load CLAUDE.md/AGENTS.md and repo structure | :construction: [PR #175](https://github.com/pydantic/pydantic-ai-harness/pull/175) | [pydantic-deep](https://github.com/vstorm-co/pydantic-deepagents) (vstorm‑co) |
|
||||
| | **Verification loop** | Run tests after edits, auto-fix failures | :construction: [PR #169](https://github.com/pydantic/pydantic-ai-harness/pull/169) | |
|
||||
| **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) | |
|
||||
| **Context management** | **Sliding window** | Trim conversation history to stay within token limits | :construction: [PR #191](https://github.com/pydantic/pydantic-ai-harness/pull/191) | [summarization-pydantic-ai](https://github.com/vstorm-co/summarization-pydantic-ai) (vstorm‑co) |
|
||||
| | **Context compaction** | LLM-powered summarization of older messages | :construction: [PR #191](https://github.com/pydantic/pydantic-ai-harness/pull/191) | [summarization-pydantic-ai](https://github.com/vstorm-co/summarization-pydantic-ai) (vstorm‑co) |
|
||||
| | **Limit warnings** | Warn agent before hitting context/iteration limits | :construction: [PR #191](https://github.com/pydantic/pydantic-ai-harness/pull/191) | [summarization-pydantic-ai](https://github.com/vstorm-co/summarization-pydantic-ai) (vstorm‑co) |
|
||||
|
||||
@@ -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/)
|
||||
@@ -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',
|
||||
]
|
||||
@@ -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
|
||||
@@ -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)
|
||||
@@ -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
|
||||
@@ -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))
|
||||
@@ -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)
|
||||
@@ -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]}
|
||||
@@ -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,
|
||||
)
|
||||
)
|
||||
@@ -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])
|
||||
@@ -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)
|
||||
@@ -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),
|
||||
|
||||
@@ -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),
|
||||
|
||||
@@ -50,6 +50,10 @@ logfire = [
|
||||
'logfire>=4.31.0',
|
||||
"pydantic-ai-slim[spec]>=2.1.0",
|
||||
]
|
||||
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'
|
||||
|
||||
@@ -0,0 +1,194 @@
|
||||
"""Shared in-memory ACP `Client` scaffolding for the adapter and toolset tests.
|
||||
|
||||
`Client` has no abstract methods at runtime, but pyright treats them as abstract (the SDK marks
|
||||
them so), so a test client must define the whole surface. `RecordingClientBase` defines that surface
|
||||
once -- it records the `session/update`s it receives and stubs every request method -- so concrete
|
||||
clients subclass it and override only the capabilities they exercise. `RecordingClient` drives
|
||||
filesystem and terminal state; `FakeClient` (in `test_acp.py`) answers permission requests; and
|
||||
`WireClient` (in `_wire.py`) records updates only.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
|
||||
from acp import Client, RequestError, schema
|
||||
|
||||
|
||||
class RecordingClientBase(Client):
|
||||
"""Records `session/update`s; every request method is an unused stub that raises if called.
|
||||
|
||||
Subclasses override only the capabilities they exercise. A stub that fires means the adapter
|
||||
called a client capability the test did not expect, which should surface loudly.
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.updates: list[object] = []
|
||||
|
||||
async def session_update(self, session_id: str, update: object, **kwargs: object) -> None:
|
||||
self.updates.append(update)
|
||||
|
||||
def on_connect(self, conn: object) -> None:
|
||||
return None # pragma: no cover - unused
|
||||
|
||||
async def request_permission(
|
||||
self,
|
||||
session_id: str,
|
||||
tool_call: schema.ToolCallUpdate,
|
||||
options: list[schema.PermissionOption],
|
||||
**kwargs: object,
|
||||
) -> schema.RequestPermissionResponse:
|
||||
raise NotImplementedError # pragma: no cover - unused
|
||||
|
||||
async def read_text_file(
|
||||
self, session_id: str, path: str, line: int | None = None, limit: int | None = None, **kwargs: object
|
||||
) -> schema.ReadTextFileResponse:
|
||||
raise NotImplementedError # pragma: no cover - unused
|
||||
|
||||
async def write_text_file(
|
||||
self, session_id: str, path: str, content: str, **kwargs: object
|
||||
) -> schema.WriteTextFileResponse | None:
|
||||
raise NotImplementedError # pragma: no cover - unused
|
||||
|
||||
async def create_terminal(
|
||||
self,
|
||||
session_id: str,
|
||||
command: str,
|
||||
args: list[str] | None = None,
|
||||
env: list[schema.EnvVariable] | None = None,
|
||||
cwd: str | None = None,
|
||||
output_byte_limit: int | None = None,
|
||||
**kwargs: object,
|
||||
) -> schema.CreateTerminalResponse:
|
||||
raise NotImplementedError # pragma: no cover - unused
|
||||
|
||||
async def wait_for_terminal_exit(
|
||||
self, session_id: str, terminal_id: str, **kwargs: object
|
||||
) -> schema.WaitForTerminalExitResponse:
|
||||
raise NotImplementedError # pragma: no cover - unused
|
||||
|
||||
async def terminal_output(
|
||||
self, session_id: str, terminal_id: str, **kwargs: object
|
||||
) -> schema.TerminalOutputResponse:
|
||||
raise NotImplementedError # pragma: no cover - unused
|
||||
|
||||
async def kill_terminal(
|
||||
self, session_id: str, terminal_id: str, **kwargs: object
|
||||
) -> schema.KillTerminalResponse | None:
|
||||
raise NotImplementedError # pragma: no cover - unused
|
||||
|
||||
async def release_terminal(
|
||||
self, session_id: str, terminal_id: str, **kwargs: object
|
||||
) -> schema.ReleaseTerminalResponse | None:
|
||||
raise NotImplementedError # pragma: no cover - unused
|
||||
|
||||
async def create_elicitation(
|
||||
self, message: str, mode: schema.ElicitationMode, **kwargs: object
|
||||
) -> schema.CreateElicitationResponse:
|
||||
raise NotImplementedError # pragma: no cover - unused
|
||||
|
||||
async def complete_elicitation(self, elicitation_id: str, **kwargs: object) -> None:
|
||||
raise NotImplementedError # pragma: no cover - unused
|
||||
|
||||
async def ext_method(self, method: str, params: dict[str, object]) -> dict[str, object]:
|
||||
raise RequestError.method_not_found(method) # pragma: no cover - unused
|
||||
|
||||
async def ext_notification(self, method: str, params: dict[str, object]) -> None:
|
||||
return None # pragma: no cover - unused
|
||||
|
||||
|
||||
class RecordingClient(RecordingClientBase):
|
||||
"""An ACP client driving in-memory filesystem/terminal state and recording session updates."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
files: dict[str, str] | None = None,
|
||||
*,
|
||||
output: str = '',
|
||||
truncated: bool = False,
|
||||
exit_code: int | None = 0,
|
||||
signal: str | None = None,
|
||||
no_exit_status: bool = False,
|
||||
block_exit: bool = False,
|
||||
block_create: bool = False,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.files: dict[str, str] = dict(files or {})
|
||||
self.reads: list[tuple[str, str]] = []
|
||||
self.writes: list[tuple[str, str, str]] = []
|
||||
self._output = output
|
||||
self._truncated = truncated
|
||||
self._exit_code = exit_code
|
||||
self._signal = signal
|
||||
self._no_exit_status = no_exit_status
|
||||
self._block_exit = block_exit
|
||||
self._block_create = block_create
|
||||
self.release_create = asyncio.Event()
|
||||
self.exit_event = asyncio.Event()
|
||||
self.created: list[tuple[str, str | None]] = []
|
||||
self.killed: list[str] = []
|
||||
self.released: list[str] = []
|
||||
self._terminals = 0
|
||||
self.create_event = asyncio.Event()
|
||||
|
||||
# --- filesystem -----------------------------------------------------------------------
|
||||
|
||||
async def read_text_file(
|
||||
self, session_id: str, path: str, line: int | None = None, limit: int | None = None, **kwargs: object
|
||||
) -> schema.ReadTextFileResponse:
|
||||
self.reads.append((path, session_id))
|
||||
return schema.ReadTextFileResponse(content=self.files[path])
|
||||
|
||||
async def write_text_file(
|
||||
self, session_id: str, path: str, content: str, **kwargs: object
|
||||
) -> schema.WriteTextFileResponse | None:
|
||||
self.files[path] = content
|
||||
self.writes.append((path, content, session_id))
|
||||
return None
|
||||
|
||||
# --- terminal -------------------------------------------------------------------------
|
||||
|
||||
async def create_terminal(
|
||||
self,
|
||||
session_id: str,
|
||||
command: str,
|
||||
args: list[str] | None = None,
|
||||
env: list[schema.EnvVariable] | None = None,
|
||||
cwd: str | None = None,
|
||||
output_byte_limit: int | None = None,
|
||||
**kwargs: object,
|
||||
) -> schema.CreateTerminalResponse:
|
||||
self._terminals += 1
|
||||
self.created.append((command, cwd))
|
||||
self.create_event.set()
|
||||
if self._block_create:
|
||||
await self.release_create.wait()
|
||||
return schema.CreateTerminalResponse(terminal_id=f'term-{self._terminals}')
|
||||
|
||||
async def wait_for_terminal_exit(
|
||||
self, session_id: str, terminal_id: str, **kwargs: object
|
||||
) -> schema.WaitForTerminalExitResponse:
|
||||
self.exit_event.set()
|
||||
if self._block_exit:
|
||||
await asyncio.Event().wait() # block until cancelled
|
||||
return schema.WaitForTerminalExitResponse(exit_code=self._exit_code, signal=self._signal)
|
||||
|
||||
async def terminal_output(
|
||||
self, session_id: str, terminal_id: str, **kwargs: object
|
||||
) -> schema.TerminalOutputResponse:
|
||||
status = (
|
||||
None if self._no_exit_status else schema.TerminalExitStatus(exit_code=self._exit_code, signal=self._signal)
|
||||
)
|
||||
return schema.TerminalOutputResponse(output=self._output, truncated=self._truncated, exit_status=status)
|
||||
|
||||
async def kill_terminal(
|
||||
self, session_id: str, terminal_id: str, **kwargs: object
|
||||
) -> schema.KillTerminalResponse | None:
|
||||
self.killed.append(terminal_id)
|
||||
return None
|
||||
|
||||
async def release_terminal(
|
||||
self, session_id: str, terminal_id: str, **kwargs: object
|
||||
) -> schema.ReleaseTerminalResponse | None:
|
||||
self.released.append(terminal_id)
|
||||
return None
|
||||
@@ -0,0 +1,24 @@
|
||||
"""A minimal ACP agent process used by the end-to-end stdio test."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
|
||||
from pydantic_ai import Agent
|
||||
from pydantic_ai.models.test import TestModel
|
||||
|
||||
from pydantic_ai_harness.experimental.acp import run_acp_stdio
|
||||
|
||||
|
||||
def build_agent() -> Agent[None, str]:
|
||||
agent = Agent(TestModel())
|
||||
|
||||
@agent.tool_plain
|
||||
def get_weather(city: str) -> str:
|
||||
return f'Sunny in {city}'
|
||||
|
||||
return agent
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
asyncio.run(run_acp_stdio(build_agent()))
|
||||
@@ -0,0 +1,22 @@
|
||||
"""An ACP agent with an approval-gated tool, used by the end-to-end permission test."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from pydantic_ai import Agent
|
||||
from pydantic_ai.models.test import TestModel
|
||||
|
||||
from pydantic_ai_harness.experimental.acp import run_acp_stdio_sync
|
||||
|
||||
|
||||
def build_agent() -> Agent[None, str]:
|
||||
agent = Agent(TestModel())
|
||||
|
||||
@agent.tool_plain(requires_approval=True)
|
||||
def delete_file(path: str) -> str:
|
||||
return f'deleted {path}'
|
||||
|
||||
return agent
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
run_acp_stdio_sync(build_agent())
|
||||
@@ -0,0 +1,13 @@
|
||||
"""An ACP agent that streams a large message, used by the large-payload stdio test."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from pydantic_ai import Agent
|
||||
from pydantic_ai.models.test import TestModel
|
||||
|
||||
from pydantic_ai_harness.experimental.acp import run_acp_stdio_sync
|
||||
|
||||
LARGE_OUTPUT = 'x' * 200_000
|
||||
|
||||
if __name__ == '__main__':
|
||||
run_acp_stdio_sync(Agent(TestModel(custom_output_text=LARGE_OUTPUT)))
|
||||
@@ -0,0 +1,11 @@
|
||||
"""An ACP agent advertising switchable models, used by the stdio model-config test."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from pydantic_ai import Agent
|
||||
from pydantic_ai.models.test import TestModel
|
||||
|
||||
from pydantic_ai_harness.experimental.acp import run_acp_stdio_sync
|
||||
|
||||
if __name__ == '__main__':
|
||||
run_acp_stdio_sync(Agent(TestModel(custom_output_text='hi')), models=['test', 'openai:gpt-4o'])
|
||||
@@ -0,0 +1,41 @@
|
||||
"""An ACP agent mounting client-backed fs/terminal toolsets, for the native-over-stdio test.
|
||||
|
||||
The model calls `read_file` then `run_command` on the first turn; the session config mounts the
|
||||
ACP-client-backed toolsets, so those tool calls must route over the wire to the test's client.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from collections.abc import AsyncIterator
|
||||
|
||||
from pydantic_ai import Agent
|
||||
from pydantic_ai.messages import ModelMessage, ToolReturnPart
|
||||
from pydantic_ai.models.function import AgentInfo, DeltaToolCall, FunctionModel
|
||||
|
||||
from pydantic_ai_harness.experimental.acp import (
|
||||
AcpSession,
|
||||
AcpSessionConfig,
|
||||
acp_filesystem,
|
||||
acp_terminal,
|
||||
run_acp_stdio_sync,
|
||||
)
|
||||
|
||||
|
||||
async def stream(messages: list[ModelMessage], info: AgentInfo) -> AsyncIterator[str | dict[int, DeltaToolCall]]:
|
||||
if any(isinstance(part, ToolReturnPart) for message in messages for part in getattr(message, 'parts', [])):
|
||||
yield 'done'
|
||||
return
|
||||
yield {
|
||||
0: DeltaToolCall(name='read_file', json_args=json.dumps({'path': 'notes.txt'})),
|
||||
1: DeltaToolCall(name='run_command', json_args=json.dumps({'command': 'echo hi'})),
|
||||
}
|
||||
|
||||
|
||||
def session_config(session: AcpSession) -> AcpSessionConfig[None]:
|
||||
toolsets = [toolset for toolset in (acp_filesystem(session), acp_terminal(session)) if toolset is not None]
|
||||
return AcpSessionConfig(deps=None, toolsets=toolsets)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
run_acp_stdio_sync(Agent(FunctionModel(stream_function=stream)), session_config=session_config)
|
||||
@@ -0,0 +1,25 @@
|
||||
"""An ACP agent with a slow tool, used by the over-the-wire cancellation test."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
|
||||
from pydantic_ai import Agent
|
||||
from pydantic_ai.models.test import TestModel
|
||||
|
||||
from pydantic_ai_harness.experimental.acp import run_acp_stdio_sync
|
||||
|
||||
|
||||
def build_agent() -> Agent[None, str]:
|
||||
agent = Agent(TestModel())
|
||||
|
||||
@agent.tool_plain
|
||||
async def slow() -> str:
|
||||
await asyncio.sleep(30)
|
||||
return 'done'
|
||||
|
||||
return agent
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
run_acp_stdio_sync(build_agent())
|
||||
@@ -0,0 +1,99 @@
|
||||
"""In-process ACP *wire* harness: drive a `PydanticAIACPAgent` through the real SDK router and codec.
|
||||
|
||||
The direct-call tests invoke adapter methods straight -- below the JSON-RPC router and JSON
|
||||
serialization. That boundary cannot see two whole classes of bug: a method the SDK router rejects
|
||||
before it reaches the adapter (e.g. an unstable method without `use_unstable_protocol`), and a
|
||||
frame whose *serialized bytes* differ from the Python object the adapter handed over (e.g. an
|
||||
`ensure_ascii` expansion that overruns the client's read buffer).
|
||||
|
||||
`wire_agent` closes that gap without a subprocess: it connects a real `ClientSideConnection` to
|
||||
`acp.run_agent` over an in-memory `socket.socketpair()`, in one event loop. Every request crosses
|
||||
the SDK's method routing and serialization, and every `session/update` arrives as bytes the client
|
||||
re-parses -- exactly as an editor would drive it. The client-side `StreamReader` keeps asyncio's
|
||||
default 64 KiB line limit, so an oversized frame fails the read just as it would over stdio.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import contextlib
|
||||
import socket
|
||||
from types import TracebackType
|
||||
|
||||
import acp
|
||||
from acp.client.connection import ClientSideConnection
|
||||
|
||||
from tests.experimental.acp._acp_clients import RecordingClientBase # pyright: ignore[reportMissingTypeStubs]
|
||||
|
||||
|
||||
class WireClient(RecordingClientBase):
|
||||
"""A real ACP client reached over the wire, recording the `session/update`s it receives.
|
||||
|
||||
The protocol-level conformance tests in this batch do not drive permission, filesystem, or
|
||||
terminal methods; those stay stubbed by `RecordingClientBase`.
|
||||
"""
|
||||
|
||||
def texts(self) -> str:
|
||||
"""The concatenated agent message text received over the wire (other update kinds skipped)."""
|
||||
out: list[str] = []
|
||||
for update in self.updates:
|
||||
if getattr(update, 'session_update', '') == 'agent_message_chunk':
|
||||
out.append(getattr(getattr(update, 'content', None), 'text', ''))
|
||||
return ''.join(out)
|
||||
|
||||
|
||||
class wire_agent:
|
||||
"""Serve `adapter` over an in-memory socket pair; `async with` yields a connected client connection.
|
||||
|
||||
The agent runs as a background task for the lifetime of the context; on exit the task is
|
||||
cancelled and both stream ends are closed. Set `unstable=False` to drive the agent without
|
||||
`use_unstable_protocol`, the configuration in which the SDK router rejects unstable methods.
|
||||
|
||||
A class rather than an `@asynccontextmanager` generator only because the bundled type stubs
|
||||
flag that decorator as deprecated under strict checking.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
adapter: acp.Agent,
|
||||
client: WireClient | None = None,
|
||||
*,
|
||||
unstable: bool = True,
|
||||
) -> None:
|
||||
self._adapter = adapter
|
||||
self._client = client or WireClient()
|
||||
self._unstable = unstable
|
||||
self._server: asyncio.Task[None] | None = None
|
||||
self._writers: tuple[asyncio.StreamWriter, asyncio.StreamWriter] | None = None
|
||||
|
||||
async def __aenter__(self) -> tuple[ClientSideConnection, WireClient]:
|
||||
agent_sock, client_sock = socket.socketpair()
|
||||
# Each side's `input_stream` is the writer it sends to the peer with; `output_stream` is
|
||||
# the reader it receives on -- the order both SDK connection constructors require.
|
||||
agent_reader, agent_writer = await asyncio.open_connection(sock=agent_sock)
|
||||
client_reader, client_writer = await asyncio.open_connection(sock=client_sock)
|
||||
self._writers = (client_writer, agent_writer)
|
||||
self._server = asyncio.ensure_future(
|
||||
acp.run_agent(
|
||||
self._adapter,
|
||||
input_stream=agent_writer,
|
||||
output_stream=agent_reader,
|
||||
use_unstable_protocol=self._unstable,
|
||||
)
|
||||
)
|
||||
conn = acp.connect_to_agent(
|
||||
self._client, input_stream=client_writer, output_stream=client_reader, use_unstable_protocol=self._unstable
|
||||
)
|
||||
return conn, self._client
|
||||
|
||||
async def __aexit__(
|
||||
self, exc_type: type[BaseException] | None, exc: BaseException | None, tb: TracebackType | None
|
||||
) -> None:
|
||||
assert self._server is not None and self._writers is not None
|
||||
self._server.cancel()
|
||||
with contextlib.suppress(asyncio.CancelledError):
|
||||
await self._server
|
||||
for writer in self._writers:
|
||||
writer.close()
|
||||
with contextlib.suppress(Exception):
|
||||
await writer.wait_closed()
|
||||
@@ -0,0 +1,30 @@
|
||||
"""Shared fixtures for the ACP adapter tests."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import importlib.util
|
||||
|
||||
import pytest
|
||||
|
||||
# The `agent-client-protocol` dependency is gated on the `acp` extra, so slim CI runs
|
||||
# (no extras) can't import these modules. Ignore them at collection; `test_packaging.py`
|
||||
# stays collected because it checks package metadata, which holds on base installs too.
|
||||
# A conditional expression rather than an `if` statement: branch coverage traces statement
|
||||
# arcs, and no single environment can take both arms of an install-dependent branch.
|
||||
collect_ignore = (
|
||||
[
|
||||
'test_acp.py',
|
||||
'test_conformance.py',
|
||||
'test_content.py',
|
||||
'test_models.py',
|
||||
'test_client_toolsets.py',
|
||||
'test_persistence.py',
|
||||
]
|
||||
if importlib.util.find_spec('acp') is None
|
||||
else []
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def anyio_backend() -> str:
|
||||
return 'asyncio'
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,288 @@
|
||||
"""Tests for the ACP-client-backed filesystem and terminal toolsets (`_client_toolsets.py`)."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from pathlib import Path
|
||||
|
||||
import anyio
|
||||
import pytest
|
||||
from acp import Client, schema
|
||||
from pydantic_ai import RunContext
|
||||
from pydantic_ai.models.test import TestModel
|
||||
from pydantic_ai.usage import RunUsage
|
||||
|
||||
from pydantic_ai_harness.experimental.acp import (
|
||||
AcpFileSystemToolset,
|
||||
AcpSession,
|
||||
AcpTerminalToolset,
|
||||
acp_filesystem,
|
||||
acp_terminal,
|
||||
)
|
||||
from tests.experimental.acp._acp_clients import RecordingClient # pyright: ignore[reportMissingTypeStubs]
|
||||
|
||||
pytestmark = pytest.mark.anyio
|
||||
|
||||
|
||||
def _ctx() -> RunContext[None]:
|
||||
return RunContext[None](deps=None, model=TestModel(), usage=RunUsage(), prompt=None, messages=[], run_step=1)
|
||||
|
||||
|
||||
def _session(client: Client, capabilities: schema.ClientCapabilities | None) -> AcpSession:
|
||||
return AcpSession(
|
||||
cwd='/ws',
|
||||
mcp_servers=[],
|
||||
client_capabilities=capabilities,
|
||||
client=client,
|
||||
session_id='sid',
|
||||
)
|
||||
|
||||
|
||||
def _fs_caps(*, read: bool, write: bool) -> schema.ClientCapabilities:
|
||||
return schema.ClientCapabilities(fs=schema.FileSystemCapabilities(read_text_file=read, write_text_file=write))
|
||||
|
||||
|
||||
# --- Filesystem toolset ----------------------------------------------------
|
||||
|
||||
|
||||
async def test_read_file_reads_through_the_client() -> None:
|
||||
client = RecordingClient({'/ws/a.py': 'hello'})
|
||||
ts = AcpFileSystemToolset[None](client=client, session_id='sid')
|
||||
assert await ts.read_file('/ws/a.py') == 'hello'
|
||||
assert client.reads == [('/ws/a.py', 'sid')] # path and session id reached the client unchanged
|
||||
|
||||
|
||||
async def test_write_file_writes_through_the_client() -> None:
|
||||
client = RecordingClient()
|
||||
ts = AcpFileSystemToolset[None](client=client, session_id='sid')
|
||||
result = await ts.write_file('/ws/b.py', 'data')
|
||||
assert client.writes == [('/ws/b.py', 'data', 'sid')]
|
||||
assert client.files['/ws/b.py'] == 'data'
|
||||
assert '/ws/b.py' in result # confirmation names the path so the model knows the write landed
|
||||
|
||||
|
||||
async def test_relative_paths_resolve_against_the_session_cwd() -> None:
|
||||
# ACP requires absolute paths on the wire, but a model routinely emits workspace-relative
|
||||
# ones (the local FileSystem tools take them); with a cwd the toolset resolves them.
|
||||
client = RecordingClient({'/ws/src/a.py': 'code'})
|
||||
ts = AcpFileSystemToolset[None](client=client, session_id='sid', cwd='/ws')
|
||||
assert await ts.read_file('src/a.py') == 'code'
|
||||
await ts.write_file('src/b.py', 'new')
|
||||
assert client.reads == [('/ws/src/a.py', 'sid')]
|
||||
assert client.writes == [('/ws/src/b.py', 'new', 'sid')]
|
||||
|
||||
|
||||
async def test_absolute_paths_and_cwdless_toolsets_pass_paths_through() -> None:
|
||||
client = RecordingClient({'/elsewhere/a.py': 'x', 'raw.txt': 'y'})
|
||||
with_cwd = AcpFileSystemToolset[None](client=client, session_id='sid', cwd='/ws')
|
||||
assert await with_cwd.read_file('/elsewhere/a.py') == 'x' # absolute paths are not rewritten
|
||||
without_cwd = AcpFileSystemToolset[None](client=client, session_id='sid')
|
||||
assert await without_cwd.read_file('raw.txt') == 'y' # no cwd: passed through unchanged
|
||||
assert client.reads == [('/elsewhere/a.py', 'sid'), ('raw.txt', 'sid')]
|
||||
|
||||
|
||||
async def test_filesystem_registers_read_file_and_write_file_tools() -> None:
|
||||
# The tool names match the local FileSystem capability so the default presenter renders them.
|
||||
ts = AcpFileSystemToolset[None](client=RecordingClient(), session_id='sid')
|
||||
assert set(await ts.get_tools(_ctx())) == {'read_file', 'write_file'}
|
||||
|
||||
|
||||
async def test_acp_filesystem_builds_a_working_toolset_when_fs_is_advertised() -> None:
|
||||
client = RecordingClient({'/ws/a.py': 'hi'})
|
||||
toolset = acp_filesystem(_session(client, _fs_caps(read=True, write=True)))
|
||||
assert isinstance(toolset, AcpFileSystemToolset)
|
||||
assert await toolset.read_file('/ws/a.py') == 'hi' # the built toolset routes through the same client
|
||||
|
||||
|
||||
async def test_acp_filesystem_read_only_client_reads_via_acp_and_writes_locally(tmp_path: Path) -> None:
|
||||
# A read-only client keeps editor-native reads, but writes go to the local workspace disk
|
||||
# rather than the client (coherent only when the agent shares that disk -- see the helper docs).
|
||||
client = RecordingClient({str(tmp_path / 'notes.txt'): 'hello'})
|
||||
session = _session(client, _fs_caps(read=True, write=False))
|
||||
session = AcpSession(
|
||||
cwd=str(tmp_path),
|
||||
mcp_servers=session.mcp_servers,
|
||||
client_capabilities=session.client_capabilities,
|
||||
client=client,
|
||||
session_id=session.session_id,
|
||||
)
|
||||
toolset = acp_filesystem(session)
|
||||
assert isinstance(toolset, AcpFileSystemToolset)
|
||||
|
||||
assert await toolset.read_file('notes.txt') == 'hello'
|
||||
assert client.reads == [(str(tmp_path / 'notes.txt'), 'sid')] # the read routed through the editor
|
||||
await toolset.write_file('out.txt', 'data')
|
||||
assert client.writes == [] # the client was never asked to write
|
||||
assert (tmp_path / 'out.txt').read_text() == 'data' # the write landed on local disk
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
'capabilities',
|
||||
[
|
||||
pytest.param(None, id='no-capabilities'),
|
||||
pytest.param(schema.ClientCapabilities(), id='no-fs'),
|
||||
pytest.param(_fs_caps(read=False, write=True), id='no-read'),
|
||||
],
|
||||
)
|
||||
def test_acp_filesystem_returns_none_when_no_readable_filesystem(
|
||||
capabilities: schema.ClientCapabilities | None,
|
||||
) -> None:
|
||||
assert acp_filesystem(_session(RecordingClient(), capabilities)) is None
|
||||
|
||||
|
||||
# --- Terminal toolset ------------------------------------------------------
|
||||
|
||||
|
||||
async def test_run_command_creates_a_terminal_in_the_session_cwd() -> None:
|
||||
client = RecordingClient(output='ok')
|
||||
await AcpTerminalToolset[None](client=client, session_id='sid', cwd='/ws').run_command('ls')
|
||||
assert client.created == [('ls', '/ws')]
|
||||
|
||||
|
||||
async def test_terminal_registers_run_command_tool() -> None:
|
||||
ts = AcpTerminalToolset[None](client=RecordingClient(), session_id='sid')
|
||||
assert set(await ts.get_tools(_ctx())) == {'run_command'}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
'kwargs, expected',
|
||||
[
|
||||
pytest.param({'output': 'hello', 'exit_code': 0}, 'hello', id='success'),
|
||||
pytest.param({'output': 'boom', 'exit_code': 2}, 'boom\n[exited with code 2]', id='nonzero-exit'),
|
||||
pytest.param(
|
||||
{'output': 'gone', 'exit_code': None, 'signal': 'SIGKILL'},
|
||||
'gone\n[terminated by signal SIGKILL]',
|
||||
id='signal',
|
||||
),
|
||||
pytest.param(
|
||||
{'output': 'partial', 'truncated': True, 'no_exit_status': True},
|
||||
'partial\n[output truncated]',
|
||||
id='truncated-no-status',
|
||||
),
|
||||
],
|
||||
)
|
||||
async def test_run_command_formats_output_and_releases(kwargs: dict[str, object], expected: str) -> None:
|
||||
client = RecordingClient(**kwargs) # pyright: ignore[reportArgumentType]
|
||||
result = await AcpTerminalToolset[None](client=client, session_id='sid', cwd='/ws').run_command('cmd')
|
||||
assert result == expected
|
||||
assert client.released == ['term-1'] # the terminal is always released
|
||||
|
||||
|
||||
async def test_run_command_kills_and_releases_the_terminal_on_cancel() -> None:
|
||||
client = RecordingClient(block_exit=True)
|
||||
ts = AcpTerminalToolset[None](client=client, session_id='sid', cwd='/ws')
|
||||
async with anyio.create_task_group() as tg:
|
||||
tg.start_soon(ts.run_command, 'sleep 100')
|
||||
await client.exit_event.wait() # the terminal exists and the command is running
|
||||
tg.cancel_scope.cancel()
|
||||
assert client.killed == ['term-1'] # killed before unwinding
|
||||
assert client.released == ['term-1'] # and released so it is not left behind
|
||||
|
||||
|
||||
async def test_run_command_cancelled_during_create_still_kills_the_terminal() -> None:
|
||||
# A raw `task.cancel()` -- how the adapter and pydantic-ai actually deliver cancellation, and
|
||||
# which pierces anyio shields -- lands while the create is in flight. The request was already
|
||||
# on the wire, so the client started the command regardless; the late-learned terminal must
|
||||
# still be killed and released, not leaked running in the editor.
|
||||
client = RecordingClient(block_create=True)
|
||||
ts = AcpTerminalToolset[None](client=client, session_id='sid', cwd='/ws')
|
||||
task = asyncio.ensure_future(ts.run_command('sleep 100'))
|
||||
await asyncio.wait_for(client.create_event.wait(), timeout=5)
|
||||
task.cancel()
|
||||
client.release_create.set() # the client answers the create only after the cancellation
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await asyncio.wait_for(task, timeout=5)
|
||||
assert client.killed == ['term-1']
|
||||
assert client.released == ['term-1']
|
||||
|
||||
|
||||
async def test_run_command_cancelled_during_a_failing_create_has_nothing_to_clean_up() -> None:
|
||||
class _CreateFails(RecordingClient):
|
||||
async def create_terminal(
|
||||
self,
|
||||
session_id: str,
|
||||
command: str,
|
||||
args: list[str] | None = None,
|
||||
env: list[schema.EnvVariable] | None = None,
|
||||
cwd: str | None = None,
|
||||
output_byte_limit: int | None = None,
|
||||
**kwargs: object,
|
||||
) -> schema.CreateTerminalResponse:
|
||||
self.create_event.set()
|
||||
await self.release_create.wait()
|
||||
raise RuntimeError('client could not create a terminal')
|
||||
|
||||
client = _CreateFails()
|
||||
ts = AcpTerminalToolset[None](client=client, session_id='sid', cwd='/ws')
|
||||
task = asyncio.ensure_future(ts.run_command('sleep 100'))
|
||||
await asyncio.wait_for(client.create_event.wait(), timeout=5)
|
||||
task.cancel()
|
||||
client.release_create.set()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await asyncio.wait_for(task, timeout=5)
|
||||
# No terminal ever existed: nothing to kill or release, and the create failure must not
|
||||
# replace the cancellation already unwinding.
|
||||
assert client.killed == []
|
||||
assert client.released == []
|
||||
|
||||
|
||||
async def test_run_command_create_failure_propagates() -> None:
|
||||
class _CreateRaises(RecordingClient):
|
||||
async def create_terminal(
|
||||
self,
|
||||
session_id: str,
|
||||
command: str,
|
||||
args: list[str] | None = None,
|
||||
env: list[schema.EnvVariable] | None = None,
|
||||
cwd: str | None = None,
|
||||
output_byte_limit: int | None = None,
|
||||
**kwargs: object,
|
||||
) -> schema.CreateTerminalResponse:
|
||||
raise RuntimeError('no terminal for you')
|
||||
|
||||
client = _CreateRaises()
|
||||
ts = AcpTerminalToolset[None](client=client, session_id='sid', cwd='/ws')
|
||||
with pytest.raises(RuntimeError, match='no terminal'):
|
||||
await ts.run_command('ls')
|
||||
assert client.released == [] # nothing came into existence, so nothing is released
|
||||
|
||||
|
||||
async def test_run_command_cancel_survives_a_failing_kill() -> None:
|
||||
class _KillRaises(RecordingClient):
|
||||
async def kill_terminal(
|
||||
self, session_id: str, terminal_id: str, **kwargs: object
|
||||
) -> schema.KillTerminalResponse | None:
|
||||
self.killed.append(terminal_id)
|
||||
raise RuntimeError('client kill failed')
|
||||
|
||||
client = _KillRaises(block_exit=True)
|
||||
ts = AcpTerminalToolset[None](client=client, session_id='sid', cwd='/ws')
|
||||
async with anyio.create_task_group() as tg:
|
||||
tg.start_soon(ts.run_command, 'sleep 100')
|
||||
await client.exit_event.wait()
|
||||
tg.cancel_scope.cancel()
|
||||
# The kill failure is suppressed, so cancellation still unwinds cleanly and the terminal is
|
||||
# still released rather than leaked.
|
||||
assert client.killed == ['term-1']
|
||||
assert client.released == ['term-1']
|
||||
|
||||
|
||||
async def test_acp_terminal_builds_a_toolset_when_terminal_is_advertised() -> None:
|
||||
client = RecordingClient(output='hi')
|
||||
toolset = acp_terminal(_session(client, schema.ClientCapabilities(terminal=True)))
|
||||
assert isinstance(toolset, AcpTerminalToolset)
|
||||
assert await toolset.run_command('echo hi') == 'hi' # the built toolset routes through the same client
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
'capabilities',
|
||||
[
|
||||
pytest.param(None, id='no-capabilities'),
|
||||
pytest.param(schema.ClientCapabilities(terminal=False), id='terminal-false'),
|
||||
pytest.param(schema.ClientCapabilities(), id='terminal-unset'),
|
||||
],
|
||||
)
|
||||
def test_acp_terminal_returns_none_when_terminal_is_unsupported(
|
||||
capabilities: schema.ClientCapabilities | None,
|
||||
) -> None:
|
||||
assert acp_terminal(_session(RecordingClient(), capabilities)) is None
|
||||
@@ -0,0 +1,217 @@
|
||||
"""ACP spec-conformance tests, driven over a real in-memory wire (`tests/experimental/acp/_wire.py`).
|
||||
|
||||
These tests are organized by spec clause, not by adapter method, and each pins the *protocol*
|
||||
invariant with an oracle constructed independently of the adapter -- not the adapter's own output
|
||||
compared back against itself. They run through the SDK's JSON-RPC router and serialization, so they
|
||||
cover what direct-call tests cannot: a method the router gates before the adapter sees it, and a
|
||||
frame whose serialized bytes overrun the client's read buffer.
|
||||
|
||||
Each test names the spec clause it enforces; the package README's feature sections and
|
||||
"Cancellation and limitations" summarize the adapter's supported surface.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import acp
|
||||
import pytest
|
||||
from acp import RequestError, schema
|
||||
from pydantic_ai import Agent
|
||||
from pydantic_ai.models.test import TestModel
|
||||
|
||||
from pydantic_ai_harness.experimental.acp import InMemorySessionStore, PydanticAIACPAgent
|
||||
from tests.experimental.acp._wire import WireClient, wire_agent # pyright: ignore[reportMissingTypeStubs]
|
||||
|
||||
pytestmark = pytest.mark.anyio
|
||||
|
||||
|
||||
def _agent(text: str = 'hello') -> PydanticAIACPAgent[None, str]:
|
||||
return PydanticAIACPAgent(Agent(TestModel(custom_output_text=text)))
|
||||
|
||||
|
||||
def _method_not_found_code() -> int:
|
||||
return RequestError.method_not_found('probe').code
|
||||
|
||||
|
||||
def _invalid_params_code() -> int:
|
||||
return RequestError.invalid_params({}).code
|
||||
|
||||
|
||||
class TestVersionNegotiation:
|
||||
"""initialization.md: the Agent echoes a supported version, else returns its latest."""
|
||||
|
||||
async def test_supported_version_is_echoed_back(self) -> None:
|
||||
# Clause (MUST): "If the Agent supports the requested version, it MUST respond with the same
|
||||
# version." Oracle: the response echoes the *input*, asserted for every in-range version
|
||||
# rather than the literal 1 -- so a future protocol bump keeps the mid-range case covered.
|
||||
async with wire_agent(_agent()) as (conn, _client):
|
||||
for version in range(1, acp.PROTOCOL_VERSION + 1):
|
||||
response = await conn.initialize(protocol_version=version)
|
||||
assert response.protocol_version == version
|
||||
|
||||
async def test_unsupported_version_negotiates_down_to_latest(self) -> None:
|
||||
# Clause (MUST): "Otherwise, the Agent MUST respond with the latest version it supports."
|
||||
async with wire_agent(_agent()) as (conn, _client):
|
||||
too_new = await conn.initialize(protocol_version=acp.PROTOCOL_VERSION + 1)
|
||||
too_old = await conn.initialize(protocol_version=0)
|
||||
assert too_new.protocol_version == acp.PROTOCOL_VERSION
|
||||
assert too_old.protocol_version == acp.PROTOCOL_VERSION
|
||||
|
||||
|
||||
class TestCapabilityAdvertisement:
|
||||
"""initialization.md: the Agent advertises exactly the optional methods it supports."""
|
||||
|
||||
async def test_unsupported_methods_are_advertised_off(self) -> None:
|
||||
# Clause (MUST): a capability omitted/absent means the method is unsupported. The adapter
|
||||
# rejects list/fork/resume/set_mode/set_config; the advertisement side must match. Oracle:
|
||||
# read the advertised capabilities back off the wire, independent of the reject handlers.
|
||||
async with wire_agent(_agent()) as (conn, _client):
|
||||
caps = (await conn.initialize(protocol_version=1)).agent_capabilities
|
||||
assert caps is not None
|
||||
session_caps = caps.session_capabilities
|
||||
assert session_caps is not None
|
||||
assert session_caps.close is not None # the one session method the adapter supports
|
||||
assert session_caps.list is None and session_caps.fork is None and session_caps.resume is None
|
||||
assert caps.load_session is False # no store given
|
||||
assert caps.mcp_capabilities == schema.McpCapabilities(http=False, sse=False)
|
||||
|
||||
async def test_no_authentication_is_advertised(self) -> None:
|
||||
# authentication.md: an agent advertises auth options in `authMethods`; none here means
|
||||
# no-auth. Oracle: the advertised list is empty, asserted directly off the wire.
|
||||
async with wire_agent(_agent()) as (conn, _client):
|
||||
response = await conn.initialize(protocol_version=1)
|
||||
assert response.auth_methods == []
|
||||
|
||||
async def test_load_session_advertised_only_with_a_store(self) -> None:
|
||||
adapter = PydanticAIACPAgent(Agent(TestModel()), session_store=InMemorySessionStore())
|
||||
async with wire_agent(adapter) as (conn, _client):
|
||||
caps = (await conn.initialize(protocol_version=1)).agent_capabilities
|
||||
assert caps is not None and caps.load_session is True
|
||||
|
||||
async def test_session_setup_without_models_advertises_no_modes_or_config_options(self) -> None:
|
||||
# session-modes.md / session-config-options.md: both are advertised per-session. The adapter
|
||||
# supports no modes and only advertises config options when a model switch-set is configured.
|
||||
async with wire_agent(_agent()) as (conn, _client):
|
||||
await conn.initialize(protocol_version=1)
|
||||
session = await conn.new_session(cwd='/ws', mcp_servers=[])
|
||||
assert session.modes is None
|
||||
assert session.config_options is None
|
||||
|
||||
|
||||
class TestModelConfigRouting:
|
||||
"""session-config-options.md: model switching is a stable config-option update."""
|
||||
|
||||
async def test_set_config_option_routes_without_unstable_protocol(self) -> None:
|
||||
adapter: PydanticAIACPAgent[None, str] = PydanticAIACPAgent(Agent(TestModel()), models=['test'])
|
||||
async with wire_agent(adapter, unstable=False) as (conn, _client):
|
||||
await conn.initialize(protocol_version=1)
|
||||
session = await conn.new_session(cwd='/ws', mcp_servers=[])
|
||||
response = await conn.set_config_option(config_id='model', value='test', session_id=session.session_id)
|
||||
[option] = response.config_options
|
||||
assert isinstance(option, schema.SessionConfigOptionSelect)
|
||||
assert option.current_value == 'test'
|
||||
|
||||
|
||||
class TestUnstableMethodRouting:
|
||||
"""transports/schema: an UNSTABLE method is reachable only with `use_unstable_protocol`."""
|
||||
|
||||
async def test_unstable_close_reaches_the_adapter_when_enabled(self) -> None:
|
||||
# session/close is UNSTABLE in the SDK router; with the flag on it must route through to
|
||||
# the adapter. (Direct-call tests bypass the router and cannot see this.)
|
||||
adapter: PydanticAIACPAgent[None, str] = PydanticAIACPAgent(Agent(TestModel()))
|
||||
async with wire_agent(adapter, unstable=True) as (conn, _client):
|
||||
await conn.initialize(protocol_version=1)
|
||||
session = await conn.new_session(cwd='/ws', mcp_servers=[])
|
||||
closed = await conn.close_session(session_id=session.session_id)
|
||||
assert closed == schema.CloseSessionResponse()
|
||||
|
||||
# The SDK router emits a UserWarning before rejecting an unstable method; this suite promotes
|
||||
# warnings to errors (so the warning would otherwise surface as a wrapped Internal error in the
|
||||
# agent task). Letting it stay a warning lets the real production error code -- method_not_found
|
||||
# -- reach the client, which is the contract this test pins.
|
||||
@pytest.mark.filterwarnings('ignore::UserWarning')
|
||||
async def test_unstable_close_is_gated_off_without_the_flag(self) -> None:
|
||||
# The load-bearing reason `run_acp_stdio` passes use_unstable_protocol=True: without it the
|
||||
# SDK router rejects close with method_not_found before the adapter.
|
||||
adapter: PydanticAIACPAgent[None, str] = PydanticAIACPAgent(Agent(TestModel()))
|
||||
async with wire_agent(adapter, unstable=False) as (conn, _client):
|
||||
await conn.initialize(protocol_version=1)
|
||||
session = await conn.new_session(cwd='/ws', mcp_servers=[])
|
||||
with pytest.raises(RequestError) as exc:
|
||||
await conn.close_session(session_id=session.session_id)
|
||||
assert exc.value.code == _method_not_found_code()
|
||||
|
||||
|
||||
class TestErrorCodes:
|
||||
"""The JSON-RPC error *code* distinguishes unsupported from malformed -- not just that it raised."""
|
||||
|
||||
async def test_unsupported_method_uses_method_not_found(self) -> None:
|
||||
async with wire_agent(_agent()) as (conn, _client):
|
||||
await conn.initialize(protocol_version=1)
|
||||
session = await conn.new_session(cwd='/ws', mcp_servers=[])
|
||||
with pytest.raises(RequestError) as exc:
|
||||
await conn.set_session_mode(mode_id='whatever', session_id=session.session_id)
|
||||
assert exc.value.code == _method_not_found_code()
|
||||
|
||||
async def test_unknown_model_uses_invalid_params(self) -> None:
|
||||
adapter: PydanticAIACPAgent[None, str] = PydanticAIACPAgent(Agent(TestModel()), models=['test'])
|
||||
async with wire_agent(adapter) as (conn, _client):
|
||||
await conn.initialize(protocol_version=1)
|
||||
session = await conn.new_session(cwd='/ws', mcp_servers=[])
|
||||
with pytest.raises(RequestError) as exc:
|
||||
# The model config option exists here, so a bad model id is malformed input.
|
||||
await conn.set_config_option(config_id='model', value='not-a-model', session_id=session.session_id)
|
||||
assert exc.value.code == _invalid_params_code()
|
||||
|
||||
async def test_unknown_config_option_uses_invalid_params(self) -> None:
|
||||
adapter: PydanticAIACPAgent[None, str] = PydanticAIACPAgent(Agent(TestModel()), models=['test'])
|
||||
async with wire_agent(adapter) as (conn, _client):
|
||||
await conn.initialize(protocol_version=1)
|
||||
session = await conn.new_session(cwd='/ws', mcp_servers=[])
|
||||
with pytest.raises(RequestError) as exc:
|
||||
await conn.set_config_option(config_id='theme', value='dark', session_id=session.session_id)
|
||||
assert exc.value.code == _invalid_params_code()
|
||||
|
||||
|
||||
class TestSessionLoadReplay:
|
||||
"""session-setup.md (MUST): `session/load` replays the *entire* conversation, user turns included."""
|
||||
|
||||
async def test_load_replays_the_user_turn_over_the_wire(self) -> None:
|
||||
store = InMemorySessionStore()
|
||||
adapter = PydanticAIACPAgent(Agent(TestModel(custom_output_text='hi there')), session_store=store)
|
||||
async with wire_agent(adapter) as (conn, _client):
|
||||
await conn.initialize(protocol_version=1)
|
||||
session = await conn.new_session(cwd='/ws', mcp_servers=[])
|
||||
await conn.prompt(prompt=[acp.text_block('remember this')], session_id=session.session_id)
|
||||
|
||||
# Reopen on a fresh connection, as an editor would after a restart.
|
||||
reopened = WireClient()
|
||||
async with wire_agent(adapter, reopened) as (conn, _client):
|
||||
await conn.initialize(protocol_version=1)
|
||||
await conn.load_session(cwd='/ws', session_id=session.session_id, mcp_servers=[])
|
||||
|
||||
# The replay must begin with the user's own message -- the half a transcript-only replay
|
||||
# would silently drop. Oracle built by hand from the known prompt, not from adapter output.
|
||||
kinds = [getattr(u, 'session_update', '') for u in reopened.updates]
|
||||
assert kinds[0] == 'user_message_chunk'
|
||||
first = reopened.updates[0]
|
||||
assert getattr(getattr(first, 'content', None), 'text', '') == 'remember this'
|
||||
# The agent's reply is replayed too (and `texts()` skips the leading user_message_chunk).
|
||||
assert 'agent_message_chunk' in kinds
|
||||
assert reopened.texts() == 'hi there'
|
||||
|
||||
|
||||
class TestStreamedFrameBytes:
|
||||
"""transports.md: a streamed frame must fit the client's read buffer -- bound by bytes, not chars."""
|
||||
|
||||
async def test_large_non_ascii_output_reassembles_intact(self) -> None:
|
||||
# The client StreamReader keeps asyncio's default 64 KiB line limit. Under `ensure_ascii`
|
||||
# each emoji serializes to a 12-byte surrogate-pair escape, so a char-count chunker would
|
||||
# emit ~98 KiB frames and the read would overrun. Oracle: the client reassembles the exact
|
||||
# payload -- only possible if every frame stayed under the buffer.
|
||||
payload = '\U0001f600' * 20_000 # 20k astral chars -> ~240 KiB of escaped JSON bytes
|
||||
async with wire_agent(_agent(payload)) as (conn, client):
|
||||
await conn.initialize(protocol_version=1)
|
||||
session = await conn.new_session(cwd='/ws', mcp_servers=[])
|
||||
response = await conn.prompt(prompt=[acp.text_block('go')], session_id=session.session_id)
|
||||
assert response.stop_reason == 'end_turn'
|
||||
assert client.texts() == payload
|
||||
@@ -0,0 +1,93 @@
|
||||
"""Tests for ACP content-block to Pydantic AI user-content translation."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
|
||||
import acp
|
||||
from pydantic_ai.messages import BinaryContent, ImageUrl
|
||||
|
||||
from pydantic_ai_harness.experimental.acp._content import prompt_blocks_to_user_content
|
||||
|
||||
|
||||
def test_text_block_becomes_string() -> None:
|
||||
content = prompt_blocks_to_user_content([acp.text_block('hello')])
|
||||
assert content == ['hello']
|
||||
|
||||
|
||||
def test_image_with_data_becomes_binary_content() -> None:
|
||||
raw = b'\x89PNG\r\n'
|
||||
block = acp.image_block(base64.b64encode(raw).decode(), 'image/png')
|
||||
[item] = prompt_blocks_to_user_content([block])
|
||||
assert isinstance(item, BinaryContent)
|
||||
assert item.data == raw
|
||||
assert item.media_type == 'image/png'
|
||||
|
||||
|
||||
def test_image_with_uri_becomes_image_url() -> None:
|
||||
block = acp.image_block('', 'image/png', uri='https://example.com/a.png')
|
||||
[item] = prompt_blocks_to_user_content([block])
|
||||
assert isinstance(item, ImageUrl)
|
||||
assert item.url == 'https://example.com/a.png'
|
||||
|
||||
|
||||
def test_image_with_both_data_and_uri_prefers_inline_data() -> None:
|
||||
# `data` is required and authoritative; `uri` is only a source reference. A block carrying both
|
||||
# must keep the inline bytes rather than be replaced by a link the model may be unable to fetch.
|
||||
raw = b'\x89PNG\r\n'
|
||||
block = acp.image_block(base64.b64encode(raw).decode(), 'image/png', uri='https://example.com/a.png')
|
||||
[item] = prompt_blocks_to_user_content([block])
|
||||
assert isinstance(item, BinaryContent)
|
||||
assert item.data == raw
|
||||
|
||||
|
||||
def test_audio_block_becomes_binary_content() -> None:
|
||||
block = acp.audio_block(base64.b64encode(b'snd').decode(), 'audio/wav')
|
||||
[item] = prompt_blocks_to_user_content([block])
|
||||
assert isinstance(item, BinaryContent)
|
||||
assert item.media_type == 'audio/wav'
|
||||
|
||||
|
||||
def test_embedded_text_resource_becomes_string() -> None:
|
||||
block = acp.resource_block(acp.embedded_text_resource('file:///a.txt', 'from file'))
|
||||
assert prompt_blocks_to_user_content([block]) == ['from file']
|
||||
|
||||
|
||||
def test_embedded_blob_resource_becomes_binary_content() -> None:
|
||||
block = acp.resource_block(
|
||||
acp.embedded_blob_resource('file:///a.bin', base64.b64encode(b'xx').decode(), mime_type='application/x-thing')
|
||||
)
|
||||
[item] = prompt_blocks_to_user_content([block])
|
||||
assert isinstance(item, BinaryContent)
|
||||
assert item.data == b'xx'
|
||||
assert item.media_type == 'application/x-thing'
|
||||
|
||||
|
||||
def test_embedded_blob_resource_without_mime_type_uses_octet_stream() -> None:
|
||||
block = acp.resource_block(acp.embedded_blob_resource('file:///a.bin', base64.b64encode(b'xx').decode()))
|
||||
[item] = prompt_blocks_to_user_content([block])
|
||||
assert isinstance(item, BinaryContent)
|
||||
# Falls back to the standard "unknown binary" type rather than an invalid empty string.
|
||||
assert item.media_type == 'application/octet-stream'
|
||||
|
||||
|
||||
def test_resource_link_labels_uri_with_name() -> None:
|
||||
block = acp.resource_link_block('a', 'file:///a.txt')
|
||||
assert prompt_blocks_to_user_content([block]) == ['a (file:///a.txt)']
|
||||
|
||||
|
||||
def test_resource_link_prefers_title_over_name() -> None:
|
||||
block = acp.resource_link_block('a.txt', 'file:///a.txt', title='My File')
|
||||
assert prompt_blocks_to_user_content([block]) == ['My File (file:///a.txt)']
|
||||
|
||||
|
||||
def test_resource_link_without_label_falls_back_to_uri() -> None:
|
||||
block = acp.resource_link_block('', 'file:///a.txt')
|
||||
assert prompt_blocks_to_user_content([block]) == ['file:///a.txt']
|
||||
|
||||
|
||||
def test_mixed_blocks_preserve_order() -> None:
|
||||
blocks = [acp.text_block('look at'), acp.image_block('', 'image/png', uri='https://example.com/a.png')]
|
||||
content = prompt_blocks_to_user_content(blocks)
|
||||
assert content[0] == 'look at'
|
||||
assert isinstance(content[1], ImageUrl)
|
||||
@@ -0,0 +1,211 @@
|
||||
"""Tests for ACP model selection through session config options."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import AsyncIterator, Sequence
|
||||
from typing import Literal
|
||||
|
||||
import acp
|
||||
import pytest
|
||||
from acp import schema
|
||||
from pydantic_ai import Agent
|
||||
from pydantic_ai.messages import ModelMessage
|
||||
from pydantic_ai.models import KnownModelName
|
||||
from pydantic_ai.models.function import AgentInfo, FunctionModel
|
||||
from pydantic_ai.models.test import TestModel
|
||||
|
||||
from pydantic_ai_harness.experimental.acp import InMemorySessionStore, PydanticAIACPAgent
|
||||
from pydantic_ai_harness.experimental.acp._adapter import _all_known_model_names
|
||||
from tests.experimental.acp._acp_clients import RecordingClient # pyright: ignore[reportMissingTypeStubs]
|
||||
|
||||
pytestmark = pytest.mark.anyio
|
||||
|
||||
|
||||
def test_all_known_model_names_are_strings() -> None:
|
||||
# Guards `models='all'` advertising only string model ids from Pydantic AI's public
|
||||
# enumeration API.
|
||||
names = _all_known_model_names()
|
||||
assert len(names) > 100
|
||||
assert all(isinstance(name, str) for name in names)
|
||||
|
||||
|
||||
def _adapter(
|
||||
*, models: Sequence[KnownModelName | str] | Literal['all'] | None = None, store: InMemorySessionStore | None = None
|
||||
) -> PydanticAIACPAgent[None, str]:
|
||||
return PydanticAIACPAgent(Agent(TestModel(custom_output_text='hi')), models=models, session_store=store)
|
||||
|
||||
|
||||
def _text_from(client: RecordingClient) -> str:
|
||||
return ''.join(
|
||||
str(getattr(getattr(update, 'content', None), 'text', ''))
|
||||
for update in client.updates
|
||||
if getattr(update, 'session_update', '') == 'agent_message_chunk'
|
||||
)
|
||||
|
||||
|
||||
def _model_option(
|
||||
response: schema.NewSessionResponse | schema.LoadSessionResponse | schema.SetSessionConfigOptionResponse,
|
||||
) -> schema.SessionConfigOptionSelect:
|
||||
options = response.config_options
|
||||
assert options is not None
|
||||
[option] = options
|
||||
assert isinstance(option, schema.SessionConfigOptionSelect)
|
||||
return option
|
||||
|
||||
|
||||
def _select_options(option: schema.SessionConfigOptionSelect) -> list[schema.SessionConfigSelectOption]:
|
||||
options: list[schema.SessionConfigSelectOption] = []
|
||||
for item in option.options:
|
||||
assert isinstance(item, schema.SessionConfigSelectOption)
|
||||
options.append(item)
|
||||
return options
|
||||
|
||||
|
||||
async def _started(adapter: PydanticAIACPAgent[None, str]) -> str:
|
||||
adapter.on_connect(RecordingClient())
|
||||
await adapter.initialize(protocol_version=1)
|
||||
return (await adapter.new_session(cwd='/ws')).session_id
|
||||
|
||||
|
||||
async def test_models_all_advertises_every_known_model() -> None:
|
||||
adapter = _adapter(models='all')
|
||||
adapter.on_connect(RecordingClient())
|
||||
await adapter.initialize(protocol_version=1)
|
||||
response = await adapter.new_session(cwd='/ws')
|
||||
option = _model_option(response)
|
||||
ids = [model.value for model in _select_options(option)]
|
||||
assert len(ids) > 100 # the whole known set, not a curated handful
|
||||
assert 'openai:gpt-4o' in ids
|
||||
assert option.current_value == ids[0] # first known model is the default
|
||||
|
||||
|
||||
async def test_new_session_advertises_configured_models() -> None:
|
||||
adapter = _adapter(models=['openai:gpt-4o', 'test'])
|
||||
adapter.on_connect(RecordingClient())
|
||||
await adapter.initialize(protocol_version=1)
|
||||
response = await adapter.new_session(cwd='/ws')
|
||||
option = _model_option(response)
|
||||
assert option.id == 'model'
|
||||
assert option.name == 'Model'
|
||||
assert [model.value for model in _select_options(option)] == ['openai:gpt-4o', 'test']
|
||||
assert option.current_value == 'openai:gpt-4o' # the first configured model is the default
|
||||
|
||||
|
||||
async def test_new_session_without_models_advertises_none() -> None:
|
||||
adapter = _adapter()
|
||||
assert (await adapter.new_session(cwd='/ws')).config_options is None
|
||||
|
||||
|
||||
async def test_set_model_config_updates_the_session_and_persists() -> None:
|
||||
store = InMemorySessionStore()
|
||||
adapter = _adapter(models=['openai:gpt-4o', 'test'], store=store)
|
||||
session_id = await _started(adapter)
|
||||
response = await adapter.set_config_option(config_id='model', value='test', session_id=session_id)
|
||||
assert response is not None
|
||||
option = _model_option(response)
|
||||
assert option.current_value == 'test'
|
||||
assert adapter._sessions[session_id].model == 'test' # pyright: ignore[reportPrivateUsage]
|
||||
stored = await store.load(session_id)
|
||||
assert stored is not None and stored.model == 'test'
|
||||
|
||||
|
||||
async def test_selected_model_applies_to_a_run() -> None:
|
||||
# The agent's own model is canned to answer 'hi'; the 'test' override resolves to a default
|
||||
# TestModel whose canned answer differs, so the override observably reached the run.
|
||||
client = RecordingClient()
|
||||
adapter = _adapter(models=['test'])
|
||||
adapter.on_connect(client)
|
||||
await adapter.initialize(protocol_version=1)
|
||||
session_id = (await adapter.new_session(cwd='/ws')).session_id
|
||||
response = await adapter.prompt(prompt=[acp.text_block('hi')], session_id=session_id)
|
||||
assert response.stop_reason == 'end_turn'
|
||||
assert _text_from(client) == 'success (no tool calls)' # TestModel's default, not the agent model's 'hi'
|
||||
|
||||
|
||||
async def test_model_resolver_applies_selected_model_to_a_run() -> None:
|
||||
async def stream(messages: list[ModelMessage], info: AgentInfo) -> AsyncIterator[str]:
|
||||
yield 'resolved model'
|
||||
|
||||
seen: list[str] = []
|
||||
resolved_model = FunctionModel(stream_function=stream)
|
||||
|
||||
def resolve(model_id: str) -> FunctionModel:
|
||||
seen.append(model_id)
|
||||
return resolved_model
|
||||
|
||||
client = RecordingClient()
|
||||
adapter: PydanticAIACPAgent[None, str] = PydanticAIACPAgent(
|
||||
Agent(TestModel(custom_output_text='agent model')),
|
||||
models=['test', 'host:gpt-custom'],
|
||||
model_resolver=resolve,
|
||||
)
|
||||
adapter.on_connect(client)
|
||||
await adapter.initialize(protocol_version=1)
|
||||
session_id = (await adapter.new_session(cwd='/ws')).session_id
|
||||
await adapter.set_config_option(config_id='model', value='host:gpt-custom', session_id=session_id)
|
||||
|
||||
response = await adapter.prompt(prompt=[acp.text_block('hi')], session_id=session_id)
|
||||
|
||||
assert response.stop_reason == 'end_turn'
|
||||
assert seen == ['host:gpt-custom']
|
||||
assert _text_from(client) == 'resolved model'
|
||||
|
||||
|
||||
async def test_selected_model_survives_reload() -> None:
|
||||
store = InMemorySessionStore()
|
||||
adapter = _adapter(models=['openai:gpt-4o', 'test'], store=store)
|
||||
session_id = await _started(adapter)
|
||||
await adapter.set_config_option(config_id='model', value='test', session_id=session_id)
|
||||
response = await adapter.load_session(cwd='/ws', session_id=session_id)
|
||||
assert adapter._sessions[session_id].model == 'test' # pyright: ignore[reportPrivateUsage]
|
||||
assert response is not None
|
||||
option = _model_option(response)
|
||||
assert option.current_value == 'test'
|
||||
|
||||
|
||||
async def test_session_model_option_returns_available_and_current_models() -> None:
|
||||
adapter = _adapter(models=['openai:gpt-4o', 'test'])
|
||||
session_id = await _started(adapter)
|
||||
await adapter.set_config_option(config_id='model', value='test', session_id=session_id)
|
||||
|
||||
state = adapter.session_model_option(session_id)
|
||||
|
||||
assert state is not None
|
||||
assert [model.value for model in _select_options(state)] == ['openai:gpt-4o', 'test']
|
||||
assert state.current_value == 'test'
|
||||
assert adapter.session_model_option('no-such-session') is None
|
||||
|
||||
|
||||
async def test_session_model_option_without_models_returns_none() -> None:
|
||||
adapter = _adapter()
|
||||
session_id = await _started(adapter)
|
||||
|
||||
assert adapter.session_model_option(session_id) is None
|
||||
|
||||
|
||||
async def test_set_unknown_model_is_rejected() -> None:
|
||||
adapter = _adapter(models=['test'])
|
||||
session_id = await _started(adapter)
|
||||
with pytest.raises(acp.RequestError):
|
||||
await adapter.set_config_option(config_id='model', value='not-a-model', session_id=session_id)
|
||||
|
||||
|
||||
async def test_set_model_config_for_unknown_session_is_rejected() -> None:
|
||||
adapter = _adapter(models=['test'])
|
||||
await _started(adapter)
|
||||
with pytest.raises(acp.RequestError):
|
||||
await adapter.set_config_option(config_id='model', value='test', session_id='no-such-session')
|
||||
|
||||
|
||||
async def test_unknown_config_option_is_rejected() -> None:
|
||||
adapter = _adapter(models=['test'])
|
||||
session_id = await _started(adapter)
|
||||
with pytest.raises(acp.RequestError):
|
||||
await adapter.set_config_option(config_id='theme', value='test', session_id=session_id)
|
||||
|
||||
|
||||
async def test_set_model_config_without_configured_models_is_rejected() -> None:
|
||||
adapter = _adapter()
|
||||
session_id = await _started(adapter)
|
||||
with pytest.raises(acp.RequestError):
|
||||
await adapter.set_config_option(config_id='model', value='test', session_id=session_id)
|
||||
@@ -0,0 +1,31 @@
|
||||
"""The ACP integration ships as the optional `acp` extra: `pip install pydantic-ai-harness[acp]`.
|
||||
|
||||
Needed because a clean install during review could not import `pydantic_ai_harness.experimental.acp` at all:
|
||||
`agent-client-protocol` was missing from the package metadata, which only an isolated install
|
||||
(not the dev environment) reveals. It is declared as an optional extra so base installs stay
|
||||
lean; these checks keep it in the metadata. The full clean-install import
|
||||
(`uv run --isolated --with '.[acp]' ...`) is verified manually until a CI job covers it.
|
||||
|
||||
Only metadata checks belong here: this file stays collected on base (no-extras) installs, so it
|
||||
must not import `pydantic_ai_harness.experimental.acp` or the `acp` SDK.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import importlib.metadata
|
||||
|
||||
|
||||
def _requires_dist() -> list[str]:
|
||||
return importlib.metadata.metadata('pydantic-ai-harness').get_all('Requires-Dist') or []
|
||||
|
||||
|
||||
def test_acp_extra_is_advertised() -> None:
|
||||
provides = importlib.metadata.metadata('pydantic-ai-harness').get_all('Provides-Extra') or []
|
||||
assert 'acp' in provides
|
||||
|
||||
|
||||
def test_agent_client_protocol_is_an_optional_acp_dependency() -> None:
|
||||
acp_requirements = [req for req in _requires_dist() if req.startswith('agent-client-protocol')]
|
||||
assert acp_requirements, 'agent-client-protocol must be declared'
|
||||
# Gated on the `acp` extra (an optional dependency), so base installs stay lean.
|
||||
assert all('extra ==' in req and 'acp' in req for req in acp_requirements)
|
||||
@@ -0,0 +1,298 @@
|
||||
"""Tests for ACP session persistence (`session/load` via a `SessionStore`)."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
|
||||
import acp
|
||||
import pytest
|
||||
from pydantic import TypeAdapter
|
||||
from pydantic_ai import Agent
|
||||
from pydantic_ai.messages import ModelRequest, ModelResponse, TextPart, UserPromptPart
|
||||
from pydantic_ai.models.test import TestModel
|
||||
|
||||
from pydantic_ai_harness.experimental.acp import InMemorySessionStore, PydanticAIACPAgent, StoredSession
|
||||
from tests.experimental.acp._acp_clients import RecordingClient # pyright: ignore[reportMissingTypeStubs]
|
||||
|
||||
pytestmark = pytest.mark.anyio
|
||||
|
||||
|
||||
def _adapter(store: InMemorySessionStore | None) -> PydanticAIACPAgent[None, str]:
|
||||
return PydanticAIACPAgent(Agent(TestModel(custom_output_text='hello')), session_store=store)
|
||||
|
||||
|
||||
def test_stored_session_round_trips_through_pydantic() -> None:
|
||||
# A durable store serializes `StoredSession` with Pydantic; the whole `SessionUpdate` union
|
||||
# (not just the variants a turn happens to produce) must survive a JSON round-trip.
|
||||
original = StoredSession(
|
||||
messages=[
|
||||
ModelRequest(parts=[UserPromptPart(content='hi')]),
|
||||
ModelResponse(parts=[TextPart(content='yo')]),
|
||||
],
|
||||
updates=[
|
||||
acp.update_user_message_text('hi'),
|
||||
acp.update_agent_message_text('yo'),
|
||||
acp.update_agent_thought_text('thinking'),
|
||||
],
|
||||
model='openai:gpt-4o',
|
||||
)
|
||||
adapter = TypeAdapter(StoredSession)
|
||||
assert adapter.validate_json(adapter.dump_json(original)) == original
|
||||
|
||||
|
||||
async def test_initialize_advertises_load_session_only_with_a_store() -> None:
|
||||
with_store = (await _adapter(InMemorySessionStore()).initialize(protocol_version=1)).agent_capabilities
|
||||
without_store = (await _adapter(None).initialize(protocol_version=1)).agent_capabilities
|
||||
assert with_store is not None and with_store.load_session is True
|
||||
assert without_store is not None and without_store.load_session is False
|
||||
|
||||
|
||||
async def test_new_session_persists_an_empty_session() -> None:
|
||||
store = InMemorySessionStore()
|
||||
adapter = _adapter(store)
|
||||
adapter.on_connect(RecordingClient())
|
||||
await adapter.initialize(protocol_version=1)
|
||||
session = await adapter.new_session(cwd='/ws')
|
||||
# The session is stored on creation, so it can be reopened before its first turn.
|
||||
stored = await store.load(session.session_id)
|
||||
assert stored == StoredSession(messages=[], updates=[])
|
||||
|
||||
|
||||
async def test_turn_persists_history_and_transcript() -> None:
|
||||
store = InMemorySessionStore()
|
||||
adapter = _adapter(store)
|
||||
client = RecordingClient()
|
||||
adapter.on_connect(client)
|
||||
await adapter.initialize(protocol_version=1)
|
||||
session = await adapter.new_session(cwd='/ws')
|
||||
await adapter.prompt(prompt=[acp.text_block('hi')], session_id=session.session_id)
|
||||
|
||||
stored = await store.load(session.session_id)
|
||||
assert stored is not None
|
||||
assert len(stored.messages) > 0 # the model exchange was persisted
|
||||
# The transcript is the user's prompt (recorded for replay, never sent live -- the client
|
||||
# renders its own prompt) followed by exactly what the client was shown this turn.
|
||||
assert stored.updates == [acp.update_user_message_text('hi'), *client.updates]
|
||||
|
||||
|
||||
async def test_load_session_restores_history_and_replays_transcript() -> None:
|
||||
store = InMemorySessionStore()
|
||||
adapter = _adapter(store)
|
||||
first = RecordingClient()
|
||||
adapter.on_connect(first)
|
||||
await adapter.initialize(protocol_version=1)
|
||||
session = await adapter.new_session(cwd='/ws')
|
||||
await adapter.prompt(prompt=[acp.text_block('hi')], session_id=session.session_id)
|
||||
shown = list(first.updates)
|
||||
|
||||
# Reopen the session on a fresh connection, as an editor would after a restart.
|
||||
reopened = RecordingClient()
|
||||
adapter.on_connect(reopened)
|
||||
await adapter.load_session(cwd='/ws', session_id=session.session_id)
|
||||
|
||||
# The whole conversation is replayed to the new client: the user's turn (which the live
|
||||
# client rendered itself, so it was never sent as an update) followed by what was shown.
|
||||
assert reopened.updates == [acp.update_user_message_text('hi'), *shown]
|
||||
# ...and the model history is restored so the next turn continues the conversation.
|
||||
stored = await store.load(session.session_id)
|
||||
assert stored is not None
|
||||
assert adapter._sessions[session.session_id].history == stored.messages # pyright: ignore[reportPrivateUsage]
|
||||
|
||||
|
||||
async def test_load_over_an_active_session_cancels_the_in_flight_turn() -> None:
|
||||
started = asyncio.Event()
|
||||
release = asyncio.Event()
|
||||
unwound = asyncio.Event()
|
||||
agent = Agent(TestModel())
|
||||
|
||||
@agent.tool_plain
|
||||
async def slow_tool() -> str:
|
||||
started.set()
|
||||
try:
|
||||
await release.wait() # released only by cancellation
|
||||
return 'done' # pragma: no cover - cancelled by the load
|
||||
finally:
|
||||
unwound.set()
|
||||
|
||||
store = InMemorySessionStore()
|
||||
adapter: PydanticAIACPAgent[None, str] = PydanticAIACPAgent(agent, session_store=store)
|
||||
adapter.on_connect(RecordingClient())
|
||||
await adapter.initialize(protocol_version=1)
|
||||
session = await adapter.new_session(cwd='/ws')
|
||||
turn = asyncio.ensure_future(adapter.prompt(prompt=[acp.text_block('go')], session_id=session.session_id))
|
||||
await asyncio.wait_for(started.wait(), timeout=5)
|
||||
|
||||
# Reopen the still-open session; its in-flight turn must be torn down, not orphaned to later
|
||||
# persist stale state over the restored transcript.
|
||||
await adapter.load_session(cwd='/ws', session_id=session.session_id)
|
||||
|
||||
assert unwound.is_set()
|
||||
response = await asyncio.wait_for(turn, timeout=5)
|
||||
assert response.stop_reason == 'cancelled'
|
||||
# The freshly loaded state replaced the old one and carries no leftover in-flight turn.
|
||||
assert adapter._sessions[session.session_id].active_turn is None # pyright: ignore[reportPrivateUsage]
|
||||
|
||||
|
||||
async def test_queued_prompt_does_not_clobber_a_reloaded_session() -> None:
|
||||
started = asyncio.Event()
|
||||
release = asyncio.Event()
|
||||
agent = Agent(TestModel())
|
||||
|
||||
@agent.tool_plain
|
||||
async def slow_tool() -> str:
|
||||
started.set()
|
||||
await release.wait()
|
||||
return 'done' # pragma: no cover - cancelled by the load
|
||||
|
||||
store = InMemorySessionStore()
|
||||
adapter: PydanticAIACPAgent[None, str] = PydanticAIACPAgent(agent, session_store=store)
|
||||
adapter.on_connect(RecordingClient())
|
||||
await adapter.initialize(protocol_version=1)
|
||||
session = await adapter.new_session(cwd='/ws')
|
||||
snapshot = await store.load(session.session_id)
|
||||
|
||||
first = asyncio.ensure_future(adapter.prompt(prompt=[acp.text_block('one')], session_id=session.session_id))
|
||||
await asyncio.wait_for(started.wait(), timeout=5)
|
||||
# A second prompt queues on the old state's turn lock; the load then replaces that state.
|
||||
queued = asyncio.ensure_future(adapter.prompt(prompt=[acp.text_block('two')], session_id=session.session_id))
|
||||
await asyncio.sleep(0)
|
||||
await adapter.load_session(cwd='/ws', session_id=session.session_id)
|
||||
|
||||
assert (await asyncio.wait_for(first, timeout=5)).stop_reason == 'cancelled'
|
||||
# The queued prompt held the *replaced* state: letting it run would commit and persist that
|
||||
# orphaned history over the session just restored from the store.
|
||||
with pytest.raises(acp.RequestError):
|
||||
await asyncio.wait_for(queued, timeout=5)
|
||||
assert await store.load(session.session_id) == snapshot
|
||||
|
||||
|
||||
async def test_load_session_dropped_from_memory_restores_from_the_store() -> None:
|
||||
store = InMemorySessionStore()
|
||||
adapter = _adapter(store)
|
||||
adapter.on_connect(RecordingClient())
|
||||
await adapter.initialize(protocol_version=1)
|
||||
session = await adapter.new_session(cwd='/ws')
|
||||
await adapter.prompt(prompt=[acp.text_block('hi')], session_id=session.session_id)
|
||||
|
||||
# Drop the live session (as a close, or a process restart, would) so only the stored copy
|
||||
# remains -- exercising the load path where no in-memory session is being replaced.
|
||||
await adapter.close_session(session_id=session.session_id)
|
||||
assert session.session_id not in adapter._sessions # pyright: ignore[reportPrivateUsage]
|
||||
|
||||
adapter.on_connect(RecordingClient())
|
||||
await adapter.load_session(cwd='/ws', session_id=session.session_id)
|
||||
|
||||
assert session.session_id in adapter._sessions # pyright: ignore[reportPrivateUsage]
|
||||
|
||||
|
||||
async def test_load_unknown_session_is_rejected() -> None:
|
||||
adapter = _adapter(InMemorySessionStore())
|
||||
adapter.on_connect(RecordingClient())
|
||||
await adapter.initialize(protocol_version=1)
|
||||
with pytest.raises(acp.RequestError):
|
||||
await adapter.load_session(cwd='/ws', session_id='does-not-exist')
|
||||
|
||||
|
||||
async def test_load_session_without_a_store_is_method_not_found() -> None:
|
||||
adapter = _adapter(None)
|
||||
adapter.on_connect(RecordingClient())
|
||||
await adapter.initialize(protocol_version=1)
|
||||
with pytest.raises(acp.RequestError):
|
||||
await adapter.load_session(cwd='/ws', session_id='whatever')
|
||||
|
||||
|
||||
class _BlockingSaveStore(InMemorySessionStore):
|
||||
"""A store whose next save suspends until released, modeling a durable backend mid-write."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self.saving = asyncio.Event()
|
||||
self.block_next_save = False
|
||||
|
||||
async def save(self, session_id: str, session: StoredSession) -> None:
|
||||
if self.block_next_save:
|
||||
self.block_next_save = False
|
||||
self.saving.set()
|
||||
await asyncio.Event().wait() # suspend until cancelled
|
||||
await super().save(session_id, session)
|
||||
|
||||
|
||||
async def test_cancel_landing_in_the_post_commit_save_commits_but_answers_cancelled() -> None:
|
||||
store = _BlockingSaveStore()
|
||||
adapter: PydanticAIACPAgent[None, str] = PydanticAIACPAgent(
|
||||
Agent(TestModel(custom_output_text='done')), session_store=store
|
||||
)
|
||||
adapter.on_connect(RecordingClient())
|
||||
await adapter.initialize(protocol_version=1)
|
||||
session = await adapter.new_session(cwd='/ws')
|
||||
store.block_next_save = True
|
||||
|
||||
turn = asyncio.ensure_future(
|
||||
adapter.prompt(prompt=[acp.text_block('hi')], session_id=session.session_id, message_id='m1')
|
||||
)
|
||||
await asyncio.wait_for(store.saving.wait(), timeout=5)
|
||||
# The turn is fully committed in memory and is suspended inside the store's save: a cancel
|
||||
# arriving now came too late to roll anything back. The spec still requires the prompt to
|
||||
# answer `cancelled`, but the committed signals must survive: usage is reported, and the
|
||||
# session history keeps the turn.
|
||||
await adapter.cancel(session_id=session.session_id)
|
||||
response = await asyncio.wait_for(turn, timeout=5)
|
||||
|
||||
assert response.stop_reason == 'cancelled'
|
||||
assert response.usage is not None
|
||||
assert len(adapter._sessions[session.session_id].history) >= 2 # pyright: ignore[reportPrivateUsage]
|
||||
|
||||
|
||||
class _FailingStore:
|
||||
"""A SessionStore whose operations always raise, standing in for a broken durable backend."""
|
||||
|
||||
async def save(self, session_id: str, session: StoredSession) -> None:
|
||||
raise OSError('disk full')
|
||||
|
||||
async def load(self, session_id: str) -> StoredSession | None:
|
||||
raise OSError('unreadable')
|
||||
|
||||
|
||||
async def test_save_failure_is_logged_and_does_not_fail_the_turn(caplog: pytest.LogCaptureFixture) -> None:
|
||||
adapter: PydanticAIACPAgent[None, str] = PydanticAIACPAgent(
|
||||
Agent(TestModel(custom_output_text='ok')), session_store=_FailingStore()
|
||||
)
|
||||
adapter.on_connect(RecordingClient())
|
||||
await adapter.initialize(protocol_version=1)
|
||||
with caplog.at_level(logging.ERROR):
|
||||
session = await adapter.new_session(cwd='/ws') # persisting the empty session fails...
|
||||
response = await adapter.prompt(prompt=[acp.text_block('hi')], session_id=session.session_id)
|
||||
|
||||
# ...as does persisting the turn, yet the turn the user already saw stream still succeeds; the
|
||||
# durable-write failure is only logged.
|
||||
assert response.stop_reason == 'end_turn'
|
||||
assert 'failed to persist' in caplog.text
|
||||
|
||||
|
||||
async def test_set_model_config_save_failure_is_logged_and_does_not_fail_the_request(
|
||||
caplog: pytest.LogCaptureFixture,
|
||||
) -> None:
|
||||
adapter: PydanticAIACPAgent[None, str] = PydanticAIACPAgent(
|
||||
Agent(TestModel()), models=['test'], session_store=_FailingStore()
|
||||
)
|
||||
adapter.on_connect(RecordingClient())
|
||||
await adapter.initialize(protocol_version=1)
|
||||
session = await adapter.new_session(cwd='/ws')
|
||||
with caplog.at_level(logging.ERROR):
|
||||
response = await adapter.set_config_option(config_id='model', value='test', session_id=session.session_id)
|
||||
# The selection took effect in memory; only the durable copy is behind, which is logged.
|
||||
assert response is not None
|
||||
assert adapter._sessions[session.session_id].model == 'test' # pyright: ignore[reportPrivateUsage]
|
||||
assert 'failed to persist' in caplog.text
|
||||
|
||||
|
||||
async def test_load_read_failure_is_reported_as_a_clean_error() -> None:
|
||||
adapter: PydanticAIACPAgent[None, str] = PydanticAIACPAgent(Agent(TestModel()), session_store=_FailingStore())
|
||||
adapter.on_connect(RecordingClient())
|
||||
await adapter.initialize(protocol_version=1)
|
||||
with pytest.raises(acp.RequestError) as excinfo:
|
||||
await adapter.load_session(cwd='/ws', session_id='whatever')
|
||||
# A read/deserialize failure surfaces as a purpose-built internal error, not a leaked exception.
|
||||
assert excinfo.value.code == -32603
|
||||
assert 'could not be read' in str(excinfo.value.data)
|
||||
@@ -15,6 +15,18 @@ pydantic-graph = false
|
||||
pydantic-ai = false
|
||||
pydantic-evals = false
|
||||
|
||||
[[package]]
|
||||
name = "agent-client-protocol"
|
||||
version = "0.11.0"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "pydantic" },
|
||||
]
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/6c/93/396d02c91b387b2678f23081869ca47bd495fc6c78789022c683180cf5f1/agent_client_protocol-0.11.0.tar.gz", hash = "sha256:8920a3b27bdcc852fd7c73066f47ab53257dbfbe57b505e20a40c69f80a5391c", size = 87402, upload-time = "2026-07-05T17:07:56.803Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/a0/48/a1367dded7765f5489f9840389949d5aeac53a73aeb1b4e6c0ab61e1617a/agent_client_protocol-0.11.0-py3-none-any.whl", hash = "sha256:2d8570cd4911f8af9fbab4808f7b05f09204a756f346c7409ecdcae3f37cdb2f", size = 68089, upload-time = "2026-07-05T17:07:55.518Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "annotated-doc"
|
||||
version = "0.0.4"
|
||||
@@ -1019,6 +1031,9 @@ dependencies = [
|
||||
]
|
||||
|
||||
[package.optional-dependencies]
|
||||
acp = [
|
||||
{ name = "agent-client-protocol" },
|
||||
]
|
||||
code-mode = [
|
||||
{ name = "pydantic-monty" },
|
||||
]
|
||||
@@ -1057,6 +1072,7 @@ lint = [
|
||||
|
||||
[package.metadata]
|
||||
requires-dist = [
|
||||
{ name = "agent-client-protocol", marker = "extra == 'acp'", specifier = ">=0.11,<0.12" },
|
||||
{ name = "httpx", specifier = ">=0.28.1" },
|
||||
{ name = "logfire", marker = "extra == 'logfire'", specifier = ">=4.31.0" },
|
||||
{ name = "pydantic-ai-slim", specifier = ">=2.1.0" },
|
||||
@@ -1066,7 +1082,7 @@ requires-dist = [
|
||||
{ name = "pydantic-monty", marker = "extra == 'code-mode'", specifier = ">=0.0.16" },
|
||||
{ name = "pydantic-monty", marker = "extra == 'codemode'", specifier = ">=0.0.16" },
|
||||
]
|
||||
provides-extras = ["code-mode", "codemode", "dbos", "logfire", "temporal"]
|
||||
provides-extras = ["acp", "code-mode", "codemode", "dbos", "logfire", "temporal"]
|
||||
|
||||
[package.metadata.requires-dev]
|
||||
dev = [
|
||||
|
||||
Reference in New Issue
Block a user