feat(sdk): add PiAxonConnection and transport - #155
Conversation
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>
| : 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; |
There was a problem hiding this comment.
[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.
There was a problem hiding this comment.
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.
| // receiveTurn() even though send() already surfaced it as an error. | ||
| if (settled && !(ack.command === PROMPT_COMMAND && !ack.success)) return; | ||
| } | ||
| this.messageQueue.push(frame); |
There was a problem hiding this comment.
[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.
There was a problem hiding this comment.
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.
| ); | ||
| 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)}`; |
There was a problem hiding this comment.
[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.
There was a problem hiding this comment.
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
left a comment
There was a problem hiding this comment.
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 localcheck → typecheck → build → testrun 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 changeci.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 atagent_settled, neveragent_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
PendingRequestMapand 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.
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>
|
Re the [TEST] major on missing CI: agreed, and addressed two ways.
Once #158 merges, this PR and #156/#157 will get build/typecheck/test results on their next push, so real required checks precede landing. Local gate after the three fixes in 05b03a0, on Node 22 + bun: |
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.ts—PiAxonTransportover the sharedAxonFrameTransport<PiFrame>. Outbound: the frame'stypebecomes the Axonevent_type, exceptprompt→turn/startandabort→cancel(broker control events); ids beginning withbroker-are refused. InboundAGENT_EVENTs are parsed back into frames verbatim; unparseable payloads are skipped.sdk/src/pi/connection.ts—PiAxonConnection:connect()(no handshake),send/steer/followUp/interrupt,getState/newSession/switchSession, acommand()escape hatch,receiveAgentEvents()/receiveTurn()/receiveTimelineEvents(), listeners,publish(),abortStream(),disconnect(), andsessionId/sessionFilegetters. Commands are id-stamped and correlated against{"type":"response","command":…,"success":…}acks; a rejected command rejects withPiCommandError. 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 intransport.test.ts, the caller that omits them.sdk/src/pi/index.ts,sdk/src/index.ts,sdk/package.json—@runloop/remote-agents-sdk/pisubpath export and apinamespace on the root export.sdk/AGENTS.md,sdk/README.md— module tables, thepi_jsonbroker protocol, quickstart,PiAxonConnectionreference, 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:
agent_settled, notagent_end. Pi can emitagent_endwithwillRetry: trueand keep working, soreceiveTurn()terminates onagent_settled(or on a rejectedpromptack, which ends the turn immediately).initialize(). The broker'sInitializere-issuesget_statewith noid, so its ack is uncorrelatable;getState()sends its own id-stamped Control frame instead.No
BaseAxonConnectionor other new abstraction — the connection reuses the existing shared primitives (AxonFrameTransport,runConnectionReadLoop,PendingRequestMap,AsyncMessageQueue,ListenerSet,createClassifier) directly. The SDK never sets--mode rpcor--session-dir; those are broker-owned.CI status — read before approving
GitHub ran only
conventional-commiton this PR. There is no build, typecheck or test result here, becauseci.ymlrestricts itspull_requesttrigger tobranches: [main]and this PR is based onpi/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 checkclean,typecheckandbuildexit 0,bun run test29 files / 665 tests pass,src/picoverage 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
<type>(<scope>): <description>format (see above)bun run checkpasses (lint + format)bun run buildpassesbun run testpasses (659 tests;src/picoverage 97.8% statements / 92% branches)