From 66840c0099a5e9a596b5c350f4c31ffca4eea3af Mon Sep 17 00:00:00 2001 From: mpfaffenberger Date: Mon, 26 Oct 2026 09:20:44 -0700 Subject: [PATCH 1/2] feat(agents): add the Pydantic AI agent Ship a Pydantic AI agent in the agents bundle, beside Claude Code and the others, at @local/agentenv-framework/agents/pydantic-ai (alias `pydantic-ai`). The agent-env protocol SDK serves the A2A endpoint and its extensions; each turn runs a Pydantic AI Agent in process with the Pydantic AI Harness Coder tools (shell, file tools, sub-agents) in a LocalWorkspace, the task's MCP servers as MCPToolsets and installed skills through the harness's Skills. A context's message history is persisted, so later turns continue it, its messages are the live trajectory, and snapshots carry it with the files attached in it. The model is called through the [model] endpoint as an OpenAI-compatible one; a `vllm:` prefix picks Pydantic AI's vLLM provider. Pydantic AI's own Agent.to_a2a is not used: it serves plain A2A, without the agent-env extensions the validator checks. Also: the fake OpenAI endpoint answers streamed requests with server-sent chunks, which Pydantic AI always asks for. --- .github/dependabot.yml | 1 + .github/scripts/infra_build.py | 5 +- CONTRIBUTING.md | 2 +- src/agent_env/agents/README.md | 18 +- .../agents/agents/pydantic-ai/Dockerfile | 13 + .../agents/agents/pydantic-ai/README.md | 70 ++ .../agents/agents/pydantic-ai/a2a_server.py | 24 + .../agents/agents/pydantic-ai/install.sh | 32 + .../agents/pydantic-ai/pydantic_ai_agent.py | 540 +++++++++++++++ .../agents/pydantic-ai/pydantic_ai_runtime.py | 166 +++++ .../agents/pydantic-ai/requirements.txt | 6 + .../agents/agents/pydantic-ai/start.sh | 11 + src/agent_env/agents/aliases.toml | 1 + .../agents/tasks/pydantic-ai-smoke.json | 7 + .../cli/run_bundle_pydantic_ai_test.py | 245 +++++++ tst/unit/agents/pydantic_ai/__init__.py | 0 tst/unit/agents/pydantic_ai/agent.py | 14 + .../pydantic_ai/test_pydantic_ai_agent.py | 621 ++++++++++++++++++ tst/unit/agents/test_installed_workspace.py | 3 +- tst/unit/bundle/shipped_bundles_test.py | 3 +- tst/unit/test_infra_images_script.py | 2 +- tst/util/fake_openai.py | 37 +- 22 files changed, 1809 insertions(+), 12 deletions(-) create mode 100644 src/agent_env/agents/agents/pydantic-ai/Dockerfile create mode 100644 src/agent_env/agents/agents/pydantic-ai/README.md create mode 100644 src/agent_env/agents/agents/pydantic-ai/a2a_server.py create mode 100644 src/agent_env/agents/agents/pydantic-ai/install.sh create mode 100644 src/agent_env/agents/agents/pydantic-ai/pydantic_ai_agent.py create mode 100644 src/agent_env/agents/agents/pydantic-ai/pydantic_ai_runtime.py create mode 100644 src/agent_env/agents/agents/pydantic-ai/requirements.txt create mode 100644 src/agent_env/agents/agents/pydantic-ai/start.sh create mode 100644 src/agent_env/agents/tasks/pydantic-ai-smoke.json create mode 100644 tst/integration/cli/run_bundle_pydantic_ai_test.py create mode 100644 tst/unit/agents/pydantic_ai/__init__.py create mode 100644 tst/unit/agents/pydantic_ai/agent.py create mode 100644 tst/unit/agents/pydantic_ai/test_pydantic_ai_agent.py diff --git a/.github/dependabot.yml b/.github/dependabot.yml index cf07a1f7..557f8554 100644 --- a/.github/dependabot.yml +++ b/.github/dependabot.yml @@ -47,6 +47,7 @@ updates: - "/src/agent_env/agents/agents/openhands" - "/src/agent_env/agents/agents/pi" - "/src/agent_env/agents/agents/proxy" + - "/src/agent_env/agents/agents/pydantic-ai" schedule: interval: "weekly" open-pull-requests-limit: 5 diff --git a/.github/scripts/infra_build.py b/.github/scripts/infra_build.py index 7906d238..92d54bc4 100644 --- a/.github/scripts/infra_build.py +++ b/.github/scripts/infra_build.py @@ -38,6 +38,7 @@ "agentenv-openhands": "src/agent_env/agents/agents/openhands/Dockerfile", "agentenv-pi": "src/agent_env/agents/agents/pi/Dockerfile", "agentenv-proxy": "src/agent_env/agents/agents/proxy/Dockerfile", + "agentenv-pydantic-ai": "src/agent_env/agents/agents/pydantic-ai/Dockerfile", } _TAG = re.compile(r"^src-[0-9a-f]{64}$") _DIGEST = re.compile(r"^sha256:[0-9a-f]{64}$") @@ -56,7 +57,8 @@ # Each image's check, beyond its build: the gateway imports its server, Postgres turns healthy, the sidecars and the # browser answer their own commands, Claude Code's, Codex's, Gemini's, OpenCode's and pi's agents run their CLIs and -# import their servers, OpenHands' imports its server, which loads the OpenHands SDK, and the proxy imports its server. +# import their servers, OpenHands' imports its server, which loads the OpenHands SDK, Pydantic AI's imports its +# server, which loads Pydantic AI, and the proxy imports its server. _SMOKES = { "agentenv-gateway": ["docker", "run", "--rm", "--entrypoint", "python", "{image}", "-c", "import gateway.entrypoint"], "agentenv-db-web": ["docker", "run", "--rm", "--entrypoint", "/usr/bin/pgweb", "{image}", "--version"], @@ -76,6 +78,7 @@ "agentenv-pi": ["docker", "run", "--rm", "--entrypoint", "sh", "{image}", "-c", 'pi --version && python3 -c "import a2a_server"'], "agentenv-proxy": ["docker", "run", "--rm", "--entrypoint", "sh", "{image}", "-c", 'python -c "import a2a_server"'], + "agentenv-pydantic-ai": ["docker", "run", "--rm", "--entrypoint", "sh", "{image}", "-c", 'python -c "import a2a_server"'], } diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 4593fbb9..b04b9636 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -90,7 +90,7 @@ the previous release, so a clear pull request title is also the release note. A A release also publishes the infra images agent-env ships, once the `PUBLISH_INFRA_IMAGES` repository variable is `true`: one per Dockerfile under `src/agent_env/env`, and one per shipped agent (`agentenv-claude-code`, its CLI -included, `agentenv-codex`, `agentenv-gemini`, `agentenv-opencode`, `agentenv-pi` and `agentenv-openhands`), at `ghcr.io/scaleapi/agentenv-:src-` for linux/amd64 +included, `agentenv-codex`, `agentenv-gemini`, `agentenv-opencode`, `agentenv-pi`, `agentenv-openhands` and `agentenv-pydantic-ai`), at `ghcr.io/scaleapi/agentenv-:src-` for linux/amd64 and linux/arm64, where the tag is the build context's source digest, so an unchanged Dockerfile is not built again. The agents' images install the newest `agentenv-framework-protocol` at or above the floor in their `requirements.txt`, which a release doesn't pin, so each is built again only when its folder changes: once the protocol diff --git a/src/agent_env/agents/README.md b/src/agent_env/agents/README.md index d9a8cd42..5d99cdd3 100644 --- a/src/agent_env/agents/README.md +++ b/src/agent_env/agents/README.md @@ -69,10 +69,20 @@ calls the `[model]` endpoint through LiteLLM and takes its model from the prompt `[model] default` (or the `agent` role) set. It declares `install/v1` too, for Debian or Ubuntu targets. Its folder's `README.md` covers the extensions it implements. -A release publishes Codex's, Gemini's, OpenCode's, pi's and OpenHands' images too, at +`agents/pydantic-ai/` runs a [Pydantic AI](https://ai.pydantic.dev) agent with the +[Pydantic AI Harness](https://pydantic.dev/docs/ai/harness/) `Coder` tools (shell, file tools, sub-agents) behind the +protocol's A2A agent framework, in process rather than a CLI: each turn is one run of the agent continuing the +context's persisted message history, the task env's MCP servers are its toolsets, and its messages are the turn's +trajectory. Name it as `pydantic-ai`, its alias, or by its id, `@local/agentenv-framework/agents/pydantic-ai`. It calls +the `[model]` endpoint as an OpenAI-compatible one, `/chat/completions`, and takes its model from the prompt, so +its smoke task needs `[model] default` (or the `agent` role) set; a `vllm:` prefix, as in `vllm:qwen3.8-27b`, fits its +requests to the model's family as a vLLM server needs. It declares `install/v1` too, for Debian or Ubuntu targets. Its +folder's `README.md` covers the extensions it implements. + +A release publishes Codex's, Gemini's, OpenCode's, pi's, OpenHands' and Pydantic AI's images too, at `ghcr.io/scaleapi/agentenv-codex:src-`, `ghcr.io/scaleapi/agentenv-gemini:src-`, -`ghcr.io/scaleapi/agentenv-opencode:src-`, `ghcr.io/scaleapi/agentenv-pi:src-` and -`ghcr.io/scaleapi/agentenv-openhands:src-`, each built again only when its folder changes, and a run deploys +`ghcr.io/scaleapi/agentenv-opencode:src-`, `ghcr.io/scaleapi/agentenv-pi:src-`, +`ghcr.io/scaleapi/agentenv-openhands:src-` and `ghcr.io/scaleapi/agentenv-pydantic-ai:src-`, each built again only when its folder changes, and a run deploys them as it does Claude Code's: it pulls the image published for the sources installed, by its digest, and builds it from the agent's folder only when it can't. @@ -88,4 +98,4 @@ others'. It has no alias and no smoke task: with nobody driving it, a prompt onl `agent-env run agents` runs every agent's smoke task, pulling or building each image; `--task claude-code-smoke` runs Claude Code's alone, `--task codex-smoke` Codex's, `--task gemini-smoke` Gemini's, `--task opencode-smoke` OpenCode's, -`--task pi-smoke` pi's, `--task openhands-smoke` OpenHands'. +`--task pi-smoke` pi's, `--task openhands-smoke` OpenHands', `--task pydantic-ai-smoke` Pydantic AI's. diff --git a/src/agent_env/agents/agents/pydantic-ai/Dockerfile b/src/agent_env/agents/agents/pydantic-ai/Dockerfile new file mode 100644 index 00000000..05c795a6 --- /dev/null +++ b/src/agent_env/agents/agents/pydantic-ai/Dockerfile @@ -0,0 +1,13 @@ +# docker build -t pydantic-ai-a2a src/agent_env/agents/agents/pydantic-ai +FROM mirror.gcr.io/library/ubuntu:24.04@sha256:534baea6a22c03a63003dbc8dbe78fe34bc0d7e595d9a9dc9834884ff530eb55 +WORKDIR /app +# Named sources: agent-env packs the build context install/v1 copies from the COPY sources. +COPY a2a_server.py install.sh pydantic_ai_agent.py pydantic_ai_runtime.py requirements.txt start.sh ./ +RUN sh install.sh \ + && useradd --create-home agent \ + && mkdir /workspace \ + && chown agent /app /workspace +ENV PATH="/opt/pydantic-ai-a2a-venv/bin:${PATH}" +USER agent +EXPOSE 8000 +CMD ["python", "a2a_server.py"] diff --git a/src/agent_env/agents/agents/pydantic-ai/README.md b/src/agent_env/agents/agents/pydantic-ai/README.md new file mode 100644 index 00000000..82c43887 --- /dev/null +++ b/src/agent_env/agents/agents/pydantic-ai/README.md @@ -0,0 +1,70 @@ +# Pydantic AI A2A agent + +A [Pydantic AI](https://ai.pydantic.dev) agent wrapped with the `agentenv_protocol.a2a_agent` framework, shipped in +the `agents` bundle as `@local/agentenv-framework/agents/pydantic-ai`. This directory is the whole build context: the +protocol package, Pydantic AI and the [Pydantic AI Harness](https://pydantic.dev/docs/ai/harness/) come from PyPI +(`requirements.txt`), and `install.sh` installs them in a Python of uv's own, with git for the agent's shell, both for +the image and for `install/v1`. + +Instead of wrapping a CLI, the agent runs Pydantic AI in process. `pydantic_ai_agent.py` is the A2A side and imports +no Pydantic AI; `pydantic_ai_runtime.py` runs the agent. Each task is one run of an `Agent` with the harness's `Coder` +capability, the coding tools CLAI is built on (`shell`, `read_file`, `write_file`, `edit_file`, `list_files`, `grep`, +`delegate_task`), working in a `LocalWorkspace` on the workspace. Its shell commands get only `PATH`, `HOME` and the +locale from the agent's environment, never the model key, and the harness keeps their logs under +`.pydantic-ai-harness/` in the workspace, with a `.gitignore` of its own. The agent runs in a sandbox of its own, so the file tools +are not confined to the workspace and reach the files attached outside it. Pydantic AI's own A2A server +(`Agent.to_a2a`) is not used: it speaks plain A2A, while agent-env needs the protocol's extensions, which the framework +serves. + +The model is called through the `[model]` endpoint as an OpenAI-compatible one, `/chat/completions`, with +`LITELLM_EXTRA_HEADERS` sent on every call but kept out of the persisted messages, which a snapshot uploads. The model +name is passed through as is; a `vllm:` prefix, as in `vllm:qwen3.8-27b`, picks Pydantic AI's vLLM provider, which +fits the requests to the model's family (Qwen, DeepSeek, GLM, Llama, ...) as a vLLM server needs, and `openai:` names +the default. + +A context maps to one conversation, whose id is derived from the context id, its message history persisted in +`~/.pydantic-ai-a2a/conversations//messages.json` once a turn succeeds, so follow-up tasks in the context continue +it, across restarts too. Images are sent to the model as images; every file part is also saved in the conversation's +directory under `inputs//` and its path added to the prompt, for the agent's tools, and the conversation's +snapshots carry it. + +| Extension | Implementation | +|---|---| +| `agent-config/v1` | `PydanticAIConfig`: `model`, `effort` (Pydantic AI's `thinking` setting), `system_prompt` (added to the agent's instructions), `request_limit` (model requests per turn, default 500), `timeout_seconds` | +| `mcp-config/v1` | SDK default; each server becomes an `MCPToolset` on its URL, with its headers, its tools named as the server names them | +| `skill-config/v1` | inline and bundle skills installed as `SKILL.md` directories under `~/.pydantic-ai-a2a/skills`, offered to the model through the harness's `Skills`, which it loads with `load_capability` | +| `trajectory/v1` | live: the task's Pydantic AI messages (`ModelMessagesTypeAdapter`, format `pydantic-ai-messages/v1`), appended as each request and response completes | +| `snapshot/v1` | `save` uploads the context's conversation directory, its messages and attached files, as a gzipped tar and, when asked, the workspace as another; `load` extracts both aside, checks the messages and points the paths of the attached files at the target's directory before it replaces anything, then replaces the workspace's contents and installs the conversation under the target context's id, so an archive that fails to extract, or holds no conversation, changes nothing | +| `peer-agents/v1` | SDK default: the `peer_agents` MCP server reaches the agent like any other | +| `install/v1` | copies this directory into the task container as `/opt/pydantic-ai-a2a`, runs `install.sh` (Debian or Ubuntu, as root, with network access), writes the LiteLLM credentials and the lines that trust the local transfer CA to a 0600 `agent.env` through a heredoc and starts `start.sh`, which applies them | +| `triggers/v1` | SDK default | + +The snapshot changelog (`enable_changelog`/`apply_changelog`) is not implemented, nor is `attribution-probe/v1`. + +`PYDANTIC_AI_A2A_MODEL` sets the model when agent-config sets none, as for a peer that other agents message directly, +and `PYDANTIC_AI_A2A_WORKSPACE` the workspace (default `/workspace`, and the task container's working directory when +it's installed there). `a2a_server.py` reads these and the `LITELLM_*` variables once, at startup: a +`LITELLM_EXTRA_HEADERS` that is not a JSON object stops the agent there, while a missing `LITELLM_BASE_URL` fails each +task as `pydantic_ai.no_endpoint`. + +Snapshot `load` answers 409 while a task is running, and 400 for an archive that holds no conversation. A snapshot +`save` waits for its context's running task to end; one that saves the workspace also waits for the tasks running in +other contexts, which work in the same workspace. + +Usage is the run's: input, output and cached tokens, model requests and tool calls; no cost. A failed run's error code +says whose failure it was: `pydantic_ai.model_http_` and `pydantic_ai.model_api` for the model endpoint and +`pydantic_ai.error` for anything else that stopped the run, such as an MCP server that can't be reached, are +`infra_error`s; `pydantic_ai.usage_limit` (the request limit) and `pydantic_ai.unexpected_model_behavior` are +`agent_error`s, and `pydantic_ai.history`, a persisted conversation that can't be read, an `infra_error`. A failed +turn leaves the conversation as it was. A timeout (`pydantic_ai.timeout`) counts reading the +task's file parts too and reports what the run had spent. A timeout, a `tasks/cancel` and the server shutting down each +cancel the run, which kills the commands its shell started. Errors are answered as `{"detail": ...}`. + +```bash +docker build -t pydantic-ai-a2a src/agent_env/agents/agents/pydantic-ai +docker run -p 8000:8000 -e LITELLM_BASE_URL=... -e LITELLM_API_KEY=... pydantic-ai-a2a +pytest tst/unit/agents/pydantic_ai +``` + +Pydantic AI's version is pinned in `requirements.txt` and `PYDANTIC_AI_VERSION`, and the harness's beside it; a new +release is a pull request that bumps them. diff --git a/src/agent_env/agents/agents/pydantic-ai/a2a_server.py b/src/agent_env/agents/agents/pydantic-ai/a2a_server.py new file mode 100644 index 00000000..4b1754cc --- /dev/null +++ b/src/agent_env/agents/agents/pydantic-ai/a2a_server.py @@ -0,0 +1,24 @@ +"""Serves the Pydantic AI agent, configured from its environment. ``AGENT_CARD`` is what agent-env's install/v1 reads +from ``/app/a2a_server.py``.""" + +import os +from pathlib import Path +from typing import Final + +from a2a.types import AgentCard +from starlette.applications import Starlette + +from pydantic_ai_agent import DEFAULT_MODEL_ENV, Endpoint, PydanticAIAgent +from pydantic_ai_runtime import PydanticAIRuntime + +agent: Final[PydanticAIAgent] = PydanticAIAgent( + PydanticAIRuntime(), + endpoint=Endpoint.from_environ(os.environ), + default_model=os.environ.get(DEFAULT_MODEL_ENV), + workspace=Path(os.environ.get("PYDANTIC_AI_A2A_WORKSPACE", "/workspace")), +) +app: Final[Starlette] = agent.create_app() +AGENT_CARD: Final[AgentCard] = app.state.agentenv_a2a.card + +if __name__ == "__main__": + agent.serve() diff --git a/src/agent_env/agents/agents/pydantic-ai/install.sh b/src/agent_env/agents/agents/pydantic-ai/install.sh new file mode 100644 index 00000000..ad9e9a75 --- /dev/null +++ b/src/agent_env/agents/agents/pydantic-ai/install.sh @@ -0,0 +1,32 @@ +#!/bin/sh +# Installs git for the agent's shell and this agent's Python environment, in uv's own Python so the target's is never +# used, on Debian or Ubuntu, as root. The image build and install/v1 both run it, so an install into a task container +# matches the image. +set -eu +UV_VERSION=0.13.1 +UV_SHA256_X86_64=6c86040a8f045ab9adc9e3ec4cdafa03fd00e9a988ed2bf28971c7fc2d0debd0 +UV_SHA256_AARCH64=345a8308e4d1835cdfaadf692836fd3137a6d45864773f05ad7bf59442b1ba16 +PYTHON_VERSION=3.12 +here=$(cd "$(dirname "$0")" && pwd) + +export DEBIAN_FRONTEND=noninteractive +apt-get update +apt-get install -y --no-install-recommends ca-certificates curl git +rm -rf /var/lib/apt/lists/* + +case "$(uname -m)" in + x86_64) arch=x86_64 sha256=$UV_SHA256_X86_64 ;; + aarch64 | arm64) arch=aarch64 sha256=$UV_SHA256_AARCH64 ;; + *) echo "unsupported architecture: $(uname -m)" >&2; exit 1 ;; +esac +stage=$(mktemp -d) +trap 'rm -rf "$stage"' EXIT +archive=uv-$arch-unknown-linux-musl.tar.gz +curl -fsSL --retry 3 "https://github.com/astral-sh/uv/releases/download/$UV_VERSION/$archive" -o "$stage/$archive" +echo "$sha256 $stage/$archive" | sha256sum -c - +tar -xzf "$stage/$archive" -C "$stage" +uv="$stage/uv-$arch-unknown-linux-musl/uv" + +export UV_PYTHON_INSTALL_DIR=/opt/pydantic-ai-a2a-python UV_PYTHON_PREFERENCE=only-managed UV_CACHE_DIR="$stage/cache" +"$uv" venv --quiet --clear --python "$PYTHON_VERSION" /opt/pydantic-ai-a2a-venv +"$uv" pip install --quiet --python /opt/pydantic-ai-a2a-venv/bin/python -r "$here/requirements.txt" diff --git a/src/agent_env/agents/agents/pydantic-ai/pydantic_ai_agent.py b/src/agent_env/agents/agents/pydantic-ai/pydantic_ai_agent.py new file mode 100644 index 00000000..446de1d2 --- /dev/null +++ b/src/agent_env/agents/agents/pydantic-ai/pydantic_ai_agent.py @@ -0,0 +1,540 @@ +"""Pydantic AI (https://ai.pydantic.dev) as an AgentEnv A2A agent. + +Each task is one run of a Pydantic AI agent with the Pydantic AI Harness coding tools in the workspace. A context maps +to one conversation, its message history persisted on disk, so later tasks in the context continue it. This module is +the A2A side and imports no Pydantic AI; ``pydantic_ai_runtime`` runs the agent. +""" + +import asyncio +import base64 +import io +import json +import mimetypes +import shutil +import tarfile +import tempfile +import uuid +from collections.abc import AsyncIterator, Callable, Mapping +from contextlib import asynccontextmanager, suppress +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any, Final, Literal, Protocol, Self +from urllib.parse import urlsplit + +import httpx +from agentenv_protocol import ROLE_HEADER +from agentenv_protocol.a2a_agent import ( + INSTALL_V1, + MCP_CONFIG_V1, + PEER_AGENTS_V1, + SKILL_CONFIG_V1, + SNAPSHOT_V1, + TRAJECTORY_V1, + TRIGGERS_V1, + AgentConfig, + AgentEnvAgent, + AgentIdentity, + BundleSkillRequest, + DataPart, + FilePart, + InlineSkillRequest, + ObjectSnapshotLoadRequest, + ObjectSnapshotSaveRequest, + TaskRequest, + TaskResult, + TextPart, + Usage, + a2a_agent, + download, + enable, + extension, + upload, +) +from starlette.applications import Starlette +from starlette.exceptions import HTTPException +from starlette.requests import Request +from starlette.responses import JSONResponse, Response + +PYDANTIC_AI_VERSION: Final[str] = "2.61.0" +TRAJECTORY_FORMAT: Final[str] = "pydantic-ai-messages/v1" +# The model when agent-config sets none, as for a peer that other agents message directly. +DEFAULT_MODEL_ENV: Final[str] = "PYDANTIC_AI_A2A_MODEL" +A2A_PORT: Final[int] = 8000 +INSTALL_DIR: Final[str] = "/opt/pydantic-ai-a2a" +ENV_FILE: Final[str] = "agent.env" +# A conversation's directory holds its message history and, under INPUTS, the files attached in it. +MESSAGES: Final[str] = "messages.json" +INPUTS: Final[str] = "inputs" +# Written into a conversation's archive: the directory it was saved from, under which its messages name attached files. +SOURCE: Final[str] = ".pydantic-ai-a2a-source.json" +FETCH_TIMEOUT_SECONDS: Final[int] = 120 +MAX_ATTACHMENT_BYTES: Final[int] = 256 * 1024 * 1024 +# A conversation id is derived from its context id, so a context finds its conversation across restarts. +_CONVERSATIONS: Final[uuid.UUID] = uuid.uuid5(uuid.NAMESPACE_URL, "urn:agentenv:agents:pydantic-ai") +# Images go to the model as images as well as to disk; any other file is left on disk for the agent's tools. +IMAGE_TYPES: Final[tuple[str, ...]] = ("image/png", "image/jpeg", "image/gif", "image/webp") +_NO_BODY_STATUSES: Final[frozenset[int]] = frozenset({204, 205, 304}) + +# sudo only where it needs no password: VM sandboxes have it, a local sandbox runs docker as the user. +_DOCKER: Final[str] = "$(sudo -n true 2>/dev/null && echo sudo) docker" +INSTALL_COMMANDS: Final[list[str]] = [ + f"{_DOCKER} cp {{agent_ctx_tar}} {{container}}:/tmp/pydantic-ai-a2a-ctx.tar.gz", + f"{_DOCKER} exec {{container}} sh -c 'mkdir -p {INSTALL_DIR} && tar -xf /tmp/pydantic-ai-a2a-ctx.tar.gz" + f" -C {INSTALL_DIR} && rm /tmp/pydantic-ai-a2a-ctx.tar.gz'", + f"{_DOCKER} exec {{container}} sh {INSTALL_DIR}/install.sh", + # The credentials travel on a heredoc into a 0600 file: the installer logs only a command's first line, + # and neither a command line nor the container's config ever holds them. + f"{_DOCKER} exec -i {{container}} sh -c 'umask 077 && cat > {INSTALL_DIR}/{ENV_FILE}' <<'PYDANTIC_AI_A2A_ENV'\n" + "LITELLM_API_KEY={litellm_api_key}\n" + "LITELLM_BASE_URL={litellm_base_url}\n" + "LITELLM_EXTRA_HEADERS={litellm_extra_headers}\n" + # The KEY='VALUE' lines pointing TLS clients at the local transfer CA, which start.sh applies. + "AGENT_TRUST_ENV={trust_env}\n" + "PYDANTIC_AI_A2A_ENV", + f"{_DOCKER} exec -d -e A2A_PORT={{a2a_port}} {{container}} sh {INSTALL_DIR}/start.sh", +] + + +class PydanticAIConfig(AgentConfig): + model: str | None = None + system_prompt: str | None = None + effort: Literal["minimal", "low", "medium", "high", "xhigh", "max"] | None = None + request_limit: int = 500 + timeout_seconds: int = 1800 + + +@dataclass(frozen=True, slots=True) +class Endpoint: + """The OpenAI-compatible endpoint the model is called through, with the headers every call carries.""" + + base_url: str + api_key: str | None = None + extra_headers: Mapping[str, str] = field(default_factory=dict) + + @classmethod + def from_environ(cls, environ: Mapping[str, str]) -> Self | None: + """The endpoint the ``LITELLM_*`` variables name; None without a base URL, which fails each task instead.""" + base_url = environ.get("LITELLM_BASE_URL") + if not base_url: + return None + try: + headers = json.loads(environ.get("LITELLM_EXTRA_HEADERS") or "{}") + except json.JSONDecodeError as exc: + raise ValueError(f"LITELLM_EXTRA_HEADERS is not valid JSON: {exc}") from exc + if not isinstance(headers, dict): + raise ValueError("LITELLM_EXTRA_HEADERS must be a JSON object of header names to values") + return cls(base_url, environ.get("LITELLM_API_KEY"), headers) + + +@dataclass(frozen=True, slots=True) +class Image: + path: Path + media_type: str + + +@dataclass(frozen=True, slots=True) +class Turn: + """One task, as the runtime runs it.""" + + conversation_dir: Path + workspace: Path + text: str + images: tuple[Image, ...] + mcp_servers: Mapping[str, Any] + skills_dir: Path + skill_names: tuple[str, ...] + config: PydanticAIConfig + endpoint: Endpoint + on_event: Callable[[dict[str, Any]], None] + # What a run that was cancelled had spent by then. + on_usage: Callable[[Usage], None] + + +@dataclass(frozen=True, slots=True) +class Outcome: + text: str + usage: Usage + error_code: str | None = None + error_type: Literal["agent_error", "infra_error"] = "agent_error" + + +@dataclass(slots=True) +class _ConversationLock: + lock: asyncio.Lock = field(default_factory=asyncio.Lock) + holders: int = 0 + + +class Runtime(Protocol): + async def run(self, turn: Turn) -> Outcome: + """Run ``turn`` to its end, saving the conversation's messages in ``turn.conversation_dir`` once it succeeds. + Cancelled, it stops the run and the commands it started, and re-raises.""" + ... + + async def stop(self) -> None: + """Cancel every turn still running: the server is shutting down.""" + ... + + +async def _detail_body(_request: Request, exc: HTTPException) -> Response: + """An ``HTTPException`` as the ``{"detail": ...}`` body agent-env reads errors from.""" + if exc.status_code < 200 or exc.status_code in _NO_BODY_STATUSES: + return Response(status_code=exc.status_code, headers=exc.headers) + return JSONResponse({"detail": exc.detail}, status_code=exc.status_code, headers=exc.headers) + + +@asynccontextmanager +async def _turns_stop_with_the_server(app: Starlette) -> AsyncIterator[None]: + """The server's lifespan: as it shuts down, turns still running are cancelled, which kills the commands they + started, so none outlives the server.""" + try: + yield + finally: + await app.state.agentenv_a2a.agent.runtime.stop() + + +def conversation_id(context_id: str) -> str: + return uuid.uuid5(_CONVERSATIONS, context_id).hex + + +def _with_role(mcp_servers: Mapping[str, Any], role: str | None) -> dict[str, Any]: + """The servers with the role header, which the env gateway filters tools by and which wins over a stored role.""" + if not role: + return dict(mcp_servers) + return { + name: {**registration, "headers": {**(registration.get("headers") or {}), ROLE_HEADER: role}} + for name, registration in mcp_servers.items() + } + + +def _media_type(part: FilePart) -> str: + return part.mime_type or mimetypes.guess_type(part.name or "")[0] or "application/octet-stream" + + +async def _materialize(part: FilePart, index: int, directory: Path) -> Path: + """Write a file part to disk; inline bytes and HTTP(S) URLs only.""" + name = Path(part.name).name if part.name else "" + destination = directory / f"{index}-{name or 'file'}" + if part.bytes is not None: + destination.write_bytes(base64.b64decode(part.bytes)) + return destination + if urlsplit(part.uri).scheme not in {"http", "https"}: + raise ValueError(f"unsupported file URI scheme for {name or 'file'}") + received = 0 + async with ( + httpx.AsyncClient(timeout=FETCH_TIMEOUT_SECONDS, follow_redirects=True) as client, + client.stream("GET", part.uri) as response, + ): + response.raise_for_status() + with destination.open("wb") as handle: + async for chunk in response.aiter_bytes(): + received += len(chunk) + if received > MAX_ATTACHMENT_BYTES: + raise ValueError(f"{name or 'file'} is larger than {MAX_ATTACHMENT_BYTES} bytes") + handle.write(chunk) + return destination + + +def _pack(directory: Path, archive: Path, *, record_source: bool = False) -> None: + """``directory``'s contents as a gzipped tar and, for a conversation, the directory it was saved from.""" + with tarfile.open(archive, "w:gz") as tar: + for child in sorted(directory.iterdir()): + tar.add(child, arcname=child.name) + if record_source: + source = json.dumps({"directory": str(directory)}).encode() + info = tarfile.TarInfo(SOURCE) + info.size = len(source) + tar.addfile(info, io.BytesIO(source)) + + +def _extract(archive: Path, directory: Path) -> Path: + """The archive extracted into a new staging directory inside ``directory``.""" + staging = directory / f".pydantic-ai-a2a-restore-{uuid.uuid4().hex}" + # An empty workspace's archive has no members, and extracting it creates nothing. + staging.mkdir() + try: + with tarfile.open(archive) as tar: + tar.extractall(staging, filter="data") + except BaseException: + shutil.rmtree(staging, ignore_errors=True) + raise + return staging + + +def _replace_contents(directory: Path, staging: Path) -> None: + """Replace ``directory``'s contents with those of ``staging``, a directory inside it.""" + for child in directory.iterdir(): + if child == staging: + continue + if child.is_dir() and not child.is_symlink(): + shutil.rmtree(child) + else: + child.unlink() + for child in staging.iterdir(): + child.rename(directory / child.name) + staging.rmdir() + + +def _stage_conversation(archive: Path, conversations_dir: Path, target: Path) -> Path: + """The archived conversation extracted beside the others, its messages naming the files attached in it under + ``target``, where it goes, rather than where it was saved from. Raises ValueError for an archive that holds no + conversation.""" + staging = _extract(archive, conversations_dir) + try: + if not isinstance(json.loads((staging / MESSAGES).read_text()), list): + raise ValueError("its messages are not a list") + source = staging / SOURCE + saved_from = json.loads(source.read_text())["directory"] if source.is_file() else None + except (OSError, ValueError, KeyError, TypeError) as exc: + shutil.rmtree(staging, ignore_errors=True) + raise ValueError(f"the trajectory object holds no Pydantic AI conversation: {type(exc).__name__}: {exc}") from exc + if saved_from is not None: + source.unlink() + # A path is written as is in a message's text and JSON-escaped in the file. + before, after = json.dumps(f"{saved_from}/")[1:-1], json.dumps(f"{target}/")[1:-1] + messages = staging / MESSAGES + messages.write_text(messages.read_text().replace(before, after)) + return staging + + +@a2a_agent( + identity=AgentIdentity( + name="Pydantic AI", + description=( + "A Pydantic AI agent with the Pydantic AI Harness coding tools: reads, edits and runs code in its " + "workspace with MCP tools and skills." + ), + version=PYDANTIC_AI_VERSION, + input_modes=("text", *IMAGE_TYPES), + ), + config=PydanticAIConfig, + config_description="Configure the Pydantic AI agent: model, effort, system prompt and request limit.", + lifespan=_turns_stop_with_the_server, + extensions=( + MCP_CONFIG_V1, + enable(TRAJECTORY_V1, live=True, description="The task's Pydantic AI messages, also while it runs."), + enable( + SNAPSHOT_V1, + description=( + "The trajectory object is the context's persisted Pydantic AI conversation, the workspace object a " + "gzipped tar." + ), + ), + PEER_AGENTS_V1, + enable( + INSTALL_V1, + install_commands=INSTALL_COMMANDS, + required_params=[ + "container", "agent_ctx_tar", "a2a_port", "litellm_api_key", "litellm_base_url", + "litellm_extra_headers", "trust_env", + ], + a2a_port=A2A_PORT, + ), + TRIGGERS_V1, + ), +) +class PydanticAIAgent(AgentEnvAgent): + def __init__( + self, + runtime: Runtime, + *, + endpoint: Endpoint | None, + default_model: str | None = None, + home: Path | None = None, + workspace: Path | None = None, + ) -> None: + self.runtime: Runtime = runtime + self.endpoint: Endpoint | None = endpoint + self.default_model: str | None = default_model + self.home: Path = home or Path.home() / ".pydantic-ai-a2a" + self.workspace: Path = workspace or Path("/workspace") + self.conversations_dir: Path = self.home / "conversations" + self.skills_dir: Path = self.home / "skills" + # The framework already runs a context's tasks one at a time; this lets a snapshot save wait for the task + # running in its context. An entry lives only while it is held, so contexts never seen again leave nothing. + self._conversation_locks: dict[str, _ConversationLock] = {} + # A task starts once no restore or workspace save holds this. A restore is refused while tasks run, and a + # workspace save waits for the tasks running in other contexts to end. + self._workspace_lock: asyncio.Lock = asyncio.Lock() + self._running: int = 0 + self._idle: asyncio.Event = asyncio.Event() + self._idle.set() + + def create_app(self) -> Starlette: + app = super().create_app() + app.add_exception_handler(HTTPException, _detail_body) + return app + + async def run(self, request: TaskRequest[PydanticAIConfig]) -> TaskResult: + config = request.config + if not config.model and self.default_model: + config = config.model_copy(update={"model": self.default_model}) + if not config.model: + return TaskResult.failure("pydantic_ai.no_model", "no model is configured; set 'model' through agent-config") + if self.endpoint is None: + return TaskResult.failure("pydantic_ai.no_endpoint", "LITELLM_BASE_URL is not set", error_type="infra_error") + workspace = request.workspace or self.workspace + workspace.mkdir(parents=True, exist_ok=True) + request.trajectory.set_format(TRAJECTORY_FORMAT) + + def on_event(event: dict[str, Any]) -> None: + # A cancelled task's log ends before its run has stopped, and refuses what the run records after that. + with suppress(RuntimeError): + request.trajectory.append(event) + + conversation = conversation_id(request.context_id) + # Bound before the run starts, so a live trajectory read or a mid-turn snapshot finds the conversation. + self.set_session_ref_for_context(request.context_id, conversation) + spent: list[Usage] = [] + async with self._holding(conversation), self._in_workspace(): + try: + # Reading the file parts is part of the task, and so counts toward its time. + async with asyncio.timeout(config.timeout_seconds): + try: + text, images = await self._prompt(request, conversation) + except (ValueError, OSError, httpx.HTTPError, httpx.InvalidURL) as exc: + return TaskResult.failure("pydantic_ai.attachment", f"could not read a file part: {exc}") + outcome = await self.runtime.run(Turn( + conversation_dir=self.conversations_dir / conversation, + workspace=workspace, + text=text, + images=images, + mcp_servers=_with_role(request.mcp_servers, config.role), + skills_dir=self.skills_dir, + skill_names=tuple(str(skill["name"]) for skill in request.skills), + config=config, + endpoint=self.endpoint, + on_event=on_event, + on_usage=spent.append, + )) + except TimeoutError: + message = f"Pydantic AI did not finish within {config.timeout_seconds}s" + timed_out = TaskResult.builder().failed("pydantic_ai.timeout", message).add_text(message) + return (timed_out.usage(spent[-1]) if spent else timed_out).session_ref(conversation).build() + builder = TaskResult.builder().session_ref(conversation).usage(outcome.usage) + if outcome.error_code is None: + return builder.succeeded().add_text(outcome.text).build() + return builder.failed(outcome.error_code, outcome.text, error_type=outcome.error_type).add_text(outcome.text).build() + + async def _prompt(self, request: TaskRequest[PydanticAIConfig], conversation: str) -> tuple[str, tuple[Image, ...]]: + """The message text and its images. Every file is kept on disk, in the conversation's own directory, for its + later turns and its snapshots, and its path added to the text.""" + files_dir = self.conversations_dir / conversation / INPUTS / request.task_id + chunks: list[str] = [] + images: list[Image] = [] + for index, part in enumerate(request.parts): + if isinstance(part, TextPart): + chunks.append(part.text) + elif isinstance(part, DataPart): + chunks.append(json.dumps(part.data, indent=2, default=str)) + elif isinstance(part, FilePart): + files_dir.mkdir(parents=True, exist_ok=True) + path = await _materialize(part, index, files_dir) + media_type = _media_type(part) + if media_type in IMAGE_TYPES: + images.append(Image(path, media_type)) + chunks.append(f"[Attached file ({media_type}) saved at {path}]") + return "\n\n".join(chunks), tuple(images) + + @asynccontextmanager + async def _holding(self, conversation: str) -> AsyncIterator[None]: + entry = self._conversation_locks.setdefault(conversation, _ConversationLock()) + entry.holders += 1 + try: + async with entry.lock: + yield + finally: + entry.holders -= 1 + if not entry.holders: + del self._conversation_locks[conversation] + + @asynccontextmanager + async def _in_workspace(self) -> AsyncIterator[None]: + async with self._workspace_lock: + self._running += 1 + self._idle.clear() + try: + yield + finally: + self._running -= 1 + if not self._running: + self._idle.set() + + @extension(SNAPSHOT_V1.save) + async def save_snapshot(self, request: ObjectSnapshotSaveRequest) -> dict[str, Any]: + conversation = conversation_id(request.context_id) + directory = self.conversations_dir / conversation + with tempfile.TemporaryDirectory() as scratch: + # The conversation and the workspace are copied together, so a restore cannot fall between them. + conversation_archive = Path(scratch) / "conversation.tar.gz" + workspace_archive = Path(scratch) / "workspace.tar.gz" + async with self._holding(conversation), self._workspace_lock: + if not (directory / MESSAGES).is_file(): + raise HTTPException(status_code=404, detail=f"No conversation for context '{request.context_id}'") + await asyncio.to_thread(_pack, directory, conversation_archive, record_source=True) + if request.objects.workspace is not None: + # Tasks in other contexts work in the same workspace: those running end first, and new ones wait + # for the lock this holds. + await self._idle.wait() + self.workspace.mkdir(parents=True, exist_ok=True) + await asyncio.to_thread(_pack, self.workspace, workspace_archive) + objects = {"trajectory": (await upload(request.objects.trajectory, conversation_archive)).model_dump(mode="json")} + if request.objects.workspace is not None: + objects["workspace"] = (await upload(request.objects.workspace, workspace_archive)).model_dump(mode="json") + return {"context_id": request.context_id, "objects": objects} + + @extension(SNAPSHOT_V1.load) + async def load_snapshot(self, request: ObjectSnapshotLoadRequest) -> dict[str, Any]: + context_id = request.target_context_id or uuid.uuid4().hex + with tempfile.TemporaryDirectory() as scratch: + conversation_archive = Path(scratch) / "conversation.tar.gz" + await download(request.objects.trajectory, conversation_archive) + workspace_archive = None + if request.objects.workspace is not None: + workspace_archive = Path(scratch) / "workspace.tar.gz" + await download(request.objects.workspace, workspace_archive) + async with self._workspace_lock: + if self._running: + raise HTTPException(status_code=409, detail="a task is running; restore the workspace once it finishes") + target = self.conversations_dir / conversation_id(context_id) + try: + await asyncio.to_thread(self._restore, conversation_archive, workspace_archive, target) + except ValueError as exc: + raise HTTPException(status_code=400, detail=str(exc)) from exc + return {"context_id": context_id} + + def _restore(self, conversation_archive: Path, workspace_archive: Path | None, target: Path) -> None: + """Extract both archives and check the conversation before putting either in place: one that fails to + extract, or holds no conversation, changes nothing.""" + self.conversations_dir.mkdir(parents=True, exist_ok=True) + conversation = _stage_conversation(conversation_archive, self.conversations_dir, target) + staged = [conversation] + try: + if workspace_archive is not None: + self.workspace.mkdir(parents=True, exist_ok=True) + staged.append(_extract(workspace_archive, self.workspace)) + _replace_contents(self.workspace, staged[-1]) + shutil.rmtree(target, ignore_errors=True) + conversation.rename(target) + finally: + for path in staged: + shutil.rmtree(path, ignore_errors=True) + + @extension(SKILL_CONFIG_V1.add.inline) + async def add_inline_skill(self, request: InlineSkillRequest) -> dict[str, Any]: + skill_dir = self.skills_dir / request.name + skill_dir.mkdir(parents=True, exist_ok=False) + (skill_dir / "SKILL.md").write_text(request.skill_md) + return {"name": request.name} + + @extension(SKILL_CONFIG_V1.add.bundle) + async def add_bundle_skill(self, request: BundleSkillRequest) -> dict[str, Any]: + self.skills_dir.mkdir(parents=True, exist_ok=True) + staging = self.skills_dir / f".{request.name}.{uuid.uuid4().hex}" + try: + # The protocol admits only normalized relative paths, so none leaves the staging directory. + for item in request.skill_bundle.files: + await download(item.object, staging / item.path) + staging.rename(self.skills_dir / request.name) + finally: + shutil.rmtree(staging, ignore_errors=True) + return {"name": request.name} diff --git a/src/agent_env/agents/agents/pydantic-ai/pydantic_ai_runtime.py b/src/agent_env/agents/agents/pydantic-ai/pydantic_ai_runtime.py new file mode 100644 index 00000000..b53166ee --- /dev/null +++ b/src/agent_env/agents/agents/pydantic-ai/pydantic_ai_runtime.py @@ -0,0 +1,166 @@ +"""The Pydantic AI side of the agent: a task is one run of a Pydantic AI agent with the Pydantic AI Harness ``Coder`` +tools in the workspace, the task's MCP servers and its skills, continuing the context's message history.""" + +import asyncio +from collections.abc import Sequence +from dataclasses import dataclass, field +from typing import Any, Final + +from agentenv_protocol.a2a_agent import Usage +from pydantic_ai import Agent +from pydantic_ai.capabilities import LocalWorkspace +from pydantic_ai.exceptions import ModelAPIError, ModelHTTPError, UnexpectedModelBehavior, UsageLimitExceeded +from pydantic_ai.mcp import MCPToolset +from pydantic_ai.messages import BinaryContent, ModelMessage, ModelMessagesTypeAdapter, UserContent +from pydantic_ai.models.openai import OpenAIChatModel +from pydantic_ai.providers.openai import OpenAIProvider +from pydantic_ai.providers.vllm import VLLMProvider +from pydantic_ai.settings import ModelSettings +from pydantic_ai.usage import RunUsage, UsageLimits +from pydantic_ai_harness import Skills +from pydantic_ai_harness.coder import Coder + +from pydantic_ai_agent import MESSAGES, Outcome, Turn + +# Consecutive retries a tool gets for arguments it refuses, as many as Coder gives its file tools, so a model that +# fumbles a command twice doesn't end the task. +_TOOL_RETRIES: Final[int] = 5 +# Without a key the OpenAI client refuses to start; an endpoint that needs none ignores it. +_NO_KEY: Final[str] = "unset" +# The Pydantic AI providers a model name's prefix picks for the endpoint, as in ``vllm:qwen3.8-27b``: each speaks +# OpenAI Chat Completions, and ``vllm`` fits its requests to the model's family, as a vLLM server needs. +_PROVIDERS: Final = {"openai": OpenAIProvider, "vllm": VLLMProvider} + + +def _model(turn: Turn) -> OpenAIChatModel: + prefix, _, name = turn.config.model.partition(":") + provider = _PROVIDERS.get(prefix) if name else None + if provider is None: + provider, name = OpenAIProvider, turn.config.model + return OpenAIChatModel(name, provider=provider(base_url=turn.endpoint.base_url, + api_key=turn.endpoint.api_key or _NO_KEY)) + + +def _settings(turn: Turn) -> ModelSettings: + # Sent on every call and kept out of the messages, which a snapshot uploads. + settings = ModelSettings(extra_headers=dict(turn.endpoint.extra_headers)) + if turn.config.effort is not None: + settings["thinking"] = turn.config.effort + return settings + + +def _agent(turn: Turn) -> Agent: + # The agent runs in a sandbox of its own, so its file tools reach the files attached outside the workspace. + capabilities: list[Any] = [LocalWorkspace(turn.workspace), Coder(unrestricted_filesystem=True)] + if turn.skill_names: + capabilities.append(Skills(turn.skills_dir, include=turn.skill_names)) + toolsets = [ + MCPToolset(registration["url"], id=name, headers=dict(registration.get("headers") or {})) + for name, registration in turn.mcp_servers.items() + ] + return Agent( + _model(turn), + instructions=turn.config.system_prompt, + model_settings=_settings(turn), + retries={"tools": _TOOL_RETRIES}, + capabilities=capabilities, + toolsets=toolsets, + ) + + +def _load(turn: Turn) -> list[ModelMessage]: + path = turn.conversation_dir / MESSAGES + return ModelMessagesTypeAdapter.validate_json(path.read_bytes()) if path.is_file() else [] + + +def _save(turn: Turn, messages: Sequence[ModelMessage]) -> None: + """The conversation's messages, written aside and moved into place, so a snapshot never reads half of them.""" + turn.conversation_dir.mkdir(parents=True, exist_ok=True) + staged = turn.conversation_dir / f".{MESSAGES}.tmp" + staged.write_bytes(ModelMessagesTypeAdapter.dump_json(list(messages))) + staged.replace(turn.conversation_dir / MESSAGES) + + +def _events(messages: Sequence[ModelMessage]) -> list[dict[str, Any]]: + return ModelMessagesTypeAdapter.dump_python(list(messages), mode="json") + + +def _usage(usage: RunUsage) -> Usage: + return Usage( + tool_call_count=usage.tool_calls, + input_tokens=usage.input_tokens, + output_tokens=usage.output_tokens, + total_tokens=usage.input_tokens + usage.output_tokens, + provider_details={ + "requests": usage.requests, + "cache_read_tokens": usage.cache_read_tokens, + "cache_write_tokens": usage.cache_write_tokens, + }, + ) + + +def _failure(exc: Exception, usage: Usage) -> Outcome: + """A run that raised, typed by whose failure it was: the model endpoint's, or the agent's own.""" + if isinstance(exc, UsageLimitExceeded): + return Outcome(f"Pydantic AI stopped: {exc}", usage, "pydantic_ai.usage_limit") + if isinstance(exc, ModelHTTPError): + return Outcome(f"The model endpoint answered {exc.status_code}: {exc.body}", usage, + f"pydantic_ai.model_http_{exc.status_code}", "infra_error") + if isinstance(exc, ModelAPIError): + return Outcome(f"The model endpoint failed: {exc}", usage, "pydantic_ai.model_api", "infra_error") + if isinstance(exc, UnexpectedModelBehavior): + return Outcome(f"The model misbehaved: {exc}", usage, "pydantic_ai.unexpected_model_behavior") + # An MCP server or a tool failing to start, for one. + return Outcome(f"Pydantic AI failed: {type(exc).__name__}: {exc}", usage, "pydantic_ai.error", "infra_error") + + +@dataclass(slots=True) +class PydanticAIRuntime: + _runs: set[asyncio.Task[Any]] = field(default_factory=set, init=False, repr=False) + + async def stop(self) -> None: + runs = list(self._runs) + for run in runs: + run.cancel() + await asyncio.gather(*runs, return_exceptions=True) + + async def run(self, turn: Turn) -> Outcome: + task = asyncio.current_task() + self._runs.add(task) + try: + return await self._run(turn) + finally: + self._runs.discard(task) + + async def _run(self, turn: Turn) -> Outcome: + try: + history = await asyncio.to_thread(_load, turn) + except (OSError, ValueError) as exc: + return Outcome(f"Pydantic AI could not read the conversation: {exc}", Usage(), "pydantic_ai.history", + "infra_error") + prompt: list[UserContent] = [turn.text] + for image in turn.images: + prompt.append(BinaryContent(data=await asyncio.to_thread(image.path.read_bytes), media_type=image.media_type)) + agent = _agent(turn) + recorded = len(history) + usage = RunUsage() + try: + async with agent.iter(prompt, message_history=history, + usage_limits=UsageLimits(request_limit=turn.config.request_limit)) as run: + async for _node in run: + messages = run.all_messages() + for event in _events(messages[recorded:]): + turn.on_event(event) + recorded = len(messages) + usage = run.usage + result = run.result + except asyncio.CancelledError: + turn.on_usage(_usage(usage)) + raise + except Exception as exc: # every failure is the task's, typed by whose it was + return _failure(exc, _usage(usage)) + messages = result.all_messages() + for event in _events(messages[recorded:]): + turn.on_event(event) + await asyncio.to_thread(_save, turn, messages) + return Outcome(result.output.strip() or "The agent finished without a reply.", _usage(result.usage)) diff --git a/src/agent_env/agents/agents/pydantic-ai/requirements.txt b/src/agent_env/agents/agents/pydantic-ai/requirements.txt new file mode 100644 index 00000000..cb95b795 --- /dev/null +++ b/src/agent_env/agents/agents/pydantic-ai/requirements.txt @@ -0,0 +1,6 @@ +# The image installs the newest agentenv-framework-protocol at or above this floor; a release doesn't pin it, so the +# image is built again only when this folder changes. The floor is the first that serves peer-agents/v1, which the +# agent leaves to the SDK; raise it when the agent needs a newer protocol. +agentenv-framework-protocol[agent]>=0.1.351 +pydantic-ai-slim[mcp,openai]==2.61.0 +pydantic-ai-harness[coder,skills]==0.61.0 diff --git a/src/agent_env/agents/agents/pydantic-ai/start.sh b/src/agent_env/agents/agents/pydantic-ai/start.sh new file mode 100644 index 00000000..41caa80b --- /dev/null +++ b/src/agent_env/agents/agents/pydantic-ai/start.sh @@ -0,0 +1,11 @@ +#!/bin/sh +# Starts the installed agent (install/v1) with the credentials install wrote to agent.env; the caller sets A2A_PORT. +here=$(cd "$(dirname "$0")" && pwd) +# The agent works where it's started, the task container's working directory, unless that's the root. +[ "$(pwd)" = / ] || export PYDANTIC_AI_A2A_WORKSPACE="$(pwd)" +set -a +. "$here/agent.env" +eval "${AGENT_TRUST_ENV:-}" +set +a +export PATH="/opt/pydantic-ai-a2a-venv/bin:$PATH" +cd "$here" && exec python a2a_server.py >/tmp/pydantic-ai-a2a.log 2>&1 diff --git a/src/agent_env/agents/aliases.toml b/src/agent_env/agents/aliases.toml index 02d350ed..7392d7d0 100644 --- a/src/agent_env/agents/aliases.toml +++ b/src/agent_env/agents/aliases.toml @@ -6,3 +6,4 @@ gemini = ["gemini"] opencode = ["opencode"] pi = ["pi"] openhands = ["openhands"] +pydantic-ai = ["pydantic-ai"] diff --git a/src/agent_env/agents/tasks/pydantic-ai-smoke.json b/src/agent_env/agents/tasks/pydantic-ai-smoke.json new file mode 100644 index 00000000..299c77dd --- /dev/null +++ b/src/agent_env/agents/tasks/pydantic-ai-smoke.json @@ -0,0 +1,7 @@ +[ + {"id": "deploy", "type": "deploy_agent", "a2a_agent_id": "pydantic-ai", "env_ids": [], "sandbox_type": "local"}, + {"id": "ask", "type": "prompt_agent", "prompt_id": "ask", "depends_on": ["deploy"], + "prompt": "Run the shell command `echo $((6 * 7))` and reply with only the number it prints."}, + {"id": "check", "type": "agent_prompt_response_verifier", "prompt_id": "ask", "verifier_id": "answer", + "criteria": [{"type": "response_contains", "needles": ["42"]}], "depends_on": ["ask"]} +] diff --git a/tst/integration/cli/run_bundle_pydantic_ai_test.py b/tst/integration/cli/run_bundle_pydantic_ai_test.py new file mode 100644 index 00000000..9a9ac397 --- /dev/null +++ b/tst/integration/cli/run_bundle_pydantic_ai_test.py @@ -0,0 +1,245 @@ +"""``agent-env run`` and the task steps on the agents bundle's Pydantic AI agent: its image built from the shipped folder +with real docker, the real Pydantic AI and harness, and a fake OpenAI Chat Completions endpoint as the configured model +endpoint, so no model key is needed. Needs a Docker daemon and network access for the image's downloads.""" + +import asyncio +import json +import logging +import time +import uuid +from typing import Final + +import httpx +import pytest +from click.testing import CliRunner + +from agent_env.a2a_agent import protocol +from agent_env.artifact.store import reset_artifact_store +from agent_env.cli import cli +from agent_env.config import configure, reset_config +from agent_env.providers.sandbox_providers.local_sandbox import _host_ips +from agent_env.store.object_store.local.grant_server import grant_server +from agent_env.store.routing import namespace_routing +from agent_env.task.teardown import teardown_run +from agent_env.task_step.context import TaskStepContext +from agent_env.task_step.task_steps.deploy_agent import DeployAgentTaskStep +from agent_env.task_step.task_steps.prompt_agent import PromptAgentTaskStep +from tst.util.fake_openai import Fault, FakeOpenAI, Text, ToolCall, answer, tool_result_text, tool_then_answer + +pytestmark = [pytest.mark.integration, pytest.mark.int_test_slow] + +AGENT: Final[str] = "@local/agentenv-framework/agents/pydantic-ai" +MODEL: Final[str] = "fake-model-1" +ECHO_ENV: Final[str] = "@local/agentenv-framework/echo-mcp/echo" + + +@pytest.fixture +def endpoint(): + """The fake endpoint, where this machine's containers reach it: on Linux, the Docker bridge address.""" + with FakeOpenAI(host=_host_ips()[-1]) as fake: + yield fake + + +@pytest.fixture +def state(monkeypatch, tmp_path, endpoint): + """Local stores under this test's folder, and [model] pointed at the fake endpoint with a header of its own. + Whatever ran, no sandbox work folder may be left behind.""" + sandboxes = tmp_path / "sandboxes" + config = tmp_path / "config.toml" + config.write_text(f'[model]\nbase_url = "{endpoint.url}"\napi_key = "fake-key"\ndefault = "{MODEL}"\n\n' + '[model.headers]\nX-Agentenv-Test = "header-check"\n') + monkeypatch.chdir(tmp_path) + monkeypatch.setenv("AGENT_ENV_CONFIG", str(config)) + monkeypatch.setenv("XDG_STATE_HOME", str(tmp_path / "state")) + monkeypatch.setenv("AGENT_ENV_DOCUMENT_STORE", "local") + monkeypatch.setenv("AGENT_ENV_OBJECT_STORE", "local") + monkeypatch.setenv("AGENT_ENV_LOCAL_SANDBOX_DIR", str(sandboxes)) + for name in ("LITELLM_API_KEY", "LITELLM_BASE_URL"): + monkeypatch.delenv(name, raising=False) + configure() + reset_artifact_store() + logging.disable(logging.CRITICAL) # pytest's live logging would take CliRunner's stdout + try: + yield tmp_path + assert not sandboxes.exists() or not any(sandboxes.iterdir()), "a run left a sandbox work folder" + finally: + logging.disable(logging.NOTSET) + grant_server(None, None).close() + reset_artifact_store() + reset_config() + + +def _bundle(root, name, steps): + (root / "tasks").mkdir(parents=True) + (root / f"tasks/{name}.json").write_text(json.dumps(steps)) + return root + + +def _ask(prompt, needle, env_ids=(), depends_on=()): + return [ + {"id": "agent", "type": "deploy_agent", "a2a_agent_id": AGENT, "env_ids": list(env_ids), + "sandbox_type": "local", **({"depends_on": list(depends_on)} if depends_on else {})}, + {"id": "ask", "type": "prompt_agent", "prompt_id": "ask", "prompt": prompt, "depends_on": ["agent"]}, + {"id": "check", "type": "agent_prompt_response_verifier", "prompt_id": "ask", "verifier_id": "answer", + "criteria": [{"type": "response_contains", "needles": [needle]}], "depends_on": ["ask"]}, + ] + + +def _run(root): + return CliRunner().invoke(cli, ["run", str(root)]) + + +def _with_agent(state, steps): + """Run ``steps`` against the shipped agent, after a bundle run has written it and its image to this test's + stores, and tear down what they deployed.""" + assert _run(_bundle(state / "setup", "ask", _ask("Say pong.", "pong"))).exit_code == 0 + context = TaskStepContext() + + async def run(): + try: + await steps(context) + finally: + await teardown_run(context) + + with namespace_routing(): + asyncio.run(run()) + + +async def _deployed(context) -> str: + await DeployAgentTaskStep(id="agent", version=None, sandbox_type="local", a2a_agent_id=AGENT).execute(context) + (agent,) = context.deployed_agents + base_url = agent.a2a_url or agent.api_url + await protocol.post_agent_config(f"{base_url}/ext/agent-config", {"model": MODEL}) + return base_url + + +async def _send(base_url: str, context_id: str, text: str) -> dict: + async with httpx.AsyncClient(timeout=300) as client: + response = await client.post(f"{base_url}/a2a", json={ + "jsonrpc": "2.0", "id": "1", "method": "message/send", "params": { + "message": {"kind": "message", "messageId": uuid.uuid4().hex, "role": "user", + "contextId": context_id, "parts": [{"kind": "text", "text": text}]}, + "configuration": {"blocking": True}}}) + task = response.json()["result"] + parts = task["status"]["message"]["parts"] + return {"state": task["status"]["state"], + "text": "".join(part.get("text", "") for part in parts if part["kind"] == "text"), + "data": next((part["data"] for part in parts if part["kind"] == "data"), {})} + + +def test_the_shipped_pydantic_ai_answers_through_the_configured_endpoint(state, endpoint): + endpoint.respond = answer("pong") + + result = _run(_bundle(state / "one", "ask", _ask("Say pong.", "pong"))) + + assert result.exit_code == 0, result.output + assert "tasks/ask.json v1: passed" in result.output, result.output + calls = endpoint.completions() + assert calls, "the agent never called the endpoint" + assert {call.headers["authorization"] for call in calls} == {"Bearer fake-key"} + assert {call.headers.get("x-agentenv-test") for call in calls} == {"header-check"} + assert {call.body.get("model") for call in calls} == {MODEL} + + +def test_the_shipped_pydantic_ai_runs_a_shell_command_with_its_shell_tool(state, endpoint): + endpoint.respond = tool_then_answer("shell", {"command": "echo $((6 * 7))"}, prefix="the shell said: ") + + result = _run(_bundle(state / "shell", "ask", _ask("What is six times seven? Use the shell.", "42"))) + + seen = (result.output, endpoint.tool_results()) + assert result.exit_code == 0 and "tasks/ask.json v1: passed" in result.output, seen + # The fake never computes, so 42 reached the answer through the container's shell. + assert [result.startswith("42") for result in endpoint.tool_results()] == [True], seen + + +def test_the_shipped_pydantic_ai_calls_the_tools_of_the_env_it_is_deployed_with(state, endpoint): + endpoint.respond = tool_then_answer("echo", {"text": "ping-from-tool"}, prefix="the tool said: ") + root = _bundle(state / "tools", "echo", [ + {"id": "env", "type": "deploy_env", "env_id": ECHO_ENV, "sandbox_type": "local"}, + *_ask("Call the echo tool with the text ping-from-tool.", "ping-from-tool", env_ids=[ECHO_ENV], + depends_on=["env"]), + ]) + + result = _run(root) + + offered = sorted({tool["function"]["name"] for call in endpoint.completions() for tool in call.body.get("tools") or []}) + seen = (result.output, offered, endpoint.tool_results()) + assert result.exit_code == 0 and "tasks/echo.json v1: passed" in result.output, seen + assert "echo" in offered, seen + assert ["ping-from-tool" in result for result in endpoint.tool_results()] == [True], seen + + +def test_a_context_continues_its_conversation_which_holds_neither_the_key_nor_the_headers(state, endpoint): + outcomes = {} + # Spelled so the command, which the conversation records, holds neither value itself. + command = ("cd ~/.pydantic-ai-a2a/conversations && " + "(grep -rlF -e \"$(printf 'header-%s' check)\" -e \"$(printf 'fake-%s' key)\" . || echo no-secrets) && " + "ls */") + + def respond(body): + result = tool_result_text(body) + if result is not None: + return Text(f"seen: {result}") + if "Look at your saved conversation." in json.dumps(body["messages"][-1]): + return ToolCall("shell", {"command": command}) + return Text("pong") + + endpoint.respond = answer("pong") + + async def steps(context): + base_url = await _deployed(context) + endpoint.respond = respond + outcomes["first"] = await _send(base_url, "ctx", "Say pong.") + outcomes["second"] = await _send(base_url, "ctx", "Look at your saved conversation.") + + _with_agent(state, steps) + + assert outcomes["first"]["state"] == outcomes["second"]["state"] == "completed", outcomes + assert "no-secrets" in outcomes["second"]["text"] and "messages.json" in outcomes["second"]["text"], outcomes + # The second turn carried the first. + last = json.dumps(endpoint.completions()[-1].body["messages"]) + assert "Say pong." in last and "Look at your saved conversation." in last, last + assert {call.headers.get("x-agentenv-test") for call in endpoint.completions()} == {"header-check"} + + +def test_a_model_endpoint_that_refuses_the_key_fails_the_task_as_an_infra_error(state, endpoint): + endpoint.respond = answer("pong") + + async def steps(context): + await DeployAgentTaskStep(id="agent", version=None, sandbox_type="local", a2a_agent_id=AGENT).execute(context) + endpoint.fail_next(Fault("status", 401), times=100) + with pytest.raises(Exception, match="error_type=infra_error, error_code=pydantic_ai.model_http_401"): + await PromptAgentTaskStep(id="ask", version=None, prompt_id="ask", prompt="Say pong.").execute(context) + + _with_agent(state, steps) + + +def test_a_run_past_its_timeout_mid_command_leaves_no_process_behind(state, endpoint): + outcomes = {} + + def respond(body): + result = tool_result_text(body) + if result is not None: + return Text(f"processes: {result}") + if "List the processes." in json.dumps(body["messages"][-1]): + return ToolCall("shell", {"command": "ps -eo args"}) + return ToolCall("shell", {"command": "sleep 120"}) + + endpoint.respond = answer("pong") + + async def steps(context): + base_url = await _deployed(context) + endpoint.respond = respond + await protocol.post_agent_config(f"{base_url}/ext/agent-config", {"timeout_seconds": 5}) + started = time.monotonic() + outcomes["timed_out"] = await _send(base_url, "ctx", "Sleep.") + outcomes["seconds"] = time.monotonic() - started + await protocol.post_agent_config(f"{base_url}/ext/agent-config", {"timeout_seconds": 120}) + outcomes["after"] = await _send(base_url, "ctx", "List the processes.") + + _with_agent(state, steps) + + assert outcomes["timed_out"]["data"]["error_code"] == "pydantic_ai.timeout", outcomes + assert outcomes["seconds"] < 30, outcomes + assert outcomes["after"]["state"] == "completed", outcomes + assert "ps -eo args" in outcomes["after"]["text"] and "sleep 120" not in outcomes["after"]["text"], outcomes diff --git a/tst/unit/agents/pydantic_ai/__init__.py b/tst/unit/agents/pydantic_ai/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tst/unit/agents/pydantic_ai/agent.py b/tst/unit/agents/pydantic_ai/agent.py new file mode 100644 index 00000000..aa1bdd01 --- /dev/null +++ b/tst/unit/agents/pydantic_ai/agent.py @@ -0,0 +1,14 @@ +"""The Pydantic AI agent's A2A module, imported from its folder as its image runs it: by its own name, from the folder on +the path. Its runtime module needs Pydantic AI, which only the agent's image installs, so tests give the agent a fake +runtime instead.""" + +import importlib +import sys +from pathlib import Path +from typing import Final + +import agent_env + +AGENT_DIR: Final[Path] = Path(agent_env.__file__).parent / "agents" / "agents" / "pydantic-ai" +sys.path.insert(0, str(AGENT_DIR)) +pydantic_ai_agent = importlib.import_module("pydantic_ai_agent") diff --git a/tst/unit/agents/pydantic_ai/test_pydantic_ai_agent.py b/tst/unit/agents/pydantic_ai/test_pydantic_ai_agent.py new file mode 100644 index 00000000..32a260dc --- /dev/null +++ b/tst/unit/agents/pydantic_ai/test_pydantic_ai_agent.py @@ -0,0 +1,621 @@ +import asyncio +import base64 +import io +import json +import os +import shlex +import subprocess +import sys +import tarfile +import threading +import time +import uuid +from collections import Counter +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any, Final + +import httpx +import pytest +from agentenv_protocol.a2a_agent import Usage +from starlette.testclient import TestClient + +from tst.unit.agents.pydantic_ai.agent import AGENT_DIR, pydantic_ai_agent + +PydanticAIAgent = pydantic_ai_agent.PydanticAIAgent +Endpoint = pydantic_ai_agent.Endpoint +Outcome = pydantic_ai_agent.Outcome + +# What a run cancelled midway had spent by then. +SPENT: Final[Usage] = Usage(tool_call_count=1, input_tokens=90, output_tokens=10, total_tokens=100) +GRANT: Final[dict[str, str]] = {"kind": "http-put", "url": "https://objects.test/put", "expires_at": "2099-01-01T00:00:00Z"} +READ: Final[dict[str, str]] = {"kind": "http-get", "url": "https://objects.test/get", "expires_at": "2099-01-01T00:00:00Z"} +ENDPOINT: Final[Endpoint] = Endpoint("http://litellm.test", "sk-test", {"x-project": "p1"}) + + +@dataclass(slots=True) +class FakeRuntime: + """Keeps each conversation's messages in its directory, as the real runtime does once a run succeeds, and records + the turns it ran. With ``gate`` set, a run waits for it to open.""" + + turns: list[Any] = field(default_factory=list) + outcome: Outcome | None = None + hang: bool = False + cancelled: bool = False + gate: threading.Event | None = None + active: Counter = field(default_factory=Counter) + stopped: bool = False + + async def stop(self) -> None: + self.stopped = True + + async def run(self, turn: Any) -> Outcome: + self.turns.append(turn) + turn.on_event({"kind": "request", "parts": [{"part_kind": "user-prompt", "content": turn.text}]}) + turn.on_event({"kind": "response", "parts": [{"part_kind": "tool-call", "tool_name": "shell"}]}) + self.active[turn.conversation_dir] += 1 + try: + while self.gate is not None and not self.gate.is_set(): + await asyncio.sleep(0.01) + if self.hang: + # Bounded, so an agent that applies no timeout fails the test instead of hanging it. + await asyncio.sleep(10) + except asyncio.CancelledError: + self.cancelled = True + turn.on_usage(SPENT) + raise + finally: + self.active[turn.conversation_dir] -= 1 + messages = history(turn.conversation_dir) + turn.conversation_dir.mkdir(parents=True, exist_ok=True) + (turn.conversation_dir / "messages.json").write_text(json.dumps([*messages, turn.text])) + return self.outcome or Outcome("Done.", Usage(tool_call_count=1, input_tokens=200, output_tokens=40, + total_tokens=240)) + + +def history(conversation_dir: Path) -> list[str]: + path = conversation_dir / "messages.json" + return json.loads(path.read_text()) if path.exists() else [] + + +def conversation(home: Path, context_id: str) -> Path: + return home / "conversations" / pydantic_ai_agent.conversation_id(context_id) + + +@pytest.fixture +def runtime() -> FakeRuntime: + return FakeRuntime() + + +@pytest.fixture +def agent(tmp_path: Path, runtime: FakeRuntime) -> PydanticAIAgent: + return PydanticAIAgent(runtime, endpoint=ENDPOINT, home=tmp_path / "home", workspace=tmp_path / "ws") + + +@pytest.fixture +def client(agent: PydanticAIAgent): + with TestClient(agent.create_app(), client=("127.0.0.1", 50000)) as client: + client.post("/ext/agent-config", json={"model": "vllm:qwen3.8-27b", "effort": "high"}) + yield client + + +def _send(client: TestClient, parts: list[dict[str, Any]], context_id: str = "ctx-1") -> dict[str, Any]: + response = client.post("/a2a", json={ + "jsonrpc": "2.0", + "id": "1", + "method": "message/send", + "params": { + "message": {"kind": "message", "messageId": uuid.uuid4().hex, "role": "user", + "contextId": context_id, "parts": parts}, + "configuration": {"blocking": True}, + }, + }) + return response.json()["result"] + + +def _reply(task: dict[str, Any]) -> tuple[str, dict[str, Any]]: + parts = task["status"]["message"]["parts"] + text = "".join(part.get("text", "") for part in parts if part["kind"] == "text") + data = next((part["data"] for part in parts if part["kind"] == "data"), {}) + return text, data + + +def _wait_until(condition, timeout: float = 5.0) -> None: + deadline = time.monotonic() + timeout + while not condition(): + assert time.monotonic() < deadline, "condition not reached" + time.sleep(0.01) + + +def _send_in_background(client: TestClient, text: str, context_id: str) -> threading.Thread: + thread = threading.Thread(target=_send, args=(client, [{"kind": "text", "text": text}], context_id)) + thread.start() + return thread + + +def test_run_returns_the_reply_usage_and_live_trajectory(client: TestClient, runtime: FakeRuntime) -> None: + task = _send(client, [{"kind": "text", "text": "fix the bug"}, {"kind": "data", "data": {"k": 1}}]) + + text, data = _reply(task) + assert task["status"]["state"] == "completed" + assert text == "Done." + assert data["usage"] == {"tool_call_count": 1, "input_tokens": 200, "output_tokens": 40, "total_tokens": 240} + turn = runtime.turns[0] + assert turn.text == 'fix the bug\n\n{\n "k": 1\n}' + assert (turn.config.model, turn.config.effort) == ("vllm:qwen3.8-27b", "high") + assert turn.endpoint == ENDPOINT + trajectory = client.post("/ext/trajectory", json={"task_id": task["id"]}).json()["trajectory"] + assert [event["kind"] for event in trajectory] == ["request", "response"] + + +def test_a_context_keeps_one_conversation(client: TestClient, runtime: FakeRuntime, tmp_path: Path) -> None: + _send(client, [{"kind": "text", "text": "one"}], context_id="ctx/a b") + _send(client, [{"kind": "text", "text": "two"}], context_id="ctx/a b") + _send(client, [{"kind": "text", "text": "three"}], context_id="ctx-2") + + first, second, other = (turn.conversation_dir for turn in runtime.turns) + assert first == second == conversation(tmp_path / "home", "ctx/a b") + assert other != first + assert history(first) == ["one", "two"] + + +def test_mcp_servers_and_the_peer_tools_reach_the_runtime(client: TestClient, runtime: FakeRuntime) -> None: + client.post("/ext/mcp-config", json={"name": "files", "url": "http://mcp.test/mcp", + "headers": {"Authorization": "Bearer t"}}) + + _send(client, [{"kind": "text", "text": "go"}]) + + servers = runtime.turns[0].mcp_servers + assert servers["files"] == {"url": "http://mcp.test/mcp", "headers": {"Authorization": "Bearer t"}} + assert servers["peer_agents"]["url"].endswith("/ext/peer-agents/mcp") + + +def test_the_configured_role_is_sent_to_every_mcp_server(client: TestClient, runtime: FakeRuntime) -> None: + client.post("/ext/mcp-config", json={"name": "files", "url": "http://mcp.test/mcp", + "headers": {"Authorization": "Bearer t", "AgentEnv-Role": "stale"}}) + client.post("/ext/agent-config", json={"role": "buyer"}) + + _send(client, [{"kind": "text", "text": "go"}]) + + servers = runtime.turns[0].mcp_servers + assert servers["files"]["headers"] == {"Authorization": "Bearer t", "AgentEnv-Role": "buyer"} + assert servers["peer_agents"]["headers"] == {"AgentEnv-Role": "buyer"} + + +def test_images_are_attached_and_every_file_is_left_in_the_conversation( + client: TestClient, runtime: FakeRuntime, tmp_path: Path +) -> None: + png = base64.b64encode(b"\x89PNG").decode() + mp3 = base64.b64encode(b"ID3").decode() + task = _send(client, [{"kind": "text", "text": "look"}, + {"kind": "file", "file": {"name": "../red.png", "mimeType": "image/png", "bytes": png}}, + {"kind": "file", "file": {"name": "clip.mp3", "mimeType": "audio/mpeg", "bytes": mp3}}]) + + turn = runtime.turns[0] + saved = [_attached(line) for line in turn.text.split("\n\n")[1:]] + assert saved == [turn.conversation_dir / "inputs" / task["id"] / name for name in ("1-red.png", "2-clip.mp3")] + assert [(image.path, image.media_type) for image in turn.images] == [(saved[0], "image/png")] + assert saved[1].read_bytes() == b"ID3" + + +def test_inline_and_bundle_skills_are_installed_for_the_runtime( + client: TestClient, runtime: FakeRuntime, tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + async def fake_download(source, destination: Path) -> None: + destination.parent.mkdir(parents=True, exist_ok=True) + destination.write_text(f"from {source.read.url}") + + monkeypatch.setattr(pydantic_ai_agent, "download", fake_download) + grant = {"media_type": "text/markdown", "max_bytes": 64, "size_bytes": 8, "read": READ} + client.post("/ext/skill-config", json={"name": "greet", "description": "Greet", "skill_md": "# Greet\nSay hi."}) + client.post("/ext/skill-config", json={ + "name": "lint", "description": "Lint", + "skill_bundle": {"max_total_bytes": 128, "files": [ + {"path": "SKILL.md", "object": grant}, {"path": "scripts/run.sh", "object": grant}]}, + }) + + _send(client, [{"kind": "text", "text": "go"}]) + + skills = tmp_path / "home" / "skills" + assert (skills / "greet" / "SKILL.md").read_text() == "# Greet\nSay hi." + assert (skills / "lint" / "scripts" / "run.sh").read_text() == f"from {READ['url']}" + assert sorted(path.name for path in skills.iterdir()) == ["greet", "lint"] + assert runtime.turns[0].skills_dir == skills + assert sorted(runtime.turns[0].skill_names) == ["greet", "lint"] + + +def test_a_bundle_skill_cannot_write_outside_its_directory( + client: TestClient, tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + async def fake_download(source, destination: Path) -> None: + destination.parent.mkdir(parents=True, exist_ok=True) + destination.write_text("payload") + + monkeypatch.setattr(pydantic_ai_agent, "download", fake_download) + grant = {"media_type": "text/plain", "max_bytes": 64, "size_bytes": 8, "read": READ} + responses = [client.post("/ext/skill-config", json={ + "name": "evil", "description": "Evil", + "skill_bundle": {"max_total_bytes": 128, "files": [ + {"path": "SKILL.md", "object": grant}, {"path": path, "object": grant}]}, + }) for path in ("../../escaped.sh", "/tmp/escaped.sh", "scripts/../../escaped.sh")] + + assert [response.status_code for response in responses] == [400, 400, 400] + assert not (tmp_path / "home" / "escaped.sh").exists() + assert not (tmp_path / "home" / "skills").exists() + + +def test_a_runtime_failure_keeps_its_code_and_type(client: TestClient, runtime: FakeRuntime) -> None: + runtime.outcome = Outcome("The model endpoint answered 429: slow down", Usage(), "pydantic_ai.model_http_429", + "infra_error") + + task = _send(client, [{"kind": "text", "text": "go"}]) + + text, data = _reply(task) + assert task["status"]["state"] == "failed" + assert text == "The model endpoint answered 429: slow down" + assert (data["error_code"], data["error_type"]) == ("pydantic_ai.model_http_429", "infra_error") + + +def test_a_timed_out_task_cancels_the_runtime_and_reports_what_it_had_spent( + client: TestClient, runtime: FakeRuntime +) -> None: + runtime.hang = True + client.post("/ext/agent-config", json={"timeout_seconds": 1}) + + task = _send(client, [{"kind": "text", "text": "go"}]) + + _, data = _reply(task) + assert task["status"]["state"] == "failed" + assert data["error_code"] == "pydantic_ai.timeout" + assert data["usage"] == {"tool_call_count": 1, "input_tokens": 90, "output_tokens": 10, "total_tokens": 100} + assert runtime.cancelled + + +def test_reading_a_file_part_counts_toward_the_task_timeout( + client: TestClient, runtime: FakeRuntime, monkeypatch: pytest.MonkeyPatch +) -> None: + real_client = httpx.AsyncClient + + async def stall(request: httpx.Request) -> httpx.Response: + await asyncio.sleep(30) + return httpx.Response(200, content=b"late") + + monkeypatch.setattr(pydantic_ai_agent.httpx, "AsyncClient", lambda **kwargs: real_client( + transport=httpx.MockTransport(stall), **kwargs)) + client.post("/ext/agent-config", json={"timeout_seconds": 1}) + started = time.monotonic() + + task = _send(client, [{"kind": "text", "text": "read it"}, + {"kind": "file", "file": {"name": "slow.txt", "uri": "https://files.test/slow.txt"}}]) + + assert _reply(task)[1]["error_code"] == "pydantic_ai.timeout" + assert time.monotonic() - started < 10 + assert runtime.turns == [] + + +def test_a_cancelled_task_cancels_the_runtime(client: TestClient, runtime: FakeRuntime) -> None: + runtime.hang = True + sent = client.post("/a2a", json={"jsonrpc": "2.0", "id": "1", "method": "message/send", "params": { + "message": {"kind": "message", "messageId": uuid.uuid4().hex, "role": "user", "contextId": "ctx", + "parts": [{"kind": "text", "text": "go"}]}, + "configuration": {"blocking": False}}}).json()["result"] + _wait_until(lambda: sum(runtime.active.values()) == 1) + + client.post("/a2a", json={"jsonrpc": "2.0", "id": "2", "method": "tasks/cancel", "params": {"id": sent["id"]}}) + _wait_until(lambda: runtime.cancelled) + + task = client.post("/a2a", json={"jsonrpc": "2.0", "id": "3", "method": "tasks/get", + "params": {"id": sent["id"]}}).json()["result"] + assert task["status"]["state"] == "canceled" + + +def test_the_server_shutting_down_stops_the_runtime(agent: PydanticAIAgent, runtime: FakeRuntime) -> None: + with TestClient(agent.create_app()): + assert not runtime.stopped + + assert runtime.stopped + + +def test_missing_model_fails_without_running(agent: PydanticAIAgent, runtime: FakeRuntime) -> None: + with TestClient(agent.create_app()) as client: + task = _send(client, [{"kind": "text", "text": "go"}]) + + assert _reply(task)[1]["error_code"] == "pydantic_ai.no_model" + assert runtime.turns == [] + + +def test_the_default_model_is_the_fallback(tmp_path: Path, runtime: FakeRuntime) -> None: + agent = PydanticAIAgent(runtime, endpoint=ENDPOINT, default_model="fallback-model", home=tmp_path / "home", + workspace=tmp_path / "ws") + with TestClient(agent.create_app()) as client: + task = _send(client, [{"kind": "text", "text": "go"}]) + + assert task["status"]["state"] == "completed" + assert runtime.turns[0].config.model == "fallback-model" + + +def test_missing_endpoint_is_an_infra_failure(tmp_path: Path, runtime: FakeRuntime) -> None: + agent = PydanticAIAgent(runtime, endpoint=None, default_model="m", home=tmp_path / "home", workspace=tmp_path / "ws") + with TestClient(agent.create_app()) as client: + task = _send(client, [{"kind": "text", "text": "go"}]) + + _, data = _reply(task) + assert (data["error_code"], data["error_type"]) == ("pydantic_ai.no_endpoint", "infra_error") + assert runtime.turns == [] + + +def test_the_endpoint_is_read_from_the_litellm_variables() -> None: + environ = {"LITELLM_BASE_URL": "http://litellm.test", "LITELLM_API_KEY": "sk-test", + "LITELLM_EXTRA_HEADERS": '{"x-project": "p1"}'} + without_headers = {"LITELLM_BASE_URL": "http://litellm.test", "LITELLM_EXTRA_HEADERS": ""} + + assert Endpoint.from_environ(environ) == ENDPOINT + assert Endpoint.from_environ(without_headers) == Endpoint("http://litellm.test") + assert Endpoint.from_environ({"LITELLM_API_KEY": "sk-test"}) is None + + +@pytest.mark.parametrize("headers", ["{not json", '["x-project"]']) +def test_malformed_extra_headers_are_refused_at_startup(headers: str) -> None: + with pytest.raises(ValueError, match="LITELLM_EXTRA_HEADERS"): + Endpoint.from_environ({"LITELLM_BASE_URL": "http://litellm.test", "LITELLM_EXTRA_HEADERS": headers}) + + +def test_only_http_file_uris_are_fetched(client: TestClient, runtime: FakeRuntime) -> None: + task = _send(client, [{"kind": "text", "text": "read it"}, + {"kind": "file", "file": {"name": "passwd", "uri": "file:///etc/passwd"}}]) + + text, data = _reply(task) + assert task["status"]["state"] == "failed" + assert data["error_code"] == "pydantic_ai.attachment" + assert "unsupported file URI scheme" in text + assert runtime.turns == [] + + +def test_an_oversized_attachment_is_refused( + client: TestClient, runtime: FakeRuntime, monkeypatch: pytest.MonkeyPatch +) -> None: + real_client = httpx.AsyncClient + monkeypatch.setattr(pydantic_ai_agent, "MAX_ATTACHMENT_BYTES", 8) + monkeypatch.setattr(pydantic_ai_agent.httpx, "AsyncClient", lambda **kwargs: real_client( + transport=httpx.MockTransport(lambda request: httpx.Response(200, content=b"x" * 64)), **kwargs)) + + task = _send(client, [{"kind": "text", "text": "read it"}, + {"kind": "file", "file": {"name": "big.txt", "uri": "https://files.test/big.txt"}}]) + + assert _reply(task)[1]["error_code"] == "pydantic_ai.attachment" + assert runtime.turns == [] + + +def _write() -> dict[str, Any]: + return {"media_type": "application/octet-stream", "max_bytes": 1 << 20, "write": GRANT} + + +def _read(name: str) -> dict[str, Any]: + return {"media_type": "application/octet-stream", "max_bytes": 1 << 20, + "read": {**READ, "url": f"https://objects.test/{name}"}} + + +def _attached(text: str) -> Path: + return Path(text.split(" saved at ")[1].rstrip("]")) + + +@pytest.fixture +def objects(monkeypatch: pytest.MonkeyPatch) -> dict[str, bytes]: + """An in-memory object store behind the transfer helpers the agent calls.""" + stored: dict[str, bytes] = {} + names = iter(["trajectory", "workspace"]) + + class FakeUploaded(dict): + def model_dump(self, mode: str = "json") -> dict: + return dict(self) + + async def fake_upload(target, source: Path): + body = source.read_bytes() + stored[next(names)] = body + return FakeUploaded(size_bytes=len(body)) + + async def fake_download(source, destination: Path) -> None: + destination.parent.mkdir(parents=True, exist_ok=True) + destination.write_bytes(stored[source.read.url.rsplit("/", 1)[-1]]) + + monkeypatch.setattr(pydantic_ai_agent, "upload", fake_upload) + monkeypatch.setattr(pydantic_ai_agent, "download", fake_download) + return stored + + +def test_snapshot_round_trip_restores_the_conversation_and_workspace_into_the_target_context( + client: TestClient, tmp_path: Path, objects: dict[str, bytes] +) -> None: + _send(client, [{"kind": "text", "text": "remember TOKEN-1"}], context_id="source") + (tmp_path / "ws" / "notes.txt").write_text("remember TOKEN-1") + saved = client.post("/ext/snapshot", json={"context_id": "source", "objects": { + "trajectory": _write(), "workspace": _write()}}) + assert saved.status_code == 200, saved.text + (tmp_path / "ws" / "notes.txt").unlink() + (tmp_path / "ws" / "after.txt").write_text("written after the snapshot") + + loaded = client.put("/ext/snapshot", json={"objects": { + "trajectory": _read("trajectory"), "workspace": _read("workspace")}, "target_context_id": "target"}) + task = _send(client, [{"kind": "text", "text": "what was the token?"}], context_id="target") + + target = conversation(tmp_path / "home", "target") + assert loaded.json() == {"context_id": "target"} + assert (tmp_path / "ws" / "notes.txt").read_text() == "remember TOKEN-1" + assert not (tmp_path / "ws" / "after.txt").exists() + assert task["status"]["state"] == "completed" + assert history(target) == ["remember TOKEN-1", "what was the token?"] + assert sorted(path.name for path in target.iterdir()) == ["messages.json"] + + +def test_a_snapshot_loaded_on_an_agent_with_another_home_names_its_attached_files_there( + client: TestClient, runtime: FakeRuntime, tmp_path: Path, objects: dict[str, bytes] +) -> None: + notes = base64.b64encode(b"TOKEN-8").decode() + _send(client, [{"kind": "text", "text": "keep this"}, + {"kind": "file", "file": {"name": "notes.txt", "mimeType": "text/plain", "bytes": notes}}], + context_id="source") + attached = _attached(runtime.turns[0].text) + client.post("/ext/snapshot", json={"context_id": "source", "objects": {"trajectory": _write()}}) + attached.write_bytes(b"changed after the snapshot") + other = PydanticAIAgent(runtime, endpoint=ENDPOINT, home=tmp_path / "other", workspace=tmp_path / "other-ws") + + with TestClient(other.create_app()) as other_client: + loaded = other_client.put("/ext/snapshot", json={"objects": {"trajectory": _read("trajectory")}, + "target_context_id": "target"}) + + target = conversation(tmp_path / "other", "target") + (message,) = history(target) + assert loaded.status_code == 200 + assert _attached(message) == target / attached.relative_to(conversation(tmp_path / "home", "source")) + assert _attached(message).read_bytes() == b"TOKEN-8" + + +@pytest.mark.parametrize("member, content", [ + ("README", b""), + ("messages.json", b"{not json"), + ("messages.json", b'{"kind": "request"}'), +], ids=["no-messages", "unparseable", "not-a-list"]) +def test_a_snapshot_that_holds_no_conversation_changes_nothing( + client: TestClient, tmp_path: Path, objects: dict[str, bytes], member: str, content: bytes +) -> None: + _send(client, [{"kind": "text", "text": "remember TOKEN-1"}], context_id="source") + client.post("/ext/snapshot", json={"context_id": "source", "objects": {"trajectory": _write(), "workspace": _write()}}) + (tmp_path / "ws" / "notes.txt").write_text("current") + not_a_conversation = io.BytesIO() + with tarfile.open(fileobj=not_a_conversation, mode="w:gz") as tar: + info = tarfile.TarInfo(member) + info.size = len(content) + tar.addfile(info, io.BytesIO(content)) + objects["trajectory"] = not_a_conversation.getvalue() + source = conversation(tmp_path / "home", "source") + + loaded = client.put("/ext/snapshot", json={"objects": { + "trajectory": _read("trajectory"), "workspace": _read("workspace")}, "target_context_id": "source"}) + + assert loaded.status_code == 400 + assert "holds no Pydantic AI conversation" in loaded.json()["detail"] + assert [path.name for path in (tmp_path / "ws").iterdir()] == ["notes.txt"] + assert [path.name for path in source.parent.iterdir()] == [source.name] + assert history(source) == ["remember TOKEN-1"] + + +def test_a_snapshot_of_an_unknown_context_is_a_json_404(client: TestClient, objects: dict[str, bytes]) -> None: + response = client.post("/ext/snapshot", json={"context_id": "nope", "objects": {"trajectory": _write()}}) + + assert response.status_code == 404 + assert response.json() == {"detail": "No conversation for context 'nope'"} + + +def test_a_snapshot_load_is_refused_while_a_task_runs( + client: TestClient, runtime: FakeRuntime, objects: dict[str, bytes] +) -> None: + _send(client, [{"kind": "text", "text": "remember TOKEN-1"}], context_id="source") + client.post("/ext/snapshot", json={"context_id": "source", "objects": {"trajectory": _write()}}) + runtime.gate = threading.Event() + running = _send_in_background(client, "keep working", "busy") + _wait_until(lambda: sum(runtime.active.values()) == 1) + + busy = client.put("/ext/snapshot", json={"objects": {"trajectory": _read("trajectory")}, "target_context_id": "t"}) + runtime.gate.set() + running.join() + idle = client.put("/ext/snapshot", json={"objects": {"trajectory": _read("trajectory")}, "target_context_id": "t"}) + + assert (busy.status_code, idle.status_code) == (409, 200) + + +def test_a_snapshot_save_waits_for_its_context_running_task( + client: TestClient, runtime: FakeRuntime, objects: dict[str, bytes] +) -> None: + _send(client, [{"kind": "text", "text": "plant"}], context_id="ctx-a") + runtime.gate = threading.Event() + running = _send_in_background(client, "more", "ctx-a") + _wait_until(lambda: sum(runtime.active.values()) == 1) + saved: dict[str, httpx.Response] = {} + saver = threading.Thread(target=lambda: saved.update(response=client.post( + "/ext/snapshot", json={"context_id": "ctx-a", "objects": {"trajectory": _write()}}))) + saver.start() + time.sleep(0.3) + finished_while_running = "response" in saved + runtime.gate.set() + running.join() + saver.join() + + with tarfile.open(fileobj=io.BytesIO(objects["trajectory"])) as tar: + messages = json.load(tar.extractfile("messages.json")) + assert not finished_while_running + assert saved["response"].status_code == 200 + assert messages == ["plant", "more"] + + +def test_conversation_locks_are_dropped_once_released( + agent: PydanticAIAgent, client: TestClient, objects: dict[str, bytes] +) -> None: + for context in ("one", "two", "three"): + _send(client, [{"kind": "text", "text": "go"}], context_id=context) + client.post("/ext/snapshot", json={"context_id": "one", "objects": {"trajectory": _write()}}) + + assert agent._conversation_locks == {} + + +def test_the_card_advertises_the_extensions_agent_env_validates(client: TestClient) -> None: + card = client.get("/.well-known/agent.json").json() + uris = {extension["uri"] for extension in card["capabilities"]["extensions"]} + config = next(e for e in card["capabilities"]["extensions"] if e["uri"] == "urn:agentenv:agent-config/v1") + + assert {f"urn:agentenv:{name}/v1" for name in ( + "agent-config", "mcp-config", "skill-config", "trajectory", "snapshot", "peer-agents", "install", "triggers", + )} <= uris + assert {"model", "system_prompt", "role", "timeout_seconds"} <= set( + config["params"]["methods"]["set"]["request"]["supported"]) + assert card["defaultInputModes"] == ["text", "image/png", "image/jpeg", "image/gif", "image/webp"] + + +def test_install_extension_is_declared_on_the_card(client: TestClient) -> None: + card = client.get("/.well-known/agent.json").json() + install = next(e for e in card["capabilities"]["extensions"] if e["uri"] == "urn:agentenv:install/v1")["params"] + + assert install["a2a_port"] == 8000 + assert set(install["required_params"]) == {"container", "agent_ctx_tar", "a2a_port", "litellm_api_key", + "litellm_base_url", "litellm_extra_headers", "trust_env"} + formatted = [command.format(container="c", agent_ctx_tar="/t.tar", a2a_port="8000", litellm_api_key="'sk-secret'", + litellm_base_url="u", litellm_extra_headers="'{}'", trust_env="''") + for command in install["install_commands"]] + assert formatted[-1].endswith("c sh /opt/pydantic-ai-a2a/start.sh") + assert all("sk-secret" not in command.splitlines()[0] for command in formatted) + assert "LITELLM_API_KEY='sk-secret'" in formatted[-2].splitlines() + + +def test_the_installed_agent_starts_with_the_endpoint_and_trust_lines_install_wrote(tmp_path: Path) -> None: + here = tmp_path / "pydantic-ai-a2a" + here.mkdir() + (here / "start.sh").write_text((AGENT_DIR / "start.sh").read_text().replace("/tmp/", f"{tmp_path}/")) + trust = "SSL_CERT_FILE='/etc/agentenv/ca-bundle.pem'\nREQUESTS_CA_BUNDLE='/etc/agentenv/ca-bundle.pem'\n" + (here / "agent.env").write_text(f"LITELLM_API_KEY={shlex.quote('sk x')}\nAGENT_TRUST_ENV={shlex.quote(trust)}\n") + # In place of the server: what it would be started with, written beside it. + (here / "a2a_server.py").write_text( + "import json, os, pathlib\npathlib.Path(__file__).with_name('seen').write_text(json.dumps(" + "[os.environ.get(name) for name in ('LITELLM_API_KEY', 'SSL_CERT_FILE', 'REQUESTS_CA_BUNDLE')]))\n") + python = tmp_path / "bin" / "python" + python.parent.mkdir() + python.symlink_to(sys.executable) + + subprocess.run(["sh", str(here / "start.sh")], cwd="/", env={"PATH": f"{python.parent}{os.pathsep}/usr/bin:/bin"}, + check=True, timeout=60) + + assert json.loads((here / "seen").read_text()) == ["sk x", "/etc/agentenv/ca-bundle.pem", + "/etc/agentenv/ca-bundle.pem"] + + +def test_the_card_version_is_the_pinned_pydantic_ai() -> None: + requirements = (AGENT_DIR / "requirements.txt").read_text().splitlines() + version = pydantic_ai_agent.PYDANTIC_AI_VERSION + + assert f"pydantic-ai-slim[mcp,openai]=={version}" in requirements + assert any(line.startswith("pydantic-ai-harness[") for line in requirements) + + +def test_the_image_copies_every_module_the_agent_imports() -> None: + copied = next(line for line in (AGENT_DIR / "Dockerfile").read_text().splitlines() if line.startswith("COPY ")) + sources = {path.name for path in AGENT_DIR.iterdir() if path.suffix in {".py", ".sh", ".txt"}} + + assert set(copied.split()[1:-1]) == sources diff --git a/tst/unit/agents/test_installed_workspace.py b/tst/unit/agents/test_installed_workspace.py index 168f36d0..a86f17e0 100644 --- a/tst/unit/agents/test_installed_workspace.py +++ b/tst/unit/agents/test_installed_workspace.py @@ -13,7 +13,8 @@ AGENTS = Path(agent_env.__file__).parent / "agents" / "agents" -@pytest.mark.parametrize("agent, variable", [("pi", "PI_A2A_WORKSPACE"), ("openhands", "OPENHANDS_A2A_WORKSPACE")]) +@pytest.mark.parametrize("agent, variable", [ + ("pi", "PI_A2A_WORKSPACE"), ("openhands", "OPENHANDS_A2A_WORKSPACE"), ("pydantic-ai", "PYDANTIC_AI_A2A_WORKSPACE")]) def test_an_installed_agent_works_in_the_directory_it_is_started_in_unless_that_is_the_root(tmp_path, agent, variable): here = tmp_path / f"{agent}-a2a" here.mkdir() diff --git a/tst/unit/bundle/shipped_bundles_test.py b/tst/unit/bundle/shipped_bundles_test.py index f918f524..ee02c4af 100644 --- a/tst/unit/bundle/shipped_bundles_test.py +++ b/tst/unit/bundle/shipped_bundles_test.py @@ -45,7 +45,8 @@ def test_the_shipped_aliases_are_agent_envs_alone_and_claude_code_names_the_defa assert {alias: named.target_id for alias, named in table.aliases.items()} == { "claude-code": "@local/agentenv-framework/agents/claude-code", "pi": "@local/agentenv-framework/agents/pi", "openhands": "@local/agentenv-framework/agents/openhands", "codex": "@local/agentenv-framework/agents/codex", - "gemini": "@local/agentenv-framework/agents/gemini", "opencode": "@local/agentenv-framework/agents/opencode"} + "gemini": "@local/agentenv-framework/agents/gemini", "opencode": "@local/agentenv-framework/agents/opencode", + "pydantic-ai": "@local/agentenv-framework/agents/pydantic-ai"} assert table.target("claude-code").target_id == get_config().get_default_a2a_agent_id() diff --git a/tst/unit/test_infra_images_script.py b/tst/unit/test_infra_images_script.py index 53ad217e..21c5d150 100644 --- a/tst/unit/test_infra_images_script.py +++ b/tst/unit/test_infra_images_script.py @@ -62,7 +62,7 @@ def test_plan_only_lists_the_infra_images_or_the_agents_images(script, monkeypat "agents": agents} assert "agentenv-gateway" in planned["infra"] and agents == [ "agentenv-claude-code", "agentenv-codex", "agentenv-gemini", "agentenv-opencode", "agentenv-openhands", - "agentenv-pi", "agentenv-proxy"] + "agentenv-pi", "agentenv-proxy", "agentenv-pydantic-ai"] def test_plan_missing_lists_the_images_the_registry_lacks_or_refuses_to_show(script, monkeypatch, capsys): diff --git a/tst/util/fake_openai.py b/tst/util/fake_openai.py index 142d41a1..4f47fbc6 100644 --- a/tst/util/fake_openai.py +++ b/tst/util/fake_openai.py @@ -1,6 +1,7 @@ """A fake OpenAI Chat Completions endpoint, so an agent that calls the model through LiteLLM runs with no model key: -``POST .../chat/completions`` is answered by a responder, as one JSON completion, and every request is recorded. A -fault queued with ``fail_next`` replaces the next answer: an HTTP error, or an answer that stalls.""" +``POST .../chat/completions`` is answered by a responder, as one JSON completion or, to a request that asks for a +stream, as server-sent chunks, and every request is recorded. A fault queued with ``fail_next`` replaces the next +answer: an HTTP error, or an answer that stalls.""" import json import threading @@ -117,7 +118,22 @@ def do_POST(self) -> None: return if fault is not None and fault.kind == "stall": time.sleep(fault.seconds) - self._send(200, _completion(body, fake.respond(body))) + completion = _completion(body, fake.respond(body)) + if body.get("stream"): + self._stream(completion) + else: + self._send(200, completion) + + def _stream(self, completion: dict) -> None: + self.send_response(200) + self.send_header("content-type", "text/event-stream") + self.end_headers() + try: + for chunk in _chunks(completion): + self.wfile.write(f"data: {json.dumps(chunk)}\n\n".encode()) + self.wfile.write(b"data: [DONE]\n\n") + except (BrokenPipeError, ConnectionResetError): + pass def _send(self, status: int, payload: dict) -> None: data = json.dumps(payload).encode() @@ -140,6 +156,21 @@ def __exit__(self, *exc) -> None: self._server.server_close() +def _chunks(completion: dict) -> list[dict]: + """``completion`` as the chunks a stream sends: its message whole as one delta, its finish, then its usage.""" + (choice,) = completion["choices"] + message = choice["message"] + delta = {"role": "assistant", "content": message["content"]} + if "tool_calls" in message: + delta["tool_calls"] = [{"index": index, **call} for index, call in enumerate(message["tool_calls"])] + head = {key: completion[key] for key in ("id", "created", "model")} | {"object": "chat.completion.chunk"} + return [ + {**head, "choices": [{"index": 0, "delta": delta, "finish_reason": None}]}, + {**head, "choices": [{"index": 0, "delta": {}, "finish_reason": choice["finish_reason"]}]}, + {**head, "choices": [], "usage": completion["usage"]}, + ] + + def _completion(body: dict, reply: Text | ToolCall) -> dict: if isinstance(reply, Text): message, finish = {"role": "assistant", "content": reply.text}, "stop" From 174e7ecf7e639659a945827e3dea224e1aa8efe9 Mon Sep 17 00:00:00 2001 From: mpfaffenberger Date: Mon, 26 Oct 2026 15:34:08 -0700 Subject: [PATCH 2/2] fix(agents): validate a snapshot's Pydantic AI messages, and finish threaded writes before a cancel lets go of the locks A snapshot whose messages.json was JSON but not Pydantic AI messages replaced the conversation and workspace, then failed the next turn; the runtime now validates the messages before anything is replaced. A turn or snapshot cancelled while a thread wrote its files released its locks before the write ended; to_completion waits for the thread before re-raising. --- .../agents/agents/pydantic-ai/README.md | 2 +- .../agents/pydantic-ai/pydantic_ai_agent.py | 35 ++++++++++++++----- .../agents/pydantic-ai/pydantic_ai_runtime.py | 8 +++-- .../pydantic_ai/test_pydantic_ai_agent.py | 31 +++++++++++++++- 4 files changed, 63 insertions(+), 13 deletions(-) diff --git a/src/agent_env/agents/agents/pydantic-ai/README.md b/src/agent_env/agents/agents/pydantic-ai/README.md index 82c43887..c319f807 100644 --- a/src/agent_env/agents/agents/pydantic-ai/README.md +++ b/src/agent_env/agents/agents/pydantic-ai/README.md @@ -34,7 +34,7 @@ snapshots carry it. | `mcp-config/v1` | SDK default; each server becomes an `MCPToolset` on its URL, with its headers, its tools named as the server names them | | `skill-config/v1` | inline and bundle skills installed as `SKILL.md` directories under `~/.pydantic-ai-a2a/skills`, offered to the model through the harness's `Skills`, which it loads with `load_capability` | | `trajectory/v1` | live: the task's Pydantic AI messages (`ModelMessagesTypeAdapter`, format `pydantic-ai-messages/v1`), appended as each request and response completes | -| `snapshot/v1` | `save` uploads the context's conversation directory, its messages and attached files, as a gzipped tar and, when asked, the workspace as another; `load` extracts both aside, checks the messages and points the paths of the attached files at the target's directory before it replaces anything, then replaces the workspace's contents and installs the conversation under the target context's id, so an archive that fails to extract, or holds no conversation, changes nothing | +| `snapshot/v1` | `save` uploads the context's conversation directory, its messages and attached files, as a gzipped tar and, when asked, the workspace as another; `load` extracts both aside, validates the messages as Pydantic AI's and points the paths of the attached files at the target's directory before it replaces anything, then replaces the workspace's contents and installs the conversation under the target context's id, so an archive that fails to extract, or holds no conversation, changes nothing | | `peer-agents/v1` | SDK default: the `peer_agents` MCP server reaches the agent like any other | | `install/v1` | copies this directory into the task container as `/opt/pydantic-ai-a2a`, runs `install.sh` (Debian or Ubuntu, as root, with network access), writes the LiteLLM credentials and the lines that trust the local transfer CA to a 0600 `agent.env` through a heredoc and starts `start.sh`, which applies them | | `triggers/v1` | SDK default | diff --git a/src/agent_env/agents/agents/pydantic-ai/pydantic_ai_agent.py b/src/agent_env/agents/agents/pydantic-ai/pydantic_ai_agent.py index 446de1d2..74d31ad5 100644 --- a/src/agent_env/agents/agents/pydantic-ai/pydantic_ai_agent.py +++ b/src/agent_env/agents/agents/pydantic-ai/pydantic_ai_agent.py @@ -174,6 +174,22 @@ async def stop(self) -> None: """Cancel every turn still running: the server is shutting down.""" ... + def check_messages(self, data: bytes) -> None: + """Raise ValueError unless ``data`` is a conversation's messages as the runtime saves them, so a snapshot that + holds anything else is refused before it replaces a conversation.""" + ... + + +async def to_completion(function: Callable[..., Any], *args: Any, **kwargs: Any) -> Any: + """``function(*args, **kwargs)`` in a thread, run to its end: cancelled meanwhile, this waits for the thread before it + re-raises, so the locks its caller holds stay held until the files the thread writes are written.""" + work = asyncio.ensure_future(asyncio.to_thread(function, *args, **kwargs)) + try: + return await asyncio.shield(work) + except asyncio.CancelledError: + await asyncio.wait([work]) + raise + async def _detail_body(_request: Request, exc: HTTPException) -> Response: """An ``HTTPException`` as the ``{"detail": ...}`` body agent-env reads errors from.""" @@ -274,14 +290,15 @@ def _replace_contents(directory: Path, staging: Path) -> None: staging.rmdir() -def _stage_conversation(archive: Path, conversations_dir: Path, target: Path) -> Path: +def _stage_conversation( + archive: Path, conversations_dir: Path, target: Path, check: Callable[[bytes], None] +) -> Path: """The archived conversation extracted beside the others, its messages naming the files attached in it under - ``target``, where it goes, rather than where it was saved from. Raises ValueError for an archive that holds no - conversation.""" + ``target``, where it goes, rather than where it was saved from. Raises ValueError for an archive whose messages + ``check`` refuses.""" staging = _extract(archive, conversations_dir) try: - if not isinstance(json.loads((staging / MESSAGES).read_text()), list): - raise ValueError("its messages are not a list") + check((staging / MESSAGES).read_bytes()) source = staging / SOURCE saved_from = json.loads(source.read_text())["directory"] if source.is_file() else None except (OSError, ValueError, KeyError, TypeError) as exc: @@ -470,13 +487,13 @@ async def save_snapshot(self, request: ObjectSnapshotSaveRequest) -> dict[str, A async with self._holding(conversation), self._workspace_lock: if not (directory / MESSAGES).is_file(): raise HTTPException(status_code=404, detail=f"No conversation for context '{request.context_id}'") - await asyncio.to_thread(_pack, directory, conversation_archive, record_source=True) + await to_completion(_pack, directory, conversation_archive, record_source=True) if request.objects.workspace is not None: # Tasks in other contexts work in the same workspace: those running end first, and new ones wait # for the lock this holds. await self._idle.wait() self.workspace.mkdir(parents=True, exist_ok=True) - await asyncio.to_thread(_pack, self.workspace, workspace_archive) + await to_completion(_pack, self.workspace, workspace_archive) objects = {"trajectory": (await upload(request.objects.trajectory, conversation_archive)).model_dump(mode="json")} if request.objects.workspace is not None: objects["workspace"] = (await upload(request.objects.workspace, workspace_archive)).model_dump(mode="json") @@ -497,7 +514,7 @@ async def load_snapshot(self, request: ObjectSnapshotLoadRequest) -> dict[str, A raise HTTPException(status_code=409, detail="a task is running; restore the workspace once it finishes") target = self.conversations_dir / conversation_id(context_id) try: - await asyncio.to_thread(self._restore, conversation_archive, workspace_archive, target) + await to_completion(self._restore, conversation_archive, workspace_archive, target) except ValueError as exc: raise HTTPException(status_code=400, detail=str(exc)) from exc return {"context_id": context_id} @@ -506,7 +523,7 @@ def _restore(self, conversation_archive: Path, workspace_archive: Path | None, t """Extract both archives and check the conversation before putting either in place: one that fails to extract, or holds no conversation, changes nothing.""" self.conversations_dir.mkdir(parents=True, exist_ok=True) - conversation = _stage_conversation(conversation_archive, self.conversations_dir, target) + conversation = _stage_conversation(conversation_archive, self.conversations_dir, target, self.runtime.check_messages) staged = [conversation] try: if workspace_archive is not None: diff --git a/src/agent_env/agents/agents/pydantic-ai/pydantic_ai_runtime.py b/src/agent_env/agents/agents/pydantic-ai/pydantic_ai_runtime.py index b53166ee..d96e065f 100644 --- a/src/agent_env/agents/agents/pydantic-ai/pydantic_ai_runtime.py +++ b/src/agent_env/agents/agents/pydantic-ai/pydantic_ai_runtime.py @@ -20,7 +20,7 @@ from pydantic_ai_harness import Skills from pydantic_ai_harness.coder import Coder -from pydantic_ai_agent import MESSAGES, Outcome, Turn +from pydantic_ai_agent import MESSAGES, Outcome, Turn, to_completion # Consecutive retries a tool gets for arguments it refuses, as many as Coder gives its file tools, so a model that # fumbles a command twice doesn't end the task. @@ -124,6 +124,9 @@ async def stop(self) -> None: run.cancel() await asyncio.gather(*runs, return_exceptions=True) + def check_messages(self, data: bytes) -> None: + ModelMessagesTypeAdapter.validate_json(data) + async def run(self, turn: Turn) -> Outcome: task = asyncio.current_task() self._runs.add(task) @@ -162,5 +165,6 @@ async def _run(self, turn: Turn) -> Outcome: messages = result.all_messages() for event in _events(messages[recorded:]): turn.on_event(event) - await asyncio.to_thread(_save, turn, messages) + # A turn whose run finished is saved whole even if it is cancelled meanwhile, before its locks are let go. + await to_completion(_save, turn, messages) return Outcome(result.output.strip() or "The agent finished without a reply.", _usage(result.usage)) diff --git a/tst/unit/agents/pydantic_ai/test_pydantic_ai_agent.py b/tst/unit/agents/pydantic_ai/test_pydantic_ai_agent.py index 32a260dc..80ebcf1e 100644 --- a/tst/unit/agents/pydantic_ai/test_pydantic_ai_agent.py +++ b/tst/unit/agents/pydantic_ai/test_pydantic_ai_agent.py @@ -49,6 +49,11 @@ class FakeRuntime: async def stop(self) -> None: self.stopped = True + def check_messages(self, data: bytes) -> None: + messages = json.loads(data) + if not isinstance(messages, list) or not all(isinstance(message, str) for message in messages): + raise ValueError("not this runtime's messages") + async def run(self, turn: Any) -> Outcome: self.turns.append(turn) turn.on_event({"kind": "request", "parts": [{"part_kind": "user-prompt", "content": turn.text}]}) @@ -474,7 +479,8 @@ def test_a_snapshot_loaded_on_an_agent_with_another_home_names_its_attached_file ("README", b""), ("messages.json", b"{not json"), ("messages.json", b'{"kind": "request"}'), -], ids=["no-messages", "unparseable", "not-a-list"]) + ("messages.json", b"[{}]"), +], ids=["no-messages", "unparseable", "not-a-list", "not-the-runtimes-messages"]) def test_a_snapshot_that_holds_no_conversation_changes_nothing( client: TestClient, tmp_path: Path, objects: dict[str, bytes], member: str, content: bytes ) -> None: @@ -619,3 +625,26 @@ def test_the_image_copies_every_module_the_agent_imports() -> None: sources = {path.name for path in AGENT_DIR.iterdir() if path.suffix in {".py", ".sh", ".txt"}} assert set(copied.split()[1:-1]) == sources + + +def test_a_threaded_write_cancelled_midway_finishes_before_the_cancel_reaches_its_caller() -> None: + started, release, finished = threading.Event(), threading.Event(), threading.Event() + + def write() -> None: + started.set() + release.wait(5) + finished.set() + + async def cancel_midway() -> bool: + task = asyncio.ensure_future(pydantic_ai_agent.to_completion(write)) + await asyncio.to_thread(started.wait, 5) + task.cancel() + await asyncio.sleep(0.1) + still_waiting = not task.done() + release.set() + with pytest.raises(asyncio.CancelledError): + await task + return still_waiting + + assert asyncio.run(cancel_midway()) + assert finished.is_set()