diff --git a/sdk/AGENTS.md b/sdk/AGENTS.md index a650704..de328ee 100644 --- a/sdk/AGENTS.md +++ b/sdk/AGENTS.md @@ -4,8 +4,8 @@ ## What this package does -Connects applications to Runloop-hosted remote agents (Claude Code, Codex, -OpenCode, etc.) via the Axon event bus. Three protocol modules plus shared +Connects applications to Runloop-hosted remote agents (Claude Code, Codex, Pi, +OpenCode, etc.) via the Axon event bus. Four protocol modules plus shared utilities. ## Choose your module @@ -15,10 +15,12 @@ utilities. | **ACP** | `@runloop/remote-agents-sdk/acp` | Any ACP-compatible agent (OpenCode, Claude via ACP) | | **Claude** | `@runloop/remote-agents-sdk/claude` | Claude Code with native SDK message types | | **Codex** | `@runloop/remote-agents-sdk/codex` | Codex CLI with native app-server message types | +| **Pi** | `@runloop/remote-agents-sdk/pi` | Pi coding agent with native JSONL RPC message types | | **Shared** | `@runloop/remote-agents-sdk/shared` | Common types (`BaseConnectionOptions`, `AxonEventView`, `AxonEventListener`) and utilities | Each module pairs with a devbox `broker_mount` `protocol`: `"acp"` for the ACP -module, `"claude_json"` for the Claude module, `"codex_json"` for the Codex module. +module, `"claude_json"` for the Claude module, `"codex_json"` for the Codex +module, `"pi_json"` for the Pi module. ## Required dependencies @@ -29,7 +31,7 @@ npm install @runloop/remote-agents-sdk @runloop/api-client # Only for the Claude module npm install @anthropic-ai/claude-agent-sdk -# The Codex module needs no extra dependency (protocol types are vendored) +# The Codex and Pi modules need no extra dependency (protocol types ship with the SDK) ``` ## ACP module — quick start @@ -249,6 +251,67 @@ await conn.disconnect(); | `abortStream()` | Abort the SSE stream without clearing listeners | | `disconnect()` | Close transport + run `onDisconnect` callback | +## Pi module — quick start + +```typescript +import { PiAxonConnection } from "@runloop/remote-agents-sdk/pi"; +import { RunloopSDK } from "@runloop/api-client"; + +const sdk = new RunloopSDK({ bearerToken: process.env.RUNLOOP_API_KEY }); + +// 1. Provision infrastructure — the broker spawns `pi --mode rpc`. The broker +// owns `--mode rpc` and `--session-dir`; never pass either yourself, or +// session persistence moves out of the durable root and resume breaks. +const axon = await sdk.axon.create({ name: "pi-transport" }); +const devbox = await sdk.devbox.create({ + mounts: [{ + type: "broker_mount", + axon_id: axon.id, + protocol: "pi_json", + agent_binary: "pi", + }], +}); + +// 2. Connect — there is no handshake, so no initialize() step +const conn = new PiAxonConnection(axon, devbox); +await conn.connect(); + +// 3. Send and receive. send() resolves when Pi *accepts* the prompt; the turn +// ends later at `agent_settled`. +await conn.send("What files are in this directory?"); +for await (const frame of conn.receiveTurn()) { + console.log(frame.type, frame); +} + +// 4. Persist the session path if you want to resume it later +const { sessionFile } = await conn.getState(); + +// 5. Clean up +await conn.disconnect(); +``` + +### Pi — key methods on `PiAxonConnection` + +| Method | Purpose | +|--------|---------| +| `connect()` | Open transport and start the read loop; **no `initialize()` follows** | +| `send(message, options?)` | Start a turn with a `prompt`; resolves on *acceptance*, throws `PiCommandError` if Pi rejects it | +| `steer(message, images?)` | Steer the in-flight turn without reopening a broker turn | +| `followUp(message, images?)` | Queue a follow-up message without reopening a broker turn | +| `interrupt()` | Abort the in-flight turn (`abort`) | +| `receiveTurn()` | Async iterator yielding Pi frames until `agent_settled` (**not** `agent_end`) | +| `receiveAgentEvents()` | Async iterator yielding all agent frames indefinitely | +| `getState()` | Read Pi's session state — the supported way to get `sessionFile` | +| `newSession(parentSession?)` | Start a fresh session, optionally branching from an existing one | +| `switchSession(sessionPath)` | Restore a persisted session from its transcript path | +| `command(frame)` | Escape hatch for any unwrapped Pi command (`set_model`, `compact`, `bash`, …) | +| `sessionId` / `sessionFile` | Captured from `get_state` acks, including replayed ones | +| `onAxonEvent(listener)` | Subscribe to all Axon events (returns unsubscribe fn) | +| `onTimelineEvent(listener)` | Subscribe to classified timeline events (returns unsubscribe fn) | +| `receiveTimelineEvents()` | Async generator yielding classified timeline events | +| `abortStream()` | Abort the SSE stream without clearing listeners | +| `disconnect()` | Close transport + run `onDisconnect` callback | + ## Timeline Events All protocol modules provide a unified timeline event stream that classifies @@ -263,6 +326,7 @@ events. | `acp_protocol` | `SessionUpdate \| unknown` | Known ACP protocol event (agent or client method) | | `claude_protocol` | `SDKMessage` | Known Claude protocol event | | `codex_protocol` | Typed app-server frame | Known Codex app-server event (narrow with `event.eventType` or the `isCodex*` guards) | +| `pi_protocol` | Typed Pi event or ack | Known Pi RPC event (narrow with `event.eventType` or the `isPi*` guards) | | `system` | `SystemEvent` | Broker system event (`turn.started`, `turn.completed`, `turn.failed`, `broker.error`) | | `unknown` | `null` | Anything else — inspect `axonEvent` for details | @@ -386,8 +450,11 @@ create a new instance. - **ACP permissions default to auto-approve** (`allow_always` > `allow_once` > first option). Pass `requestPermission` to customize. - **Claude permissions also auto-approve** all tool use. Register a `"can_use_tool"` handler via `onControlRequest()` to customize. - **Codex approvals also auto-approve** by default. Register handlers via `onApprovalRequest()` to customize, or mount with `launch_args: ["-c", "approval_policy=never"]` for headless full-auto (no approval traffic at all). -- **Explicit `connect()` required:** All connections require `await conn.connect()` first, followed by `initialize()` — ACP, Claude, and Codex alike. +- **Explicit `connect()` required:** All connections require `await conn.connect()` first, followed by `initialize()` — ACP, Claude, and Codex alike. **Pi is the exception:** it has no handshake and no `initialize()`; `connect()` is enough. +- **Pi's `send()` resolves on acceptance, not completion.** Pi's ack means only that the prompt was accepted. Await `receiveTurn()` (which ends at `agent_settled`) or the `turn.completed` system event for the response. +- **Pi's `agent_end` is not the end of a turn.** Pi may follow it with an automatic retry (`willRetry: true`). Only `agent_settled` ends an accepted turn. +- **Pi streams one `message_update` per token**, each carrying a full `partial` assistant message. The push surfaces (`onTimelineEvent` / `onAxonEvent`) see every frame. The pull surfaces (`receiveAgentEvents()` / `receiveTurn()`) buffer at most `maxQueuedFrames` (default 1000) and then discard the oldest frame, warning once — so a listener-only application keeps bounded memory, and a pull consumer must keep up. - **Node >= 22** required. - **`@runloop/api-client`** is a peer dep — you must install it yourself. -- **`@anthropic-ai/claude-agent-sdk`** is an optional peer dep — only needed for the Claude module. The Codex module has no extra dependency. +- **`@anthropic-ai/claude-agent-sdk`** is an optional peer dep — only needed for the Claude module. The Codex and Pi modules have no extra dependency. - **`prompt()` resolves before all session updates arrive.** The broker sends the prompt response and `turn.completed` system event *before* flushing thought/message chunks as `session/update` notifications. Use `onAxonEvent` to watch for `turn.started` / `turn.completed` system events to accurately bracket turn content. See the SDK README for details. diff --git a/sdk/README.md b/sdk/README.md index 14cc5d1..d7e405c 100644 --- a/sdk/README.md +++ b/sdk/README.md @@ -4,13 +4,14 @@ TypeScript client for connecting applications to Runloop-hosted remote agents via the Axon event bus. -This package provides three protocol modules and a shared utilities module: +This package provides four protocol modules and a shared utilities module: | Module | Import path | Protocol | Use case | |--------|-------------|----------|----------| | **ACP** | `@runloop/remote-agents-sdk/acp` | [Agent Client Protocol](https://agentclientprotocol.com) (JSON-RPC 2.0) | Any ACP-compatible agent (OpenCode, Claude via ACP, etc.) | | **Claude** | `@runloop/remote-agents-sdk/claude` | Claude Code SDK wire format | Claude Code with native SDK message types | | **Codex** | `@runloop/remote-agents-sdk/codex` | OpenAI Codex app-server protocol (JSON-RPC over stdio) | Codex CLI with native app-server message types | +| **Pi** | `@runloop/remote-agents-sdk/pi` | Pi RPC protocol (newline-delimited JSON over stdio) | Pi coding agent with native RPC message types | | **Shared** | `@runloop/remote-agents-sdk/shared` | — | Common types (`BaseConnectionOptions`, `AxonEventView`, `AxonEventListener`) and utilities | All protocol modules communicate over Runloop Axon channels. Pick the one that matches your agent's protocol. Shared types are also re-exported from each protocol module for convenience. @@ -22,6 +23,7 @@ Each module pairs with a devbox `broker_mount` whose `protocol` field tells the | `acp` | ACP | Any ACP agent binary (e.g. `opencode acp`) | | `claude_json` | Claude | `claude` (Claude Code CLI) | | `codex_json` | Codex | `codex` (the broker spawns `codex app-server`) | +| `pi_json` | Pi | `pi` (the broker spawns `pi --mode rpc`) | ## Installation @@ -39,6 +41,8 @@ npm install @anthropic-ai/claude-agent-sdk The Codex module needs no extra dependency — the app-server protocol types are vendored into this package (generated by `codex app-server generate-ts`, pinned to `@openai/codex` 0.144.x). +The Pi module needs no extra dependency either — its wire types are written directly into this package from Pi `0.82.1`, so none of Pi's own packages (and none of the provider SDKs they pull in) are required. + ## Imports ```typescript @@ -46,10 +50,11 @@ The Codex module needs no extra dependency — the app-server protocol types are import { ACPAxonConnection, PROTOCOL_VERSION } from "@runloop/remote-agents-sdk/acp"; import { ClaudeAxonConnection } from "@runloop/remote-agents-sdk/claude"; import { CodexAxonConnection } from "@runloop/remote-agents-sdk/codex"; +import { PiAxonConnection } from "@runloop/remote-agents-sdk/pi"; import type { BaseConnectionOptions, AxonEventView } from "@runloop/remote-agents-sdk/shared"; // Namespaced root import (all modules at once) -import { acp, claude, codex, shared } from "@runloop/remote-agents-sdk"; +import { acp, claude, codex, pi, shared } from "@runloop/remote-agents-sdk"; ``` ## Getting Started @@ -250,6 +255,60 @@ for await (const frame of conn.receiveTurn()) { await conn.disconnect(); ``` +### Pi Agent + +```typescript +import { PiAxonConnection, isPiAssistantTextDeltaEvent } from "@runloop/remote-agents-sdk/pi"; +import { RunloopSDK } from "@runloop/api-client"; + +const sdk = new RunloopSDK({ bearerToken: process.env.RUNLOOP_API_KEY }); + +// Create an Axon channel and a devbox with a Pi broker mount. The broker owns +// `--mode rpc` and `--session-dir`; passing either yourself moves session +// persistence out of the durable root and breaks resume across snapshots. +const axon = await sdk.axon.create({ name: "pi-transport" }); +const devbox = await sdk.devbox.create({ + mounts: [{ + type: "broker_mount", + axon_id: axon.id, + protocol: "pi_json", + agent_binary: "pi", + }], +}); + +// Pi has no handshake — connect() is the whole setup. +const conn = new PiAxonConnection(axon, devbox); +await conn.connect(); + +conn.onTimelineEvent((event) => { + switch (event.kind) { + case "pi_protocol": + // Pi emits one message_update per token, each carrying a cumulative + // `partial` assistant message alongside the delta. + if (isPiAssistantTextDeltaEvent(event)) { + process.stdout.write(event.data.assistantMessageEvent.delta); + } + break; + case "system": + break; + case "unknown": + break; + } +}); + +// send() resolves when Pi *accepts* the prompt, not when the turn finishes. +// receiveTurn() ends at `agent_settled`. +await conn.send("What files are in this directory?"); +for await (const frame of conn.receiveTurn()) { + console.log(frame.type, frame); +} + +// Persist sessionFile to resume this session later with switchSession(). +const { sessionFile } = await conn.getState(); + +await conn.disconnect(); +``` + --- ## ACP Module @@ -666,6 +725,102 @@ Lower-level transport used internally by `CodexAxonConnection` but available for --- +## Pi Module + +### Pi Protocol Types + +The Pi JSONL RPC wire types are hand-written into this package (transcribed from the `pi-codes` crate for Pi 0.82.1) and exported directly — no extra dependency: + +```typescript +import type { + PiCommand, + PiEvent, + PiResponse, + PiSessionState, + AssistantMessageEvent, + Message, + StopReason, + // ... etc. +} from "@runloop/remote-agents-sdk/pi"; +``` + +### `PiAxonConnection` + +Bidirectional, interactive client for the Pi CLI via Axon. The broker spawns Pi in RPC mode in the devbox and proxies its newline-delimited JSON frames over the Axon channel — messages are yielded as raw Pi frames (`PiFrame`). + +Pi has **no handshake**: `connect()` is the whole setup and there is no `initialize()`. Pi also issues no server-initiated requests, so there is no approval flow. Commands are correlated by an SDK-generated `id` and acknowledged with `{"type":"response","command":…,"success":…}`. + +**Constructor**: `new PiAxonConnection(axon, devbox, options?)` + +| Parameter | Type | Description | +|-----------|------|-------------| +| `axon` | `Axon` | Axon channel from `@runloop/api-client` | +| `devbox` | `Devbox` | Runloop devbox from `@runloop/api-client` | + +**Options** (`PiAxonConnectionOptions`): + +| Field | Type | Description | +|-------|------|-------------| +| `verbose` | `boolean` | Emit verbose logs to stderr | +| `replay` | `boolean` | Replay the channel's history on connect | +| `afterSequence` | `number` | Replay only events after this sequence number | +| `requestTimeoutMs` | `number` | Timeout for command acknowledgement (default `60000`) | +| `maxQueuedFrames` | `number` | Cap on frames buffered for the pull surfaces; the oldest is discarded past it (default `1000`) | +| `onError` | `(error: unknown) => void` | Error callback (defaults to `console.error`) | +| `onDisconnect` | `() => void \| Promise` | Teardown callback invoked by `disconnect()` (e.g. devbox shutdown) | + +**Methods**: + +| Method | Description | +|--------|-------------| +| `connect()` | Subscribe to the channel and start the read loop | +| `send(message, options?)` | Send a `prompt`; resolves when Pi **accepts** it, rejects with `PiCommandError` when Pi rejects it. `options` carries `images` and `streamingBehavior` | +| `steer(message, images?)` | Redirect the in-flight turn | +| `followUp(message, images?)` | Queue a message for after the current turn | +| `interrupt()` | Send `abort` (published as the `cancel` control event) | +| `getState()` | Request `get_state`; returns `PiSessionState` and updates `sessionId` / `sessionFile` | +| `newSession(parentSession?)` | Start a new session; returns `SessionChange` | +| `switchSession(sessionPath)` | Resume a persisted session by its `sessionFile` path | +| `command(frame)` | Escape hatch for any Pi command not wrapped above | +| `receiveAgentEvents()` | Async generator over every inbound frame | +| `receiveTurn()` | Async generator that ends at `agent_settled` (or at the rejection ack of the prompt it belongs to) | +| `receiveTimelineEvents()` | Async generator over classified `PiTimelineEvent`s | +| `onAxonEvent(listener)` / `onTimelineEvent(listener)` | Register listeners; returns an unsubscribe function | +| `publish(params)` | Publish a custom event to the channel | +| `abortStream()` | Abort the inbound SSE stream without tearing down | +| `disconnect()` | Stop the read loop, close the transport, run `onDisconnect` | + +**Getters**: `sessionId`, `sessionFile`, `isConnected`, `isDisconnected`. + +A turn completes at `agent_settled`, **not** at `agent_end` — Pi can emit `agent_end` with `willRetry: true` and then keep working. To resume a session across devboxes, persist `sessionFile` from `getState()` and restore it with `switchSession(sessionFile)`. + +### Pi Timeline Event Type Guards + +```typescript +import { + isPiAssistantTextDeltaEvent, + isPiToolExecutionEndEvent, + isPiAgentSettledEvent, + // ... +} from "@runloop/remote-agents-sdk/pi"; + +conn.onTimelineEvent((event) => { + if (isPiAssistantTextDeltaEvent(event)) { + process.stdout.write(event.data.assistantMessageEvent.delta); + } else if (isPiAgentSettledEvent(event)) { + console.log("turn settled"); + } +}); +``` + +Available guards: `isPiProtocolEvent`, `isPiAgentStartEvent`, `isPiMessageStartEvent`, `isPiMessageUpdateEvent`, `isPiMessageEndEvent`, `isPiToolExecutionStartEvent`, `isPiToolExecutionUpdateEvent`, `isPiToolExecutionEndEvent`, `isPiTurnStartEvent`, `isPiTurnEndEvent`, `isPiAgentEndEvent`, `isPiAgentSettledEvent`, `isPiResponseEvent`, `isPiAssistantTextDeltaEvent`, `isPiAssistantThinkingDeltaEvent` (plus the shared system guards, e.g. `isTurnCompletedEvent`, `isBrokerErrorEvent`). + +### `PiAxonTransport` + +Lower-level transport used internally by `PiAxonConnection` but available for custom integrations. It implements the same `connect`/`write`/`readMessages`/`abortStream`/`reconnect`/`close`/`isReady` surface as the Claude `AxonTransport`. Outbound frames are published with the frame's `type` as the Axon `event_type`, except `prompt` (published as `turn/start`) and `abort` (published as `cancel`) which map to broker control events; inbound `AGENT_EVENT`s are parsed back into frames verbatim. + +--- + ## Timeline Events All protocol modules provide a unified timeline event stream that classifies every Axon event into a typed discriminated union. This is the recommended way to build chat UIs that interleave protocol events, system events (turn start/end), and custom events in a single chronological view. @@ -676,8 +831,8 @@ Every timeline event has three fields: | Field | Type | Description | |-------|------|-------------| -| `kind` | `string` | Discriminant: `"acp_protocol"`, `"claude_protocol"`, `"codex_protocol"`, `"system"`, or `"unknown"` | -| `data` | varies | Parsed typed payload (`SessionUpdate`, `SDKMessage`, app-server frame, `SystemEvent`, or `null`) | +| `kind` | `string` | Discriminant: `"acp_protocol"`, `"claude_protocol"`, `"codex_protocol"`, `"pi_protocol"`, `"system"`, or `"unknown"` | +| `data` | varies | Parsed typed payload (`SessionUpdate`, `SDKMessage`, app-server frame, Pi event frame, `SystemEvent`, or `null`) | | `axonEvent` | `AxonEventView` | The raw Axon event with full metadata (origin, event_type, payload, sequence) | ### ACP timeline events (`ACPTimelineEvent`) @@ -753,6 +908,33 @@ conn.onTimelineEvent((event: CodexTimelineEvent) => { }); ``` +### Pi timeline events (`PiTimelineEvent`) + +```typescript +import type { PiTimelineEvent } from "@runloop/remote-agents-sdk/pi"; +import { tryParseTimelinePayload } from "@runloop/remote-agents-sdk/shared"; + +conn.onTimelineEvent((event: PiTimelineEvent) => { + switch (event.kind) { + case "pi_protocol": + // event.eventType narrows the data type (e.g. "message_update" -> the + // streaming delta frame, "response" -> a correlated acknowledgement) + break; + case "system": + // event.data is SystemEvent (turn.started / turn.completed / turn.failed / broker.error) + break; + case "unknown": + // Pi frames this SDK does not model yet (queue_update, auto_retry_start, …) + // arrive here with their payload intact + const payload = tryParseTimelinePayload<{ type: string }>(event); + if (payload) console.log(payload.type); + break; + } +}); +``` + +Available guards: `isPiProtocolEvent`, `isPiAgentStartEvent`, `isPiMessageStartEvent`, `isPiMessageUpdateEvent`, `isPiMessageEndEvent`, `isPiToolExecutionStartEvent`, `isPiToolExecutionUpdateEvent`, `isPiToolExecutionEndEvent`, `isPiTurnStartEvent`, `isPiTurnEndEvent`, `isPiAgentEndEvent`, `isPiAgentSettledEvent`, `isPiResponseEvent`, plus `isPiAssistantTextDeltaEvent` / `isPiAssistantThinkingDeltaEvent` for the nested streaming delta (and the shared system guards). + ### Async generator pattern All connections also provide `receiveTimelineEvents()` for pull-based consumption: @@ -765,7 +947,7 @@ for await (const event of conn.receiveTimelineEvents()) { ### Custom events via `publish()` and `tryParseTimelinePayload` -`ACPAxonConnection`, `ClaudeAxonConnection`, and `CodexAxonConnection` expose a `publish()` method for pushing custom events to the Axon channel. These arrive in the timeline as `kind: "unknown"` events that you can match on `event_type` and parse with `tryParseTimelinePayload`. +`ACPAxonConnection`, `ClaudeAxonConnection`, `CodexAxonConnection`, and `PiAxonConnection` expose a `publish()` method for pushing custom events to the Axon channel. These arrive in the timeline as `kind: "unknown"` events that you can match on `event_type` and parse with `tryParseTimelinePayload`. **Publishing a custom event:** @@ -853,12 +1035,12 @@ ACP Module Claude Module (in devbox) (in devbox) ``` -| | ACP Module | Claude Module | Codex Module | -|---|---|---|---| -| Wire format | JSON-RPC 2.0 via Axon events | Claude SDK messages via Axon events | App-server JSON-RPC frames via Axon events (verbatim) | -| Transport | `@runloop/api-client` Axon SDK | `@runloop/api-client` Axon SDK | `@runloop/api-client` Axon SDK | -| Agent protocol | `@agentclientprotocol/sdk` | `@anthropic-ai/claude-agent-sdk` | Vendored `codex app-server` protocol types | -| ID tracking | Synthetic (transport maps IDs) | Native (SDK handles correlation) | Native JSON-RPC ids (SDK correlates responses) | +| | ACP Module | Claude Module | Codex Module | Pi Module | +|---|---|---|---|---| +| Wire format | JSON-RPC 2.0 via Axon events | Claude SDK messages via Axon events | App-server JSON-RPC frames via Axon events (verbatim) | Pi JSONL frames via Axon events (verbatim) | +| Transport | `@runloop/api-client` Axon SDK | `@runloop/api-client` Axon SDK | `@runloop/api-client` Axon SDK | `@runloop/api-client` Axon SDK | +| Agent protocol | `@agentclientprotocol/sdk` | `@anthropic-ai/claude-agent-sdk` | Vendored `codex app-server` protocol types | Hand-written Pi wire types | +| ID tracking | Synthetic (transport maps IDs) | Native (SDK handles correlation) | Native JSON-RPC ids (SDK correlates responses) | Optional `id` on commands (SDK stamps and correlates acks) | ## Shared Types @@ -866,7 +1048,7 @@ Shared types are available from `@runloop/remote-agents-sdk/shared` or re-export ### `BaseConnectionOptions` -Common options accepted by `ACPAxonConnection`, `ClaudeAxonConnection`, and `CodexAxonConnection`: +Common options accepted by `ACPAxonConnection`, `ClaudeAxonConnection`, `CodexAxonConnection`, and `PiAxonConnection`: | Field | Type | Description | |-------|------|-------------| @@ -929,10 +1111,12 @@ type WireData = Record; ## Known Limitations -- **Explicit `connect()` required**: All connections require an explicit `await conn.connect()` call followed by `initialize()` — ACP, Claude, and Codex alike. The constructor is lightweight and synchronous. +- **Explicit `connect()` required**: All connections require an explicit `await conn.connect()` call, followed by `initialize()` for ACP, Claude, and Codex. Pi has no handshake and exposes no `initialize()`. The constructor is lightweight and synchronous. - **Automatic reconnection (single retry)**: If an SSE stream drops unexpectedly, the SDK re-subscribes once and logs a `console.warn`. If the retry also fails, the connection is terminal — create a new instance. - **Permission handling** (Claude): The `ClaudeAxonConnection` auto-approves all tool use by default. Register a `"can_use_tool"` handler via `onControlRequest()` to customize. - **Approval handling** (Codex): The `CodexAxonConnection` auto-approves server-initiated approval requests by default. Register handlers via `onApprovalRequest()` to customize, or mount with `launch_args: ["-c", "approval_policy=never"]` to skip approval traffic entirely. +- **Turn completion** (Pi): `send()` resolves when Pi *accepts* the prompt, not when the turn ends, and `agent_end` can be followed by a retry. Use `receiveTurn()` (or watch for `agent_settled`) to detect the real end of a turn. +- **Streaming volume** (Pi): Pi emits one `message_update` per token, each carrying a cumulative `partial` assistant message. `onTimelineEvent()`/`onAxonEvent()` listeners see every frame; the pull surfaces (`receiveAgentEvents()`/`receiveTurn()`) buffer at most `maxQueuedFrames` (default 1000) and then discard the oldest, warning once. Applications that consume only through listeners therefore keep bounded memory; a pull consumer must keep up or raise the cap. ### ACP: `prompt()` resolves before all session updates arrive diff --git a/sdk/package.json b/sdk/package.json index b8bbae7..9342f92 100644 --- a/sdk/package.json +++ b/sdk/package.json @@ -46,6 +46,11 @@ "import": "./dist/codex/index.js", "default": "./dist/codex/index.js" }, + "./pi": { + "types": "./dist/pi/index.d.ts", + "import": "./dist/pi/index.js", + "default": "./dist/pi/index.js" + }, "./shared": { "types": "./dist/shared/index.d.ts", "import": "./dist/shared/index.js", diff --git a/sdk/src/index.test.ts b/sdk/src/index.test.ts index 23871aa..56ca90c 100644 --- a/sdk/src/index.test.ts +++ b/sdk/src/index.test.ts @@ -29,6 +29,21 @@ describe("root exports", () => { expect(SDK.claude.AxonTransport).toBeDefined(); }); + it("exports a pi namespace", () => { + expect(SDK.pi).toBeDefined(); + expect(typeof SDK.pi).toBe("object"); + }); + + it("pi namespace contains PiAxonConnection and PiAxonTransport", () => { + expect(SDK.pi.PiAxonConnection).toBeDefined(); + expect(SDK.pi.PiAxonTransport).toBeDefined(); + }); + + it("pi namespace contains protocol constants and guards", () => { + expect(SDK.pi.PI_TURN_START_EVENT_TYPE).toBe("turn/start"); + expect(typeof SDK.pi.isPiAgentSettledEvent).toBe("function"); + }); + it("exports a shared namespace", () => { expect(SDK.shared).toBeDefined(); expect(typeof SDK.shared).toBe("object"); diff --git a/sdk/src/index.ts b/sdk/src/index.ts index 1aaf43a..ad2767d 100644 --- a/sdk/src/index.ts +++ b/sdk/src/index.ts @@ -1,4 +1,5 @@ export * as acp from "./acp/index.js"; export * as claude from "./claude/index.js"; export * as codex from "./codex/index.js"; +export * as pi from "./pi/index.js"; export * as shared from "./shared/index.js"; diff --git a/sdk/src/pi/connection.test.ts b/sdk/src/pi/connection.test.ts new file mode 100644 index 0000000..8aebb19 --- /dev/null +++ b/sdk/src/pi/connection.test.ts @@ -0,0 +1,449 @@ +import { describe, expect, it, vi } from "vitest"; +import { + createControllableStream, + createMockAxon, + makeAgentEvent, + makeSystemEventWithRawPayload, +} from "../__test-utils__/mock-axon.js"; +import type { ConnectionStateError } from "../shared/errors/connection-state-error.js"; +import { SystemError } from "../shared/errors/system-error.js"; +import { PiAxonConnection, PiCommandError } from "./connection.js"; +import type { PiSessionState } from "./protocol/index.js"; + +const tick = () => new Promise((resolve) => setTimeout(resolve, 0)); + +const SESSION_STATE: PiSessionState = { + model: null, + thinkingLevel: "medium", + isStreaming: false, + isCompacting: false, + steeringMode: "all", + followUpMode: "one-at-a-time", + sessionFile: "/sessions/current.jsonl", + sessionId: "session-1", + autoCompactionEnabled: true, + messageCount: 5, + pendingMessageCount: 0, +}; + +function setup(options: Record = {}) { + const ctrl = createControllableStream(true); + const mock = createMockAxon(ctrl); + const conn = new PiAxonConnection(mock.axon as never, { id: "dbx-test" } as never, { + replay: false, + ...options, + }); + const frames = () => mock.axon.publish.mock.calls.map(([event]) => JSON.parse(event.payload)); + return { ctrl, mock, conn, frames }; +} + +/** Acks every published command, mirroring its `type` as the ack `command`. */ +function ackAll( + { ctrl, mock }: Pick, "ctrl" | "mock">, + ack: (type: string) => { success?: boolean; error?: string; data?: unknown } = () => ({}), +) { + mock.axon.publish.mockImplementation(async (event: { payload: string }) => { + const frame = JSON.parse(event.payload) as { type: string; id: string }; + const { success = true, error, data } = ack(frame.type); + ctrl.push( + makeAgentEvent("response", { + type: "response", + id: frame.id, + command: frame.type, + success, + ...(error != null ? { error } : {}), + ...(data !== undefined ? { data } : {}), + }), + ); + }); +} + +describe("PiAxonConnection", () => { + it("connects without a handshake, rejects a duplicate connect, and disconnects", async () => { + const { conn } = setup(); + await conn.connect(); + expect(conn.isConnected).toBe(true); + await expect(conn.connect()).rejects.toMatchObject({ code: "already_connected" }); + await conn.disconnect(); + expect(conn.isDisconnected).toBe(true); + }); + + it("rejects commands before connect", async () => { + const { conn } = setup(); + await expect(conn.getState()).rejects.toMatchObject({ code: "not_connected" }); + }); + + it("publishes exactly one turn/start prompt frame and resolves on its ack", async () => { + const ctx = setup(); + ackAll(ctx); + await ctx.conn.connect(); + await ctx.conn.send("hi"); + expect(ctx.mock.axon.publish.mock.calls).toHaveLength(1); + const [event] = ctx.mock.axon.publish.mock.calls[0] as [Record]; + expect(event).toMatchObject({ + event_type: "turn/start", + origin: "USER_EVENT", + source: "pi-sdk-client", + }); + const frame = JSON.parse(event.payload as string); + expect(frame).toEqual({ type: "prompt", id: expect.stringMatching(/^pi-sdk-/), message: "hi" }); + }); + + it("passes prompt images and streamingBehavior through", async () => { + const ctx = setup(); + ackAll(ctx); + await ctx.conn.connect(); + await ctx.conn.send("look", { + images: [{ type: "image", data: "abc", mimeType: "image/png" }], + streamingBehavior: "steer", + }); + expect(ctx.frames()[0]).toMatchObject({ + type: "prompt", + message: "look", + images: [{ type: "image", data: "abc", mimeType: "image/png" }], + streamingBehavior: "steer", + }); + }); + + it("rejects send with a PiCommandError carrying Pi's error string", async () => { + const ctx = setup(); + ackAll(ctx, () => ({ success: false, error: "agent is streaming" })); + await ctx.conn.connect(); + const error = await ctx.conn.send("hi").catch((e: unknown) => e); + expect(error).toBeInstanceOf(PiCommandError); + expect(error).toMatchObject({ + command: "prompt", + error: "agent is streaming", + message: "agent is streaming", + }); + }); + + it("falls back to a synthesized message when a rejection carries no error", async () => { + const ctx = setup(); + ackAll(ctx, () => ({ success: false })); + await ctx.conn.connect(); + await expect(ctx.conn.getState()).rejects.toThrow("Pi rejected the get_state command"); + }); + + it("routes interrupt, steer and followUp to their own event types", async () => { + const ctx = setup(); + ackAll(ctx); + await ctx.conn.connect(); + await ctx.conn.interrupt(); + await ctx.conn.steer("actually, stop"); + await ctx.conn.followUp("and then this", [ + { type: "image", data: "abc", mimeType: "image/png" }, + ]); + const events = ctx.mock.axon.publish.mock.calls.map(([event]) => event.event_type); + expect(events).toEqual(["cancel", "steer", "follow_up"]); + expect(events).not.toContain("turn/start"); + expect(ctx.frames()).toEqual([ + { type: "abort", id: expect.any(String) }, + { type: "steer", id: expect.any(String), message: "actually, stop" }, + { + type: "follow_up", + id: expect.any(String), + message: "and then this", + images: [{ type: "image", data: "abc", mimeType: "image/png" }], + }, + ]); + }); + + it("refuses to publish a broker-reserved command id", async () => { + const ctx = setup(); + ackAll(ctx); + await ctx.conn.connect(); + await expect(ctx.conn.command({ type: "prompt", id: "broker-1" })).rejects.toThrow("reserved"); + }); + + it("terminates receiveTurn at agent_settled, not at a retrying agent_end", async () => { + const ctx = setup(); + ackAll(ctx); + await ctx.conn.connect(); + await ctx.conn.send("hi"); + ctx.ctrl.push(makeAgentEvent("turn_start", { type: "turn_start" })); + ctx.ctrl.push( + makeAgentEvent("agent_end", { type: "agent_end", messages: [], willRetry: true }), + ); + ctx.ctrl.push(makeAgentEvent("turn_start", { type: "turn_start" })); + ctx.ctrl.push( + makeAgentEvent("agent_end", { type: "agent_end", messages: [], willRetry: false }), + ); + ctx.ctrl.push(makeAgentEvent("agent_settled", { type: "agent_settled" })); + ctx.ctrl.push(makeAgentEvent("turn_start", { type: "turn_start" })); + const frames = []; + for await (const frame of ctx.conn.receiveTurn()) frames.push(frame.type); + expect(frames).toEqual(["turn_start", "agent_end", "turn_start", "agent_end", "agent_settled"]); + }); + + it("terminates receiveTurn on a rejected prompt ack", async () => { + const ctx = setup(); + ackAll(ctx, (type) => (type === "prompt" ? { success: false, error: "busy" } : {})); + await ctx.conn.connect(); + await expect(ctx.conn.send("hi")).rejects.toBeInstanceOf(PiCommandError); + const frames = []; + for await (const frame of ctx.conn.receiveTurn()) frames.push(frame); + expect(frames).toEqual([ + { + type: "response", + id: expect.any(String), + command: "prompt", + success: false, + error: "busy", + }, + ]); + }); + + it("does not let an undrained rejection terminate a later accepted turn", async () => { + const ctx = setup(); + let rejectPrompt = true; + ackAll(ctx, (type) => + type === "prompt" && rejectPrompt ? { success: false, error: "busy" } : {}, + ); + await ctx.conn.connect(); + // The first prompt is rejected and its caller never drains the turn, so + // the rejection ack is left sitting in the queue. + await expect(ctx.conn.send("first")).rejects.toBeInstanceOf(PiCommandError); + rejectPrompt = false; + await ctx.conn.send("second"); + ctx.ctrl.push(makeAgentEvent("turn_start", { type: "turn_start" })); + ctx.ctrl.push(makeAgentEvent("agent_settled", { type: "agent_settled" })); + const frames = []; + for await (const frame of ctx.conn.receiveTurn()) frames.push(frame); + // The stale rejection is skipped; the accepted turn runs to agent_settled. + expect(frames.map((frame) => frame.type)).toEqual(["turn_start", "agent_settled"]); + }); + + it("bounds the pull queue when frames are consumed only through listeners", async () => { + const ctx = setup({ maxQueuedFrames: 10, onError: () => {} }); + const seen: string[] = []; + ctx.conn.onTimelineEvent(() => seen.push("event")); + await ctx.conn.connect(); + for (let index = 0; index < 50; index++) + ctx.ctrl.push( + makeAgentEvent("message_update", { + type: "message_update", + assistantMessageEvent: { type: "text_delta", contentIndex: 0, delta: `${index}` }, + }), + ); + ctx.ctrl.push(makeAgentEvent("agent_settled", { type: "agent_settled" })); + await tick(); + // Every frame reached the listener, but the pull queue kept only the cap. + expect(seen).toHaveLength(51); + const frames = []; + for await (const frame of ctx.conn.receiveTurn()) frames.push(frame); + expect(frames).toHaveLength(10); + expect(frames.at(-1)?.type).toBe("agent_settled"); + }); + + it("generates an id for a blank caller-supplied id", async () => { + const ctx = setup(); + ackAll(ctx); + await ctx.conn.connect(); + await ctx.conn.command({ type: "prompt", id: " ", message: "hi" }); + expect(ctx.frames()).toEqual([ + { type: "prompt", id: expect.stringMatching(/^pi-sdk-/), message: "hi" }, + ]); + }); + + it("resolves getState and captures session identity from the ack", async () => { + const ctx = setup(); + ackAll(ctx, () => ({ data: SESSION_STATE })); + await ctx.conn.connect(); + expect(ctx.conn.sessionId).toBeUndefined(); + expect(await ctx.conn.getState()).toEqual(SESSION_STATE); + expect(ctx.conn.sessionId).toBe("session-1"); + expect(ctx.conn.sessionFile).toBe("/sessions/current.jsonl"); + }); + + it("captures session identity from acks seen during replay", async () => { + const ctrl = createControllableStream(true); + ctrl.push( + makeAgentEvent( + "response", + { + type: "response", + id: "broker-1", + command: "get_state", + success: true, + data: SESSION_STATE, + }, + 1, + ), + ); + const mock = createMockAxon(ctrl); + Object.assign(mock.axon, { + client: { get: vi.fn().mockResolvedValue({ events: [], has_more: false, total_count: 1 }) }, + }); + const conn = new PiAxonConnection(mock.axon as never, { id: "dbx" } as never); + await conn.connect(); + await tick(); + expect(conn.sessionId).toBe("session-1"); + expect(conn.sessionFile).toBe("/sessions/current.jsonl"); + }); + + it("queues acks it did not ask for rather than swallowing them", async () => { + const ctx = setup(); + await ctx.conn.connect(); + // The adapter's own `get_state`, both id-less and `broker-N` stamped. + ctx.ctrl.push( + makeAgentEvent("response", { type: "response", command: "get_state", success: true }), + ); + ctx.ctrl.push( + makeAgentEvent("response", { + type: "response", + id: "broker-2", + command: "get_state", + success: true, + }), + ); + ctx.ctrl.push(makeAgentEvent("agent_settled", { type: "agent_settled" })); + const frames = []; + for await (const frame of ctx.conn.receiveTurn()) frames.push([frame.type, frame.id]); + expect(frames).toEqual([ + ["response", undefined], + ["response", "broker-2"], + ["agent_settled", undefined], + ]); + }); + + it("returns SessionChange from newSession and switchSession", async () => { + const ctx = setup(); + ackAll(ctx, () => ({ data: { cancelled: false } })); + await ctx.conn.connect(); + expect(await ctx.conn.newSession("/sessions/parent.jsonl")).toEqual({ cancelled: false }); + expect(await ctx.conn.switchSession("/sessions/other.jsonl")).toEqual({ cancelled: false }); + expect(ctx.frames()).toEqual([ + { type: "new_session", id: expect.any(String), parentSession: "/sessions/parent.jsonl" }, + { type: "switch_session", id: expect.any(String), sessionPath: "/sessions/other.jsonl" }, + ]); + }); + + it("sends unwrapped commands through the escape hatch", async () => { + const ctx = setup(); + ackAll(ctx, () => ({ data: { model: "glm-5.2" } })); + await ctx.conn.connect(); + expect(await ctx.conn.command({ type: "set_model", model: "glm-5.2" })).toEqual({ + model: "glm-5.2", + }); + expect(ctx.frames()[0]).toMatchObject({ type: "set_model", model: "glm-5.2" }); + }); + + it("emits raw and classified events to listeners", async () => { + const ctx = setup(); + const axonEvents = vi.fn(); + const timelineEvents = vi.fn(); + ctx.conn.onAxonEvent(axonEvents); + ctx.conn.onTimelineEvent(timelineEvents); + await ctx.conn.connect(); + ctx.ctrl.push(makeAgentEvent("agent_settled", { type: "agent_settled" })); + await tick(); + expect(axonEvents).toHaveBeenCalledWith( + expect.objectContaining({ event_type: "agent_settled" }), + ); + expect(timelineEvents).toHaveBeenCalledWith( + expect.objectContaining({ kind: "pi_protocol", eventType: "agent_settled" }), + ); + }); + + it("streams classified timeline events until the stream ends", async () => { + const ctx = setup(); + await ctx.conn.connect(); + const events = ctx.conn.receiveTimelineEvents(); + ctx.ctrl.push(makeAgentEvent("agent_settled", { type: "agent_settled" })); + const first = await events.next(); + expect(first.value).toMatchObject({ eventType: "agent_settled" }); + await ctx.conn.disconnect(); + }); + + it("treats broker errors as fatal and refuses to reconnect", async () => { + const onError = vi.fn(); + const { ctrl, conn } = setup({ onError }); + await conn.connect(); + ctrl.push(makeSystemEventWithRawPayload("broker.error", "boom", 1)); + await tick(); + expect(onError).toHaveBeenCalledWith(expect.any(SystemError)); + await expect(conn.getState()).rejects.toEqual( + expect.objectContaining>({ code: "terminated" }), + ); + await expect(conn.connect()).rejects.toMatchObject({ code: "terminated" }); + }); + + it("rejects a pending command when a broker error kills the read loop", async () => { + const onError = vi.fn(); + const { ctrl, mock, conn } = setup({ onError }); + mock.axon.publish.mockImplementation(async () => { + ctrl.push(makeSystemEventWithRawPayload("broker.error", "boom", 1)); + }); + await conn.connect(); + await expect(conn.getState()).rejects.toBeInstanceOf(SystemError); + }); + + it("keeps frames buffered before a fatal broker error drainable", async () => { + const onError = vi.fn(); + const { ctrl, conn } = setup({ onError }); + await conn.connect(); + ctrl.push(makeAgentEvent("turn_start", { type: "turn_start" }, 1)); + ctrl.push( + makeAgentEvent("agent_end", { type: "agent_end", messages: [], willRetry: false }, 2), + ); + ctrl.push(makeSystemEventWithRawPayload("broker.error", "boom", 3)); + await tick(); + const frames = []; + for await (const frame of conn.receiveAgentEvents()) frames.push(frame.type); + expect(frames).toEqual(["turn_start", "agent_end"]); + }); + + it("auto-reconnects once after a transient stream error", async () => { + const second = createControllableStream(true); + const failing = { + controller: { abort: vi.fn() }, + [Symbol.asyncIterator]() { + return { + next: async () => { + throw new Error("network blip"); + }, + }; + }, + }; + const axon = { + id: "axon", + publish: vi.fn(), + subscribeSse: vi.fn().mockResolvedValueOnce(failing).mockResolvedValueOnce(second.stream), + }; + const conn = new PiAxonConnection(axon as never, { id: "dbx" } as never, { + replay: false, + onError: vi.fn(), + }); + await conn.connect(); + await tick(); + expect(axon.subscribeSse).toHaveBeenCalledTimes(2); + second.end(); + }); + + it("aborts only the current stream and publishes raw events", async () => { + const ctx = setup(); + await ctx.conn.connect(); + await ctx.conn.publish({ event_type: "custom", origin: "USER_EVENT", payload: "{}" } as never); + expect(ctx.mock.published[0]).toMatchObject({ event_type: "custom" }); + ctx.conn.abortStream(); + await ctx.conn.disconnect(); + expect(ctx.conn.isConnected).toBe(false); + }); + + it("times out a command that is never acknowledged", async () => { + const { conn, mock } = setup({ requestTimeoutMs: 10 }); + mock.axon.publish.mockImplementation(async () => {}); + await conn.connect(); + await expect(conn.getState()).rejects.toThrow("Command timeout: get_state"); + }); + + it("drops the pending entry when publishing fails", async () => { + const { conn, mock } = setup(); + mock.axon.publish.mockRejectedValue(new Error("publish failed")); + await conn.connect(); + await expect(conn.getState()).rejects.toThrow("publish failed"); + // A dropped entry means the same generated id space stays usable. + await expect(conn.getState()).rejects.toThrow("publish failed"); + }); +}); diff --git a/sdk/src/pi/connection.ts b/sdk/src/pi/connection.ts new file mode 100644 index 0000000..ba250fc --- /dev/null +++ b/sdk/src/pi/connection.ts @@ -0,0 +1,412 @@ +import type { AxonPublishParams, PublishResultView } from "@runloop/api-client/resources/axons"; +import type { Axon, Devbox } from "@runloop/api-client/sdk"; +import { AsyncMessageQueue } from "../shared/async-message-queue.js"; +import { resolveReplayTarget } from "../shared/connect-guards.js"; +import { runConnectionReadLoop } from "../shared/connection-read-loop.js"; +import { ConnectionStateError } from "../shared/errors/connection-state-error.js"; +import { runDisconnectHook } from "../shared/lifecycle.js"; +import { ListenerSet } from "../shared/listener-set.js"; +import { makeDefaultOnError, makeLogger } from "../shared/logging.js"; +import { isFromAgent } from "../shared/origin-guards.js"; +import { PendingRequestMap } from "../shared/pending-request-map.js"; +import { timelineEventGenerator } from "../shared/timeline-generator.js"; +import type { + AxonEventListener, + BaseConnectionOptions, + TimelineEventListener, +} from "../shared/types.js"; +import { classifyPiAxonEvent } from "./classify-pi-axon-event.js"; +import type { + ImageContent, + PiCommand, + PiCommandFrame, + PiResponse, + PiSessionState, + SessionChange, + StreamingBehavior, +} from "./protocol/index.js"; +import { PI_RESPONSE_EVENT_TYPE } from "./protocol/index.js"; +import { PiAxonTransport, type PiFrame, type PiTransport } from "./transport.js"; +import type { PiTimelineEvent } from "./types.js"; + +/** The Pi event that ends an accepted turn. */ +const AGENT_SETTLED = "agent_settled"; +const GET_STATE_COMMAND = "get_state"; +const PROMPT_COMMAND = "prompt"; + +/** Per-prompt options for {@link PiAxonConnection.send}. @category Configuration */ +export interface PiSendOptions { + images?: ImageContent[]; + /** + * Required by Pi to accept a prompt while a turn is already streaming. + * Prefer {@link PiAxonConnection.steer} or + * {@link PiAxonConnection.followUp}, which do not reopen a broker turn. + */ + streamingBehavior?: StreamingBehavior; +} + +/** + * A Pi command that came back with `success: false`. + * Carries the command name and Pi's own `error` string. + * @category Errors + */ +export class PiCommandError extends Error { + constructor( + message: string, + readonly command: string, + readonly error?: string, + ) { + super(message); + this.name = "PiCommandError"; + } +} + +/** Default cap on frames buffered for {@link PiAxonConnection.receiveAgentEvents}. */ +const DEFAULT_MAX_QUEUED_FRAMES = 1000; + +/** Options for a native Pi connection. @category Configuration */ +export interface PiAxonConnectionOptions extends BaseConnectionOptions { + requestTimeoutMs?: number; + /** + * Cap on frames buffered for the pull surfaces + * ({@link PiAxonConnection.receiveAgentEvents} and + * {@link PiAxonConnection.receiveTurn}). Once reached the oldest frame is + * discarded, so an application that consumes only through + * {@link PiAxonConnection.onTimelineEvent} keeps bounded memory. Defaults to + * 1000. + */ + maxQueuedFrames?: number; +} + +/** + * Native Pi RPC connection over an Axon channel. + * + * Unlike ACP, Claude and Codex, Pi has **no handshake**: there is no + * `initialize()` to call. {@link connect} is enough. Read session identity + * with {@link getState} instead. + * + * @example + * ```ts + * await connection.connect(); + * await connection.send("Explain this repository"); + * for await (const frame of connection.receiveTurn()) console.log(frame); + * ``` + * @category Connection + */ +export class PiAxonConnection { + readonly axonId: string; + readonly devboxId: string; + private _sessionId: string | undefined; + private _sessionFile: string | undefined; + private transport?: PiTransport; + private running = false; + private closed = false; + private fatal = false; + private everConnected = false; + private aborted = false; + private suppressAutoReconnect = false; + private counter = 0; + /** Id of the most recent `prompt`, used to correlate a rejection to its turn. */ + private latestPromptId: string | undefined; + private pending = new PendingRequestMap(); + private messageQueue: AsyncMessageQueue; + private abortController = new AbortController(); + private readonly axonListeners: ListenerSet; + private readonly timelineListeners: ListenerSet>; + private readonly handleError: (error: unknown) => void; + private readonly log; + constructor( + private readonly axon: Axon, + devbox: Devbox, + private readonly options: PiAxonConnectionOptions = {}, + ) { + this.axonId = axon.id; + this.devboxId = devbox.id; + this.handleError = options.onError ?? makeDefaultOnError("PiAxonConnection"); + this.log = makeLogger("pi-sdk", options.verbose ?? false); + this.axonListeners = new ListenerSet(this.handleError); + this.timelineListeners = new ListenerSet(this.handleError); + const maxQueuedFrames = options.maxQueuedFrames ?? DEFAULT_MAX_QUEUED_FRAMES; + this.messageQueue = new AsyncMessageQueue( + maxQueuedFrames, + (size) => + this.handleError( + `[PiAxonConnection] Message queue is full at ${size} buffered frames; ` + + "the oldest are being discarded. Consume frames via receiveAgentEvents() " + + "or receiveTurn(), or raise maxQueuedFrames.", + ), + maxQueuedFrames, + ); + } + get isConnected(): boolean { + return this.running && !this.closed; + } + get isDisconnected(): boolean { + return this.everConnected && !this.running; + } + /** + * Pi's session id, captured from `get_state` acknowledgements. The broker + * issues one after every turn, so this populates without an explicit + * {@link getState} call — including from replayed history. + */ + get sessionId(): string | undefined { + return this._sessionId; + } + /** + * The session transcript path Pi writes to, captured from `get_state` + * acknowledgements. Pass it to {@link switchSession} to restore a session. + */ + get sessionFile(): string | undefined { + return this._sessionFile; + } + /** + * Opens the SSE transport and starts its read loop. There is no handshake + * to follow it — send a prompt straight away. + * @throws {@link ConnectionStateError} with `terminated` or `already_connected`. + */ + async connect(): Promise { + if (this.fatal) + throw new ConnectionStateError( + "terminated", + "This connection hit a fatal broker error and cannot be reused. Create a new instance.", + ); + if (this.running) + throw new ConnectionStateError( + "already_connected", + "Already connected. Call disconnect() before reconnecting.", + ); + this.closed = false; + this.aborted = false; + this.suppressAutoReconnect = false; + this.abortController = new AbortController(); + this.messageQueue.reopen(); + const replayTargetSequence = await resolveReplayTarget(this.axon, this.options, this.log); + this.transport = new PiAxonTransport(this.axon, { + verbose: this.options.verbose, + afterSequence: this.options.afterSequence, + replayTargetSequence, + onAxonEvent: (event) => { + this.axonListeners.emit(event); + this.timelineListeners.emit(classifyPiAxonEvent(event)); + // Session identity is recovered here rather than in route() so that + // acknowledgements inside the replay window count too. + if (isFromAgent(event) && event.event_type === PI_RESPONSE_EVENT_TYPE && event.payload) + try { + this.captureSessionState(JSON.parse(event.payload) as PiFrame); + } catch { + // Classification reports malformed events through the timeline. + } + }, + }); + await this.transport.connect(); + this.running = true; + this.everConnected = true; + this.readLoop(); + } + /** Aborts only the current SSE stream, preserving listeners. */ + abortStream(): void { + this.aborted = true; + this.transport?.abortStream(); + } + /** Gracefully closes the transport and rejects pending commands. Idempotent. */ + async disconnect(): Promise { + if (!this.transport && !this.running) return; + this.suppressAutoReconnect = true; + this.closed = true; + this.abortController.abort(); + this.pending.rejectAll(new Error("Client disconnected")); + this.messageQueue.close(); + await this.transport?.close(); + this.transport = undefined; + this.running = false; + await runDisconnectHook(this.options.onDisconnect, this.log, this.handleError); + this.closed = false; + } + /** Registers a raw Axon event listener. */ + onAxonEvent(listener: AxonEventListener): () => void { + return this.axonListeners.add(listener); + } + /** Registers a classified Pi timeline listener. */ + onTimelineEvent(listener: TimelineEventListener): () => void { + return this.timelineListeners.add(listener); + } + /** Pull-based classified timeline event stream. */ + async *receiveTimelineEvents(): AsyncGenerator { + yield* timelineEventGenerator( + (listener) => this.onTimelineEvent(listener), + this.abortController.signal, + ); + } + private readLoop(): void { + const transport = this.transport; + if (!transport) return; + void runConnectionReadLoop({ + transport, + route: (frame) => this.route(frame), + isClosed: () => this.closed, + isReconnectSuppressed: () => this.suppressAutoReconnect, + isStreamAborted: () => this.aborted, + isCurrent: () => transport === this.transport, + onError: this.handleError, + onFatal: (error) => { + this.fatal = true; + this.closed = true; + this.pending.rejectAll(error); + }, + onTerminalError: (error) => this.pending.rejectAll(error), + onFinished: () => { + this.running = false; + this.abortController.abort(); + this.messageQueue.close(false); + }, + log: this.log, + }); + } + private captureSessionState(frame: PiFrame): void { + const ack = frame as Partial; + if (ack.type !== PI_RESPONSE_EVENT_TYPE || ack.command !== GET_STATE_COMMAND || !ack.success) + return; + const state = ack.data as Partial | undefined; + if (typeof state?.sessionId === "string") this._sessionId = state.sessionId; + if (typeof state?.sessionFile === "string") this._sessionFile = state.sessionFile; + } + private route(frame: PiFrame): void { + if (frame.type === PI_RESPONSE_EVENT_TYPE) { + const ack = frame as Partial; + this.captureSessionState(frame); + // Acks the SDK did not ask for — the adapter's own `broker-N` commands, + // or an id-less `get_state` — fall through to the queue rather than + // vanishing. + const settled = + typeof ack.id === "string" && + (ack.success + ? this.pending.resolve(ack.id, ack.data) + : this.pending.reject(ack.id, toCommandError(ack))); + // A rejected prompt completes the broker turn, so it stays visible to + // receiveTurn() even though send() already surfaced it as an error. + // receiveTurn() correlates it on `id` before acting on it. + if (settled && !(ack.command === PROMPT_COMMAND && !ack.success)) return; + } + this.messageQueue.push(frame); + } + /** + * Publishes a Pi command and awaits its acknowledgement, correlated on an + * SDK-stamped `id`. + * @throws {@link ConnectionStateError} with `terminated` or `not_connected`. + * @throws {@link PiCommandError} If Pi answers `success: false`. + */ + private async request( + command: PiCommand | PiCommandFrame, + timeoutMs = this.options.requestTimeoutMs ?? 60_000, + ): Promise { + if (this.fatal) + throw new ConnectionStateError( + "terminated", + "This connection hit a fatal broker error and cannot be reused. Create a new instance.", + ); + if (!this.transport?.isReady()) + throw new ConnectionStateError("not_connected", "Not connected. Call connect() first."); + // A blank id is what the broker treats as absent: it would stamp its own + // `broker-N` id and the ack could never resolve a pending `""` entry. + const explicitId = typeof command.id === "string" && command.id.trim() !== "" ? command.id : ""; + const id = explicitId || `pi-sdk-${++this.counter}-${Math.random().toString(36).slice(2, 10)}`; + if (command.type === PROMPT_COMMAND) this.latestPromptId = id; + const promise = this.pending.create(id, timeoutMs, `Command timeout: ${command.type}`); + try { + await this.transport.write({ ...command, id }); + } catch (error) { + this.pending.delete(id); + throw error; + } + return (await promise) as T; + } + /** + * Starts a turn with a `prompt` command. + * + * **Resolves on acceptance, not completion.** Pi's ack means only that the + * prompt was accepted; the turn ends later at `agent_settled`. Await + * {@link receiveTurn} (or the `turn.completed` system event) for the + * response. Pi rejects a prompt sent while it is already streaming unless + * {@link PiSendOptions.streamingBehavior} is set. + * @throws {@link PiCommandError} If Pi rejects the prompt. + */ + async send(message: string, options?: PiSendOptions): Promise { + await this.request({ type: PROMPT_COMMAND, message, ...options }); + } + /** Steers the in-flight turn without reopening a broker turn. */ + async steer(message: string, images?: ImageContent[]): Promise { + await this.request({ type: "steer", message, ...(images ? { images } : {}) }); + } + /** Queues a follow-up message without reopening a broker turn. */ + async followUp(message: string, images?: ImageContent[]): Promise { + await this.request({ type: "follow_up", message, ...(images ? { images } : {}) }); + } + /** Aborts the in-flight turn. */ + async interrupt(): Promise { + await this.request({ type: "abort" }); + } + /** Reads Pi's session state — the supported way to get `sessionFile`. */ + async getState(): Promise { + return this.request({ type: GET_STATE_COMMAND }); + } + /** Starts a fresh Pi session, optionally branching from an existing one. */ + async newSession(parentSession?: string): Promise { + return this.request({ + type: "new_session", + ...(parentSession != null ? { parentSession } : {}), + }); + } + /** Restores a persisted session from its transcript path. */ + async switchSession(sessionPath: string): Promise { + return this.request({ type: "switch_session", sessionPath }); + } + /** + * Sends any Pi command this class does not wrap (`set_model`, `compact`, + * `bash`, `get_messages`, `export_html`, …) and returns its ack `data`. + */ + async command(frame: PiCommandFrame): Promise { + return this.request(frame); + } + /** Yields agent frames until the connection closes. */ + async *receiveAgentEvents(): AsyncGenerator { + while (true) { + // Delegate closed-state handling to the queue so frames buffered before + // a fatal error or stream end remain drainable; disconnect() clears them. + const value = await this.messageQueue.next(); + if (!value) return; + yield value; + } + } + /** + * Yields one turn and terminates at `agent_settled` — **not** at + * `agent_end`, which Pi may follow with an automatic retry + * (`willRetry: true`). Also terminates on a rejected `prompt`, which + * completes the turn immediately. + * + * A rejection ends only the turn it belongs to: one left over from an + * earlier {@link send} whose caller caught the error without draining is + * skipped, so it cannot cut a later accepted turn short. The broker reports + * rejections of its own prompts through the `turn.failed` system event. + */ + async *receiveTurn(): AsyncGenerator { + for await (const frame of this.receiveAgentEvents()) { + const ack = frame as Partial; + if (ack.type === PI_RESPONSE_EVENT_TYPE && ack.command === PROMPT_COMMAND && !ack.success) { + if (ack.id !== undefined && ack.id !== this.latestPromptId) continue; + this.latestPromptId = undefined; + yield frame; + return; + } + yield frame; + if (frame.type === AGENT_SETTLED) return; + } + } + async publish(params: AxonPublishParams): Promise { + return this.axon.publish(params); + } +} + +function toCommandError(ack: Partial): PiCommandError { + const command = ack.command ?? "unknown"; + return new PiCommandError(ack.error ?? `Pi rejected the ${command} command`, command, ack.error); +} diff --git a/sdk/src/pi/index.ts b/sdk/src/pi/index.ts new file mode 100644 index 0000000..0c8e135 --- /dev/null +++ b/sdk/src/pi/index.ts @@ -0,0 +1,6 @@ +export { classifyPiAxonEvent, isPiProtocolEventType } from "./classify-pi-axon-event.js"; +export * from "./connection.js"; +export * from "./protocol/index.js"; +export * from "./timeline-event-guards.js"; +export * from "./transport.js"; +export type * from "./types.js"; diff --git a/sdk/src/pi/transport.test.ts b/sdk/src/pi/transport.test.ts new file mode 100644 index 0000000..df61c3f --- /dev/null +++ b/sdk/src/pi/transport.test.ts @@ -0,0 +1,95 @@ +import { describe, expect, it } from "vitest"; +import { + createControllableStream, + createMockAxon, + makeAgentEvent, + makeUserEvent, +} from "../__test-utils__/mock-axon.js"; +import { PiAxonTransport } from "./transport.js"; + +function setup(replayTargetSequence?: number) { + const ctrl = createControllableStream(true); + const { axon, published } = createMockAxon(ctrl); + const transport = new PiAxonTransport( + axon as never, + replayTargetSequence != null ? { replayTargetSequence } : {}, + ); + return { ctrl, published, transport }; +} + +describe("PiAxonTransport", () => { + it("publishes a prompt as turn/start with the frame verbatim", async () => { + const { published, transport } = setup(); + await transport.connect(); + await transport.write({ type: "prompt", id: "pi-sdk-1", message: "hi" }); + expect(published[0]).toMatchObject({ + event_type: "turn/start", + origin: "USER_EVENT", + source: "pi-sdk-client", + }); + expect(JSON.parse(published[0]?.payload ?? "null")).toEqual({ + type: "prompt", + id: "pi-sdk-1", + message: "hi", + }); + }); + + it("maps abort to cancel and leaves every other command as itself", async () => { + const { published, transport } = setup(); + await transport.connect(); + await transport.write({ type: "abort", id: "pi-sdk-1" }); + await transport.write({ type: "steer", id: "pi-sdk-2", message: "actually" }); + await transport.write({ type: "follow_up", id: "pi-sdk-3", message: "and then" }); + await transport.write({ type: "get_state", id: "pi-sdk-4" }); + expect(published.map((call) => call.event_type)).toEqual([ + "cancel", + "steer", + "follow_up", + "get_state", + ]); + }); + + it("falls back to unknown for a frame with no type", async () => { + const { published, transport } = setup(); + await transport.connect(); + await transport.write({ id: "pi-sdk-1" }); + expect(published[0]?.event_type).toBe("unknown"); + }); + + it("rejects broker-reserved command ids", async () => { + const { transport } = setup(); + await transport.connect(); + await expect( + transport.write({ type: "prompt", id: "broker-1", message: "hi" }), + ).rejects.toThrow("reserved"); + }); + + // Pi has no server-initiated requests, so PiAxonTransport omits the four + // replay-request callbacks and the shared transport buffers nothing. + it("buffers nothing across the replay window", async () => { + const { ctrl, transport } = setup(2); + await transport.connect(); + ctrl.push(makeAgentEvent("agent_start", { type: "agent_start" }, 1)); + ctrl.push(makeUserEvent("turn/start", { type: "prompt", id: "pi-sdk-1", message: "hi" }, 2)); + ctrl.push(makeAgentEvent("agent_settled", { type: "agent_settled" }, 3)); + ctrl.end(); + const frames = []; + for await (const frame of transport.readMessages()) frames.push(frame); + expect(frames).toEqual([{ type: "agent_settled" }]); + }); + + it("skips unparseable payloads and reports readiness", async () => { + const { ctrl, transport } = setup(); + expect(transport.isReady()).toBe(false); + await transport.connect(); + expect(transport.isReady()).toBe(true); + ctrl.push({ event_type: "agent_start", payload: "not json", origin: "AGENT_EVENT" }); + ctrl.push(makeAgentEvent("agent_settled", { type: "agent_settled" })); + ctrl.end(); + const frames = []; + for await (const frame of transport.readMessages()) frames.push(frame); + expect(frames).toEqual([{ type: "agent_settled" }]); + await transport.close(); + expect(transport.isReady()).toBe(false); + }); +}); diff --git a/sdk/src/pi/transport.ts b/sdk/src/pi/transport.ts new file mode 100644 index 0000000..77422e4 --- /dev/null +++ b/sdk/src/pi/transport.ts @@ -0,0 +1,97 @@ +import type { AxonEventView } from "@runloop/api-client/resources/axons"; +import type { Axon } from "@runloop/api-client/sdk"; +import { AxonFrameTransport } from "../shared/axon-frame-transport.js"; +import { + PI_CANCEL_EVENT_TYPE, + PI_TURN_START_EVENT_TYPE, + RESERVED_REQUEST_ID_PREFIX, +} from "./protocol/index.js"; + +/** A Pi JSONL frame in either direction: a command, an event, or an ack. */ +export type PiFrame = { + type?: string; + id?: string; + [key: string]: unknown; +}; + +export interface PiAxonTransportOptions { + verbose?: boolean; + onAxonEvent?: (event: AxonEventView) => void; + afterSequence?: number; + replayTargetSequence?: number; +} + +export interface PiTransport { + connect(): Promise; + reconnect(): Promise; + write(frame: PiFrame | string): Promise; + readMessages(): AsyncIterable; + close(): Promise; + abortStream(): void; + isReady(): boolean; +} + +/** + * Whole-frame Pi JSONL transport over Axon publish/SSE. + * + * The Pi broker adapter is a translating proxy: the published `event_type` + * selects its behaviour, and the payload is the raw Pi command frame. + * @category Transport + */ +export class PiAxonTransport implements PiTransport { + private readonly inner: AxonFrameTransport; + constructor(axon: Axon, options: PiAxonTransportOptions = {}) { + const parseFrame = (payload: string): PiFrame | undefined => { + try { + const frame: unknown = JSON.parse(payload); + return typeof frame === "object" && frame !== null ? (frame as PiFrame) : undefined; + } catch { + return undefined; + } + }; + this.inner = new AxonFrameTransport(axon, { + ...options, + source: "pi-sdk-client", + logPrefix: "pi-axon-transport", + parseFrame, + // The adapter's `classify_input` matches these two event types exactly: + // `turn/start` opens a broker turn, `cancel` aborts it. Every other + // frame is forwarded to Pi's stdin verbatim as a Control frame, which is + // what `steer` and `follow_up` need so they do not reopen a turn. + resolveEventType: (frame) => + frame?.type === "prompt" + ? PI_TURN_START_EVENT_TYPE + : frame?.type === "abort" + ? PI_CANCEL_EVENT_TYPE + : (frame?.type ?? "unknown"), + validateOutbound: (frame) => { + if (typeof frame.id === "string" && frame.id.startsWith(RESERVED_REQUEST_ID_PREFIX)) + throw new Error(`Request IDs beginning with ${RESERVED_REQUEST_ID_PREFIX} are reserved`); + }, + systemErrorsDuringReplay: false, + // Pi has no server-initiated requests, so the replay-request callbacks + // are omitted and nothing is buffered across the replay window. + }); + } + async connect(): Promise { + await this.inner.connect(); + } + async reconnect(): Promise { + await this.inner.reconnect(); + } + async write(frame: PiFrame | string): Promise { + await this.inner.write(frame); + } + async *readMessages(): AsyncGenerator { + yield* this.inner.readMessages(); + } + abortStream(): void { + this.inner.abortStream(); + } + async close(): Promise { + await this.inner.close(); + } + isReady(): boolean { + return this.inner.isReady(); + } +} diff --git a/sdk/src/shared/async-message-queue.test.ts b/sdk/src/shared/async-message-queue.test.ts index 4507fd3..b48953c 100644 --- a/sdk/src/shared/async-message-queue.test.ts +++ b/sdk/src/shared/async-message-queue.test.ts @@ -69,6 +69,19 @@ describe("AsyncMessageQueue", () => { expect(onHighWater).toHaveBeenCalledTimes(2); }); + it("discards the oldest values once maxBuffered is reached", async () => { + const queue = new AsyncMessageQueue(3, undefined, 3); + for (const value of [1, 2, 3, 4, 5]) queue.push(value); + expect(queue.size).toBe(3); + expect([await queue.next(), await queue.next(), await queue.next()]).toEqual([3, 4, 5]); + }); + + it("buffers without bound when maxBuffered is omitted", () => { + const queue = new AsyncMessageQueue(2); + for (const value of [1, 2, 3, 4, 5]) queue.push(value); + expect(queue.size).toBe(5); + }); + it("buffers values pushed after close(false) for a later drain", async () => { const queue = new AsyncMessageQueue(); queue.close(false); diff --git a/sdk/src/shared/async-message-queue.ts b/sdk/src/shared/async-message-queue.ts index c756b19..906ec87 100644 --- a/sdk/src/shared/async-message-queue.ts +++ b/sdk/src/shared/async-message-queue.ts @@ -5,16 +5,29 @@ export class AsyncMessageQueue { private closed = false; private warned = false; + /** + * @param maxBuffered Hard cap on buffered values. Once reached, pushing + * discards the oldest value so a queue with no consumer cannot grow without + * bound. Omit for an unbounded queue. + */ constructor( private readonly highWaterMark = 1000, private readonly onHighWater?: (size: number) => void, + private readonly maxBuffered?: number, ) {} + /** Number of values buffered because no consumer was waiting. */ + get size(): number { + return this.values.length; + } + push(value: T): void { const waiter = this.waiters.shift(); if (waiter) waiter(value); else { this.values.push(value); + if (this.maxBuffered !== undefined) + while (this.values.length > this.maxBuffered) this.values.shift(); if (!this.warned && this.values.length >= this.highWaterMark) { this.warned = true; this.onHighWater?.(this.values.length); diff --git a/sdk/src/shared/axon-frame-transport.ts b/sdk/src/shared/axon-frame-transport.ts index 487fea3..31ea9cf 100644 --- a/sdk/src/shared/axon-frame-transport.ts +++ b/sdk/src/shared/axon-frame-transport.ts @@ -17,10 +17,15 @@ export interface AxonFrameTransportOptions { logPrefix: string; parseFrame(payload: string): TFrame | undefined; resolveEventType(frame: TFrame | undefined, raw: string): string; - isReplayRequest(event: AxonEventView, frame: TFrame): boolean; - requestId(frame: TFrame): FrameRequestId | undefined; - isReplayAnswer(event: AxonEventView, frame: TFrame): boolean; - answerId(frame: TFrame): FrameRequestId | undefined; + /** + * Replay-window buffering of server-initiated requests that were never + * answered. Omit all four for protocols with no server-initiated requests: + * nothing is buffered and the replay window yields nothing. + */ + isReplayRequest?(event: AxonEventView, frame: TFrame): boolean; + requestId?(frame: TFrame): FrameRequestId | undefined; + isReplayAnswer?(event: AxonEventView, frame: TFrame): boolean; + answerId?(frame: TFrame): FrameRequestId | undefined; validateOutbound?(frame: TFrame): void; allowInvalidOutbound?: boolean; systemErrorsDuringReplay?: boolean; @@ -100,11 +105,11 @@ export class AxonFrameTransport { if (event.payload != null) { const frame = this.options.parseFrame(event.payload); if (frame !== undefined) { - if (this.options.isReplayRequest(event, frame)) { - const id = this.options.requestId(frame); + if (this.options.isReplayRequest?.(event, frame)) { + const id = this.options.requestId?.(frame); if (id !== undefined) replayBuffer.set(id, frame); - } else if (this.options.isReplayAnswer(event, frame)) { - const id = this.options.answerId(frame); + } else if (this.options.isReplayAnswer?.(event, frame)) { + const id = this.options.answerId?.(frame); if (id !== undefined) replayBuffer.delete(id); } }