From d5beb5001a410f6be734ca77913c833ad4e6f80b Mon Sep 17 00:00:00 2001 From: George Date: Wed, 28 Oct 2026 11:01:36 -0700 Subject: [PATCH 1/7] harnesses: pi, over `pi --mode rpc`, with an extension as its approval gate --- README.md | 4 +- .../content/docs/basics/harnesses.mdx | 114 +- docs/ai-python/content/docs/basics/index.mdx | 2 +- .../content/docs/reference/harnesses.mdx | 46 +- examples/harnesses/tui.py | 7 +- skills/ai/SKILL.md | 4 +- skills/ai/references/harnesses.md | 24 +- src/ai/harnesses/experimental/__init__.py | 5 +- src/ai/harnesses/experimental/_adapters/pi.py | 1412 +++++++++++++++++ .../experimental/_adapters/pi_extension.js | 88 + src/ai/harnesses/experimental/_factories.py | 74 + src/ai/harnesses/experimental/_handle.py | 2 +- src/ai/harnesses/experimental/_handoff.py | 130 +- tests/harnesses/experimental/conftest.py | 16 +- .../experimental/e2e/test_handoff.py | 2 +- .../experimental/e2e/test_options.py | 5 +- tests/harnesses/experimental/e2e/test_pty.py | 2 + .../experimental/e2e/test_reconnect.py | 14 +- tests/harnesses/experimental/test_extras.py | 6 + tests/harnesses/experimental/test_options.py | 9 +- .../harnesses/experimental/test_pi_adapter.py | 957 +++++++++++ 21 files changed, 2837 insertions(+), 86 deletions(-) create mode 100644 src/ai/harnesses/experimental/_adapters/pi.py create mode 100644 src/ai/harnesses/experimental/_adapters/pi_extension.js create mode 100644 tests/harnesses/experimental/test_pi_adapter.py diff --git a/README.md b/README.md index 0f35eb25..f69d7aa2 100644 --- a/README.md +++ b/README.md @@ -161,7 +161,7 @@ ai.resolve_hook("approve_send_email", {"granted": True, "reason": "approved"}) ## Harnesses (experimental) -Harnesses drive coding-agent CLIs you already have, Claude Code and Codex, +Harnesses drive coding-agent CLIs you already have, Claude Code, Codex and Pi, on this machine or in a Vercel Sandbox. A turn streams the same `ai.events` as an agent: @@ -185,7 +185,7 @@ Focused samples live in category directories under `examples/`. - `examples/media/` - image, video, speech, transcription, embeddings, reranking, and multimodal input/output - `examples/models/` - streaming, structured output, and provider examples -- `examples/harnesses/` - driving Claude Code and Codex, locally and in a +- `examples/harnesses/` - driving Claude Code, Codex and Pi, locally and in a sandbox (experimental) - `examples/apps/` - end-to-end demos diff --git a/docs/ai-python/content/docs/basics/harnesses.mdx b/docs/ai-python/content/docs/basics/harnesses.mdx index e5f673cc..6818c789 100644 --- a/docs/ai-python/content/docs/basics/harnesses.mdx +++ b/docs/ai-python/content/docs/basics/harnesses.mdx @@ -1,13 +1,14 @@ --- title: Harnesses -description: Drive Claude Code and Codex from Python. +description: Drive Claude Code, Codex, and Pi from Python. type: guide summary: Run turns on coding-agent CLIs, stream events, steer, stop, approve tool calls, resume conversations, and open the CLI's own TUI. --- -A harness drives a coding-agent CLI you already have, `claude` or `codex`. +A harness drives a coding-agent CLI you already have: `claude`, `codex`, or +`pi`. Your program sends prompts. The CLI does the agent work with its own tools, -login, and session store. The same code runs on either CLI, on this machine or +login, and session store. The same code runs on any of them, on this machine or in a [workspace](/docs/basics/workspaces) such as a Vercel Sandbox. This API is experimental and may change or be removed. @@ -35,8 +36,8 @@ if __name__ == "__main__": asyncio.run(main()) ``` -`claude_code()` and `codex()` return a `Harness` and share these arguments, -so swapping one for the other is a one-word edit: +`claude_code()`, `codex()`, and `pi()` return a `Harness` and share these +arguments, so swapping one for another is a one-word edit: ```python harness = harnesses.codex(workspace=workspaces.Local("."), effort="high") @@ -50,7 +51,7 @@ harness = harnesses.codex(workspace=workspaces.Local("."), effort="high") | `model` | Model name passed to the CLI. | `None`: the CLI's default | | `effort` | `"low"`, `"medium"`, `"high"`, or `"max"`. | `None`: the CLI's default | | `gateway` | Upstream model provider LLM Gateway. See [Workspaces](/docs/basics/workspaces#reach-a-model-through-ai-gateway). | `None`: the workspace's gateway | -| `executable` | CLI name or path. | `"claude"` or `"codex"` | +| `executable` | CLI name or path. | `"claude"`, `"codex"`, or `"pi"` | | `instructions` | Text added to the CLI's own system prompt, for every turn. | `None` | | `mcp_servers` | MCP servers to connect: `{"name": {"command": ..., "args": [...], "env": {...}}}`, or `{"name": {"url": ...}}` for a remote one. | `None` | @@ -81,8 +82,8 @@ A missing extra raises `ai.errors.InstallationError` when you call the factory. A missing CLI raises `ExecutableMissingError` when the harness opens. -Each harness was verified against one CLI version: `claude` 2.1.289 and -`codex` 0.160.1. An older CLI warns when the harness opens. A harness error +Each harness was verified against one CLI version: `claude` 2.1.289, +`codex` 0.160.1, and `pi` 1.3.0. An older CLI warns when the harness opens. A harness error carries the CLI version in its notes, so a failure on a newer CLI says so. On your own machine the SDK uses whatever CLI you installed and never @@ -137,6 +138,38 @@ It takes the [shared arguments](#run-a-prompt), plus: `fileChange`. The shell command is in the call's arguments, never the name. - `effort="max"` is sent to Codex as `"xhigh"`. +### Pi + +`harnesses.pi()` drives the `pi` CLI over `pi --mode rpc`. It needs no +extra, only the `pi` CLI (`npm install -g @earendil-works/pi-coding-agent`). + +It takes the [shared arguments](#run-a-prompt), plus Claude Code's `tools` +and `disallowed_tools`: + +| Argument | Description | Default | +|---|---|---| +| `tools` | The tools the agent has, such as `["read", "grep"]`. `[]` for none. | `None`: pi's defaults | +| `disallowed_tools` | Tools taken away, such as `["bash"]`. | `None` | + +- `model` is any model pi knows, such as `"sonnet"` or + `"anthropic/claude-sonnet-4-5"`. With a `gateway`, it is the gateway's + model, such as `"anthropic/claude-sonnet-4.5"`. Without a `model`, pi + chooses exactly as it does on its own: your settings' default, or a + default of a provider it has credentials for, which may not be the + gateway. With no credentials at all, the turn raises + `NotAuthenticatedError`. +- `effort` is pi's thinking level. pi clamps it to what the model supports. +- `instructions` are appended to pi's system prompt. +- Your own pi extensions, skills, and settings load as they always do. +- pi has no permission prompt. The SDK loads a small extension into pi that + sends every tool call to your hook. See [Approve tool calls](#approve-tool-calls). +- `pi()` also accepts `max_turns`, but raises `UnsupportedError` when it is + set: pi has no such limit. +- Tool names are pi's own, such as `bash`, `read`, `edit`, and `write`. +- Context a pi extension injects into a turn (`sendMessage`) is in + `history()`, but not in the live `session.messages` of that turn. +- `result.usage.raw["cost_usd"]` reports the turn's cost. + ## Keep a conversation `harness.run()` uses a throwaway conversation. Use a session for several turns @@ -175,11 +208,11 @@ result = await session.run( ) ``` -| File | Claude Code | Codex | -|---|---|---| -| Images (`image/*`) | yes | yes | -| PDFs (`application/pdf`) | yes | no | -| Text files (`text/*`) | yes, as text | yes, as text | +| File | Claude Code | Codex | Pi | +|---|---|---|---| +| Images (`image/*`) | yes | yes | yes | +| PDFs (`application/pdf`) | yes | no | no | +| Text files (`text/*`) | yes, as text | yes, as text | yes, as text | - A file the harness cannot take raises `UnsupportedError` when you call `run()` or `stream()`, before anything runs. Check @@ -260,7 +293,7 @@ task waits for the first. Pass `approve=` to decide which tool calls run: ```python -SHELL = {"Bash", "commandExecution"} # Claude Code's and Codex's shell tools +SHELL = {"Bash", "commandExecution", "bash"} # Claude Code, Codex, Pi async def approve(ctx: harnesses.ApprovalContext) -> harnesses.Decision: @@ -269,7 +302,7 @@ async def approve(ctx: harnesses.ApprovalContext) -> harnesses.Decision: if "git push" in command or "rm -rf" in command: return harnesses.Deny("not that command") return harnesses.Allow() - target = ctx.call.input.get("file_path") + target = ctx.call.input.get("file_path") or ctx.call.input.get("path") if target: real = await ctx.workspace.resolve(target) # follows symlinks if not ctx.workspace.contains(real): @@ -293,17 +326,17 @@ The hook receives an `ApprovalContext`: | `call` | The `ToolCall`: the CLI's own tool name (`call.name`) and arguments (`call.input`), untranslated. See [Harnesses](#harnesses) for each CLI's names. | | `history` | The conversation so far. | | `workspace` | The workspace the CLI runs in. | -| `harness` | `"claude-code"` or `"codex"`. | +| `harness` | `"claude-code"`, `"codex"`, or `"pi"`. | | `session_id` | The conversation's id. | -| `prompt` | The turn's prompt on Claude Code; `None` on Codex. | +| `prompt` | The turn's prompt on Claude Code; `None` on Codex and Pi. | It returns a decision: | Decision | Effect | |---|---| | `Allow()` | The call runs. | -| `Allow(input={...})` | The call runs with rewritten arguments. Claude Code only. | -| `Deny(reason)` | The call is blocked. On Claude Code the agent reads the reason; on Codex it only learns the call was refused. | +| `Allow(input={...})` | The call runs with rewritten arguments. Claude Code and Pi. | +| `Deny(reason)` | The call is blocked. On Claude Code and Pi the agent reads the reason; on Codex it only learns the call was refused. | | `Deny(reason, stop=True)` | The call is blocked and the turn ends. | - A decision the harness cannot honor, such as `Allow(input=...)` on Codex @@ -311,7 +344,7 @@ It returns a decision: `UnsupportedError` from the turn. - A hook that raises or runs past `approval_timeout` denies the call. The failure is listed in `result.approval_errors`. -- On Claude Code every tool call reaches the hook, reads included. On +- On Claude Code and Pi every tool call reaches the hook, reads included. On Codex the CLI's policy decides which calls to ask about; shell commands and file changes do. See `approval` in [Capabilities](#capabilities). - A shell command is a string the agent wrote. Matching on it, as above, is @@ -335,21 +368,21 @@ Each harness declares what it can do in `harness.capabilities`, before anything runs. Using a capability a harness lacks raises `UnsupportedError`; the one exception is `deny_reason`, whose absence only drops the reason. -| Capability | Description | Claude Code | Codex | -|---|---|---|---| -| `run` | Run a prompt and wait for the result: `run()`. | yes | yes | -| `stream` | Iterate a turn's events as they arrive: `stream()`. | yes | yes | -| `steer` | Send guidance to the turn that is running: `session.steer()`. | yes | yes | -| `stop` | Interrupt the running turn; it settles as cancelled: `session.stop()`. | yes | yes | -| `resume` | Continue a conversation this process did not create: `resume()`, `harness_from()`. | yes | yes | -| `fork` | Branch a conversation into a new one with the same past: `fork()`. | yes | yes | -| `history` | Read stored transcripts in full, tool calls included: `history()`, `sessions()`. | yes | yes | -| `images` | Images in a prompt. See [Send files](#send-files). | yes | yes | -| `documents` | PDFs in a prompt. | yes | no | -| `rewrite_tool_input` | Rewrite a call's arguments before it runs: `Allow(input={...})`. | yes | no | -| `deny_reason` | The reason in `Deny(reason)` reaches the agent. Without it the call is still denied, but the reason is dropped. | yes | no | -| `remember_decisions` | Keep an approval for the rest of the session: `Allow(remember=True)`. | no | no | -| `approval` | Which tool calls reach the hook. `"all"`: every call, reads included. `"policy"`: the CLI's own policy decides, so not every call may. | `"all"` | `"policy"` | +| Capability | Description | Claude Code | Codex | Pi | +|---|---|---|---|---| +| `run` | Run a prompt and wait for the result: `run()`. | yes | yes | yes | +| `stream` | Iterate a turn's events as they arrive: `stream()`. | yes | yes | yes | +| `steer` | Send guidance to the turn that is running: `session.steer()`. | yes | yes | yes | +| `stop` | Interrupt the running turn; it settles as cancelled: `session.stop()`. | yes | yes | yes | +| `resume` | Continue a conversation this process did not create: `resume()`, `harness_from()`. | yes | yes | yes | +| `fork` | Branch a conversation into a new one with the same past: `fork()`. | yes | yes | yes | +| `history` | Read stored transcripts in full, tool calls included: `history()`, `sessions()`. | yes | yes | yes | +| `images` | Images in a prompt. See [Send files](#send-files). | yes | yes | yes | +| `documents` | PDFs in a prompt. | yes | no | no | +| `rewrite_tool_input` | Rewrite a call's arguments before it runs: `Allow(input={...})`. | yes | no | yes | +| `deny_reason` | The reason in `Deny(reason)` reaches the agent. Without it the call is still denied, but the reason is dropped. | yes | no | yes | +| `remember_decisions` | Keep an approval for the rest of the session: `Allow(remember=True)`. | no | no | no | +| `approval` | Which tool calls reach the hook. `"all"`: every call, reads included. `"policy"`: the CLI's own policy decides, so not every call may. | `"all"` | `"policy"` | `"all"` | ## Resume and fork @@ -389,11 +422,12 @@ same = await harness.resume(infos[0].session_id) - `resume()` continues the original. Only one client may write a conversation. If another SDK client or a `harness.tui()` holds it, `resume()` raises `SessionBusyError`, and `SessionInfo.running` is `True`. -- A `claude` or `codex` that a person started in a terminal holds no claim, - so `resume()` does not refuse it. `sessions()` still shows it as +- A `claude`, `codex`, or `pi` that a person started in a terminal holds no + claim, so `resume()` does not refuse it. `sessions()` still shows it as `running`, with its `pid`: any `claude`, and a `codex` started as - `codex resume `. A fresh `codex` names no conversation, so it is not - seen. Fork conversations you did not start. + `codex resume `. A fresh `codex` names no conversation, and `pi` hides + its arguments from `ps`, so neither is seen. Fork conversations you did + not start. `harness_from(handle, fork=True)` forks instead of resuming. @@ -411,7 +445,7 @@ session = sandbox_agent.session(history=history) The receiving CLI stores the history natively before the first turn. Tool calls keep their original names. Two limits: -- Reasoning cannot go into Claude Code. It is dropped with a warning. +- Reasoning cannot go into Claude Code or Pi. It is dropped with a warning. - Paths point to the source machine. Copy the files and tell the agent where they are now. diff --git a/docs/ai-python/content/docs/basics/index.mdx b/docs/ai-python/content/docs/basics/index.mdx index 3d5a2b92..55630da4 100644 --- a/docs/ai-python/content/docs/basics/index.mdx +++ b/docs/ai-python/content/docs/basics/index.mdx @@ -22,7 +22,7 @@ Start with the model call, then add the pieces your app needs: - Pause tools for approval with hooks. - Run subagents as tools and stream their output. - Bridge agent streams to AI SDK UI clients over Server-Sent Events (SSE). -- Drive Claude Code and Codex from Python, locally or in a Vercel Sandbox +- Drive Claude Code, Codex, and Pi from Python, locally or in a Vercel Sandbox (experimental). ## Main primitives diff --git a/docs/ai-python/content/docs/reference/harnesses.mdx b/docs/ai-python/content/docs/reference/harnesses.mdx index b784bc91..0ca8e002 100644 --- a/docs/ai-python/content/docs/reference/harnesses.mdx +++ b/docs/ai-python/content/docs/reference/harnesses.mdx @@ -2,7 +2,7 @@ title: "ai.harnesses.experimental" description: Reference for driving coding-agent CLIs. type: reference -summary: Reference for claude_code, codex, Harness, Session, Turn, Tui, Result, Handle, approvals, capabilities, and harness errors. +summary: Reference for claude_code, codex, pi, Harness, Session, Turn, Tui, Result, Handle, approvals, capabilities, and harness errors. --- `ai.harnesses.experimental` drives coding-agent CLIs, Claude Code and Codex, @@ -31,7 +31,7 @@ Each factory returns a `Harness`. They share these arguments: | `model` | Model name passed to the CLI. | `None`: the CLI's default | | `gateway` | Upstream model provider LLM Gateway, as a `Gateway`. Overrides the workspace's gateway. | `None`: the workspace's gateway | | `effort` | `"low"`, `"medium"`, `"high"`, or `"max"`. | `None`: the CLI's default | -| `executable` | CLI name or path. | `"claude"` or `"codex"` | +| `executable` | CLI name or path. | `"claude"`, `"codex"`, or `"pi"` | | `instructions` | Text added to the CLI's own system prompt. | `None` | | `mcp_servers` | MCP servers: `{"name": {"command": ..., "args": [...], "env": {...}}}` or `{"name": {"url": ...}}`. Not stored in a `Handle`. | `None` | @@ -104,14 +104,46 @@ when one is set. `effort="max"` is sent to Codex as `"xhigh"`. +### pi + +```python +harnesses.pi( + *, + workspace, + approve=None, + approval_timeout=300.0, + model=None, + gateway=None, + effort=None, + executable="pi", + instructions=None, + mcp_servers=None, + tools=None, + disallowed_tools=None, + max_turns=None, +) +``` + +Pi over `pi --mode rpc`. Needs only the `pi` CLI. Takes the shared +arguments, plus: + +| Argument | Description | Default | +|---|---|---| +| `tools` | The tools the agent has. `[]` for none. | `None`: pi's defaults | +| `disallowed_tools` | Tools taken away. | `None` | + +`effort` is pi's thinking level. `instructions` are appended to pi's system +prompt. With a `gateway`, `model` names a gateway model. It also takes +`max_turns`, and raises `UnsupportedError` when it is set. + ## Harness -One configured CLI, open on one workspace. Build it with `claude_code()` or -`codex()`. Use it as an async context manager, or call `open()` and +One configured CLI, open on one workspace. Build it with `claude_code()`, +`codex()`, or `pi()`. Use it as an async context manager, or call `open()` and `close()`. ```python -harness.kind # "claude-code" or "codex" +harness.kind # "claude-code", "codex", or "pi" harness.name # same as kind harness.version # what the CLI reported, or None harness.capabilities @@ -278,7 +310,7 @@ One conversation in the CLI's store. | Field | Description | |---|---| -| `kind` | `"claude-code"` or `"codex"`. | +| `kind` | `"claude-code"`, `"codex"`, or `"pi"`. | | `session_id` | The CLI's own id. | | `title` | The conversation's title or first prompt, or `None`. | | `cwd` | The directory it ran in, or `None`. | @@ -292,7 +324,7 @@ A conversation as plain, serializable data. | Field | Description | |---|---| -| `kind` | `"claude-code"` or `"codex"`. | +| `kind` | `"claude-code"`, `"codex"`, or `"pi"`. | | `session_id` | The CLI's own id. | | `workspace` | `WorkspaceCoords`. | | `options` | Configuration the conversation was created with. | diff --git a/examples/harnesses/tui.py b/examples/harnesses/tui.py index 81d67f41..b7dd7759 100644 --- a/examples/harnesses/tui.py +++ b/examples/harnesses/tui.py @@ -11,6 +11,7 @@ python examples/harnesses/tui.py # claude, in a sandbox python examples/harnesses/tui.py codex # codex, in a sandbox python examples/harnesses/tui.py codex # reattach later + python examples/harnesses/tui.py pi # pi, in a sandbox # push a conversation from this directory into it: PUSH_FROM= python examples/harnesses/tui.py claude @@ -32,7 +33,11 @@ from ai.workspaces.experimental import tty NAME = "demo-tui" -HARNESSES = {"claude": harnesses.claude_code, "codex": harnesses.codex} +HARNESSES = { + "claude": harnesses.claude_code, + "codex": harnesses.codex, + "pi": harnesses.pi, +} async def main() -> None: diff --git a/skills/ai/SKILL.md b/skills/ai/SKILL.md index 9c9e2c91..68985d84 100644 --- a/skills/ai/SKILL.md +++ b/skills/ai/SKILL.md @@ -1,6 +1,6 @@ --- name: ai -description: AI SDK for Python (the `ai` package). Use for Python model calls (LLMs, image, video, speech, embedding, transcription, reranking, or evaluation), agents, model interaction tests, tool calling, subagents, approvals, durable execution, telemetry, AI SDK UI backends, custom providers, driving coding-agent CLIs (Claude Code, Codex), and sandboxes. +description: AI SDK for Python (the `ai` package). Use for Python model calls (LLMs, image, video, speech, embedding, transcription, reranking, or evaluation), agents, model interaction tests, tool calling, subagents, approvals, durable execution, telemetry, AI SDK UI backends, custom providers, driving coding-agent CLIs (Claude Code, Codex, Pi), and sandboxes. metadata: sdk-version: "0.5.0" --- @@ -84,7 +84,7 @@ read the listed local notes before writing code. | Telemetry and tracing | `basics/telemetry.md` | — | | AI SDK UI backends | `basics/ai-sdk-ui.md` | [ui.md](references/ui.md) | | Custom providers | `basics/providers.md` | [custom-provider.md](references/custom-provider.md) | -| Drive Claude Code / Codex CLIs | `basics/harnesses.md` | [harnesses.md](references/harnesses.md) | +| Drive Claude Code / Codex / Pi CLIs | `basics/harnesses.md` | [harnesses.md](references/harnesses.md) | | Run agents in a sandbox | `basics/workspaces.md` | [harnesses.md](references/harnesses.md) | For exact APIs, use `reference.md` and the relevant `reference/*.md` page. diff --git a/skills/ai/references/harnesses.md b/skills/ai/references/harnesses.md index 6082ad2c..09a3ff33 100644 --- a/skills/ai/references/harnesses.md +++ b/skills/ai/references/harnesses.md @@ -9,7 +9,7 @@ from ai.harnesses import experimental as harnesses from ai.workspaces import experimental as workspaces ``` -A harness drives a CLI (`claude`, `codex`). It is not `ai.Agent`: there is no +A harness drives a CLI (`claude`, `codex`, `pi`). It is not `ai.Agent`: there is no `ai.get_model`, and Python `@ai.tool` functions do not reach it. The CLI brings its own tools, login, and session store. @@ -18,8 +18,8 @@ brings its own tools, login, and session store. - `claude_code()` needs `ai[claude-code]`; `VercelSandbox` needs `ai[sandbox]` with a `vercel-sandbox` build that has interactive PTY support (vercel/vercel-py#415, unreleased). Both raise - `ai.errors.InstallationError` at construction. `codex()` needs only the - `codex` binary. + `ai.errors.InstallationError` at construction. `codex()` and `pi()` need + only their binary. - On `Local`, a missing CLI raises `ExecutableMissingError`. The SDK installs CLIs only into a sandbox, and there it installs or upgrades to the latest npm release on open (not while a copy of the CLI is running there). @@ -29,8 +29,8 @@ brings its own tools, login, and session store. sandbox), claude's `tools=`/`disallowed_tools=`, or a `VercelSandbox`. - Shared options: `instructions=` (added to the CLI's system prompt) and `mcp_servers=` (claude's shape: `{"name": {"command", "args", "env"}}` or - `{"name": {"url"}}`). Claude only: `tools=`, `disallowed_tools=`, - `max_turns=`. A `Handle` never stores `mcp_servers`: pass them again to + `{"name": {"url"}}`). Claude and pi: `tools=`, `disallowed_tools=`. + Claude only: `max_turns=`. A `Handle` never stores `mcp_servers`: pass them again to `harness_from`. - Open one harness per workspace and reuse it for many sessions. The CLI processes live as long as the `async with`. @@ -71,21 +71,23 @@ brings its own tools, login, and session store. - No `approve=` hook means every tool call is allowed, not "ask the user". Pass a hook to restrict. -- Claude Code's hook sees every call, reads included - (`capabilities.approval == "all"`). Codex's sees what its policy asks +- Claude Code's and pi's hooks see every call, reads included + (`capabilities.approval == "all"`). Pi has no permission prompt: the SDK + loads an extension into it that asks the hook. Codex's sees what its policy asks about (`"policy"`): shell commands and file changes. Read the event stream to see every call on any harness. - `ctx.call.name` and `ctx.call.input` are native. Claude uses `Bash`, `Write`, `Edit` with `file_path`. Codex uses `commandExecution` (the command is a string, often wrapped in `/bin/bash -lc`) and `fileChange`. - In a sandbox codex often writes files through the shell. A hook for both - harnesses must handle both dialects. Do not parse shell to judge + In a sandbox codex often writes files through the shell. Pi uses `bash`, + `read`, `edit`, `write` with `path`. A hook for several harnesses must + handle each dialect. Do not parse shell to judge containment; use `ctx.workspace.contains(await ctx.workspace.resolve(p))` so a symlink cannot lead out. For real isolation use a `VercelSandbox`. - `harnesses.Deny("reason")` and `Deny(reason="reason")` are the same. - On codex (`deny_reason=False`) the reason never reaches the agent. - `Allow(input=...)` works only on Claude. `Allow(remember=True)` works on - neither. Unsupported options raise `UnsupportedError` from the turn. + `Allow(input=...)` works on Claude and pi. `Allow(remember=True)` works + nowhere. Unsupported options raise `UnsupportedError` from the turn. - A hook that raises or exceeds `approval_timeout` (300s) denies the call and is listed in `result.approval_errors`. - Never call session methods from inside the hook; the agent is blocked diff --git a/src/ai/harnesses/experimental/__init__.py b/src/ai/harnesses/experimental/__init__.py index ab482079..918d1331 100644 --- a/src/ai/harnesses/experimental/__init__.py +++ b/src/ai/harnesses/experimental/__init__.py @@ -1,4 +1,4 @@ -"""Drive coding-agent CLIs you don't own: Claude Code and Codex. +"""Drive coding-agent CLIs you don't own: Claude Code, Codex and Pi. Experimental: not part of the stable API, may change or be removed. @@ -13,7 +13,7 @@ from . import errors from ._approval import Allow, ApprovalContext, Decision, Deny, ToolCall from ._capabilities import ApprovalCoverage, Capabilities -from ._factories import claude_code, codex, harness_from, workspace_from +from ._factories import claude_code, codex, harness_from, pi, workspace_from from ._handle import Handle from ._harness import Harness from ._result import Result @@ -39,5 +39,6 @@ "codex", "errors", "harness_from", + "pi", "workspace_from", ] diff --git a/src/ai/harnesses/experimental/_adapters/pi.py b/src/ai/harnesses/experimental/_adapters/pi.py new file mode 100644 index 00000000..9086d8a9 --- /dev/null +++ b/src/ai/harnesses/experimental/_adapters/pi.py @@ -0,0 +1,1412 @@ +"""Pi, driven through `pi --mode rpc`. + +One pi process per conversation, like claude: an RPC process holds one +active session. The id we hand pi (`--session-id`) IS its session id, so +`pi --session ` opens the same conversation in pi's own TUI. + +pi has no permission prompt. Approval is an extension this SDK loads into +every pi it launches (`pi_extension.js`): its `tool_call` handler asks +through an RPC dialog, and this adapter answers it. The protocol shapes +here were measured against a real pi (see VERIFIED_VERSION). +""" + +from __future__ import annotations + +import asyncio +import base64 +import contextlib +import datetime +import json +import pathlib +import uuid +import warnings +from typing import TYPE_CHECKING, Any + +from ....types import events as events_ +from ....types import messages as messages_ +from ....types import usage as usage_ +from ....workspaces.experimental import _base, _gateway +from ....workspaces.experimental import errors as workspace_errors +from .. import ( + _approval, + _capabilities, + _handoff, + _input, + _session, + _session_lock, + errors, +) +from . import base + +if TYPE_CHECKING: + from collections.abc import AsyncIterator + +KIND = "pi" +PACKAGE = "@earendil-works/pi-coding-agent" +INSTALL = f"npm install -g {PACKAGE}" +VERIFIED_VERSION = "1.3.0" +"""The CLI the live suite last passed against. Older ones get a warning.""" +APPROVE = "ai-python:approve" +"""The dialog title the extension asks approvals under.""" +READY = "ai-python-ready" +"""The command the extension registers last: proof that it loaded.""" +CONFIG_VAR = "AI_PYTHON_PI" +"""The environment variable that configures the extension.""" +# Internal markers delivered on a session's event queue: pi died, or this +# client closed the connection under a running turn. +CRASHED = "__crashed__" +CLOSED = "__closed__" +DIALOGS = ("select", "confirm", "input", "editor") +"""Extension UI methods that wait for an answer; the rest are notices.""" +PROVIDER = "vercel-ai-gateway" +"""pi's built-in provider for the Vercel AI Gateway.""" + +# The hook is pi's `tool_call` event, which fires for every tool call — +# built-in, extension, MCP, and the nested calls a codemode script makes — +# and can rewrite the input in place or block with a reason the agent reads. +CAPABILITIES = _capabilities.Capabilities( + steer=True, + stop=True, + resume=True, + fork=True, + history=True, + rewrite_tool_input=True, + deny_reason=True, + images=True, + approval="all", +) + + +class PiAdapter: + kind = KIND + capabilities = CAPABILITIES + verified_version: str | None = VERIFIED_VERSION + + def __init__( + self, + executable: str = "pi", + model: str | None = None, + gateway: _gateway.Gateway | None = None, + effort: str | None = None, + instructions: str | None = None, + mcp_servers: dict[str, Any] | None = None, + tools: list[str] | None = None, + disallowed_tools: list[str] | None = None, + ) -> None: + self._executable = executable + self._model = model + self._gateway = gateway + self._effort = effort + self._instructions = instructions + self._mcp_servers = dict(mcp_servers or {}) + self._tools = tools + self._disallowed_tools = list(disallowed_tools or []) + self._version: str | None = None + self._workspace: _base.Workspace | None = None + self._extension: str | None = None + self._lock: _session_lock.SessionLock | None = None + self._held: set[str] = set() + self._opening: dict[str, asyncio.Lock] = {} + self._agents: dict[str, _Agent] = {} + # Sessions whose turn WE ended: they settle as cancelled. + self._interrupted: set[str] = set() + # Sessions with a prompt accepted and no `agent_settled` yet. + self._running: set[str] = set() + # An error raised while answering an approval cannot reach anyone + # where it happens; it is parked and raised on that session's turn. + self._pending_error: dict[str, BaseException] = {} + self.approve: base.ApprovalHook | None = None + self.approval_timeout: float = 300.0 + self.on_approval_error: Any = None + self.live_history: Any = lambda session_id: [] + + @property + def version(self) -> str | None: + return self._version + + @property + def process_ids(self) -> list[int]: + return [ + pid + for agent in self._agents.values() + if isinstance(pid := getattr(agent.process, "pid", None), int) + ] + + # -- lifecycle ------------------------------------------------------------ + + async def start( + self, workspace: _base.Workspace, options: dict[str, Any] + ) -> None: + if not workspace.duplex_spawn: + raise errors.UnsupportedError( + KIND, + f"{workspace.kind} workspaces", + "pi's RPC mode speaks JSON lines on stdin, which this " + + "workspace cannot provide", + ) + # An explicit gateway wins; otherwise the workspace's. Where the + # workspace injects the credential at egress, pi only ever holds the + # placeholder: the extension sends it as `Authorization: Bearer`, + # the header the firewall rewrites. + gateway = self._gateway or workspace.gateway + if gateway is not None and workspace.injects_credentials_for( + gateway.host + ): + gateway = gateway.brokered() + self._gateway = gateway + probe = await base.provision(workspace, self._executable, PACKAGE) + if probe.exit_code != 0: + raise errors.ExecutableMissingError(KIND, self._executable, INSTALL) + self._version = probe.stdout.strip().splitlines()[0] or None + self._workspace = workspace + self._lock = _session_lock.SessionLock(workspace, KIND) + # The extension travels to wherever pi runs, as a file: pi loads + # extensions by path. + self._extension = f"{await workspace.home()}/.ai-python/pi/extension.js" + await workspace.write_text( + self._extension, + pathlib.Path(__file__).with_name("pi_extension.js").read_text(), + ) + + async def close(self) -> None: + for session_id in list(self._agents): + await self.close_session(session_id) + + async def detach(self) -> None: + """Let go of every session without terminating pi. + + Stop reading and drop the input conduit; where the workspace outlives + the client (a kept sandbox) the process runs on. The session locks + stay held: the process is still the writer. + """ + for session_id, agent in list(self._agents.items()): + await agent.close() + with contextlib.suppress(Exception): + await agent.process.detach() + self._agents.pop(session_id, None) + + def _args(self) -> list[str]: + """How this harness runs pi: headless and in the TUI alike.""" + assert self._extension is not None + args = ["-e", self._extension] + model = self._model + if ( + model + and self._gateway is not None + and not model.startswith(f"{PROVIDER}/") + ): + # With a gateway, the model is the gateway's: `anthropic/...` + # names it there, not at Anthropic. + model = f"{PROVIDER}/{model}" + if model: + args += ["--model", model] + if self._effort: + args += ["--thinking", self._effort] + if self._instructions: + args += ["--append-system-prompt", self._instructions] + if self._tools is not None: + args += ( + ["--tools", ",".join(self._tools)] + if self._tools + else ["--no-tools"] + ) + if self._disallowed_tools: + args += ["--exclude-tools", ",".join(self._disallowed_tools)] + return args + + def _env(self, *, approve: bool) -> dict[str, str]: + """Return the extension's configuration, and the gateway's key.""" + config: dict[str, Any] = {"approve": approve} + if self._mcp_servers: + config["mcp"] = mcp_config(self._mcp_servers) + env: dict[str, str] = {} + if self._gateway is not None: + base_url = self._gateway.base_url + config["gateway"] = ( + {"baseUrl": base_url} + if base_url != _gateway.DEFAULT_BASE_URL + else {} + ) + env[_gateway.KEY_VAR] = self._gateway.credential + env[CONFIG_VAR] = json.dumps(config) + return env + + def _auth_hint(self) -> str: + if self._gateway is not None and self._gateway.is_brokered: + return BROKERED_HINT.format(host=self._gateway.host) + return AUTH_HINT + + # -- sessions ------------------------------------------------------------- + + async def new_session( + self, history: list[messages_.Message] | None = None + ) -> str: + session_id = ( + await self.stage_history(history) if history else str(uuid.uuid4()) + ) + await self._acquire(session_id) + try: + await self._spawn(session_id, ["--session-id", session_id]) + except BaseException: + await self._release(session_id) + raise + return session_id + + async def resume_session(self, session_id: str) -> None: + # Checked first: `--session-id` CREATES a session it cannot find, + # so a typo would silently start an empty conversation. + await self._find(session_id) + # One launch per conversation at a time: two resumes racing through + # the hand-over would each start a pi (two writers, one untracked). + async with self._opening.setdefault(session_id, asyncio.Lock()): + await self._hand_over(session_id) + await self._acquire(session_id) + try: + await self._spawn(session_id, ["--session-id", session_id]) + except BaseException as exc: + # Whatever ended the launch — a crash, a timeout, a cancelled + # task — nothing holds the conversation now: free the claim, + # or every other client is refused until it goes stale. + await self._release(session_id) + if isinstance(exc, errors.HarnessError): + raise errors.ResumeFailedError( + session_id, str(exc) + ) from exc + raise + + async def fork(self, session_id: str) -> str: + found = await self._find(session_id) + forked = str(uuid.uuid4()) + await self._acquire(forked) + try: + # By path, not id: pi matches a partial id, a path is exact. + await self._spawn( + forked, ["--fork", found["path"], "--session-id", forked] + ) + except BaseException as exc: + await self._release(forked) # see resume_session + if isinstance(exc, errors.HarnessError): + raise errors.ResumeFailedError( + session_id, f"fork failed: {exc}" + ) from exc + raise + return forked + + async def close_session(self, session_id: str) -> None: + agent = self._agents.pop(session_id, None) + self._running.discard(session_id) + if agent is not None: + await agent.close() + with contextlib.suppress(Exception): + await agent.process.terminate() + await self._release(session_id) + + async def _spawn(self, session_id: str, args: list[str]) -> None: + """Launch pi on one conversation and wait until it answers. + + A pi that cannot start — an unknown model, a bad flag — exits at + once, and the handshake reports its last words. + """ + assert self._workspace is not None + process = await self._workspace.spawn( + [self._executable, "--mode", "rpc", *args, *self._args()], + env=self._env(approve=self.approve is not None), + ) + agent = _Agent(process, on_dialog=self._dialog_handler(session_id)) + agent.start() + try: + state = await agent.call({"type": "get_state"}, timeout=60) + reported = (state or {}).get("sessionId") + if reported and reported != session_id: + raise errors.AgentCrashedError( + KIND, f"pi opened session {reported}, not {session_id}" + ) + # The approval gate must be there before any work is accepted: + # a gate that failed to load is a gate that fails open. + listed = await agent.call({"type": "get_commands"}, timeout=60) + if READY not in { + c.get("name") for c in (listed or {}).get("commands") or [] + }: + raise errors.AgentCrashedError( + KIND, + "pi did not load the SDK's extension " + + f"({self._extension}); refusing to run without it", + ) + # Published before the session is registered: whatever fails + # here, the process goes with it, never outliving its claim. + if self._lock is not None: + pid = await process.os_pid() + await self._lock.claim_pid(session_id, pid, client=KIND) + except BaseException: + await agent.close() + with contextlib.suppress(Exception): + await process.terminate() + raise + self._agents[session_id] = agent + + async def _hand_over(self, session_id: str) -> None: + """One writer per conversation, this client included. + + The lock is per client, so it lets this client through. A second pi + on a conversation this client already has open would be a second + writer on one transcript: the one already open is ended first (the + claim stays ours), unless a turn is running on it. + """ + if session_id in self._running: + raise errors.SessionBusyError( + session_id, + None, + "a turn is running on it in this client; stop() it first", + ) + agent = self._agents.pop(session_id, None) + if agent is not None: + await agent.close() + with contextlib.suppress(Exception): + await agent.process.terminate() + + async def _acquire(self, session_id: str) -> None: + """Refuse before launching a second pi: one writer per conversation.""" + if self._lock is None or session_id in self._held: + return + await self._lock.acquire(session_id, client=KIND) + self._held.add(session_id) + + async def _release(self, session_id: str) -> None: + self._held.discard(session_id) + if self._lock is not None: + with contextlib.suppress(Exception): + await self._lock.release(session_id) + + # -- the session store ---------------------------------------------------- + + async def _store(self, op: str, session_id: str = "") -> Any: + """Ask the workspace about pi's session files: one command, one trip.""" + assert self._workspace is not None + result = await self._workspace.exec( + ["python3", "-c", _STORE, op, self._workspace.path, session_id], + timeout=60, + ) + if result.exit_code != 0: + raise errors.HarnessError( + f"could not read pi's sessions in the workspace (exit " + f"{result.exit_code}): " + + (result.stderr.strip()[-300:] or "no output") + + ". Harnesses need python3 where the CLI runs." + ) + return json.loads(result.stdout.strip().splitlines()[-1]) + + async def _find(self, session_id: str) -> dict[str, Any]: + """Return a conversation of THIS workspace, or raise ResumeFailed.""" + assert self._workspace is not None + found = await self._store("find", session_id) + if not found: + raise errors.ResumeFailedError( + session_id, + f"no such session in workspace {self._workspace.path}", + ) + return dict(found) + + async def list_sessions(self) -> list[_session.SessionInfo]: + return [ + _session.SessionInfo( + kind=KIND, + session_id=s["id"], + title=s.get("title"), + cwd=s.get("cwd"), + updated_at=s.get("updated"), + created_at=s.get("created"), + ) + for s in await self._store("list") + ] + + async def running_sessions(self) -> dict[str, int | None]: + """Conversations a client of this SDK, or its `tui()`, has open. + + A pi started in a terminal is not seen: it renames its process to + `pi`, so `ps` shows no session in its arguments, and it records no + process per conversation. + """ + return await self._lock.holders() if self._lock is not None else {} + + async def history( + self, session_id: str, *, limit: int | None = None, offset: int = 0 + ) -> list[messages_.Message]: + """Return the active branch of the session file, in full. + + Every message on it, from the root to the latest entry — including + what a compaction summarized away, which pi keeps on disk. + """ + assert self._workspace is not None + found = await self._find(session_id) + try: + text = await self._workspace.read_text(found["path"]) + except Exception as exc: + raise errors.ResumeFailedError(session_id, str(exc)) from exc + messages: list[messages_.Message] = [] + for record in _active_branch(text): + messages.extend(_stored_to_messages(record, replay=True)) + window = messages[offset:] + return window[:limit] if limit is not None else window + + async def stage_history(self, history: list[messages_.Message]) -> str: + """Write a session file with this past, and start nothing on it. + + Where pi will look for it: the session directory of this workspace, + under a fresh id that `--session-id` then opens. + """ + assert self._workspace is not None + session_id = str(uuid.uuid4()) + _handoff.warn_unsupported_files(history, KIND) + dropped = _handoff.reasoning_count(history) + if dropped: + # pi itself drops thinking that another model produced (its + # signature is only valid for that model), and replaying it as + # user text is refused by Anthropic's safeguards — pi often runs + # Claude. Say so rather than drop it silently. + warnings.warn( + f"{dropped} reasoning part(s) in the history were not carried " + + "into pi: a model only reads its own reasoning. Text, tool " + + "calls and images were carried in full.", + stacklevel=2, + ) + where = await self._store("dir") + now = datetime.datetime.now(datetime.UTC) + stamp = now.isoformat(timespec="milliseconds").replace("+00:00", "Z") + lines: list[dict[str, Any]] = [ + { + "type": "session", + "version": 3, + "id": session_id, + "timestamp": stamp, + "cwd": where["cwd"], + } + ] + parent: str | None = None + for message in _handoff.to_pi_messages( + history, timestamp=int(now.timestamp() * 1000) + ): + entry_id = uuid.uuid4().hex[:8] + lines.append( + { + "type": "message", + "id": entry_id, + "parentId": parent, + "timestamp": stamp, + "message": message, + } + ) + parent = entry_id + name = stamp.replace(":", "-").replace(".", "-") + await self._workspace.write_text( + f"{where['dir']}/{name}_{session_id}.jsonl", + "".join(json.dumps(line) + "\n" for line in lines), + ) + return session_id + + # -- the TUI -------------------------------------------------------------- + + async def prepare_tui(self) -> None: + """Nothing to do: pi's TUI has no first-run walk to get past.""" + + def tui_launch( + self, session_id: str | None + ) -> tuple[list[str], dict[str, str], str | None]: + assert self._workspace is not None + running_as = session_id or str(uuid.uuid4()) + argv = [self._executable, "--session-id", running_as, *self._args()] + # The person at the TUI is the gate: the extension asks nothing. + env = { + "TERM": "xterm-256color", + **dict(self._workspace.env), + **self._env(approve=False), + } + return argv, env, running_as + + async def claim_new_session(self, session_id: str) -> None: + await self._acquire(session_id) + + async def record_tui_pid(self, session_id: str, pid: int | None) -> None: + if self._lock is not None and pid is not None: + await self._lock.claim_pid(session_id, pid, client=KIND) + + async def release_tui_lock(self, session_id: str) -> None: + await self._release(session_id) + + async def resume_session_lock(self, session_id: str) -> None: + await self._find(session_id) + async with self._opening.setdefault(session_id, asyncio.Lock()): + await self._hand_over(session_id) + await self._acquire(session_id) + + # -- approval ------------------------------------------------------------- + + def _dialog_handler(self, session_id: str) -> Any: + async def on_dialog(agent: _Agent, request: dict[str, Any]) -> None: + if request.get("method") not in DIALOGS: + return # a notice: nothing waits for an answer + reply: dict[str, Any] = { + "type": "extension_ui_response", + "id": request.get("id"), + } + if request.get("method") != "input" or request.get("title") != ( + APPROVE + ): + # Some other extension asking a person. Nobody is there: + # dismiss it rather than leave the agent waiting. + agent.pending.pop(str(request.get("id")), None) + await agent.send({**reply, "cancelled": True}) + return + if session_id in self._interrupted: + # The turn is being stopped: a question asked now is + # withdrawn at once, or `abort` would wait on it. + agent.pending.pop(str(request.get("id")), None) + await agent.send({**reply, "cancelled": True}) + return + decision = await self._judge(session_id, request) + # Answered from here on: `dismiss` must not withdraw it too. + agent.pending.pop(str(request.get("id")), None) + await agent.send({**reply, "value": json.dumps(decision)}) + if decision.get("stop"): + # Only now: the answer must reach pi before it is aborted. + with contextlib.suppress(Exception): + await agent.call({"type": "abort"}, timeout=30) + + return on_dialog + + async def _judge( + self, session_id: str, request: dict[str, Any] + ) -> dict[str, Any]: + """Ask the hook about one call; always come back with an answer. + + A hook that fails or runs out of time denies (a gate that fails open + is not a gate) and is reported. A decision pi cannot honor denies + and is parked to raise on the turn. + """ + try: + payload = json.loads(request.get("placeholder") or "{}") + assert self._workspace is not None + if self.approve is None: + return {"decision": "allow"} + parent = payload.get("parentToolCallId") + call = _approval.ToolCall( + id=str(payload.get("toolCallId") or ""), + name=str(payload.get("toolName") or ""), + input=dict(payload.get("input") or {}), + hints={"parentToolCallId": parent} if parent else {}, + raw=payload, + ) + context = _approval.ApprovalContext( + call=call, + history=list(self.live_history(session_id)), + workspace=self._workspace, + harness=KIND, + session_id=session_id, + ) + async with asyncio.timeout(self.approval_timeout): + decision = await self.approve(context) + except Exception as exc: + self._record_error(session_id, f"{type(exc).__name__}: {exc}") + return {"decision": "deny", "reason": "approval failed; denied"} + if isinstance(decision, _approval.Allow): + if decision.remember: + self._pending_error[session_id] = errors.UnsupportedError( + KIND, + "remember", + "pi has no permission rules to remember a decision in", + ) + return {"decision": "deny", "reason": "remember is unsupported"} + return {"decision": "allow", "input": decision.input} + if isinstance(decision, _approval.Deny): + if decision.stop: + self._interrupted.add(session_id) + return { + "decision": "deny", + "reason": decision.reason or "denied by policy", + "stop": decision.stop, + } + return {"decision": "deny", "reason": "approval hook returned nothing"} + + def _record_error(self, session_id: str, detail: str) -> None: + if callable(self.on_approval_error): + self.on_approval_error(session_id, detail) + + # -- turns ---------------------------------------------------------------- + + def _agent(self, session_id: str) -> _Agent: + agent = self._agents.get(session_id) + if agent is None: + raise RuntimeError(f"unknown session {session_id}") + return agent + + async def turn( + self, session_id: str, prompt: list[_input.Part] + ) -> AsyncIterator[events_.AgentEvent]: + agent = self._agent(session_id) + if agent.dead is not None: + raise errors.AgentCrashedError(KIND, agent.dead) + # Whatever the last turn left behind (a stopped turn's tail) is not + # this turn's: `abort` answers only once pi is idle, so it has all + # arrived by now. + while not agent.events.empty(): + agent.events.get_nowait() + self._interrupted.discard(session_id) + self._pending_error.pop(session_id, None) + command = await _prompt_command(prompt) + if session_id in self._interrupted: + # Stopped while its files were being read: the abort found pi + # idle, so nothing would stop a prompt sent now. Send none. + self._interrupted.discard(session_id) + yield events_.StreamStart() + yield events_.StreamEnd( + message=messages_.Message(role="assistant", parts=[]), + finish_reason="cancelled", + ) + return + try: + accepted = await agent.call(command) + except _CommandError as exc: + raise _turn_failure(str(exc), hint=self._auth_hint()) from exc + if (accepted or {}).get("disposition") == "handled": + # An extension command consumed it: no run starts. + yield events_.StreamStart() + yield events_.StreamEnd( + message=messages_.Message(role="assistant", parts=[]), + finish_reason="stop", + ) + return + # Running BEFORE the first event: a consumer may steer on it. + self._running.add(session_id) + parts: list[Any] = [] + usage = usage_.Usage() + cost = 0.0 + started: set[str] = set() + open_blocks: dict[str, str] = {} + message_no = 0 + last: dict[str, Any] = {} + retry_error = "" + settled: dict[str, Any] = {} + try: + yield events_.StreamStart() + while True: + event = await agent.events.get() + kind = event.get("type") + if kind == CRASHED: + raise errors.AgentCrashedError( + KIND, event.get("detail") or "exited mid-turn" + ) + if kind == CLOSED: + raise errors.HarnessClosedError( + "the session was closed while its turn was running" + ) + message = event.get("message") or {} + if kind == "message_start" and message.get("role") == ( + "assistant" + ): + message_no += 1 + elif kind == "message_update": + for out in _stream_events( + event.get("assistantMessageEvent") or {}, + f"{message_no}", + open_blocks, + started, + ): + yield out + elif kind == "message_end" and message.get("role") == ( + "assistant" + ): + last = message + if message.get("stopReason") != "error": + # A failed attempt is not the answer: pi retries it + # (or the turn raises below), and its partial text + # would be glued onto the reply that replaces it. + parts += _assistant_parts(message.get("content")) + raw = message.get("usage") or {} + usage = usage + _usage(raw) + cost += float((raw.get("cost") or {}).get("total") or 0) + for out in _close_blocks(open_blocks): + yield out + elif kind == "tool_execution_end" and not event.get( + "parentToolCallId" + ): + # A nested call (a codemode script's) is not in the + # transcript: only the tool that made it reports it. + result = _tool_result( + event.get("toolCallId") or "", + event.get("toolName") or "", + event.get("result") or {}, + is_error=bool(event.get("isError")), + ) + yield events_.ToolCallResult( + message=messages_.Message(role="tool", parts=[result]), + results=[result], + ) + elif kind == "auto_retry_end" and not event.get("success"): + retry_error = str(event.get("finalError") or "") + elif kind == "agent_settled": + settled = event + break + finally: + self._running.discard(session_id) + error = self._pending_error.pop(session_id, None) + if error is not None: + raise error + interrupted = session_id in self._interrupted or bool( + settled.get("aborted") + ) + self._interrupted.discard(session_id) + stop = last.get("stopReason") + if stop == "error" and not interrupted: + raise _turn_failure( + str(last.get("errorMessage") or retry_error or "unknown error"), + hint=self._auth_hint(), + ) + yield events_.StreamEnd( + message=messages_.Message(role="assistant", parts=parts), + # Cost has no home in Usage; pi's own number rides in `raw`. + usage=usage.model_copy(update={"raw": {"cost_usd": cost}}), + finish_reason=( + "cancelled" + if interrupted or stop == "aborted" + else "length" + if stop == "length" + else "stop" + ), + ) + + async def steer(self, session_id: str, text: str) -> None: + """Queue guidance for the running turn. + + pi delivers it once the current assistant message's tool calls are + done, before its next model call. + """ + if session_id not in self._running: + raise RuntimeError("no active turn to steer") + try: + await self._agent(session_id).call( + {"type": "steer", "message": text} + ) + except _CommandError as exc: + raise errors.TurnFailedError( + KIND, f"steer rejected: {exc}" + ) from exc + + async def stop(self, session_id: str) -> None: + self._interrupted.add(session_id) + agent = self._agents.get(session_id) + if agent is None or agent.dead is not None: + return + # A tool call waiting on the hook keeps pi busy, and `abort` only + # answers once pi is idle — measured: stop() timed out and the turn + # ran on. Withdraw the question first: the call is denied, pi goes + # idle, and the abort lands at once. + await agent.dismiss() + # `abort` answers once pi is idle: everything the turn produced has + # been sent by then. + with contextlib.suppress(Exception): + await agent.call({"type": "abort"}, timeout=30) + + +class _CommandError(Exception): + """pi answered a command with `success: false`.""" + + +class _Agent: + """One `pi --mode rpc` process: its commands, its events, its dialogs. + + Responses resolve the command that asked (by id). Session events go on + `events`, in order, for whichever turn is reading. A dialog is answered + in a task of its own: an approval can take minutes, and the reader must + keep reading meanwhile — including the reply to an `abort` sent from + inside it. + """ + + def __init__(self, process: _base.Process, *, on_dialog: Any) -> None: + self.process = process + self.events: asyncio.Queue[dict[str, Any]] = asyncio.Queue() + self.dead: str | None = None + self._on_dialog = on_dialog + self._waiters: dict[str, asyncio.Future[Any]] = {} + self._pump: asyncio.Task[None] | None = None + self._dialogs: set[asyncio.Task[None]] = set() + # Dialogs pi is still waiting on, by request id: the ones `dismiss` + # withdraws. A handler takes its entry out just before it answers. + self.pending: dict[str, asyncio.Task[None]] = {} + + def start(self) -> None: + self._pump = asyncio.create_task(self._read()) + + async def call( + self, command: dict[str, Any], *, timeout: float = 120 + ) -> Any: + if self.dead is not None: + raise errors.AgentCrashedError(KIND, self.dead) + request_id = uuid.uuid4().hex + future: asyncio.Future[Any] = asyncio.get_running_loop().create_future() + self._waiters[request_id] = future + try: + await self.send({**command, "id": request_id}) + async with asyncio.timeout(timeout): + return await future + finally: + self._waiters.pop(request_id, None) + + async def send(self, record: dict[str, Any]) -> None: + try: + await self.process.write(json.dumps(record) + "\n") + except Exception as exc: + raise errors.AgentCrashedError( + KIND, f"could not write to pi: {exc}" + ) from exc + + def _track(self, request_id: str, task: asyncio.Task[None]) -> None: + self.pending[request_id] = task + + def forget(_: asyncio.Task[None]) -> None: + if self.pending.get(request_id) is task: + del self.pending[request_id] + + task.add_done_callback(forget) + + async def dismiss(self) -> None: + """Withdraw every question pi is still waiting on. + + The hook deciding it is cancelled, and pi hears `cancelled`: the + extension then blocks the call, so nothing runs on a question nobody + answered. + """ + for request_id, task in list(self.pending.items()): + self.pending.pop(request_id, None) + task.cancel() + with contextlib.suppress(asyncio.CancelledError, Exception): + await task + with contextlib.suppress(Exception): + await self.send( + { + "type": "extension_ui_response", + "id": request_id, + "cancelled": True, + } + ) + + async def close(self) -> None: + if self._pump is not None: + self._pump.cancel() + with contextlib.suppress(asyncio.CancelledError, Exception): + await self._pump + self._pump = None + for task in list(self._dialogs): + task.cancel() + for task in list(self._dialogs): + with contextlib.suppress(asyncio.CancelledError, Exception): + await task + self._dialogs.clear() + self._fail("connection closed") + # A turn reading events must hear it too: the producer is gone. + self.events.put_nowait({"type": CLOSED}) + + def _fail(self, detail: str) -> None: + self.dead = self.dead or detail + for future in self._waiters.values(): + if not future.done(): + future.set_exception(errors.AgentCrashedError(KIND, detail)) + self._waiters.clear() + + async def _read(self) -> None: + try: + while True: + line = await self.process.readline() + if not line: + detail = await self._death_note() + self._fail(detail) + # Whoever is waiting on events hears it too. + self.events.put_nowait({"type": CRASHED, "detail": detail}) + return + try: + record = json.loads(line) + except json.JSONDecodeError: + continue # stdout is protocol only; tolerate noise + if not isinstance(record, dict): + continue + kind = record.get("type") + if kind == "response": + self._resolve(record) + elif kind == "extension_ui_request": + task = asyncio.create_task(self._on_dialog(self, record)) + self._dialogs.add(task) + task.add_done_callback(self._dialogs.discard) + if record.get("method") in DIALOGS: + self._track(str(record.get("id")), task) + else: + self.events.put_nowait(record) + except asyncio.CancelledError: + raise + except Exception as exc: + self._fail(str(exc)) + self.events.put_nowait({"type": CRASHED, "detail": str(exc)}) + + def _resolve(self, record: dict[str, Any]) -> None: + future = self._waiters.get(str(record.get("id"))) + if future is None or future.done(): + return + if record.get("success") is False: + future.set_exception(_CommandError(str(record.get("error") or ""))) + else: + future.set_result(record.get("data")) + + async def _death_note(self) -> str: + """Describe how pi died and its last words: a crash to diagnose.""" + code: int | None = None + tail = "" + with contextlib.suppress(Exception): + async with asyncio.timeout(2): + code = await self.process.wait() + with contextlib.suppress(Exception): + async with asyncio.timeout(2): + tail = (await self.process.read_stderr()).strip() + how = f"exit code {code}" if code is not None else "exited" + return f"{how}; stderr: {tail[-800:]}" if tail else f"{how}; no stderr" + + +async def _prompt_command(prompt: list[_input.Part]) -> dict[str, Any]: + """Return a prompt as pi's `prompt` command: text, and images as base64. + + Text files travel inline, as text. Images are read HERE (`_input.data` + downloads a URL first): pi in a sandbox could not fetch one. + """ + texts: list[str] = [] + images: list[dict[str, Any]] = [] + for part in prompt: + if isinstance(part, messages_.TextPart): + texts.append(part.text) + elif part.media_type.startswith("text/"): + texts.append(await _input.as_text(part)) + else: + images.append( + { + "type": "image", + "data": base64.b64encode(await _input.data(part)).decode( + "ascii" + ), + "mimeType": part.media_type, + } + ) + command: dict[str, Any] = {"type": "prompt", "message": "\n\n".join(texts)} + if images: + command["images"] = images + return command + + +def _stream_events( + update: dict[str, Any], + message: str, + open_blocks: dict[str, str], + started: set[str], +) -> list[Any]: + """One `assistantMessageEvent` -> AI SDK events. + + Block ids are `.`: pi numbers content blocks per + assistant message, and a turn holds several. Tool-call arguments stream + as raw JSON fragments, but they travel here as ONE delta, the final + arguments, at `toolcall_end`: the hydrator builds a call only from its + deltas, and fragments are not guaranteed to add up to what pi settles on. + """ + kind = update.get("type") or "" + block = f"{message}.{update.get('contentIndex', 0)}" + out: list[Any] = [] + if kind == "text_start": + open_blocks[block] = "text" + out.append(events_.TextStart(block_id=block)) + elif kind == "text_delta" and update.get("delta"): + out.append(events_.TextDelta(chunk=update["delta"], block_id=block)) + elif kind == "thinking_start": + open_blocks[block] = "thinking" + out.append(events_.ReasoningStart(block_id=block)) + elif kind == "thinking_delta" and update.get("delta"): + out.append( + events_.ReasoningDelta(chunk=update["delta"], block_id=block) + ) + elif kind in ("text_end", "thinking_end"): + out += _close_blocks(open_blocks, only=block) + elif kind == "toolcall_start" and update.get("id"): + started.add(update["id"]) + out.append( + events_.ToolStart( + tool_call_id=update["id"], + tool_name=update.get("toolName") or "", + ) + ) + elif kind == "toolcall_end": + call = _tool_call(update.get("toolCall") or {}) + if call.tool_call_id not in started: + started.add(call.tool_call_id) + out.append( + events_.ToolStart( + tool_call_id=call.tool_call_id, tool_name=call.tool_name + ) + ) + out.append( + events_.ToolDelta( + tool_call_id=call.tool_call_id, chunk=call.tool_args + ) + ) + out.append( + events_.ToolEnd(tool_call_id=call.tool_call_id, tool_call=call) + ) + return out + + +def _close_blocks( + open_blocks: dict[str, str], only: str | None = None +) -> list[Any]: + """End the text and reasoning blocks still open (an aborted message's).""" + out: list[Any] = [] + for block in [only] if only is not None else list(open_blocks): + opened = open_blocks.pop(block, None) + if opened == "text": + out.append(events_.TextEnd(block_id=block)) + elif opened == "thinking": + out.append(events_.ReasoningEnd(block_id=block)) + return out + + +def _tool_call(block: dict[str, Any]) -> messages_.ToolCallPart: + return messages_.ToolCallPart( + tool_call_id=str(block.get("id") or ""), + tool_name=str(block.get("name") or ""), + tool_args=json.dumps(block.get("arguments") or {}), + ) + + +def _assistant_parts(content: Any) -> list[Any]: + parts: list[Any] = [] + for block in content if isinstance(content, list) else []: + if not isinstance(block, dict): + continue + kind = block.get("type") + if kind == "text" and block.get("text"): + parts.append(messages_.TextPart(text=block["text"])) + elif kind == "thinking" and block.get("thinking"): + parts.append(messages_.ReasoningPart(text=block["thinking"])) + elif kind == "toolCall": + parts.append(_tool_call(block)) + return parts + + +def _content_parts(content: Any) -> list[Any]: + """Return a user message's parts: from a string, or text and images.""" + if isinstance(content, str): + return [messages_.TextPart(text=content)] if content else [] + parts: list[Any] = [] + for block in content if isinstance(content, list) else []: + if not isinstance(block, dict): + continue + if block.get("type") == "text" and block.get("text"): + parts.append(messages_.TextPart(text=block["text"])) + elif block.get("type") == "image" and block.get("data"): + parts.append( + messages_.FilePart( + data=str(block["data"]), + media_type=str(block.get("mimeType") or "image/png"), + ) + ) + return parts + + +def _tool_result( + tool_call_id: str, tool_name: str, result: Any, *, is_error: bool +) -> messages_.ToolResultPart: + """Return what a tool returned: its text, when that is all it is. + + pi tools return content blocks; text-only content (nearly every call) + reads as the string it is, anything else stays the blocks themselves. + """ + content = result.get("content") if isinstance(result, dict) else result + if isinstance(content, list) and all( + isinstance(b, dict) and b.get("type") == "text" for b in content + ): + content = "".join(b.get("text") or "" for b in content) + return messages_.ToolResultPart( + tool_call_id=tool_call_id, + tool_name=tool_name, + result=content, + result_kind="error" if is_error else "json", + ) + + +def _active_branch(text: str) -> list[dict[str, Any]]: + """Return the session file's active branch: root to latest entry. + + A session is a tree (`id`/`parentId`); the leaf is the entry appended + last, and abandoned branches are not part of the conversation. + """ + entries: dict[str, dict[str, Any]] = {} + leaf: str | None = None + for line in text.splitlines(): + try: + entry = json.loads(line) + except json.JSONDecodeError: + continue + if ( + isinstance(entry, dict) + and entry.get("id") + and entry.get("type") != ("session") + ): + entries[str(entry["id"])] = entry + leaf = str(entry["id"]) + path: list[dict[str, Any]] = [] + seen: set[str] = set() + while leaf is not None and leaf in entries and leaf not in seen: + seen.add(leaf) + path.append(entries[leaf]) + parent = entries[leaf].get("parentId") + leaf = str(parent) if parent is not None else None + return list(reversed(path)) + + +def _stored_to_messages( + entry: dict[str, Any], *, replay: bool = False +) -> list[messages_.Message]: + """One session entry -> AI SDK messages, shaped like a live turn's. + + Everything pi sends the model is kept: what an extension injected + (`custom_message`) and the output of a `!` command (`bashExecution`) + reach it as user messages, and so they are user messages here; a branch + summary is the context pi carried over from an abandoned branch. + Compaction summaries are not: the messages they summarize are all here. + """ + kind = entry.get("type") + if kind == "custom_message": + parts = _content_parts(entry.get("content")) + return ( + [messages_.Message(role="user", parts=parts, replay=replay)] + if parts + else [] + ) + if kind == "branch_summary" and entry.get("summary"): + text = str(entry["summary"]) + return [ + messages_.Message( + role="user", + parts=[messages_.TextPart(text=text)], + replay=replay, + ) + ] + if kind != "message": + return [] + message = entry.get("message") or {} + role = message.get("role") + if role == "custom": + parts = _content_parts(message.get("content")) + return ( + [messages_.Message(role="user", parts=parts, replay=replay)] + if parts + else [] + ) + if role == "bashExecution" and not message.get("excludeFromContext"): + # pi's own rendering of it for the model. + text = ( + f"Ran `{message.get('command', '')}`\n```\n" + + f"{message.get('output', '')}\n```" + ) + return [ + messages_.Message( + role="user", + parts=[messages_.TextPart(text=text)], + replay=replay, + ) + ] + if role == "user": + parts = _content_parts(message.get("content")) + return ( + [messages_.Message(role="user", parts=parts, replay=replay)] + if parts + else [] + ) + if role == "assistant": + parts = _assistant_parts(message.get("content")) + return ( + [messages_.Message(role="assistant", parts=parts, replay=replay)] + if parts + else [] + ) + if role == "toolResult": + result = _tool_result( + str(message.get("toolCallId") or ""), + str(message.get("toolName") or ""), + message, + is_error=bool(message.get("isError")), + ) + return [messages_.Message(role="tool", parts=[result], replay=replay)] + return [] + + +def _usage(raw: dict[str, Any]) -> usage_.Usage: + """Return pi's per-message usage. + + pi's `input` excludes cached input; ours includes it. + """ + + def count(name: str) -> int: + value = raw.get(name) + return value if isinstance(value, int) else 0 + + reasoning = raw.get("reasoning") + return usage_.Usage( + input_tokens=count("input") + count("cacheRead") + count("cacheWrite"), + output_tokens=count("output"), + cache_read_tokens=count("cacheRead"), + cache_write_tokens=count("cacheWrite"), + reasoning_tokens=reasoning if isinstance(reasoning, int) else None, + ) + + +def mcp_config(servers: dict[str, Any]) -> dict[str, Any]: + """MCP servers, in the shape claude takes, as pi registers them. + + `{"name": {"command": ..., "args": [...], "env": {...}}}` for a local + server, `{"name": {"type": "http", "url": ...}}` for a remote one. pi + validates the config, so only the keys it knows are passed. + """ + keys = ("command", "args", "env", "cwd", "url", "headers") + return { + name: {k: server[k] for k in keys if server.get(k) is not None} + for name, server in servers.items() + } + + +AUTH_HINT = ( + "pi authenticates where it RUNS: its own login (`/login`, auth.json) or a " + "provider key in the environment. A sandbox is another machine that has " + "never seen either, so give it a way to reach a model: " + "VercelSandbox(gateway=vercel_ai_gateway()) injects a gateway key at " + "egress, or VercelSandbox(env={'ANTHROPIC_API_KEY': ...}) puts a key of " + "its own in the VM." +) +BROKERED_HINT = ( + "This workspace injects the credential into requests to {host} at egress, " + "so pi never held it and no login in the VM is involved. A 401 here means " + "the gateway rejected the injected key — check AI_GATEWAY_API_KEY on the " + "machine that opened the sandbox — or the egress policy is not rewriting " + "requests to that host." +) + + +def _turn_failure(detail: str, *, hint: str = AUTH_HINT) -> Exception: + """Rename the failure for what the caller has to fix.""" + detail = detail[:500] + if errors._looks_unauthenticated(detail): + return workspace_errors.NotAuthenticatedError(KIND, detail, hint) + return errors.TurnFailedError(KIND, detail) + + +_STORE = r""" +import json, os, sys + +op, where, sid = sys.argv[1], sys.argv[2], sys.argv[3] +cwd = os.path.realpath(where) + + +def settings(path): + try: + with open(path) as f: + value = json.load(f).get("sessionDir") + return value if isinstance(value, str) and value else None + except Exception: + return None + + +# Where pi keeps this directory's sessions, decided the way pi decides it: +# the environment, then project settings, then global ones; else one +# directory per working directory under the agent dir. +agent = os.path.expanduser( + os.environ.get("PI_CODING_AGENT_DIR") or "~/.pi/agent" +) +custom = ( + os.environ.get("PI_CODING_AGENT_SESSION_DIR") + or settings(os.path.join(cwd, ".pi", "settings.json")) + or settings(os.path.join(agent, "settings.json")) +) +if custom: + directory = os.path.join(cwd, os.path.expanduser(custom)) +else: + safe = cwd.lstrip("/\\") + for c in "/\\:": + safe = safe.replace(c, "-") + directory = os.path.join(agent, "sessions", "--" + safe + "--") + + +def iso_ms(text): + import datetime + + try: + when = datetime.datetime.strptime(text[:19], "%Y-%m-%dT%H:%M:%S") + except Exception: + return None + when = when.replace(tzinfo=datetime.timezone.utc) + return int(when.timestamp() * 1000) + + +def describe(path, full): + with open(path) as f: + header = json.loads(f.readline() or "{}") + if header.get("type") != "session" or not header.get("id"): + return None + if header.get("cwd") not in (cwd, where): + return None + info = { + "id": header["id"], + "path": path, + "cwd": header.get("cwd"), + "created": iso_ms(header.get("timestamp") or ""), + "updated": int(os.path.getmtime(path) * 1000), + } + if not full: + return info + name = first = None + for line in f: + try: + entry = json.loads(line) + except ValueError: + continue + if entry.get("type") == "session_info": + name = entry.get("name") or None + elif first is None and entry.get("type") == "message": + message = entry.get("message") or {} + if message.get("role") == "user": + content = message.get("content") + if isinstance(content, list): + content = " ".join( + b.get("text", "") for b in content + if isinstance(b, dict) and b.get("type") == "text" + ) + first = (content or "").strip()[:200] or None + info["title"] = name or first + return info + + +def sessions(full): + try: + names = sorted(os.listdir(directory)) + except OSError: + return [] + out = [] + for name in names: + if not name.endswith(".jsonl"): + continue + try: + info = describe(os.path.join(directory, name), full) + except Exception: + continue + if info is not None: + out.append(info) + return out + + +if op == "dir": + out = {"dir": directory, "cwd": cwd} +elif op == "list": + out = sorted(sessions(True), key=lambda s: -s["updated"]) +else: + out = next((s for s in sessions(False) if s["id"] == sid), None) +print(json.dumps(out)) +""" diff --git a/src/ai/harnesses/experimental/_adapters/pi_extension.js b/src/ai/harnesses/experimental/_adapters/pi_extension.js new file mode 100644 index 00000000..164712c0 --- /dev/null +++ b/src/ai/harnesses/experimental/_adapters/pi_extension.js @@ -0,0 +1,88 @@ +// Loaded into every `pi` the ai-python SDK launches (`pi -e `). +// It does nothing on its own: AI_PYTHON_PI, a JSON object in the +// environment, says which of its three jobs this process needs. +// +// - "approve": ask the SDK about every tool call. pi has no permission +// channel of its own; an extension's `tool_call` handler is the gate, and +// an RPC dialog (`ctx.ui.input`) is how it reaches the SDK. The answer is +// a JSON decision: allow (optionally with rewritten input) or deny. +// - "mcp": MCP servers, registered for this session only. +// - "gateway": the Vercel AI Gateway, reached with `Authorization: Bearer` +// rather than the `x-api-key` header pi's provider sends, so a sandbox +// firewall that rewrites `Authorization` can inject the credential. +// `baseUrl` moves the provider when the gateway is not the default one. +const APPROVE = "ai-python:approve"; +// Registered LAST: its presence in `get_commands` is how the SDK knows every +// part above loaded — above all the approval gate. No marker, no session. +const READY = "ai-python-ready"; + +export default function (pi) { + // Unparseable configuration fails the load (and so the session) rather + // than reading as "no approval asked for". + const config = JSON.parse(process.env.AI_PYTHON_PI || "{}") || {}; + + if (config.approve) { + pi.on("tool_call", async (event, ctx) => { + const answer = await ctx.ui.input( + APPROVE, + JSON.stringify({ + toolCallId: event.toolCallId, + parentToolCallId: event.parentToolCallId, + toolName: event.toolName, + input: event.input, + }), + ); + let decision = null; + try { + decision = JSON.parse(answer ?? ""); + } catch { + decision = null; + } + if (decision && decision.decision === "allow") { + if (decision.input && typeof decision.input === "object") { + // pi reads the arguments from this object after the handler: + // rewrite it in place. + for (const key of Object.keys(event.input)) delete event.input[key]; + Object.assign(event.input, decision.input); + } + return undefined; + } + return { + block: true, + // No answer: the SDK withdrew the question (the turn was stopped). + reason: + answer === undefined + ? "the approval was withdrawn" + : (decision && decision.reason) || "denied by policy", + ...(decision && decision.stop ? { terminate: true } : {}), + }; + }); + } + + for (const [name, server] of Object.entries(config.mcp || {})) { + pi.registerMcpServer(name, server); + } + + const gateway = config.gateway; + if (gateway) { + if (gateway.baseUrl) { + pi.registerProvider("vercel-ai-gateway", { baseUrl: gateway.baseUrl }); + } + pi.on("before_provider_headers", (event, ctx) => { + if (ctx.model?.provider !== "vercel-ai-gateway") return; + for (const key of Object.keys(event.headers)) { + const lower = key.toLowerCase(); + if (lower === "x-api-key" || lower === "authorization") { + event.headers[key] = null; + } + } + event.headers["authorization"] = + "Bearer " + (process.env.AI_GATEWAY_API_KEY || ""); + }); + } + + pi.registerCommand(READY, { + description: "Marks the ai-python SDK's extension as loaded", + handler: async () => {}, + }); +} diff --git a/src/ai/harnesses/experimental/_factories.py b/src/ai/harnesses/experimental/_factories.py index f5440330..0fb0c4c2 100644 --- a/src/ai/harnesses/experimental/_factories.py +++ b/src/ai/harnesses/experimental/_factories.py @@ -16,6 +16,7 @@ from . import _handle, _harness, _session, errors from ._adapters import base from ._adapters import codex as codex_ +from ._adapters import pi as pi_ if TYPE_CHECKING: from collections.abc import AsyncIterator, Callable, Mapping, Sequence @@ -173,9 +174,82 @@ def codex( ) +def pi( + *, + workspace: _base.Workspace, + approve: base.ApprovalHook | None = None, + approval_timeout: float = _harness.DEFAULT_APPROVAL_TIMEOUT, + model: str | None = None, + gateway: _gateway.Gateway | None = None, + effort: Effort | None = None, + executable: str = "pi", + instructions: str | None = None, + mcp_servers: Mapping[str, Mapping[str, Any]] | None = None, + tools: Sequence[str] | None = None, + disallowed_tools: Sequence[str] | None = None, + max_turns: int | None = None, +) -> _harness.Harness: + """Pi over `pi --mode rpc`, driving your own installation and login. + + `model` is any model pi knows (`--model`): `"sonnet"`, + `"anthropic/claude-sonnet-4-5"`, with an optional `:` suffix. + Unset, pi chooses exactly as it does on its own — gateway or not. + + `gateway` sends pi to the Vercel AI Gateway (its `vercel-ai-gateway` + provider) with that key — see `ai.workspaces.experimental. + vercel_ai_gateway`. A `model` is then the gateway's: + `"anthropic/claude-sonnet-4.5"`. + + `effort` is pi's thinking level, on the same "low" to "max" scale the + other harnesses take; pi clamps it to what the model supports. + + `instructions` are appended to pi's system prompt. `mcp_servers` takes + the same shape as on `claude_code()`. `tools` limits the tools the + agent has (`["read", "grep"]`; `[]` for none); `disallowed_tools` takes + tools away. + + pi has no permission prompt, so the approval hook is an extension this + SDK loads into pi: it sees EVERY tool call, can rewrite its input, and + a `Deny(reason)` reaches the agent. Your own pi extensions, skills and + settings load as they always do. + + `max_turns` is accepted so that every `claude_code()` argument is a + `pi()` argument, but pi has no such limit: passing one raises + UnsupportedError. + """ + if max_turns is not None: + raise errors.UnsupportedError( + pi_.KIND, "max_turns", "pi has no limit on an agent's steps" + ) + return _harness.Harness( + pi_.PiAdapter( + executable=executable, + model=model, + gateway=gateway, + effort=effort, + instructions=instructions, + mcp_servers={k: dict(v) for k, v in (mcp_servers or {}).items()}, + tools=list(tools) if tools is not None else None, + disallowed_tools=list(disallowed_tools or []), + ), + workspace=workspace, + approve=approve, + approval_timeout=approval_timeout, + options={ + "model": model, + "executable": executable, + "effort": effort, + "instructions": instructions, + "tools": list(tools) if tools is not None else None, + "disallowed_tools": list(disallowed_tools or []) or None, + }, + ) + + BUILDERS: dict[str, Callable[..., _harness.Harness]] = { "claude-code": claude_code, "codex": codex, + "pi": pi, } RUNTIME_ONLY = frozenset( {"workspace", "approve", "approval_timeout", "gateway", "mcp_servers"} diff --git a/src/ai/harnesses/experimental/_handle.py b/src/ai/harnesses/experimental/_handle.py index 57201968..10dd17ad 100644 --- a/src/ai/harnesses/experimental/_handle.py +++ b/src/ai/harnesses/experimental/_handle.py @@ -18,7 +18,7 @@ class Handle(BaseModel): model_config = ConfigDict(frozen=True) kind: str - """Which harness owns this conversation ("claude-code", "codex").""" + """Which harness owns this conversation ("claude-code", "codex", "pi").""" session_id: str """The harness's OWN conversation id — not one we minted on the side.""" workspace: _base.WorkspaceCoords diff --git a/src/ai/harnesses/experimental/_handoff.py b/src/ai/harnesses/experimental/_handoff.py index 8dc984b0..645cc858 100644 --- a/src/ai/harnesses/experimental/_handoff.py +++ b/src/ai/harnesses/experimental/_handoff.py @@ -1,4 +1,4 @@ -"""Starting a conversation FROM a history: one rendering rule, both harnesses. +"""Starting a conversation FROM a history: one rendering rule, every harness. A conversation is a `list[Message]` — what `history()` returns and what `Result.messages` carries. To move one, read it as messages and start a new @@ -23,10 +23,13 @@ hand-written notes passed, genuine Opus thinking was refused. So the Claude renderer omits reasoning, and the adapter warns with the count, because dropping context silently is the one thing this SDK never does. -- files: images go to both, as the image input each takes; PDFs to Claude - as documents. Anything else (a PDF into codex, a text file) is left out - and counted (`unsupported_files`), and the adapter warns, like it does - for reasoning into Claude. + Pi omits it too, for both reasons: pi drops thinking another model wrote, + and pi often runs Claude. +- files: images go to every harness, as the image input each takes (pi + takes bytes, so an image known only by URL becomes a note naming it); + PDFs to Claude as documents. Anything else (a PDF into codex, a text + file) is left out and counted (`unsupported_files`), and the adapter + warns, like it does for reasoning into Claude. """ from __future__ import annotations @@ -41,6 +44,7 @@ from ...types import messages as messages_ NOTE_PREFIX = "[Handoff from a previous session — earlier reasoning: " +SYSTEM_PREFIX = "[Handoff from a previous session — instructions it was given: " def _text(value: Any) -> str: @@ -244,3 +248,119 @@ def add(role: str, content: Any) -> None: # the CLI keeps them as user records with tool_result blocks. add("user" if m.role in ("user", "tool") else "assistant", blocks) return records + + +def to_pi_messages( + messages: list[messages_.Message], *, timestamp: int +) -> list[dict[str, Any]]: + """Build pi `AgentMessage`s for a session file. + + Assistant messages are attributed to a model that is not the one pi + runs (`ai-python/handoff`): pi then treats them as another model's — + it keeps their text and tool calls and normalizes the call ids for its + own provider. Consecutive assistant messages are merged, the way one + model response holds text and calls together. Reasoning is omitted: see + the module docstring. + """ + out: list[dict[str, Any]] = [] + usage = { + "input": 0, + "output": 0, + "cacheRead": 0, + "cacheWrite": 0, + "totalTokens": 0, + "cost": { + "input": 0, + "output": 0, + "cacheRead": 0, + "cacheWrite": 0, + "total": 0, + }, + } + for m in messages: + if m.role == "tool": + for p in m.parts: + if isinstance(p, messages_.ToolResultPart): + out.append( + { + "role": "toolResult", + "toolCallId": p.tool_call_id, + "toolName": p.tool_name, + "content": [ + {"type": "text", "text": _text(p.result)} + ], + "isError": p.result_kind == "error", + "timestamp": timestamp, + } + ) + continue + blocks: list[dict[str, Any]] = [] + for p in m.parts: + if isinstance(p, messages_.TextPart) and p.text.strip(): + blocks.append({"type": "text", "text": p.text}) + elif ( + isinstance(p, messages_.FilePart) + and m.role == "user" + and p.media_type.startswith("image/") + ): + if isinstance(p.data, str) and media.is_downloadable_url( + p.data + ): + blocks.append( + {"type": "text", "text": f"[image: {p.data}]"} + ) + else: + blocks.append( + { + "type": "image", + "data": media.data_to_base64(p.data), + "mimeType": p.media_type, + } + ) + elif ( + isinstance(p, messages_.ToolCallPart) and m.role == "assistant" + ): + try: + args = json.loads(p.tool_args) if p.tool_args else {} + except json.JSONDecodeError: + args = {"raw": p.tool_args} + blocks.append( + { + "type": "toolCall", + "id": p.tool_call_id or p.id, + "name": p.tool_name, + "arguments": args if isinstance(args, dict) else {}, + } + ) + if not blocks: + continue + if m.role == "system": + # Measured (pi 1.3.0): a system message in a session file is + # pi's own prompt state, and replaying one REPLACED pi's system + # prompt and its tools. So instructions travel as what they are + # to the agent — the caller's words — never as its own output. + text = " ".join(b["text"] for b in blocks if b["type"] == "text") + blocks = [{"type": "text", "text": SYSTEM_PREFIX + text + "]"}] + if m.role in ("user", "system"): + out.append( + {"role": "user", "content": blocks, "timestamp": timestamp} + ) + continue + if out and out[-1]["role"] == "assistant": + out[-1]["content"] += blocks + else: + out.append( + { + "role": "assistant", + "content": blocks, + "api": "ai-python", + "provider": "ai-python", + "model": "handoff", + "usage": usage, + "stopReason": "stop", + "timestamp": timestamp, + } + ) + if any(b["type"] == "toolCall" for b in out[-1]["content"]): + out[-1]["stopReason"] = "toolUse" + return out diff --git a/tests/harnesses/experimental/conftest.py b/tests/harnesses/experimental/conftest.py index 62f40947..1c2d65f3 100644 --- a/tests/harnesses/experimental/conftest.py +++ b/tests/harnesses/experimental/conftest.py @@ -26,7 +26,7 @@ import pytest -from ai.harnesses.experimental import Harness, claude_code, codex +from ai.harnesses.experimental import Harness, claude_code, codex, pi from ai.workspaces.experimental import Local, VercelSandbox, Workspace, copy from tests.workspaces.experimental import conftest as workspaces @@ -37,11 +37,21 @@ workspace = workspaces.workspace sandbox_credentials = workspaces.sandbox_credentials -HARNESS_CLI = {"claude": "claude", "codex": "codex"} +HARNESS_CLI = {"claude": "claude", "codex": "codex", "pi": "pi"} + + +def _pi(**options: Any) -> Harness: + """Pi on PI_TEST_MODEL when set: pi has no model variable of its own, + and the default is whatever this machine's pi settings say.""" + if os.environ.get("PI_TEST_MODEL") and "model" not in options: + options["model"] = os.environ["PI_TEST_MODEL"] + return pi(**options) + SPECS: dict[str, Callable[..., Harness]] = { "claude": claude_code, "codex": codex, + "pi": _pi, } @@ -49,7 +59,7 @@ def _cli_missing(kind: str) -> bool: return shutil.which(HARNESS_CLI[kind]) is None -@pytest.fixture(params=["claude", "codex"]) +@pytest.fixture(params=["claude", "codex", "pi"]) def harness_kind(request: pytest.FixtureRequest) -> str: kind: str = request.param if _cli_missing(kind): diff --git a/tests/harnesses/experimental/e2e/test_handoff.py b/tests/harnesses/experimental/e2e/test_handoff.py index bcd43393..542e91c4 100644 --- a/tests/harnesses/experimental/e2e/test_handoff.py +++ b/tests/harnesses/experimental/e2e/test_handoff.py @@ -27,7 +27,7 @@ pytestmark = pytest.mark.live -OTHER = {"claude": "codex", "codex": "claude"} +OTHER = {"claude": "codex", "codex": "claude", "pi": "claude"} async def _source_history( diff --git a/tests/harnesses/experimental/e2e/test_options.py b/tests/harnesses/experimental/e2e/test_options.py index ad58d034..6f71d83d 100644 --- a/tests/harnesses/experimental/e2e/test_options.py +++ b/tests/harnesses/experimental/e2e/test_options.py @@ -48,12 +48,13 @@ async def test_tools_limits_what_claude_can_do( any_make_harness: Callable[..., Any], any_workspace: Workspace, ) -> None: - if harness_kind != "claude": + read = {"claude": "Read", "pi": "read"}.get(harness_kind) + if read is None: with pytest.raises(UnsupportedError, match="tools"): async with any_make_harness(tools=["Read"]): pass return - async with any_make_harness(tools=["Read"]) as harness: + async with any_make_harness(tools=[read]) as harness: await harness.run( "Write the word NOPE into nope.txt. If you cannot, reply CANNOT." ) diff --git a/tests/harnesses/experimental/e2e/test_pty.py b/tests/harnesses/experimental/e2e/test_pty.py index 83f90f5e..04d08332 100644 --- a/tests/harnesses/experimental/e2e/test_pty.py +++ b/tests/harnesses/experimental/e2e/test_pty.py @@ -355,6 +355,8 @@ async def test_the_tui_starts_from_a_history( tui.session_id != source.session_id ), "a NEW conversation, not the source" try: + if any_workspace.owner == "provider" and agent.kind not in READY: + pytest.skip(f"{agent.kind}'s prompt marker is not measured") if any_workspace.owner == "provider": screen = await _read_until(tui.pty, READY[agent.kind], timeout=150) assert ( diff --git a/tests/harnesses/experimental/e2e/test_reconnect.py b/tests/harnesses/experimental/e2e/test_reconnect.py index 2727c9aa..9ceb97d6 100644 --- a/tests/harnesses/experimental/e2e/test_reconnect.py +++ b/tests/harnesses/experimental/e2e/test_reconnect.py @@ -168,7 +168,10 @@ async def test_a_hooked_harness_survives_and_blocks(harness_kind: str) -> None: creds, gw = _creds_or_skip(), vercel_ai_gateway() reached = asyncio.Event() - binary = {"claude": "claude", "codex": "codex"}[harness_kind] + # pi renames its process to `pi`: match the name, not a substring. + binary = {"claude": ["-f", "claude"], "codex": ["-f", "codex"]}.get( + harness_kind, ["-x", harness_kind] + ) async def approve(ctx: ApprovalContext) -> None: reached.set() @@ -196,7 +199,7 @@ async def drive() -> None: try: await asyncio.wait_for(reached.wait(), 90) # the gate was hit async with VercelSandbox(name=name) as back: - alive = await back.exec(["pgrep", "-f", binary]) + alive = await back.exec(["pgrep", *binary]) assert ( alive.exit_code == 0 ), "the hooked harness must still be running" @@ -220,10 +223,13 @@ async def test_close_ends_the_harness_and_detach_keeps_it( and leaves no FIFO or pid file behind. detach() leaves it running. """ creds, gw = _creds_or_skip(), vercel_ai_gateway() - binary = {"claude": "claude", "codex": "codex"}[harness_kind] + # pi renames its process to `pi`: match the name, not a substring. + binary = {"claude": ["-f", "claude"], "codex": ["-f", "codex"]}.get( + harness_kind, ["-x", harness_kind] + ) async def alive(ws: Workspace) -> bool: - return (await ws.exec(["pgrep", "-f", binary])).exit_code == 0 + return (await ws.exec(["pgrep", *binary])).exit_code == 0 async with VercelSandbox( **creds, keep=True, execution_time_limit=900 diff --git a/tests/harnesses/experimental/test_extras.py b/tests/harnesses/experimental/test_extras.py index 779f0f45..e3b5e4dc 100644 --- a/tests/harnesses/experimental/test_extras.py +++ b/tests/harnesses/experimental/test_extras.py @@ -28,6 +28,12 @@ def test_codex_needs_no_extra(monkeypatch: pytest.MonkeyPatch) -> None: assert harness.name == "codex" +def test_pi_needs_no_extra(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setitem(sys.modules, "claude_agent_sdk", None) + harness = harnesses.pi(workspace=workspaces.Local(".")) + assert harness.name == "pi" + + def test_import_ai_loads_no_optional_harness_package() -> None: code = ( "import sys, ai\n" diff --git a/tests/harnesses/experimental/test_options.py b/tests/harnesses/experimental/test_options.py index 1413a9cc..743a3789 100644 --- a/tests/harnesses/experimental/test_options.py +++ b/tests/harnesses/experimental/test_options.py @@ -129,8 +129,9 @@ def test_a_dict_of_both_factories_keeps_a_callable_type() -> None: `{"claude-code": claude_code, "codex": codex}` a dict of bare `function` to mypy, and every call through it an error (afk's HARNESSES, examples/harnesses/tui.py).""" - from ai.harnesses.experimental import claude_code, codex + from ai.harnesses.experimental import claude_code, codex, pi - assert set(inspect.signature(claude_code).parameters) <= set( - inspect.signature(codex).parameters - ) + for other in (codex, pi): + assert set(inspect.signature(claude_code).parameters) <= set( + inspect.signature(other).parameters + ) diff --git a/tests/harnesses/experimental/test_pi_adapter.py b/tests/harnesses/experimental/test_pi_adapter.py new file mode 100644 index 00000000..a2d99174 --- /dev/null +++ b/tests/harnesses/experimental/test_pi_adapter.py @@ -0,0 +1,957 @@ +"""The pi adapter against a scripted `pi --mode rpc`, offline. + +The records here are the shapes a real pi 1.3.0 sent when measured. +""" + +from __future__ import annotations + +import asyncio +import json +from collections.abc import Callable +from pathlib import Path +from typing import Any + +import pytest + +from ai.harnesses.experimental import Allow, ApprovalContext, Deny +from ai.harnesses.experimental import errors as harness_errors +from ai.harnesses.experimental._adapters import pi +from ai.types import events +from ai.types import messages as m +from ai.workspaces.experimental import Local +from ai.workspaces.experimental import errors as workspace_errors + +from . import fakes + +SID = "s1" + + +def _update(kind: str, index: int = 0, **fields: Any) -> dict[str, Any]: + return { + "type": "message_update", + "assistantMessageEvent": { + "type": kind, + "contentIndex": index, + **fields, + }, + } + + +def _assistant(*content: dict[str, Any], stop: str = "stop") -> dict[str, Any]: + return { + "role": "assistant", + "content": list(content), + "usage": { + "input": 10, + "output": 5, + "cacheRead": 100, + "cacheWrite": 0, + "cost": {"total": 0.01}, + }, + "stopReason": stop, + } + + +CALL = { + "type": "toolCall", + "id": "toolu_1", + "name": "bash", + "arguments": {"command": "echo hi"}, +} + +# text -> tool call -> tool result -> text: what pi sent for +# "run `echo hi`, then say done". +TOOL_TURN: list[dict[str, Any]] = [ + {"type": "agent_start"}, + {"type": "message_start", "message": {"role": "assistant"}}, + _update("thinking_start"), + _update("thinking_delta", delta="plan"), + _update("thinking_end", content="plan"), + _update("toolcall_start", 1, id="toolu_1", toolName="bash"), + _update("toolcall_delta", 1, delta='{"command": "ec'), + _update("toolcall_end", 1, toolCall=CALL), + { + "type": "message_end", + "message": _assistant( + {"type": "thinking", "thinking": "plan"}, CALL, stop="toolUse" + ), + }, + { + "type": "tool_execution_end", + "toolCallId": "toolu_1", + "toolName": "bash", + "result": {"content": [{"type": "text", "text": "hi\n"}]}, + "isError": False, + }, + {"type": "message_start", "message": {"role": "assistant"}}, + _update("text_start"), + _update("text_delta", delta="done"), + _update("text_end", content="done"), + { + "type": "message_end", + "message": _assistant({"type": "text", "text": "done"}), + }, + {"type": "agent_end", "messages": []}, + {"type": "agent_settled", "aborted": False}, +] + + +class Pi: + """Answers like `pi --mode rpc`: a prompt is accepted, then `turn` + plays out; `abort` settles the turn as aborted.""" + + def __init__( + self, + turn: list[dict[str, Any]] | None = None, + *, + prompt_error: str | None = None, + ) -> None: + self.process = fakes.FakeProcess(self.answer) + self.turn = turn if turn is not None else TOOL_TURN + self.prompt_error = prompt_error + self.on_prompt: Callable[[], None] | None = None + self.pending_abort: dict[str, Any] | None = None + self.extension_loaded = True + self.terminated = False + + def says(self, record: dict[str, Any]) -> None: + self.process.server_says(record) + + def respond(self, message: dict[str, Any], **fields: Any) -> None: + self.says( + { + "type": "response", + "id": message["id"], + "command": message["type"], + **fields, + } + ) + + def answer(self, message: dict[str, Any]) -> None: + kind = message.get("type") + if kind == "prompt": + if self.prompt_error is not None: + self.respond(message, success=False, error=self.prompt_error) + return + self.respond(message, success=True, data={"disposition": "started"}) + if self.on_prompt is not None: + self.on_prompt() + return + for record in self.turn: + self.says(record) + elif kind == "abort": + self.says({"type": "agent_settled", "aborted": True}) + self.respond(message, success=True) + elif kind == "get_state": + self.respond(message, success=True, data={"sessionId": SID}) + elif kind == "get_commands": + commands = [{"name": pi.READY}] if self.extension_loaded else [] + self.respond(message, success=True, data={"commands": commands}) + elif "id" in message: + self.respond(message, success=True) + + def sent(self, kind: str) -> list[dict[str, Any]]: + return [w for w in self.process.written if w.get("type") == kind] + + +def _adapter(server: Pi, approve: Any = None) -> pi.PiAdapter: + adapter = pi.PiAdapter() + adapter.approve = approve + adapter._workspace = fakes.FakeWorkspace() + agent = pi._Agent(server.process, on_dialog=adapter._dialog_handler(SID)) + agent.start() + adapter._agents[SID] = agent + return adapter + + +async def _turn(adapter: pi.PiAdapter, text: str = "go") -> list[Any]: + async with asyncio.timeout(5): + return [e async for e in adapter.turn(SID, [m.TextPart(text=text)])] + + +async def test_a_turn_streams_and_settles_on_agent_settled() -> None: + server = Pi() + adapter = _adapter(server) + seen = await _turn(adapter, "say done") + kinds = [type(e).__name__ for e in seen] + assert kinds == [ + "StreamStart", + "ReasoningStart", + "ReasoningDelta", + "ReasoningEnd", + "ToolStart", + "ToolDelta", + "ToolEnd", + "ToolCallResult", + "TextStart", + "TextDelta", + "TextEnd", + "StreamEnd", + ] + # The arguments travel once, whole: the hydrator builds calls from + # deltas, and pi's fragments are not what it settles on. + delta = next(e for e in seen if isinstance(e, events.ToolDelta)) + assert json.loads(delta.chunk) == {"command": "echo hi"} + result = next(e for e in seen if isinstance(e, events.ToolCallResult)) + assert result.results[0].result == "hi\n" + end = seen[-1] + assert isinstance(end, events.StreamEnd) + assert end.finish_reason == "stop" + assert [type(p).__name__ for p in end.message.parts] == [ + "ReasoningPart", + "ToolCallPart", + "TextPart", + ] + # Two model calls: cached input counts as input, cost rides in raw. + usage = end.usage + assert usage is not None + assert usage.input_tokens == 220 + assert usage.output_tokens == 10 + assert usage.cache_read_tokens == 200 + assert usage.raw == {"cost_usd": 0.02} + assert server.sent("prompt")[0]["message"] == "say done" + # Blocks of different model calls never share an id. + starts = [e.block_id for e in seen if isinstance(e, events.TextStart)] + assert starts == ["2.0"] + + +async def test_images_travel_as_base64() -> None: + server = Pi() + adapter = _adapter(server) + png = m.FilePart(data=b"\x89PNG", media_type="image/png") + async with asyncio.timeout(5): + async for _ in adapter.turn(SID, [m.TextPart(text="look"), png]): + pass + prompt = server.sent("prompt")[0] + assert prompt["images"] == [ + {"type": "image", "data": "iVBORw==", "mimeType": "image/png"} + ] + + +async def test_a_rejected_prompt_without_a_key_is_not_authenticated() -> None: + """Measured: pi refuses the prompt itself when it has no credentials.""" + server = Pi(prompt_error="No API key found for openai.\n\nUse /login") + adapter = _adapter(server) + with pytest.raises(workspace_errors.NotAuthenticatedError): + await _turn(adapter) + + +async def test_a_failed_model_call_fails_the_turn() -> None: + server = Pi( + [ + { + "type": "message_end", + "message": { + **_assistant(stop="error"), + "errorMessage": "529 overloaded", + }, + }, + {"type": "agent_settled", "aborted": False}, + ] + ) + adapter = _adapter(server) + with pytest.raises(harness_errors.TurnFailedError, match="overloaded"): + await _turn(adapter) + + +async def test_the_hook_sees_the_call_and_can_rewrite_it() -> None: + server = Pi() + seen: list[ApprovalContext] = [] + + async def approve(ctx: ApprovalContext) -> Allow: + seen.append(ctx) + return Allow(input={"command": "echo rewritten"}) + + adapter = _adapter(server, approve) + answered = asyncio.Event() + + def ask() -> None: + server.says( + { + "type": "extension_ui_request", + "id": "ui-1", + "method": "input", + "title": pi.APPROVE, + "placeholder": json.dumps( + { + "toolCallId": "toolu_1", + "toolName": "bash", + "input": {"command": "echo hi"}, + } + ), + } + ) + + def answer(message: dict[str, Any]) -> None: + if message.get("type") == "extension_ui_response": + answered.set() + for record in TOOL_TURN: + server.says(record) + else: + Pi.answer(server, message) + + server.on_prompt = ask + server.process.answer = answer + seen_events = await _turn(adapter) + assert answered.is_set() + assert seen[0].call.name == "bash" + assert seen[0].call.input == {"command": "echo hi"} + assert seen[0].harness == "pi" + reply = server.sent("extension_ui_response")[0] + assert reply["id"] == "ui-1" + assert json.loads(reply["value"]) == { + "decision": "allow", + "input": {"command": "echo rewritten"}, + } + assert isinstance(seen_events[-1], events.StreamEnd) + + +async def test_deny_and_stop_answers_then_aborts() -> None: + server = Pi() + + async def deny(ctx: ApprovalContext) -> Deny: + return Deny("not that", stop=True) + + adapter = _adapter(server, deny) + server.on_prompt = lambda: server.says( + { + "type": "extension_ui_request", + "id": "ui-1", + "method": "input", + "title": pi.APPROVE, + "placeholder": json.dumps({"toolCallId": "t", "toolName": "bash"}), + } + ) + seen = await _turn(adapter) + reply = json.loads(server.sent("extension_ui_response")[0]["value"]) + assert reply == {"decision": "deny", "reason": "not that", "stop": True} + # The answer first, then the abort: pi must hear the denial. + order = [w["type"] for w in server.process.written] + assert order.index("extension_ui_response") < order.index("abort") + end = seen[-1] + assert isinstance(end, events.StreamEnd) + assert end.finish_reason == "cancelled" + + +async def test_a_failing_hook_denies_and_is_reported() -> None: + server = Pi() + reported: list[str] = [] + + async def broken(ctx: ApprovalContext) -> Allow: + raise ValueError("boom") + + adapter = _adapter(server, broken) + adapter.on_approval_error = lambda sid, detail: reported.append(detail) + request = {"method": "input", "title": pi.APPROVE, "placeholder": "{}"} + decision = await adapter._judge(SID, request) + assert decision["decision"] == "deny" + assert reported == ["ValueError: boom"] + + +async def test_another_extensions_dialog_is_dismissed() -> None: + """Nobody is at a terminal to answer it; the agent must not wait.""" + server = Pi() + adapter = _adapter(server) + agent = adapter._agents[SID] + await adapter._dialog_handler(SID)( + agent, + {"id": "x", "method": "confirm", "title": "Clear session?"}, + ) + await adapter._dialog_handler(SID)( + agent, {"id": "y", "method": "notify", "message": "hello"} + ) + assert server.sent("extension_ui_response") == [ + {"type": "extension_ui_response", "id": "x", "cancelled": True} + ] + + +async def test_stop_settles_the_turn_as_cancelled() -> None: + server = Pi([]) + adapter = _adapter(server) + turn = adapter.turn(SID, [m.TextPart(text="go")]) + assert isinstance(await anext(turn), events.StreamStart) + reader = asyncio.ensure_future(_rest(turn)) + await adapter.stop(SID) + seen = await asyncio.wait_for(reader, 5) + assert isinstance(seen[-1], events.StreamEnd) + assert seen[-1].finish_reason == "cancelled" + + +async def _rest(turn: Any) -> list[Any]: + return [e async for e in turn] + + +async def test_a_dead_pi_crashes_the_turn() -> None: + server = Pi([]) + adapter = _adapter(server) + server.on_prompt = lambda: server.process._lines.put_nowait("") + with pytest.raises(harness_errors.AgentCrashedError): + await _turn(adapter) + + +def test_history_follows_the_active_branch_only() -> None: + lines = [ + {"type": "session", "id": "S", "cwd": "/w"}, + { + "type": "message", + "id": "a", + "parentId": None, + "message": {"role": "user", "content": "first"}, + }, + { + "type": "message", + "id": "b", + "parentId": "a", + "message": _assistant({"type": "text", "text": "abandoned"}), + }, + { + "type": "message", + "id": "c", + "parentId": "a", + "message": _assistant(CALL, stop="toolUse"), + }, + { + "type": "message", + "id": "d", + "parentId": "c", + "message": { + "role": "toolResult", + "toolCallId": "toolu_1", + "toolName": "bash", + "content": [{"type": "text", "text": "hi"}], + "isError": True, + }, + }, + {"type": "label", "id": "e", "parentId": "d"}, + ] + text = "\n".join(json.dumps(line) for line in lines) + messages = [ + msg + for entry in pi._active_branch(text) + for msg in pi._stored_to_messages(entry, replay=True) + ] + assert [msg.role for msg in messages] == ["user", "assistant", "tool"] + result = messages[2].parts[0] + assert isinstance(result, m.ToolResultPart) + assert (result.result, result.result_kind) == ("hi", "error") + assert all(msg.replay for msg in messages) + + +@pytest.fixture +def agent_dir(tmp_path: Path) -> Path: + return tmp_path / "agent" + + +async def test_a_staged_history_is_listed_and_read_back( + tmp_path: Path, agent_dir: Path +) -> None: + """The store helper, for real: write where pi looks, find it there.""" + project = tmp_path / "project" + project.mkdir() + workspace = Local(project, env={"PI_CODING_AGENT_DIR": str(agent_dir)}) + adapter = pi.PiAdapter() + adapter._workspace = workspace + history = [ + m.Message(role="user", parts=[m.TextPart(text="the word is HERON")]), + m.Message( + role="assistant", + parts=[ + m.ReasoningPart(text="hmm"), + m.TextPart(text="Reading."), + m.ToolCallPart( + tool_call_id="c1", + tool_name="read", + tool_args='{"path": "a.txt"}', + ), + ], + ), + m.Message( + role="tool", + parts=[ + m.ToolResultPart( + tool_call_id="c1", tool_name="read", result="milk" + ) + ], + ), + m.Message(role="assistant", parts=[m.TextPart(text="It says milk.")]), + ] + with pytest.warns(UserWarning, match="reasoning"): + session_id = await adapter.stage_history(history) + infos = await adapter.list_sessions() + assert [i.session_id for i in infos] == [session_id] + assert infos[0].title == "the word is HERON" + assert infos[0].kind == "pi" + files = list(agent_dir.glob("sessions/*/*.jsonl")) + assert len(files) == 1 + assert files[0].parent.name.startswith("--") and files[0].name.endswith( + f"_{session_id}.jsonl" + ) + back = await adapter.history(session_id) + assert [msg.role for msg in back] == [ + "user", + "assistant", + "tool", + "assistant", + ] + call = back[1].parts[-1] + assert isinstance(call, m.ToolCallPart) + assert json.loads(call.tool_args) == {"path": "a.txt"} + with pytest.raises(harness_errors.ResumeFailedError): + await adapter.history("not-a-session") + + +def test_handoff_merges_assistant_messages_and_keeps_tool_results() -> None: + from ai.harnesses.experimental import _handoff + + rendered = _handoff.to_pi_messages( + [ + m.Message(role="user", parts=[m.TextPart(text="hi")]), + m.Message(role="assistant", parts=[m.TextPart(text="one")]), + m.Message( + role="assistant", + parts=[ + m.ToolCallPart( + tool_call_id="c1", tool_name="Bash", tool_args="{}" + ) + ], + ), + m.Message( + role="tool", + parts=[ + m.ToolResultPart( + tool_call_id="c1", + tool_name="Bash", + result={"ok": True}, + result_kind="error", + ) + ], + ), + ], + timestamp=1, + ) + assert [r["role"] for r in rendered] == ["user", "assistant", "toolResult"] + assert [b["type"] for b in rendered[1]["content"]] == ["text", "toolCall"] + assert rendered[1]["stopReason"] == "toolUse" + assert rendered[2]["content"] == [{"type": "text", "text": '{"ok": true}'}] + assert rendered[2]["isError"] is True + + +def test_options_reach_pi_in_its_own_spelling() -> None: + from ai.workspaces.experimental import vercel_ai_gateway + + adapter = pi.PiAdapter( + model="anthropic/claude-haiku-4.5", + gateway=vercel_ai_gateway(api_key="k"), + effort="high", + instructions="be brief", + mcp_servers={"web": {"type": "http", "url": "https://mcp.example"}}, + tools=[], + disallowed_tools=["bash"], + ) + adapter._workspace = fakes.FakeWorkspace() + adapter._extension = "/home/.ai-python/pi/extension.js" + argv, env, running_as = adapter.tui_launch(None) + assert argv[argv.index("--model") + 1] == ( + "vercel-ai-gateway/anthropic/claude-haiku-4.5" + ) + assert argv[argv.index("--thinking") + 1] == "high" + assert argv[argv.index("--append-system-prompt") + 1] == "be brief" + assert "--no-tools" in argv + assert argv[argv.index("--exclude-tools") + 1] == "bash" + assert argv[argv.index("--session-id") + 1] == running_as + assert argv[argv.index("-e") + 1] == adapter._extension + assert env["AI_GATEWAY_API_KEY"] == "k" + config = json.loads(env[pi.CONFIG_VAR]) + # The person at the TUI is the gate; the type key is claude's, not pi's. + assert config == { + "approve": False, + "mcp": {"web": {"url": "https://mcp.example"}}, + "gateway": {}, + } + + +def test_pi_refuses_max_turns() -> None: + from ai.harnesses.experimental import pi as make_pi + + with pytest.raises(harness_errors.UnsupportedError, match="max_turns"): + make_pi(workspace=fakes.FakeWorkspace(), max_turns=3) + + +class _LaunchFails(fakes.FakeWorkspace): + """A workspace whose launches fail with what `exc` names.""" + + def __init__(self, exc: BaseException) -> None: + super().__init__() + self.exc = exc + + async def spawn(self, argv: Any, *, env: Any = None) -> Any: + raise self.exc + + +class _Lock: + def __init__(self) -> None: + self.held: set[str] = set() + + async def acquire(self, session_id: str, **kwargs: Any) -> None: + self.held.add(session_id) + + async def release(self, session_id: str) -> None: + self.held.discard(session_id) + + +@pytest.mark.parametrize( + "exc", [TimeoutError(), asyncio.CancelledError(), OSError("no pi")] +) +@pytest.mark.parametrize("verb", ["resume", "fork"]) +async def test_a_failed_launch_releases_its_claim( + exc: BaseException, verb: str, monkeypatch: pytest.MonkeyPatch +) -> None: + """Codex review: only a HarnessError released the lock, so a timeout, + a cancellation or a workspace error left the conversation claimed and + every other client got SessionBusyError until the claim went stale.""" + adapter = pi.PiAdapter() + adapter._workspace = _LaunchFails(exc) + adapter._extension = "/ext.js" + lock = _Lock() + monkeypatch.setattr(adapter, "_lock", lock) + + async def found(session_id: str) -> dict[str, Any]: + return {"id": session_id, "path": f"/s/{session_id}.jsonl"} + + monkeypatch.setattr(adapter, "_find", found) + # Not a harness failure: it propagates as itself, never renamed. + with pytest.raises(type(exc)): + if verb == "resume": + await adapter.resume_session("S") + else: + await adapter.fork("S") + assert lock.held == set() + assert adapter._held == set() + + +async def test_stop_withdraws_an_approval_nobody_has_answered() -> None: + """Codex review, then measured on pi 1.3.0: pi answers `abort` only once + it is idle, and a tool call waiting on the hook keeps it busy — stop() + timed out after 60s and the next run waited behind it. The question is + withdrawn first now: pi hears `cancelled`, blocks the call, goes idle.""" + server = Pi([]) + parked = asyncio.Event() + + async def never(ctx: ApprovalContext) -> Allow: + parked.set() + await asyncio.Event().wait() # a person who never answers + return Allow() + + adapter = _adapter(server, never) + asking = True + + def answer(message: dict[str, Any]) -> None: + nonlocal asking + kind = message.get("type") + if kind == "extension_ui_response": + asking = False + if server.pending_abort is not None: + Pi.answer(server, server.pending_abort) + elif kind == "abort" and asking: + server.pending_abort = message # like pi: busy, no answer yet + else: + Pi.answer(server, message) + + server.process.answer = answer + server.on_prompt = lambda: server.says( + { + "type": "extension_ui_request", + "id": "ui-1", + "method": "input", + "title": pi.APPROVE, + "placeholder": json.dumps({"toolCallId": "t", "toolName": "bash"}), + } + ) + turn = adapter.turn(SID, [m.TextPart(text="go")]) + assert isinstance(await anext(turn), events.StreamStart) + reader = asyncio.ensure_future(_rest(turn)) + await asyncio.wait_for(parked.wait(), 2) + async with asyncio.timeout(2): + await adapter.stop(SID) + seen = await asyncio.wait_for(reader, 2) + assert isinstance(seen[-1], events.StreamEnd) + assert seen[-1].finish_reason == "cancelled" + assert server.sent("extension_ui_response") == [ + {"type": "extension_ui_response", "id": "ui-1", "cancelled": True} + ] + assert adapter._agents[SID].pending == {} + + +def test_without_a_model_pi_chooses_as_it_always_does() -> None: + """No model means pi's own choice, gateway or not — measured: a fresh pi + with only a gateway key picks its own default gateway model, and with + no credentials refuses the prompt (NotAuthenticatedError). The SDK + passes nothing that would change either.""" + from ai.workspaces.experimental import vercel_ai_gateway + + for gateway in (None, vercel_ai_gateway(api_key="k")): + adapter = pi.PiAdapter(gateway=gateway) + adapter._workspace = fakes.FakeWorkspace() + adapter._extension = "/ext.js" + argv, _, _ = adapter.tui_launch(None) + assert "--model" not in argv and "--provider" not in argv + + +class _Spawns(fakes.FakeWorkspace): + """A workspace whose every launch is the scripted pi.""" + + def __init__(self, server: Pi) -> None: + super().__init__() + self.server = server + + async def spawn(self, argv: Any, *, env: Any = None) -> Any: + return _Terminable(self.server) + + +class _Terminable(fakes.FakeProcess): + """The scripted pi's process, noting whether it was terminated.""" + + def __init__(self, server: Pi) -> None: + super().__init__(server.process.answer) + self._lines = server.process._lines + self.written = server.process.written + self.server = server + + async def terminate(self) -> None: + self.server.terminated = True + await super().terminate() + + +def _launcher(server: Pi) -> pi.PiAdapter: + adapter = pi.PiAdapter() + adapter._workspace = _Spawns(server) + adapter._extension = "/ext.js" + return adapter + + +async def test_a_session_without_the_extension_is_refused() -> None: + """Codex review: a gate that did not load is a gate that fails open. + The extension's marker command must be listed, or nothing starts.""" + server = Pi() + server.extension_loaded = False + adapter = _launcher(server) + with pytest.raises(harness_errors.AgentCrashedError, match="extension"): + await adapter._spawn(SID, ["--session-id", SID]) + assert server.terminated + assert SID not in adapter._agents + + +async def test_a_process_whose_claim_fails_does_not_outlive_it( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Codex review: a failed pid claim left pi registered and running + while the caller released the lock — a writer nobody accounted for.""" + server = Pi() + adapter = _launcher(server) + + class Lock: + async def claim_pid(self, *args: Any, **kwargs: Any) -> None: + raise OSError("the workspace went away") + + monkeypatch.setattr(adapter, "_lock", Lock()) + with pytest.raises(OSError): + await adapter._spawn(SID, ["--session-id", SID]) + assert server.terminated + assert SID not in adapter._agents + + +async def test_closing_a_session_ends_the_turn_reading_it() -> None: + """Codex review: close() stopped the event producer and nobody told the + turn waiting on it, so it waited forever.""" + server = Pi([]) + adapter = _adapter(server) + turn = adapter.turn(SID, [m.TextPart(text="go")]) + assert isinstance(await anext(turn), events.StreamStart) + reader = asyncio.ensure_future(_rest(turn)) + await asyncio.sleep(0) + await adapter.close_session(SID) + with pytest.raises(harness_errors.HarnessClosedError): + await asyncio.wait_for(reader, 2) + + +async def test_reopening_a_session_hands_it_over( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Codex review: the lock lets its own holder through, so a second + resume spawned a second pi on the same transcript and orphaned the + first. Now the first is ended before the second starts — and while a + turn runs on it, the resume is refused.""" + server = Pi() + adapter = _adapter(server) + old = adapter._agents[SID] + spawned: list[Any] = [] + + async def found(session_id: str) -> dict[str, Any]: + return {"id": session_id, "path": "/s.jsonl"} + + async def spawn(session_id: str, args: list[str]) -> None: + assert session_id not in adapter._agents, "one pi at a time" + spawned.append(session_id) + + monkeypatch.setattr(adapter, "_find", found) + monkeypatch.setattr(adapter, "_spawn", spawn) + adapter._running.add(SID) + with pytest.raises(harness_errors.SessionBusyError): + await adapter.resume_session(SID) + assert adapter._agents[SID] is old and spawned == [] + adapter._running.discard(SID) + await adapter.resume_session(SID) + assert spawned == [SID] + assert old.dead is not None, "the first pi was closed" + + +async def test_a_question_asked_while_stopping_is_withdrawn_at_once() -> None: + """Codex review: an approval that arrived after stop() withdrew the + pending ones still ran the hook, and `abort` waited on it.""" + server = Pi() + asked: list[Any] = [] + + async def approve(ctx: ApprovalContext) -> Allow: + asked.append(ctx) + return Allow() + + adapter = _adapter(server, approve) + adapter._interrupted.add(SID) + await adapter._dialog_handler(SID)( + adapter._agents[SID], + {"id": "late", "method": "input", "title": pi.APPROVE}, + ) + assert asked == [] + assert server.sent("extension_ui_response") == [ + {"type": "extension_ui_response", "id": "late", "cancelled": True} + ] + + +def test_history_keeps_what_pi_showed_the_model() -> None: + """Codex review: extension-injected context was dropped from history, + so a resumed or handed-off conversation lost what the agent was told.""" + entries: list[dict[str, Any]] = [ + {"type": "custom_message", "content": "ticket #42 is P1"}, + { + "type": "message", + "message": { + "role": "bashExecution", + "command": "ls", + "output": "a.txt", + }, + }, + { + "type": "message", + "message": { + "role": "bashExecution", + "command": "env", + "output": "SECRET=1", + "excludeFromContext": True, + }, + }, + {"type": "branch_summary", "summary": "tried approach A first"}, + {"type": "compaction", "summary": "everything above"}, + {"type": "model_change", "provider": "x"}, + ] + messages = [msg for e in entries for msg in pi._stored_to_messages(e)] + assert [(msg.role, msg.text) for msg in messages] == [ + ("user", "ticket #42 is P1"), + ("user", "Ran `ls`\n```\na.txt\n```"), + ("user", "tried approach A first"), + ] + + +def test_handoff_never_makes_instructions_the_agents_own_words() -> None: + """Codex review: a system message fell through to the assistant branch + and merged into the model's reply. Measured: as a pi system message it + replaces pi's own prompt and tools, so it travels as a user note.""" + from ai.harnesses.experimental import _handoff + + rendered = _handoff.to_pi_messages( + [ + m.Message(role="assistant", parts=[m.TextPart(text="hello")]), + m.Message(role="system", parts=[m.TextPart(text="be terse")]), + m.Message(role="assistant", parts=[m.TextPart(text="ok")]), + ], + timestamp=1, + ) + assert [r["role"] for r in rendered] == ["assistant", "user", "assistant"] + assert "be terse" in rendered[1]["content"][0]["text"] + assert rendered[0]["content"] == [{"type": "text", "text": "hello"}] + + +async def test_a_turn_stopped_before_its_prompt_is_sent_sends_none( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Codex review: stop() while the prompt's files were being read found + pi idle, and the prompt went out anyway — a cancelled turn that still + called the model and ran tools.""" + server = Pi() + adapter = _adapter(server) + + async def slow(prompt: Any) -> dict[str, Any]: + await adapter.stop(SID) # the Stop button, mid-download + return {"type": "prompt", "message": "go"} + + monkeypatch.setattr(pi, "_prompt_command", slow) + seen = await _turn(adapter) + assert server.sent("prompt") == [] + end = seen[-1] + assert isinstance(end, events.StreamEnd) + assert end.finish_reason == "cancelled" + + +async def test_a_failed_attempt_is_not_part_of_the_retried_answer() -> None: + """Codex review: pi retries a failed model call, and both attempts end + in `message_end` — the half-written first one was glued to the reply.""" + failed = { + **_assistant({"type": "text", "text": "Half an ans"}, stop="error"), + "errorMessage": "529 overloaded", + } + server = Pi( + [ + {"type": "message_end", "message": failed}, + {"type": "auto_retry_start", "attempt": 1}, + { + "type": "message_end", + "message": _assistant({"type": "text", "text": "The answer"}), + }, + {"type": "auto_retry_end", "success": True}, + {"type": "agent_settled", "aborted": False}, + ] + ) + seen = await _turn(_adapter(server)) + end = seen[-1] + assert isinstance(end, events.StreamEnd) + assert end.finish_reason == "stop" + assert end.message.text == "The answer" + + +async def test_two_resumes_at_once_start_one_pi_at_a_time( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Codex review: two resumes racing through the hand-over each started + a pi, one of them untracked.""" + adapter = _adapter(Pi()) + live: list[str] = [] + most = 0 + + async def found(session_id: str) -> dict[str, Any]: + return {"id": session_id, "path": "/s.jsonl"} + + async def spawn(session_id: str, args: list[str]) -> None: + nonlocal most + live.append(session_id) + most = max(most, len(live)) + await asyncio.sleep(0.01) + adapter._agents[session_id] = pi._Agent( + fakes.FakeProcess(), on_dialog=None + ) + live.remove(session_id) + + monkeypatch.setattr(adapter, "_find", found) + monkeypatch.setattr(adapter, "_spawn", spawn) + await asyncio.gather( + adapter.resume_session(SID), adapter.resume_session(SID) + ) + assert most == 1 From 13df1155d2184f88c879a3b6d086bae575705ba3 Mon Sep 17 00:00:00 2001 From: George Date: Wed, 28 Oct 2026 12:31:46 -0700 Subject: [PATCH 2/7] harnesses: a pi conversation handed to pi stays on its model; afk drives pi --- examples/apps/afk/README.md | 9 +- examples/apps/afk/afk/cli.py | 7 +- examples/apps/afk/afk/rows.py | 10 +-- examples/apps/afk/afk/tui.py | 6 +- examples/apps/afk/tests/test_tui.py | 18 +++- src/ai/harnesses/experimental/_adapters/pi.py | 76 ++++++++++++++++- src/ai/harnesses/experimental/_handoff.py | 28 +++++-- .../harnesses/experimental/test_pi_adapter.py | 82 +++++++++++++++++++ 8 files changed, 213 insertions(+), 23 deletions(-) diff --git a/examples/apps/afk/README.md b/examples/apps/afk/README.md index fb6ea115..89d192d1 100644 --- a/examples/apps/afk/README.md +++ b/examples/apps/afk/README.md @@ -1,8 +1,8 @@ # afk — step away; your agents don't A small CLI on `ai.harnesses` for one everyday situation: you have been working -with `claude` or `codex` in a directory, and you need to leave. `afk` lists -the conversations in that directory across both harnesses, opens one in the +with `claude`, `codex` or `pi` in a directory, and you need to leave. `afk` +lists the conversations in that directory across every harness, opens one in the harness's own TUI in a sandbox, lets you detach and close the laptop, pushes work to run unattended, and comes back later to attach, watch, bring home, or stop. Nothing has to be started "the afk way": it reads the @@ -36,7 +36,7 @@ A key means the same thing on every row it appears on: | key | | |---|---| -| Enter | into the conversation: one here resumes in its harness's own TUI (`claude --resume ` or `codex resume `), a pushed one is attached. One here that is open in another terminal is taken over, after a y: that process ends first, then the TUI opens here (`afk stop && claude --resume `). Where afk cannot end what drives it, Enter peeks | +| Enter | into the conversation: one here resumes in its harness's own TUI (`claude --resume `, `codex resume ` or `pi --session `), a pushed one is attached. One here that is open in another terminal is taken over, after a y: that process ends first, then the TUI opens here (`afk stop && claude --resume `). Where afk cannot end what drives it, Enter peeks | | `p` | peek | | `s` | stop: a pushed one's sandbox (its files are kept), or the process a conversation here is open in | | `r` | rm: delete a pushed one's sandbox and snapshot | @@ -47,7 +47,8 @@ On a stopped or expired one, Enter resumes it. `stop` and `rm` ask first. A conversation here that is open in another terminal is `in use`. afk sees any `claude`, but a `codex` only when it was started as `codex resume `: -a fresh `codex` names no conversation afk can find. +a fresh `codex` names no conversation afk can find. A `pi` started in a +terminal is never seen: it shows no conversation in its process name. ## The verbs diff --git a/examples/apps/afk/afk/cli.py b/examples/apps/afk/afk/cli.py index d6d4af82..43503612 100644 --- a/examples/apps/afk/afk/cli.py +++ b/examples/apps/afk/afk/cli.py @@ -121,7 +121,8 @@ async def listing( if not here and away is None: print() print( - " no conversations here yet — start one with `claude` or `codex`" + " no conversations here yet — start one with `claude`, `codex` " + + "or `pi`" ) print() return @@ -214,8 +215,8 @@ async def away() -> list[Row]: for note in notes: print(f" {DIM}{note}{RESET}") print( - "\n no conversations here yet — start one with `claude` or " - "`codex`\n" + "\n no conversations here yet — start one with `claude`, " + "`codex` or `pi`\n" ) return None shown = here if everything else here[:LIST_CAP] diff --git a/examples/apps/afk/afk/rows.py b/examples/apps/afk/afk/rows.py index 84aac92a..a0a7b049 100644 --- a/examples/apps/afk/afk/rows.py +++ b/examples/apps/afk/afk/rows.py @@ -16,7 +16,7 @@ from pydantic import BaseModel -from ai.harnesses.experimental import Handle, claude_code, codex +from ai.harnesses.experimental import Handle, claude_code, codex, pi from ai.harnesses.experimental.errors import HarnessError from ai.workspaces.experimental import Local, VercelSandbox from ai.workspaces.experimental.errors import ( @@ -34,8 +34,8 @@ from .state import Remote, State -HARNESSES = {"claude-code": claude_code, "codex": codex} -KIND_SHORT = {"claude-code": "claude", "codex": "codex"} +HARNESSES = {"claude-code": claude_code, "codex": codex, "pi": pi} +KIND_SHORT = {"claude-code": "claude", "codex": "codex", "pi": "pi"} CONTINUE = ( CONTINUE_MARK @@ -118,7 +118,7 @@ def age(ts: int | None) -> str: async def local_rows( cwd: Path, gateway: Gateway | None ) -> tuple[list[Row], list[str], list[str]]: - """This directory's conversations, both harnesses. + """This directory's conversations, every harness. Returns rows, notes (a harness that is not installed here is a note, not an error), and the harnesses that answered. @@ -135,7 +135,7 @@ async def ask(factory: Any) -> Any: except HarnessError as exc: return exc - # Both harnesses at once: neither waits on the other's store. + # Every harness at once: none waits on another's store. results = await asyncio.gather(*(ask(f) for f in HARNESSES.values())) for kind, infos in zip(HARNESSES, results, strict=True): if isinstance(infos, HarnessError): diff --git a/examples/apps/afk/afk/tui.py b/examples/apps/afk/afk/tui.py index 7f3f80a3..4dead88b 100644 --- a/examples/apps/afk/afk/tui.py +++ b/examples/apps/afk/afk/tui.py @@ -32,7 +32,11 @@ """key -> (what the footer calls it, the commands it runs, in order, each only if the one before it succeeded: `a && b`).""" -RESUME = {"claude-code": ["claude", "--resume"], "codex": ["codex", "resume"]} +RESUME = { + "claude-code": ["claude", "--resume"], + "codex": ["codex", "resume"], + "pi": ["pi", "--session"], +} """Each harness's own command to reopen a conversation in its TUI here.""" diff --git a/examples/apps/afk/tests/test_tui.py b/examples/apps/afk/tests/test_tui.py index 4c43324f..0b25ea6e 100644 --- a/examples/apps/afk/tests/test_tui.py +++ b/examples/apps/afk/tests/test_tui.py @@ -6,7 +6,7 @@ import pytest from afk.rows import Row -from afk.tui import Picker +from afk.tui import Picker, verbs from rich.text import Text from textual.widgets import Static @@ -235,3 +235,19 @@ async def test_stop_is_offered_only_with_a_process( async def test_stop_does_nothing_without_a_process() -> None: app, _ = await _press_in_use("down", "s", "y") assert app.return_value is None + + +def test_a_pi_conversation_resumes_in_pi() -> None: + """Enter on a pi conversation here opens pi's own TUI on it: the + command pi itself takes to reopen a session (measured: `pi --session + `, the project's session first).""" + row = Row( + where="here", + kind="pi", + session_id="01a1bbbb", + title="remember the codeword", + status="idle 1h", + ) + assert verbs(row)["enter"] == ("resume", [["pi", "--session", "01a1bbbb"]]) + # pi shows no process per conversation: afk cannot name one to stop. + assert "s" not in verbs(row) diff --git a/src/ai/harnesses/experimental/_adapters/pi.py b/src/ai/harnesses/experimental/_adapters/pi.py index 9086d8a9..c11795ae 100644 --- a/src/ai/harnesses/experimental/_adapters/pi.py +++ b/src/ai/harnesses/experimental/_adapters/pi.py @@ -482,6 +482,51 @@ async def stage_history(self, history: list[messages_.Message]) -> str: } ] parent: str | None = None + # The model the conversation was on, as pi records a switch. pi + # restores a session's model from its last reply's author, which + # `to_pi_messages` keeps for replies pi wrote; this covers a past + # whose last reply is not one. Measured: without either, a pushed + # conversation ran on pi's default model, which the gateway refused. + # Only a conversation pi itself wrote says which; from another + # harness, pi's own choice applies. + written: dict[str, Any] = next( + ( + found + for m in reversed(history) + if m.role == "assistant" + and isinstance( + found := (m.provider_metadata or {}).get(KIND), dict + ) + ), + {}, + ) + changes: list[dict[str, Any]] = [] + if written.get("provider") and written.get("model"): + changes.append( + { + "type": "model_change", + "provider": written["provider"], + "modelId": written["model"], + } + ) + if written.get("thinkingLevel"): + changes.append( + { + "type": "thinking_level_change", + "thinkingLevel": written["thinkingLevel"], + } + ) + for change in changes: + entry_id = uuid.uuid4().hex[:8] + lines.append( + { + **change, + "id": entry_id, + "parentId": parent, + "timestamp": stamp, + } + ) + parent = entry_id for message in _handoff.to_pi_messages( history, timestamp=int(now.timestamp() * 1000) ): @@ -761,7 +806,11 @@ async def turn( hint=self._auth_hint(), ) yield events_.StreamEnd( - message=messages_.Message(role="assistant", parts=parts), + message=messages_.Message( + role="assistant", + parts=parts, + provider_metadata=_written_by(last), + ), # Cost has no home in Usage; pi's own number rides in `raw`. usage=usage.model_copy(update={"raw": {"cost_usd": cost}}), finish_reason=( @@ -1225,7 +1274,14 @@ def _stored_to_messages( if role == "assistant": parts = _assistant_parts(message.get("content")) return ( - [messages_.Message(role="assistant", parts=parts, replay=replay)] + [ + messages_.Message( + role="assistant", + parts=parts, + replay=replay, + provider_metadata=_written_by(message), + ) + ] if parts else [] ) @@ -1240,6 +1296,22 @@ def _stored_to_messages( return [] +def _written_by(message: dict[str, Any]) -> dict[str, Any] | None: + """Which provider and model wrote an assistant message, as pi says. + + Kept on the message (`provider_metadata["pi"]`) so a conversation + handed to another pi continues on the model it was on. + """ + written = { + k: message[k] + for k in ("api", "provider", "model", "thinkingLevel") + if isinstance(message.get(k), str) + } + if not written or written.get("provider") == "ai-python": + return None + return {KIND: written} + + def _usage(raw: dict[str, Any]) -> usage_.Usage: """Return pi's per-message usage. diff --git a/src/ai/harnesses/experimental/_handoff.py b/src/ai/harnesses/experimental/_handoff.py index 645cc858..e6baf731 100644 --- a/src/ai/harnesses/experimental/_handoff.py +++ b/src/ai/harnesses/experimental/_handoff.py @@ -44,6 +44,8 @@ from ...types import messages as messages_ NOTE_PREFIX = "[Handoff from a previous session — earlier reasoning: " +AUTHOR = ("api", "provider", "model") +"""What names a reply's author in a pi session: its API and model.""" SYSTEM_PREFIX = "[Handoff from a previous session — instructions it was given: " @@ -255,10 +257,11 @@ def to_pi_messages( ) -> list[dict[str, Any]]: """Build pi `AgentMessage`s for a session file. - Assistant messages are attributed to a model that is not the one pi - runs (`ai-python/handoff`): pi then treats them as another model's — - it keeps their text and tool calls and normalizes the call ids for its - own provider. Consecutive assistant messages are merged, the way one + Assistant messages pi wrote keep their author (`provider_metadata["pi"]`): + pi restores a session's model from it. Any other is attributed to a + model pi does not run (`ai-python/handoff`), so pi treats it as another + model's — it keeps its text and tool calls and normalizes the call ids + for its own provider. Consecutive assistant messages are merged, the way one model response holds text and calls together. Reasoning is omitted: see the module docstring. """ @@ -349,13 +352,24 @@ def to_pi_messages( if out and out[-1]["role"] == "assistant": out[-1]["content"] += blocks else: + # A reply pi wrote keeps its author: pi restores a session's model + # from its last reply's, and a model it does not know (ours) + # put the conversation on pi's default one (measured). + author = (m.provider_metadata or {}).get("pi") or {} + known = all(isinstance(author.get(k), str) for k in AUTHOR) out.append( { "role": "assistant", "content": blocks, - "api": "ai-python", - "provider": "ai-python", - "model": "handoff", + **( + {k: author[k] for k in AUTHOR} + if known + else { + "api": "ai-python", + "provider": "ai-python", + "model": "handoff", + } + ), "usage": usage, "stopReason": "stop", "timestamp": timestamp, diff --git a/tests/harnesses/experimental/test_pi_adapter.py b/tests/harnesses/experimental/test_pi_adapter.py index a2d99174..49509e9b 100644 --- a/tests/harnesses/experimental/test_pi_adapter.py +++ b/tests/harnesses/experimental/test_pi_adapter.py @@ -955,3 +955,85 @@ async def spawn(session_id: str, args: list[str]) -> None: adapter.resume_session(SID), adapter.resume_session(SID) ) assert most == 1 + + +async def test_a_handed_off_pi_conversation_stays_on_its_model( + tmp_path: Path, agent_dir: Path +) -> None: + """Measured with afk: a pushed pi conversation ran on pi's default model + in the sandbox (refused by the gateway), not the one it was on. History + now says which model wrote each reply, and staging writes it back the + way pi records a switch, so pi restores it on open.""" + project = tmp_path / "project" + project.mkdir() + workspace = Local(project, env={"PI_CODING_AGENT_DIR": str(agent_dir)}) + adapter = pi.PiAdapter() + adapter._workspace = workspace + stored = { + "type": "message", + "message": { + **_assistant({"type": "text", "text": "hi"}), + "api": "anthropic-messages", + "provider": "vercel-ai-gateway", + "model": "anthropic/claude-haiku-4.5", + "thinkingLevel": "high", + }, + } + history = [ + m.Message(role="user", parts=[m.TextPart(text="hello")]), + *pi._stored_to_messages(stored), + ] + assert history[1].provider_metadata == { + "pi": { + "api": "anthropic-messages", + "provider": "vercel-ai-gateway", + "model": "anthropic/claude-haiku-4.5", + "thinkingLevel": "high", + } + } + session_id = await adapter.stage_history(history) + (written,) = agent_dir.glob(f"sessions/*/*_{session_id}.jsonl") + entries = [json.loads(line) for line in written.read_text().splitlines()] + assert [e["type"] for e in entries] == [ + "session", + "model_change", + "thinking_level_change", + "message", + "message", + ] + assert (entries[1]["provider"], entries[1]["modelId"]) == ( + "vercel-ai-gateway", + "anthropic/claude-haiku-4.5", + ) + # pi restores the model from the last reply's author: it is the real one. + reply = entries[-1]["message"] + assert (reply["api"], reply["provider"], reply["model"]) == ( + "anthropic-messages", + "vercel-ai-gateway", + "anthropic/claude-haiku-4.5", + ) + # A chain, root first: pi reads it as one branch. + assert [e.get("parentId") for e in entries[1:]] == [ + None, + *[e["id"] for e in entries[1:-1]], + ] + + +async def test_without_pi_history_no_model_is_carried( + tmp_path: Path, agent_dir: Path +) -> None: + """From another harness there is no pi model to name: pi chooses.""" + project = tmp_path / "project" + project.mkdir() + adapter = pi.PiAdapter() + adapter._workspace = Local( + project, env={"PI_CODING_AGENT_DIR": str(agent_dir)} + ) + session_id = await adapter.stage_history( + [ + m.Message(role="user", parts=[m.TextPart(text="hello")]), + m.Message(role="assistant", parts=[m.TextPart(text="hi")]), + ] + ) + (written,) = agent_dir.glob(f"sessions/*/*_{session_id}.jsonl") + assert "model_change" not in written.read_text() From 2187a02c879c62258b9e14cab3a5ac811669800d Mon Sep 17 00:00:00 2001 From: George Date: Wed, 28 Oct 2026 12:33:25 -0700 Subject: [PATCH 3/7] examples: pi, with a hook that denies, explains and rewrites --- .../content/docs/basics/harnesses.mdx | 2 + examples/.test_scripts/run-examples.py | 6 ++- examples/harnesses/pi.py | 52 +++++++++++++++++++ 3 files changed, 58 insertions(+), 2 deletions(-) create mode 100644 examples/harnesses/pi.py diff --git a/docs/ai-python/content/docs/basics/harnesses.mdx b/docs/ai-python/content/docs/basics/harnesses.mdx index 6818c789..f499748e 100644 --- a/docs/ai-python/content/docs/basics/harnesses.mdx +++ b/docs/ai-python/content/docs/basics/harnesses.mdx @@ -702,6 +702,8 @@ One runnable file per feature: redirect a running turn. - [`examples/harnesses/tool_approval.py`](https://github.com/vercel-labs/ai-python/blob/main/examples/harnesses/tool_approval.py): an approval hook. +- [`examples/harnesses/pi.py`](https://github.com/vercel-labs/ai-python/blob/main/examples/harnesses/pi.py): + your own pi, with a hook that denies, explains and rewrites. - [`examples/harnesses/take_over.py`](https://github.com/vercel-labs/ai-python/blob/main/examples/harnesses/take_over.py): fork a conversation you did not start. - [`examples/harnesses/resume.py`](https://github.com/vercel-labs/ai-python/blob/main/examples/harnesses/resume.py): diff --git a/examples/.test_scripts/run-examples.py b/examples/.test_scripts/run-examples.py index 7b145d86..d5f08730 100755 --- a/examples/.test_scripts/run-examples.py +++ b/examples/.test_scripts/run-examples.py @@ -7,8 +7,9 @@ uv run examples/.test_scripts/run-examples.py --video # also video samples uv run examples/.test_scripts/run-examples.py --e2e # also e2e test scripts uv run examples/.test_scripts/run-examples.py --harnesses - # also harness samples: need logged-in `claude` and `codex` CLIs, and - # AI_GATEWAY_API_KEY plus Vercel credentials for the sandbox ones + # also harness samples: need logged-in `claude`, `codex` and `pi` + # CLIs, and AI_GATEWAY_API_KEY plus Vercel credentials for the + # sandbox ones uv run examples/.test_scripts/run-examples.py --all # everything uv run examples/.test_scripts/run-examples.py --parallel # in parallel uv run examples/.test_scripts/run-examples.py models/stream.py @@ -88,6 +89,7 @@ class Sample: Sample("harnesses/until_done.py", timeout=600.0), Sample("harnesses/steering.py", timeout=300.0), Sample("harnesses/tool_approval.py", timeout=300.0), + Sample("harnesses/pi.py", timeout=300.0), Sample("harnesses/take_over.py", timeout=300.0), Sample("harnesses/resume.py", timeout=300.0), Sample("harnesses/handoff.py", timeout=600.0), diff --git a/examples/harnesses/pi.py b/examples/harnesses/pi.py new file mode 100644 index 00000000..d209a0bf --- /dev/null +++ b/examples/harnesses/pi.py @@ -0,0 +1,52 @@ +"""Drive your own pi from Python: its login, settings and extensions. + +`model` is left unset, so pi runs on your own default, exactly as when you +type `pi`. pi has no permission prompt of its own: with `approve=`, the SDK +loads a small extension into pi that asks your hook about EVERY tool call, +reads included. The hook can allow a call, deny it with a reason the agent +reads and acts on, or rewrite its arguments before it runs. +""" + +import asyncio +import pathlib +import tempfile + +from ai.harnesses import experimental as harnesses +from ai.workspaces import experimental as workspaces + + +async def approve( + ctx: harnesses.ApprovalContext, +) -> harnesses.Decision: + print(f" [asking] {ctx.call.name} {ctx.call.input}") + if ctx.call.name == "bash" and "rm " in str(ctx.call.input.get("command")): + # A command is a string the agent wrote: matching on it is a speed + # bump, not isolation. For isolation, use a VercelSandbox. + return harnesses.Deny("never delete files; move them into trash/") + if ctx.call.name == "write": + content = str(ctx.call.input.get("content", "")) + if not content.endswith("\n"): + # Rewrite the call: every file the agent writes ends in a newline. + return harnesses.Allow( + input={**ctx.call.input, "content": content + "\n"} + ) + return harnesses.Allow() + + +async def main() -> None: + with tempfile.TemporaryDirectory() as tmp: + async with harnesses.pi( + workspace=workspaces.Local(tmp), approve=approve + ) as harness: + result = await harness.run( + "Write HELLO into a.txt and BYE into b.txt with your write " + + "tool, then delete a.txt with rm. Reply with the files " + + "left, in one line." + ) + print(result.text) + for path in sorted(pathlib.Path(tmp).rglob("*.txt")): + print(f"{path.relative_to(tmp)}: {path.read_text()!r}") + print(f"cost: ${(result.usage.raw or {})['cost_usd']:.4f}") + + +asyncio.run(main()) From a01262d697a4b9959258af749b3571e2d776926a Mon Sep 17 00:00:00 2001 From: George Date: Wed, 28 Oct 2026 12:42:22 -0700 Subject: [PATCH 4/7] docs: pi in the tables and lists that came with the rebase --- .../content/docs/basics/harnesses.mdx | 14 +++--- .../content/docs/reference/harnesses.mdx | 44 +++++++++---------- examples/apps/afk/README.md | 6 +-- skills/ai/references/harnesses.md | 11 ++--- 4 files changed, 38 insertions(+), 37 deletions(-) diff --git a/docs/ai-python/content/docs/basics/harnesses.mdx b/docs/ai-python/content/docs/basics/harnesses.mdx index f499748e..1b089c1e 100644 --- a/docs/ai-python/content/docs/basics/harnesses.mdx +++ b/docs/ai-python/content/docs/basics/harnesses.mdx @@ -535,13 +535,13 @@ that `ai.Agent` does not have. What the event stream holds: -| Event detail | `ai.Agent` | Claude Code | Codex | -|---|---|---|---| -| Text deltas | yes | yes | yes | -| Tool calls and their results | yes | yes | yes | -| Tool arguments while they stream | yes | no: sent whole | no: sent whole | -| One `StreamStart`/`StreamEnd` per model call | yes | no: one per turn | no: one per turn | -| Usage per model call | yes | no: per turn, with cost | no: per turn | +| Event detail | `ai.Agent` | Claude Code | Codex | Pi | +|---|---|---|---|---| +| Text deltas | yes | yes | yes | yes | +| Tool calls and their results | yes | yes | yes | yes | +| Tool arguments while they stream | yes | no: sent whole | no: sent whole | no: sent whole | +| One `StreamStart`/`StreamEnd` per model call | yes | no: one per turn | no: one per turn | no: one per turn | +| Usage per model call | yes | no: per turn, with cost | no: per turn | no: per turn, with cost | Code that reads events should not need more than this: text, tool calls and results, the final `StreamEnd`, and `messages`. The SDK does not invent diff --git a/docs/ai-python/content/docs/reference/harnesses.mdx b/docs/ai-python/content/docs/reference/harnesses.mdx index 0ca8e002..e031f34e 100644 --- a/docs/ai-python/content/docs/reference/harnesses.mdx +++ b/docs/ai-python/content/docs/reference/harnesses.mdx @@ -5,7 +5,7 @@ type: reference summary: Reference for claude_code, codex, pi, Harness, Session, Turn, Tui, Result, Handle, approvals, capabilities, and harness errors. --- -`ai.harnesses.experimental` drives coding-agent CLIs, Claude Code and Codex, +`ai.harnesses.experimental` drives coding-agent CLIs, Claude Code, Codex and Pi, on a [workspace](/docs/reference/workspaces). See [Harnesses](/docs/basics/harnesses) for a guide. @@ -315,7 +315,7 @@ One conversation in the CLI's store. | `title` | The conversation's title or first prompt, or `None`. | | `cwd` | The directory it ran in, or `None`. | | `updated_at`, `created_at` | Timestamps as the CLI reports them, or `None`. | -| `running` | `True` when a process has the conversation open. An SDK client or `harness.tui()` holds it, so `resume()` would raise `SessionBusyError`. A CLI started in a terminal shows as running too: any `claude`, and a `codex` started as `codex resume `. `resume()` does not check for those. | +| `running` | `True` when a process has the conversation open. An SDK client or `harness.tui()` holds it, so `resume()` would raise `SessionBusyError`. A CLI started in a terminal shows as running too: any `claude`, and a `codex` started as `codex resume `; a `pi` started in a terminal is never seen. `resume()` does not check for those. | | `pid` | The process that has it open, in the workspace, or `None`. | ## Handle @@ -372,9 +372,9 @@ waiting on the decision. | `call` | The `ToolCall` being decided. | | `history` | The conversation so far. | | `workspace` | The workspace the CLI runs in. | -| `harness` | `"claude-code"` or `"codex"`. | +| `harness` | `"claude-code"`, `"codex"`, or `"pi"`. | | `session_id` | The conversation's id. | -| `prompt` | The turn's prompt on Claude Code; `None` on Codex. | +| `prompt` | The turn's prompt on Claude Code; `None` on Codex and Pi. | ### ToolCall @@ -383,7 +383,7 @@ One tool call, as the CLI described it. | Field | Description | |---|---| | `id` | Call id. | -| `name` | The CLI's native tool name, such as `Bash` or `fileChange`. Codex names shell calls `commandExecution`; the command is in `input`. | +| `name` | The CLI's native tool name, such as `Bash`, `fileChange`, or `bash`. Codex names shell calls `commandExecution`; the command is in `input`. | | `input` | Native arguments. | | `title` | Display text. Do not key policy on it. | | `kind` | A hint the CLI supplied, or `None`. | @@ -398,8 +398,8 @@ harnesses.Allow(input=None, remember=False) ``` Lets the call run. `input` rewrites the arguments where -`capabilities.rewrite_tool_input` is true (Claude Code only). `remember` is -not supported on either harness. An unsupported option denies the call and +`capabilities.rewrite_tool_input` is true (Claude Code and Pi). `remember` +is not supported on any harness. An unsupported option denies the call and raises `UnsupportedError` from the turn. ### Deny @@ -430,21 +430,21 @@ contextvars the turn's caller set are visible inside it. What a harness can do, declared before anything runs. Read it from `harness.capabilities`. -| Capability | Description | Claude Code | Codex | -|---|---|---|---| -| `run` | Run a prompt and wait for the result: `run()`. | yes | yes | -| `stream` | Iterate a turn's events as they arrive: `stream()`. | yes | yes | -| `steer` | Send guidance to the turn that is running: `session.steer()`. | yes | yes | -| `stop` | Interrupt the running turn; it settles as cancelled: `session.stop()`. | yes | yes | -| `resume` | Continue a conversation this process did not create: `resume()`, `harness_from()`. | yes | yes | -| `fork` | Branch a conversation into a new one with the same past: `fork()`. | yes | yes | -| `history` | Read stored transcripts in full, tool calls included: `history()`, `sessions()`. | yes | yes | -| `images` | Images in a prompt. | yes | yes | -| `documents` | PDFs in a prompt. | yes | no | -| `rewrite_tool_input` | Rewrite a call's arguments before it runs: `Allow(input={...})`. | yes | no | -| `deny_reason` | The reason in `Deny(reason)` reaches the agent. Without it the call is still denied, but the reason is dropped. | yes | no | -| `remember_decisions` | Keep an approval for the rest of the session: `Allow(remember=True)`. | no | no | -| `approval` | Which tool calls reach the hook. `"all"`: every call, reads included. `"policy"`: the CLI's own policy decides, so not every call may. | `"all"` | `"policy"` | +| Capability | Description | Claude Code | Codex | Pi | +|---|---|---|---|---| +| `run` | Run a prompt and wait for the result: `run()`. | yes | yes | yes | +| `stream` | Iterate a turn's events as they arrive: `stream()`. | yes | yes | yes | +| `steer` | Send guidance to the turn that is running: `session.steer()`. | yes | yes | yes | +| `stop` | Interrupt the running turn; it settles as cancelled: `session.stop()`. | yes | yes | yes | +| `resume` | Continue a conversation this process did not create: `resume()`, `harness_from()`. | yes | yes | yes | +| `fork` | Branch a conversation into a new one with the same past: `fork()`. | yes | yes | yes | +| `history` | Read stored transcripts in full, tool calls included: `history()`, `sessions()`. | yes | yes | yes | +| `images` | Images in a prompt. | yes | yes | yes | +| `documents` | PDFs in a prompt. | yes | no | no | +| `rewrite_tool_input` | Rewrite a call's arguments before it runs: `Allow(input={...})`. | yes | no | yes | +| `deny_reason` | The reason in `Deny(reason)` reaches the agent. Without it the call is still denied, but the reason is dropped. | yes | no | yes | +| `remember_decisions` | Keep an approval for the rest of the session: `Allow(remember=True)`. | no | no | no | +| `approval` | Which tool calls reach the hook. `"all"`: every call, reads included. `"policy"`: the CLI's own policy decides, so not every call may. | `"all"` | `"policy"` | `"all"` | `approval` is an `ApprovalCoverage`: diff --git a/examples/apps/afk/README.md b/examples/apps/afk/README.md index 89d192d1..c896de4c 100644 --- a/examples/apps/afk/README.md +++ b/examples/apps/afk/README.md @@ -11,7 +11,7 @@ ago is listed like any other. ``` $ afk -~/proj · claude, codex +~/proj · claude, codex, pi here a3f5 claude migrate the billing tests to pytest active 2m @@ -215,8 +215,8 @@ uv tool install "git+https://github.com/vercel-labs/ai-python@harnesses-experime ``` Or from a checkout: `uv tool install ./examples/apps/afk`. Then run `afk` -in any directory where you use `claude` or `codex`. afk keeps its own files -in `~/.afk/` and writes into your project only on `afk pull --files`. +in any directory where you use `claude`, `codex` or `pi`. afk keeps its own +files in `~/.afk/` and writes into your project only on `afk pull --files`. ## Vercel, set up on first use diff --git a/skills/ai/references/harnesses.md b/skills/ai/references/harnesses.md index 09a3ff33..78590776 100644 --- a/skills/ai/references/harnesses.md +++ b/skills/ai/references/harnesses.md @@ -24,9 +24,10 @@ brings its own tools, login, and session store. CLIs only into a sandbox, and there it installs or upgrades to the latest npm release on open (not while a copy of the CLI is running there). - Every harness has full access by default (claude: no permission prompts; - codex: `sandbox="danger-full-access"`). Restrict with `approve=`, codex's - `sandbox=` (on `Local` only; other modes raise `UnsupportedError` in a - sandbox), claude's `tools=`/`disallowed_tools=`, or a `VercelSandbox`. + codex: `sandbox="danger-full-access"`; pi has no permission prompt). + Restrict with `approve=`, codex's `sandbox=` (on `Local` only; other modes + raise `UnsupportedError` in a sandbox), claude's and pi's + `tools=`/`disallowed_tools=`, or a `VercelSandbox`. - Shared options: `instructions=` (added to the CLI's system prompt) and `mcp_servers=` (claude's shape: `{"name": {"command", "args", "env"}}` or `{"name": {"url"}}`). Claude and pi: `tools=`, `disallowed_tools=`. @@ -55,8 +56,8 @@ brings its own tools, login, and session store. - `result.messages` reads in order: assistant message with the tool call, then a `"tool"` message with its result. Safe to pass as `history=`. - A prompt can be a list: `[ "describe this", ai.file_part(png) ]`. Images - on both, PDFs on claude only (`capabilities.documents`), `text/*` files on - both as text. Anything else raises `UnsupportedError` before the turn. + on every harness, PDFs on claude only (`capabilities.documents`), `text/*` + files everywhere as text. Anything else raises `UnsupportedError` before the turn. - To continue a stored conversation, `resume(id)` (its past is in `session.messages`) and run or stream as usual. There is no separate follow API; `history(id, limit=, offset=)` pages by message. From 5be941eda819568a93f78a2cf2e479595edbc70b Mon Sep 17 00:00:00 2001 From: George Date: Wed, 28 Oct 2026 13:06:43 -0700 Subject: [PATCH 5/7] harnesses: pi follows a run an extension command starts; a retried cut-off answer is not the reply --- src/ai/harnesses/experimental/_adapters/pi.py | 30 ++++- .../harnesses/experimental/test_pi_adapter.py | 124 ++++++++++++++++++ 2 files changed, 147 insertions(+), 7 deletions(-) diff --git a/src/ai/harnesses/experimental/_adapters/pi.py b/src/ai/harnesses/experimental/_adapters/pi.py index c11795ae..e6989079 100644 --- a/src/ai/harnesses/experimental/_adapters/pi.py +++ b/src/ai/harnesses/experimental/_adapters/pi.py @@ -712,16 +712,23 @@ async def turn( except _CommandError as exc: raise _turn_failure(str(exc), hint=self._auth_hint()) from exc if (accepted or {}).get("disposition") == "handled": - # An extension command consumed it: no run starts. - yield events_.StreamStart() - yield events_.StreamEnd( - message=messages_.Message(role="assistant", parts=[]), - finish_reason="stop", - ) - return + # An extension command consumed it. The command may start a run + # of its own (`pi.sendUserMessage`) — measured: then pi is + # already streaming when it answers, and that run is this turn's. + state = await agent.call({"type": "get_state"}) + if not (state or {}).get("isStreaming"): + yield events_.StreamStart() + yield events_.StreamEnd( + message=messages_.Message(role="assistant", parts=[]), + finish_reason="stop", + ) + return # Running BEFORE the first event: a consumer may steer on it. self._running.add(session_id) parts: list[Any] = [] + # Where the latest assistant message's parts are in `parts`: pi may + # abandon that attempt and retry it (see `agent_end`). + attempt = (0, 0) usage = usage_.Usage() cost = 0.0 started: set[str] = set() @@ -760,11 +767,13 @@ async def turn( "assistant" ): last = message + begun = len(parts) if message.get("stopReason") != "error": # A failed attempt is not the answer: pi retries it # (or the turn raises below), and its partial text # would be glued onto the reply that replaces it. parts += _assistant_parts(message.get("content")) + attempt = (begun, len(parts)) raw = message.get("usage") or {} usage = usage + _usage(raw) cost += float((raw.get("cost") or {}).get("total") or 0) @@ -785,6 +794,13 @@ async def turn( message=messages_.Message(role="tool", parts=[result]), results=[result], ) + elif kind == "agent_end" and event.get("willRetry"): + # pi abandons the run's last response and retries it as + # a fresh run — after a cut-off (`length`) answer, it + # compacts first. That response is not the answer. + if attempt[1] == len(parts): + del parts[attempt[0] :] + attempt = (len(parts), len(parts)) elif kind == "auto_retry_end" and not event.get("success"): retry_error = str(event.get("finalError") or "") elif kind == "agent_settled": diff --git a/tests/harnesses/experimental/test_pi_adapter.py b/tests/harnesses/experimental/test_pi_adapter.py index 49509e9b..cff8e9c9 100644 --- a/tests/harnesses/experimental/test_pi_adapter.py +++ b/tests/harnesses/experimental/test_pi_adapter.py @@ -1037,3 +1037,127 @@ async def test_without_pi_history_no_model_is_carried( ) (written,) = agent_dir.glob(f"sessions/*/*_{session_id}.jsonl") assert "model_change" not in written.read_text() + + +async def test_an_extension_command_that_starts_a_run_is_followed() -> None: + """Vercel Agent review, then measured: `/ask` answered `handled`, and the + run the command started (`pi.sendUserMessage`) followed — pi was already + streaming. The turn ended empty and left that run unread.""" + server = Pi( + [ + {"type": "agent_start"}, + {"type": "message_start", "message": {"role": "assistant"}}, + { + "type": "message_end", + "message": _assistant({"type": "text", "text": "PONG"}), + }, + {"type": "agent_settled", "aborted": False}, + ] + ) + + def answer(message: dict[str, Any]) -> None: + if message.get("type") == "prompt": + server.respond( + message, success=True, data={"disposition": "handled"} + ) + for record in server.turn: + server.says(record) + elif message.get("type") == "get_state": + server.respond(message, success=True, data={"isStreaming": True}) + else: + Pi.answer(server, message) + + server.process.answer = answer + seen = await _turn(_adapter(server), "/ask PONG") + end = seen[-1] + assert isinstance(end, events.StreamEnd) + assert end.message.text == "PONG" + + +async def test_a_command_that_starts_nothing_ends_the_turn_empty() -> None: + server = Pi() + + def answer(message: dict[str, Any]) -> None: + if message.get("type") == "prompt": + server.respond( + message, success=True, data={"disposition": "handled"} + ) + elif message.get("type") == "get_state": + server.respond(message, success=True, data={"isStreaming": False}) + else: + Pi.answer(server, message) + + server.process.answer = answer + seen = await _turn(_adapter(server), "/noop") + assert [type(e).__name__ for e in seen] == ["StreamStart", "StreamEnd"] + + +async def test_a_cut_off_answer_pi_retries_is_not_part_of_the_reply() -> None: + """Vercel Agent review: pi compacts and retries a final `length` + answer as a fresh run (its docs: "an early final stopReason: length can + select one compact-and-retry recovery attempt"); the cut-off text was + glued onto the retried reply.""" + server = Pi( + [ + { + "type": "message_end", + "message": _assistant( + {"type": "text", "text": "Half an ans"}, stop="length" + ), + }, + {"type": "agent_end", "messages": [], "willRetry": True}, + {"type": "compaction_start", "reason": "overflow"}, + {"type": "compaction_end", "reason": "overflow", "willRetry": True}, + {"type": "agent_start"}, + { + "type": "message_end", + "message": _assistant({"type": "text", "text": "The answer"}), + }, + {"type": "agent_end", "messages": [], "willRetry": False}, + {"type": "agent_settled", "aborted": False}, + ] + ) + seen = await _turn(_adapter(server)) + end = seen[-1] + assert isinstance(end, events.StreamEnd) + assert end.message.text == "The answer" + assert end.finish_reason == "stop" + + +async def test_a_retry_after_tool_work_keeps_the_tool_work() -> None: + """Only the abandoned response goes: calls made before it stay.""" + server = Pi( + [ + { + "type": "message_end", + "message": _assistant(CALL, stop="toolUse"), + }, + { + "type": "tool_execution_end", + "toolCallId": "toolu_1", + "toolName": "bash", + "result": {"content": [{"type": "text", "text": "hi"}]}, + "isError": False, + }, + { + "type": "message_end", + "message": _assistant( + {"type": "text", "text": "cut o"}, stop="length" + ), + }, + {"type": "agent_end", "messages": [], "willRetry": True}, + { + "type": "message_end", + "message": _assistant({"type": "text", "text": "done"}), + }, + {"type": "agent_settled", "aborted": False}, + ] + ) + seen = await _turn(_adapter(server)) + end = seen[-1] + assert isinstance(end, events.StreamEnd) + assert [type(p).__name__ for p in end.message.parts] == [ + "ToolCallPart", + "TextPart", + ] + assert end.message.text == "done" From 240d46bbaab614d66c32ab2da4fc917be189e872 Mon Sep 17 00:00:00 2001 From: George Date: Wed, 28 Oct 2026 14:17:31 -0700 Subject: [PATCH 6/7] harnesses: a session id is a name, never a path Security review: pi reads a project's sessionDir before trust, so a repository can commit a session file whose id is a path; resuming it made the session lock rmtree that path as a stale lock. - The session lock refuses any id that is not a plain name, in Python and again in the helper where the files are; the helper also checks on disk that the lock directory is a direct child of its root, so a planted symlink cannot lead out. - pi's session store skips files whose id breaks pi's own id rule, and resume/fork/history refuse such an id before the lock is touched. --- src/ai/harnesses/experimental/_adapters/pi.py | 19 +++++++- .../harnesses/experimental/_session_lock.py | 44 +++++++++++++++++-- .../harnesses/experimental/test_pi_adapter.py | 33 ++++++++++++++ .../experimental/test_session_lock.py | 42 ++++++++++++++++++ 4 files changed, 132 insertions(+), 6 deletions(-) diff --git a/src/ai/harnesses/experimental/_adapters/pi.py b/src/ai/harnesses/experimental/_adapters/pi.py index e6989079..fc842a09 100644 --- a/src/ai/harnesses/experimental/_adapters/pi.py +++ b/src/ai/harnesses/experimental/_adapters/pi.py @@ -18,6 +18,7 @@ import datetime import json import pathlib +import re import uuid import warnings from typing import TYPE_CHECKING, Any @@ -48,6 +49,9 @@ """The CLI the live suite last passed against. Older ones get a warning.""" APPROVE = "ai-python:approve" """The dialog title the extension asks approvals under.""" +SESSION_ID = re.compile(r"[A-Za-z0-9](?:[A-Za-z0-9._-]{0,198}[A-Za-z0-9])?\Z") +"""pi's rule for a session id (`--session-id`): letters, digits, `.`, `_` +and `-`, starting and ending with a letter or digit. Never a path.""" READY = "ai-python-ready" """The command the extension registers last: proof that it loaded.""" CONFIG_VAR = "AI_PYTHON_PI" @@ -398,6 +402,10 @@ async def _store(self, op: str, session_id: str = "") -> Any: async def _find(self, session_id: str) -> dict[str, Any]: """Return a conversation of THIS workspace, or raise ResumeFailed.""" assert self._workspace is not None + if not SESSION_ID.match(session_id): + raise errors.ResumeFailedError( + session_id, "not a pi session id (see SESSION_ID)" + ) found = await self._store("find", session_id) if not found: raise errors.ResumeFailedError( @@ -1388,10 +1396,15 @@ def _turn_failure(detail: str, *, hint: str = AUTH_HINT) -> Exception: _STORE = r""" -import json, os, sys +import json, os, re, sys op, where, sid = sys.argv[1], sys.argv[2], sys.argv[3] cwd = os.path.realpath(where) +# pi's own rule for a session id. A session file can come from anyone — a +# project's settings choose the directory, before pi asks for trust — and +# its id becomes a lock path and a CLI argument: one that breaks the rule +# is not a session. +ID = re.compile(r"[A-Za-z0-9](?:[A-Za-z0-9._-]{0,198}[A-Za-z0-9])?\Z") def settings(path): @@ -1437,7 +1450,9 @@ def iso_ms(text): def describe(path, full): with open(path) as f: header = json.loads(f.readline() or "{}") - if header.get("type") != "session" or not header.get("id"): + if header.get("type") != "session" or not ID.match( + str(header.get("id") or "") + ): return None if header.get("cwd") not in (cwd, where): return None diff --git a/src/ai/harnesses/experimental/_session_lock.py b/src/ai/harnesses/experimental/_session_lock.py index f18e7615..84d4a62f 100644 --- a/src/ai/harnesses/experimental/_session_lock.py +++ b/src/ai/harnesses/experimental/_session_lock.py @@ -41,6 +41,7 @@ from __future__ import annotations import json +import re import uuid from dataclasses import dataclass from typing import TYPE_CHECKING, Any @@ -53,11 +54,39 @@ STALE_AFTER = 120.0 """Seconds a pid-less lock is honoured before it counts as abandoned.""" +SESSION_ID = re.compile(r"[A-Za-z0-9][A-Za-z0-9._-]{0,199}\Z") +"""What a lock accepts as a session id: a name, never a path. + +Every harness's ids fit (UUIDs, pi's letters, digits, `.`, `_`, `-`). An id +can come from a file someone else wrote — a session store a repository +points to — and the lock turns it into a directory it may delete. +""" + HELPER = r""" -import fcntl, json, os, shutil, subprocess, sys, time +import fcntl, json, os, re, shutil, subprocess, sys, time op, root, args = sys.argv[1], sys.argv[2], json.loads(sys.argv[3]) STALE = float(args.get("stale_after", 120)) +SAFE = re.compile(r"[A-Za-z0-9][A-Za-z0-9._-]{0,199}\Z") + + +# The lock directory of `sid`: one entry directly in the root, or none. An +# id is a name, never a path: an absolute or `..` id would escape the root, +# and a stale lock there is deleted. Checked again on disk, so a symlink +# planted in the root cannot lead out either. +def place(sid, strict=True): + d = os.path.join(root, sid) if isinstance(sid, str) else None + ok = ( + d is not None + and SAFE.match(sid) is not None + and os.path.dirname(os.path.realpath(d)) == os.path.realpath(root) + ) + if ok: + return d + if strict: + sys.stderr.write("not a session id: %r\n" % (sid,)) + sys.exit(3) + return None def started(pid): @@ -93,7 +122,7 @@ def alive(rec): def read(sid): - d = os.path.join(root, sid) + d = place(sid) try: with open(os.path.join(d, "holder.json")) as f: rec = json.load(f) @@ -131,7 +160,7 @@ def live(sid): if op == "holders": out = {} for sid in os.listdir(root): - if not sid.startswith("."): + if not sid.startswith(".") and place(sid, strict=False): rec = live(sid) if rec is not None: out[sid] = rec.get("pid") @@ -141,7 +170,7 @@ def live(sid): with open(os.path.join(root, ".mutex"), "a") as mutex: fcntl.flock(mutex, fcntl.LOCK_EX) sid, token = args["sid"], args["token"] - d = os.path.join(root, sid) + d = place(sid) rec = read(sid) if op == "acquire": if rec is not None and alive(rec): @@ -263,6 +292,13 @@ async def holders(self) -> dict[str, int | None]: } async def _run(self, op: str, **args: Any) -> Any: + sid = args.get("sid") + if sid is not None and not SESSION_ID.match(str(sid)): + # Refused here, and again by the helper where the files are. + raise errors.HarnessError( + f"{sid!r} is not a session id: an id is a name of letters, " + + "digits, '.', '_' and '-', never a path" + ) result = await self._workspace.exec( ["python3", "-c", HELPER, op, await self.root(), json.dumps(args)], timeout=30, diff --git a/tests/harnesses/experimental/test_pi_adapter.py b/tests/harnesses/experimental/test_pi_adapter.py index cff8e9c9..ef3f336e 100644 --- a/tests/harnesses/experimental/test_pi_adapter.py +++ b/tests/harnesses/experimental/test_pi_adapter.py @@ -16,6 +16,7 @@ from ai.harnesses.experimental import Allow, ApprovalContext, Deny from ai.harnesses.experimental import errors as harness_errors from ai.harnesses.experimental._adapters import pi +from ai.harnesses.experimental._session_lock import SessionLock from ai.types import events from ai.types import messages as m from ai.workspaces.experimental import Local @@ -1161,3 +1162,35 @@ async def test_a_retry_after_tool_work_keeps_the_tool_work() -> None: "TextPart", ] assert end.message.text == "done" + + +async def test_a_planted_session_cannot_reach_outside_the_lock( + tmp_path: Path, agent_dir: Path +) -> None: + """Security review: a project's `.pi/settings.json` picks the session + directory (pi reads it before trust), a committed session file there + names an id that is a path, and resuming it deleted that path as a stale + lock. Such a file is not a session: not listed, not found, and resuming + by that id is refused before the lock is touched.""" + victim = tmp_path / "victim" + victim.mkdir() + (victim / "keep.txt").write_text("precious") + project = tmp_path / "project" + (project / ".pi" / "sessions").mkdir(parents=True) + (project / ".pi" / "settings.json").write_text( + json.dumps({"sessionDir": ".pi/sessions"}) + ) + for name, sid in (("a", str(victim)), ("b", "../../victim")): + header = {"type": "session", "id": sid, "cwd": str(project.resolve())} + (project / ".pi" / "sessions" / f"{name}.jsonl").write_text( + json.dumps(header) + "\n" + ) + workspace = fakes.LockWorkspace(project) + adapter = pi.PiAdapter() + adapter._workspace = workspace + adapter._lock = SessionLock(workspace, pi.KIND) + assert await adapter.list_sessions() == [] + for sid in (str(victim), "../../victim"): + with pytest.raises(harness_errors.ResumeFailedError): + await adapter.resume_session(sid) + assert (victim / "keep.txt").read_text() == "precious" diff --git a/tests/harnesses/experimental/test_session_lock.py b/tests/harnesses/experimental/test_session_lock.py index 7ae51d95..3225a7fb 100644 --- a/tests/harnesses/experimental/test_session_lock.py +++ b/tests/harnesses/experimental/test_session_lock.py @@ -215,3 +215,45 @@ async def test_listing_holders_is_one_command( workspace.execs.clear() assert len(await locks[0].holders()) == 5 assert len(workspace.execs) == 1 + + +@pytest.mark.parametrize( + "bad", ["/victim", "../victim", "a/../../victim", "", ".", "..", "-rf"] +) +async def test_a_session_id_is_a_name_never_a_path( + lock: SessionLock, workspace: fakes.LockWorkspace, tmp_path: Path, bad: str +) -> None: + """Security review: an id joined onto the lock root escaped it, and a + stale "lock" there was deleted with rmtree. Refused before anything + runs in the workspace.""" + from ai.harnesses.experimental.errors import HarnessError + + ran = len(workspace.execs) + with pytest.raises(HarnessError, match="not a session id"): + await lock.acquire(bad.replace("victim", str(tmp_path / "victim"))) + assert len(workspace.execs) == ran + + +async def test_the_helper_refuses_an_escaping_id_itself( + lock: SessionLock, workspace: fakes.LockWorkspace, tmp_path: Path +) -> None: + """The helper does not trust its caller: an absolute id, or a symlink + planted in the root, never reaches rmtree. Called directly, past the + check `SessionLock` makes first.""" + from ai.harnesses.experimental._session_lock import HELPER + + victim = tmp_path / "victim" + victim.mkdir() + (victim / "keep.txt").write_text("precious") + root = await lock.root() + Path(root).mkdir(parents=True, exist_ok=True) + (Path(root) / "planted").symlink_to(victim) + for sid in (str(victim), "../victim", "planted"): + args = json.dumps({"sid": sid, "token": "t", "stale_after": 0}) + result = await workspace.exec( + ["python3", "-c", HELPER, "acquire", root, args], timeout=30 + ) + assert result.exit_code != 0 + assert "not a session id" in result.stderr + assert (victim / "keep.txt").read_text() == "precious" + assert await lock.holders() == {} From c2b13a7fe02afea59bdf984f984e7fdd7df91ea5 Mon Sep 17 00:00:00 2001 From: George Date: Wed, 28 Oct 2026 14:26:13 -0700 Subject: [PATCH 7/7] harnesses: every harness refuses a session id that is a path Follow-up to the pi finding, for all harnesses: Harness.resume, fork, history and tui check the id where it enters (it may come from a Handle out of a database), before any adapter turns it into a lock directory, a transcript path or a CLI argument. The remote Claude store refuses a key whose parts would leave the store. --- .../experimental/_adapters/claude_store.py | 18 +++++++++++++ src/ai/harnesses/experimental/_harness.py | 22 +++++++++++++++- .../experimental/test_claude_adapter.py | 26 +++++++++++++++++++ tests/harnesses/experimental/test_session.py | 22 ++++++++++++++++ 4 files changed, 87 insertions(+), 1 deletion(-) diff --git a/src/ai/harnesses/experimental/_adapters/claude_store.py b/src/ai/harnesses/experimental/_adapters/claude_store.py index 17d5337b..b16447c8 100644 --- a/src/ai/harnesses/experimental/_adapters/claude_store.py +++ b/src/ai/harnesses/experimental/_adapters/claude_store.py @@ -71,6 +71,24 @@ async def discover( return cls(workspace, directory or await config_dir(workspace)) def _path(self, key: sdk_types.SessionKey) -> str: + # Every part becomes a path in the workspace: a part that is itself + # a path (absolute, `..`) would read or write outside the store. A + # project key or session id is one name; a subagent's subpath nests. + for part, nests in ( + (key["project_key"], False), + (key["session_id"], False), + (key.get("subpath"), True), + ): + if part is None: + continue + names = part.split("/") if nests else [part] + if ( + not part + or part.startswith("/") + or (not nests and "/" in part) + or any(n in ("", ".", "..") for n in names) + ): + raise ValueError(f"not a store key: {part!r}") project_dir = posixpath.join(self._projects, key["project_key"]) subpath = key.get("subpath") if subpath: diff --git a/src/ai/harnesses/experimental/_harness.py b/src/ai/harnesses/experimental/_harness.py index fd5360bf..80658412 100644 --- a/src/ai/harnesses/experimental/_harness.py +++ b/src/ai/harnesses/experimental/_harness.py @@ -7,7 +7,7 @@ import warnings from typing import TYPE_CHECKING, Any, Self, TypeVar -from . import _capabilities, _input, _result, _session +from . import _capabilities, _input, _result, _session, _session_lock from . import errors as errors_ if TYPE_CHECKING: @@ -251,6 +251,7 @@ async def history( """ self._require_open() self._require("history") + self._require_id(session_id) return await self.adapter.history( session_id, limit=limit, offset=offset ) @@ -269,6 +270,7 @@ async def resume(self, session_id: str) -> _session.Session: """ self._require_open() self._require("resume") + self._require_id(session_id) await self.adapter.resume_session(session_id) session = _session.Session(self, session_id) session._seed(await self.adapter.history(session_id)) @@ -307,6 +309,8 @@ async def tui( existing one always. """ self._require_open() + if session_id is not None: + self._require_id(session_id) if session_id is not None and history is not None: raise ValueError( "tui(): pass session_id to continue a conversation, or history " @@ -341,6 +345,7 @@ async def fork(self, session_id: str) -> _session.Session: """ self._require_open() self._require("fork") + self._require_id(session_id) forked = await self.adapter.fork(session_id) session = _session.Session(self, forked) session._seed(await self.adapter.history(forked)) @@ -353,6 +358,21 @@ def _require_open(self) -> None: "harness is not open; call open() or use `async with`" ) + def _require_id(self, session_id: str) -> None: + """Refuse an id that is not a plain name, on every harness. + + An id arrives from a caller, a `Handle` out of a database, or a + session file someone else wrote, and adapters turn it into paths + (a lock directory, a transcript file) and CLI arguments. A path or + a `-flag` is never a conversation. + """ + if not _session_lock.SESSION_ID.match(session_id): + raise errors_.ResumeFailedError( + session_id, + "not a session id: a name of letters, digits, '.', '_' and " + + "'-', never a path", + ) + def _require(self, capability: str) -> None: if not getattr(self.capabilities, capability): raise errors_.UnsupportedError(self.name, capability) diff --git a/tests/harnesses/experimental/test_claude_adapter.py b/tests/harnesses/experimental/test_claude_adapter.py index 19f68e9f..e03b03ce 100644 --- a/tests/harnesses/experimental/test_claude_adapter.py +++ b/tests/harnesses/experimental/test_claude_adapter.py @@ -115,3 +115,29 @@ def shape(infos: list[Any]) -> list[tuple[Any, ...]]: full = await claude_agent_sdk.list_sessions_from_store(slow, ws.path) assert len(quick) == 6 assert shape(quick) == shape(full) + + +async def test_the_remote_claude_store_keeps_its_paths_inside() -> None: + """A store key is names, never a path out of the store.""" + from ai.harnesses.experimental._adapters import claude_store + + store = await claude_store.WorkspaceSessionStore.discover( + fakes.FakeWorkspace(), "/c" + ) + keys: list[Any] = [ + {"project_key": "-fake", "session_id": "abc"}, + {"project_key": "-fake", "session_id": "abc", "subpath": "a/b"}, + ] + assert [store._path(k) for k in keys] == [ + "/c/projects/-fake/abc.jsonl", + "/c/projects/-fake/abc/a/b.jsonl", + ] + bad: list[Any] = [ + {"project_key": "-fake", "session_id": "../../etc/x"}, + {"project_key": "-fake", "session_id": "/etc/x"}, + {"project_key": "..", "session_id": "abc"}, + {"project_key": "-fake", "session_id": "abc", "subpath": "../x"}, + ] + for key in bad: + with pytest.raises(ValueError, match="not a store key"): + store._path(key) diff --git a/tests/harnesses/experimental/test_session.py b/tests/harnesses/experimental/test_session.py index 0b5f5b91..5ae6c98a 100644 --- a/tests/harnesses/experimental/test_session.py +++ b/tests/harnesses/experimental/test_session.py @@ -643,3 +643,25 @@ async def test_output_is_unavailable_until_the_turn_settles() -> None: turn = harness.session().stream("go") with pytest.raises(RuntimeError, match="not settled"): _ = turn.output + + +@pytest.mark.parametrize("bad", ["/etc", "../x", "a/b", "-rf", "", ".."]) +async def test_every_harness_refuses_a_session_id_that_is_a_path( + bad: str, +) -> None: + """Security review: ids become lock directories, transcript paths and + CLI arguments in the adapters. Refused where they enter, on every + harness, before an adapter sees one.""" + from ai.harnesses.experimental.errors import ResumeFailedError + + harness, adapter = await fakes.open_harness() + calls = [ + harness.resume(bad), + harness.fork(bad), + harness.history(bad), + harness.tui(bad), + ] + for call in calls: + with pytest.raises(ResumeFailedError, match="not a session id"): + await call + assert adapter.turns_started == []