Skip to content

feat(sdk): add PiAxonConnection and transport - #155

Open
ross-rl wants to merge 3 commits into
pi/sdk-wire-typesfrom
pi/sdk-connection
Open

feat(sdk): add PiAxonConnection and transport#155
ross-rl wants to merge 3 commits into
pi/sdk-wire-typesfrom
pi/sdk-connection

Conversation

@ross-rl

@ross-rl ross-rl commented Jul 29, 2026

Copy link
Copy Markdown
Collaborator

What

Stacked on #154 (pi/sdk-wire-types) — review that first; this PR's diff is the client layer on top of those types.

  • sdk/src/pi/transport.tsPiAxonTransport over the shared AxonFrameTransport<PiFrame>. Outbound: the frame's type becomes the Axon event_type, except promptturn/start and abortcancel (broker control events); ids beginning with broker- are refused. Inbound AGENT_EVENTs are parsed back into frames verbatim; unparseable payloads are skipped.
  • sdk/src/pi/connection.tsPiAxonConnection: connect() (no handshake), send/steer/followUp/interrupt, getState/newSession/switchSession, a command() escape hatch, receiveAgentEvents()/receiveTurn()/receiveTimelineEvents(), listeners, publish(), abortStream(), disconnect(), and sessionId/sessionFile getters. Commands are id-stamped and correlated against {"type":"response","command":…,"success":…} acks; a rejected command rejects with PiCommandError. Acks the SDK did not ask for (including broker-stamped ones) fall through to the message queue rather than being swallowed.
  • sdk/src/shared/axon-frame-transport.ts — the four replay-request callbacks (isReplayRequest, requestId, isReplayAnswer, answerId) become optional; with all four omitted the replay window buffers nothing. Covered by a test in transport.test.ts, the caller that omits them.
  • sdk/src/pi/index.ts, sdk/src/index.ts, sdk/package.json@runloop/remote-agents-sdk/pi subpath export and a pi namespace on the root export.
  • sdk/AGENTS.md, sdk/README.md — module tables, the pi_json broker protocol, quickstart, PiAxonConnection reference, timeline-event guards, and gotchas.

Why

Completes the Pi module: the wire types in #154 describe the protocol, this makes it usable. Two behaviours drive the design:

  • A turn ends at agent_settled, not agent_end. Pi can emit agent_end with willRetry: true and keep working, so receiveTurn() terminates on agent_settled (or on a rejected prompt ack, which ends the turn immediately).
  • There is no handshake and no initialize(). The broker's Initialize re-issues get_state with no id, so its ack is uncorrelatable; getState() sends its own id-stamped Control frame instead.

No BaseAxonConnection or other new abstraction — the connection reuses the existing shared primitives (AxonFrameTransport, runConnectionReadLoop, PendingRequestMap, AsyncMessageQueue, ListenerSet, createClassifier) directly. The SDK never sets --mode rpc or --session-dir; those are broker-owned.

CI status — read before approving

GitHub ran only conventional-commit on this PR. There is no build, typecheck or test result here, because ci.yml restricts its pull_request trigger to branches: [main] and this PR is based on pi/sdk-wire-types. The one-item check list below is not evidence that anything passed.

Verification of record is local, on Node 22 + bun: bun run check clean, typecheck and build exit 0, bun run test 29 files / 665 tests pass, src/pi coverage 98.42% statements / 92.85% branches / 98.52% functions / 99.44% lines.

#158 broadens the CI trigger to cover any PR base; once it lands, this PR gets real checks on its next push.

Checklist

  • PR title follows <type>(<scope>): <description> format (see above)
  • bun run check passes (lint + format)
  • bun run build passes
  • bun run test passes (659 tests; src/pi coverage 97.8% statements / 92% branches)
  • SDK documentation updated (if applicable)

Add the Pi client on top of the wire types: PiAxonTransport maps outbound
frames onto Axon event types (prompt -> turn/start, abort -> cancel) and
parses inbound agent events back into frames verbatim; PiAxonConnection
correlates commands with id-stamped response acks, tracks session identity
from get_state, and ends a turn at agent_settled rather than agent_end.

Pi issues no server-initiated requests, so the four replay-request callbacks
on AxonFrameTransportOptions become optional: with all four omitted the
replay window buffers nothing.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Comment thread sdk/src/pi/connection.ts
: 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.
if (settled && !(ack.command === PROMPT_COMMAND && !ack.success)) return;

@ross-rl ross-rl Jul 29, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[BUG] major — The rejected-prompt ack can be orphaned and then terminate the next turn. Sequence: (1) await send("A") rejects, and the normal caller catches it without calling receiveTurn() because no turn started; (2) this ack remains in messageQueue; (3) send("B") is accepted; (4) the subsequent receiveTurn() immediately yields A's stale rejection and returns, leaving all of B's frames queued. The current test only exercises the special case where callers drain immediately after a rejected send(). Please tie this queued termination to the corresponding receive operation/turn, or ensure a rejection cannot contaminate a later turn, and add the two-send regression sequence.

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed in 05b03a0, by correlation — the same shape as the broker-side fix.

request() records the id of the most recent prompt in latestPromptId. receiveTurn() now inspects a rejected-prompt ack before yielding it: if it carries an id that is not latestPromptId, it is skipped (not yielded, not terminal) because send() already surfaced it to its own caller as a PiCommandError. Only a rejection belonging to the latest prompt attempt ends the turn, and consuming it clears latestPromptId so it cannot fire twice.

An id-less rejection still terminates: it cannot be mis-attributed to an SDK prompt. A broker-N rejection does not — that is the broker's own prompt, and it reports the failure through the turn.failed system event. Documented on receiveTurn().

Regression test: does not let an undrained rejection terminate a later accepted turn runs your exact sequence — send("first") rejected and never drained, send("second") accepted, then receiveTurn(). It yields ["turn_start", "agent_settled"]; before the fix it returned after the stale ack with B's frames stranded.

Comment thread sdk/src/pi/connection.ts
// receiveTurn() even though send() already surfaced it as an error.
if (settled && !(ack.command === PROMPT_COMMAND && !ack.success)) return;
}
this.messageQueue.push(frame);

@ross-rl ross-rl Jul 29, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[BUG] major — Every non-consumed Pi frame is queued even when the application uses onTimelineEvent() as its sole (documented) consumption surface. A 2,000-token response produces roughly 2,000 cumulative message_update frames; the listener handles them, but this queue retains every frame (quadratic payload growth), and repeated turns grow memory without bound. PR #156's combined app uses exactly this push-only pattern and never drains receiveAgentEvents(). Either avoid retaining frames when a push consumer is configured, provide a bounded/drop policy or separate pull subscription, or make the documented/example usage run a drain loop. Please cover a listener-only turn and verify the internal pull backlog does not grow.

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed in 05b03a0. Policy: the pull queue is bounded, oldest-dropped.

PiAxonConnectionOptions.maxQueuedFrames (default 1000) caps what receiveAgentEvents()/receiveTurn() retain; past the cap the oldest frame is discarded and the existing high-water callback warns once, with a message that now says frames are being dropped. Listeners are unaffected — onTimelineEvent/onAxonEvent still see every frame — so a listener-only application has bounded memory no matter how many turns it runs.

I bounded rather than "stop retaining when a push consumer exists": the common sequence is await send(...) then receiveTurn(), and a listener registered for UI purposes must not silently disable retention for frames that arrive in the gap before the generator starts. Dropping the oldest also keeps termination correct — agent_settled sits at the tail, so a lagging puller loses cumulative deltas rather than the frame that ends its turn.

This needed a second shared-primitive change: AsyncMessageQueue gains an optional maxBuffered ctor arg and a size getter. Both are additive — existing Claude/Codex queues pass no cap and keep unbounded behaviour — and are covered by two new tests in async-message-queue.test.ts.

Coverage for the listener-only case: bounds the pull queue when frames are consumed only through listeners pushes 50 message_updates plus agent_settled with maxQueuedFrames: 10, asserts all 51 reached the listener, and asserts the backlog drained afterwards is 10 frames ending at agent_settled.

Documented in README.md (options table + Known Limitations) and AGENTS.md. This also resolves the #156 finding without the example changing — I have said so on that PR.

Comment thread sdk/src/pi/connection.ts Outdated
);
if (!this.transport?.isReady())
throw new ConnectionStateError("not_connected", "Not connected. Call connect() first.");
const id = command.id ?? `pi-sdk-${++this.counter}-${Math.random().toString(36).slice(2, 10)}`;

@ross-rl ross-rl Jul 29, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[BUG] minor — Empty caller-supplied ids bypass generation here. For command({ type: "prompt", id: "", message: "x" }), the broker's prepare_prompt treats the empty id as absent and replaces it with its own broker-N id. The response therefore cannot resolve the pending "" entry: it falls through to the queue and this request waits until the 60s timeout. Please reject empty ids or treat them as absent before creating the pending entry; add this alongside the reserved-prefix test.

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed in 05b03a0. A blank id (empty or whitespace-only) is treated as absent, so request() generates a pi-sdk-… id rather than creating a pending "" entry the broker's broker-N ack could never resolve. Test generates an id for a blank caller-supplied id sits next to the reserved-prefix test and asserts the published frame carries a generated id.

I treated it as absent rather than rejecting it: absence is exactly what the broker infers from a blank id, and a caller passing id: "" wants the default, not an error.

@ross-rl ross-rl left a comment

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Review summary

Findings: 3 major, 1 minor.

  • [BUG] major — A rejected prompt ack can remain queued and prematurely terminate the next accepted turn; see inline.
  • [BUG] major — Push-only timeline consumers still retain every frame in the pull queue, causing unbounded (and for cumulative partials effectively quadratic) memory growth; see inline.
  • [TEST] major — This stacked PR received only the conventional-commit check. The shared transport change and new connection have no GitHub build/typecheck/test result because CI is restricted to pull requests targeting main. The reported local check → typecheck → build → test run materially reduces risk, but it is not an acceptable substitute for reproducible required checks on an SDK/shared-primitive change. Please run CI on this head before landing and change ci.yml/relevant smoke trigger to cover stacked PR bases (or provide an explicit manually-triggered equivalent).
  • [BUG] minor — An empty explicit prompt id is replaced by the broker, leaving the SDK request uncorrelatable until timeout; see inline.

Focused verdicts:

  • Turn termination is otherwise correct: receiveTurn() stops at agent_settled, never agent_end, and also stops on a rejected prompt ack. A terminal read-loop closure closes the queue rather than hanging the iterator, and an absent ack has the configured timeout.
  • The rejected-prompt ack is settled once in PendingRequestMap and enqueued once; I did not find a duplicate queue delivery or an inherent unhandled-rejection path. The stale/orphaned queue entry is the problem.
  • Broker/SDK id namespaces are separated by the broker- refusal, duplicate live SDK ids are rejected, and unknown/id-less acks safely fall through to the queue.
  • Existing Claude and Codex callers still provide all four replay callbacks. For Pi, buffering nothing during replay is correct because it has no server-initiated requests. The Pi transport test exercises the actual shared default path, so its placement still gives real coverage. The optional-callback API would be stronger as an all-four-or-none union, but I do not consider that merge-blocking.

Reflex and others added 2 commits July 29, 2026 07:56
A rejected prompt ack left in the queue by a caller that caught the error
without draining could terminate a later accepted turn. receiveTurn() now
correlates a rejection to the prompt it belongs to and skips an orphaned one.

The pull queue gains a cap (maxQueuedFrames, default 1000) past which the
oldest frame is discarded, so an application that consumes only through
onTimelineEvent() no longer retains every frame of every turn.

A blank caller-supplied command id is treated as absent: the broker replaces
it with its own id, leaving the pending entry uncorrelatable until timeout.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
@ross-rl

ross-rl commented Jul 29, 2026

Copy link
Copy Markdown
Collaborator Author

Re the [TEST] major on missing CI: agreed, and addressed two ways.

  1. Stated explicitly in the PR body — there is now a "CI status — read before approving" section at the top saying that only conventional-commit ran, that the one-item check list is not evidence of a passing build, and what the local verification actually covered. An empty check list should not read as "verified".
  2. Opened ci(project): run CI on pull requests with any base #158ci(project): run CI on pull requests with any base drops the branches: [main] filter from ci.yml's pull_request trigger, so check (Node 22 and 24) and audit run on every PR whatever its base. ci(project): run CI on pull requests with any base #158 is green (check (22), check (24), audit, conventional-commit). smoke-tests.yml keeps its main-only trigger: its single job is already gated on startsWith(github.head_ref, 'release-please--branches--main') and it drives real devboxes, so broadening it adds cost without signal. Say the word if you want that broadened too.

Once #158 merges, this PR and #156/#157 will get build/typecheck/test results on their next push, so real required checks precede landing. ci.yml has no workflow_dispatch, so I cannot force a run on this head before then — #158 is the shortest path to it.

Local gate after the three fixes in 05b03a0, on Node 22 + bun: bun run check clean, typecheck 0, build 0, bun run test 29 files / 665 tests pass, src/pi coverage 98.42% statements / 92.85% branches / 98.52% functions / 99.44% lines.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant