From a4475344fb765850ec5321efe3c67e6f364ead5c Mon Sep 17 00:00:00 2001 From: Christian Klotz Date: Mon, 3 Aug 2026 19:33:22 +0300 Subject: [PATCH] fix(coding-agent): make JSON streaming output linear (#7394) Fixes #7290 Co-authored-by: Mario Zechner --- packages/coding-agent/CHANGELOG.md | 1 + packages/coding-agent/docs/json.md | 33 ++++++++------ packages/coding-agent/docs/rpc.md | 24 +++++----- packages/coding-agent/src/index.ts | 1 + packages/coding-agent/src/modes/index.ts | 1 + packages/coding-agent/src/modes/json-event.ts | 40 +++++++++++++++++ packages/coding-agent/src/modes/print-mode.ts | 14 +++++- .../coding-agent/src/modes/rpc/rpc-client.ts | 13 +++--- .../coding-agent/src/modes/rpc/rpc-mode.ts | 3 +- packages/coding-agent/test/print-mode.test.ts | 4 +- .../7290-json-stream-linear.test.ts | 44 +++++++++++++++++++ 11 files changed, 142 insertions(+), 36 deletions(-) create mode 100644 packages/coding-agent/src/modes/json-event.ts create mode 100644 packages/coding-agent/test/suite/regressions/7290-json-stream-linear.test.ts diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index debafcc97..ca625c284 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -4,6 +4,7 @@ ### Breaking Changes +- Changed JSON and RPC `message_update` events to emit only `assistantMessageEvent` deltas, removing the cumulative `message` and `assistantMessageEvent.partial` fields that caused quadratic output growth. Clients that need partial messages must assemble deltas between `message_start` and `message_end`; the latter remains authoritative ([#7290](https://github.com/earendil-works/pi/issues/7290)). - `ModelRegistry.getApiKeyAndHeaders()` now returns `ProviderHeaders` with `string | null` values and preserves `null` header-deletion markers. Extensions that inspect returned headers must handle `null`; extensions forwarding them to pi-ai streams should pass them through unchanged. This prevents placeholder OpenAI credentials from being sent through Cloudflare AI Gateway ([#7030](https://github.com/earendil-works/pi/issues/7030)). ### Added diff --git a/packages/coding-agent/docs/json.md b/packages/coding-agent/docs/json.md index a9d28cbc0..f497f3406 100644 --- a/packages/coding-agent/docs/json.md +++ b/packages/coding-agent/docs/json.md @@ -8,25 +8,25 @@ Outputs all session events as JSON lines to stdout. Useful for integrating pi in ## Event Types -Events are defined in [`AgentSessionEvent`](https://github.com/earendil-works/pi-mono/blob/main/packages/coding-agent/src/core/agent-session.ts#L102): +Wire events use `JsonAgentSessionEvent`. It matches +[`AgentSessionEvent`](https://github.com/earendil-works/pi-mono/blob/main/packages/coding-agent/src/core/agent-session.ts) +except that streaming message updates omit cumulative snapshots: ```typescript -type AgentSessionEvent = - | AgentEvent - | { type: "queue_update"; steering: readonly string[]; followUp: readonly string[] } - | { type: "compaction_start"; reason: "manual" | "threshold" | "overflow" } - | { type: "compaction_end"; reason: "manual" | "threshold" | "overflow"; result: CompactionResult | undefined; aborted: boolean; willRetry: boolean; errorMessage?: string } - | { type: "auto_retry_start"; attempt: number; maxAttempts: number; delayMs: number; errorMessage: string } - | { type: "auto_retry_end"; success: boolean; attempt: number; finalError?: string } - | { type: "summarization_retry_scheduled"; attempt: number; maxAttempts: number; delayMs: number; errorMessage: string } - | { type: "summarization_retry_attempt_start"; source: "branchSummary" } - | { type: "summarization_retry_attempt_start"; source: "compaction"; reason: "manual" | "threshold" | "overflow" } - | { type: "summarization_retry_finished" }; +type WithoutPartial = T extends { partial: unknown } ? Omit : T; + +type JsonAgentSessionEvent = + | Exclude + | { + type: "message_update"; + assistantMessageEvent: WithoutPartial; + }; ``` `queue_update` emits the full pending steering and follow-up queues whenever they change. `compaction_start` and `compaction_end` cover both manual and automatic compaction. -Base events from [`AgentEvent`](https://github.com/earendil-works/pi-mono/blob/main/packages/agent/src/types.ts#L179): +Other base events come from +[`AgentEvent`](https://github.com/earendil-works/pi-mono/blob/main/packages/agent/src/types.ts): ```typescript type AgentEvent = @@ -73,12 +73,17 @@ Followed by events as they occur: {"type":"agent_start"} {"type":"turn_start"} {"type":"message_start","message":{"role":"assistant","content":[],...}} -{"type":"message_update","message":{...},"assistantMessageEvent":{"type":"text_delta","delta":"Hello",...}} +{"type":"message_update","assistantMessageEvent":{"type":"text_delta","contentIndex":0,"delta":"Hello"}} {"type":"message_end","message":{...}} {"type":"turn_end","message":{...},"toolResults":[]} {"type":"agent_end","messages":[...]} ``` +`message_update` records are delta-only. They omit both the cumulative `message` field and +`assistantMessageEvent.partial` to keep stream size linear. Use `contentIndex` and `delta` +to assemble live text, thinking, or tool-call arguments if needed. `message_end` contains +the final authoritative message. + ## Example ```bash diff --git a/packages/coding-agent/docs/rpc.md b/packages/coding-agent/docs/rpc.md index 12f5b02ab..499ca673b 100644 --- a/packages/coding-agent/docs/rpc.md +++ b/packages/coding-agent/docs/rpc.md @@ -914,17 +914,15 @@ Emitted when a message begins and completes. The `message` field contains an `Ag ### message_update (Streaming) -Emitted during streaming of assistant messages. Contains both the partial message and a streaming delta event. +Emitted during streaming of assistant messages. Contains a delta event without a cumulative message snapshot. ```json { "type": "message_update", - "message": {...}, "assistantMessageEvent": { "type": "text_delta", "contentIndex": 0, - "delta": "Hello ", - "partial": {...} + "delta": "Hello " } } ``` @@ -933,7 +931,6 @@ The `assistantMessageEvent` field contains one of these delta types: | Type | Description | |------|-------------| -| `start` | Message generation started | | `text_start` | Text content block started | | `text_delta` | Text content chunk | | `text_end` | Text content block ended | @@ -943,17 +940,21 @@ The `assistantMessageEvent` field contains one of these delta types: | `toolcall_start` | Tool call started | | `toolcall_delta` | Tool call arguments chunk | | `toolcall_end` | Tool call ended (includes full `toolCall` object) | -| `done` | Message complete (reason: `"stop"`, `"length"`, `"toolUse"`) | -| `error` | Error occurred (reason: `"aborted"`, `"error"`) | Example streaming a text response: ```json -{"type":"message_update","message":{...},"assistantMessageEvent":{"type":"text_start","contentIndex":0,"partial":{...}}} -{"type":"message_update","message":{...},"assistantMessageEvent":{"type":"text_delta","contentIndex":0,"delta":"Hello","partial":{...}}} -{"type":"message_update","message":{...},"assistantMessageEvent":{"type":"text_delta","contentIndex":0,"delta":" world","partial":{...}}} -{"type":"message_update","message":{...},"assistantMessageEvent":{"type":"text_end","contentIndex":0,"content":"Hello world","partial":{...}}} +{"type":"message_update","assistantMessageEvent":{"type":"text_start","contentIndex":0}} +{"type":"message_update","assistantMessageEvent":{"type":"text_delta","contentIndex":0,"delta":"Hello"}} +{"type":"message_update","assistantMessageEvent":{"type":"text_delta","contentIndex":0,"delta":" world"}} +{"type":"message_update","assistantMessageEvent":{"type":"text_end","contentIndex":0,"content":"Hello world"}} ``` +`message_update` intentionally omits the former cumulative `message` field and +`assistantMessageEvent.partial`. Clients that need a live partial message must assemble it +from `message_start` and subsequent events using `contentIndex`. Treat `message_end.message` +as authoritative. For tool calls, buffer `toolcall_delta.delta`; `toolcall_end.toolCall` +contains the completed call. + ### bash_execution_update Emitted once for each output chunk from a direct `bash` command. `id` matches the command's `id`, allowing clients to associate output with the correct command. @@ -1362,6 +1363,7 @@ Source files: - [`packages/ai/src/types.ts`](../../ai/src/types.ts) - `Model`, `UserMessage`, `AssistantMessage`, `ToolResultMessage` - [`packages/agent/src/types.ts`](../../agent/src/types.ts) - `AgentMessage`, `AgentEvent` - [`src/core/messages.ts`](../src/core/messages.ts) - `BashExecutionMessage` +- [`src/modes/json-event.ts`](../src/modes/json-event.ts) - `JsonAgentSessionEvent` - [`src/modes/rpc/rpc-types.ts`](../src/modes/rpc/rpc-types.ts) - RPC command/response types, extension UI request/response types ### Model diff --git a/packages/coding-agent/src/index.ts b/packages/coding-agent/src/index.ts index 3615bccad..394851162 100644 --- a/packages/coding-agent/src/index.ts +++ b/packages/coding-agent/src/index.ts @@ -332,6 +332,7 @@ export { type MainOptions, main } from "./main.ts"; export { InteractiveMode, type InteractiveModeOptions, + type JsonAgentSessionEvent, type ModelInfo, type PrintModeOptions, RpcClient, diff --git a/packages/coding-agent/src/modes/index.ts b/packages/coding-agent/src/modes/index.ts index 732aee558..ce860b075 100644 --- a/packages/coding-agent/src/modes/index.ts +++ b/packages/coding-agent/src/modes/index.ts @@ -3,6 +3,7 @@ */ export { InteractiveMode, type InteractiveModeOptions } from "./interactive/interactive-mode.ts"; +export type { JsonAgentSessionEvent } from "./json-event.ts"; export { type PrintModeOptions, runPrintMode } from "./print-mode.ts"; export { type ModelInfo, RpcClient, type RpcClientOptions, type RpcEventListener } from "./rpc/rpc-client.ts"; export { runRpcMode } from "./rpc/rpc-mode.ts"; diff --git a/packages/coding-agent/src/modes/json-event.ts b/packages/coding-agent/src/modes/json-event.ts new file mode 100644 index 000000000..6cff8ebf2 --- /dev/null +++ b/packages/coding-agent/src/modes/json-event.ts @@ -0,0 +1,40 @@ +import type { AgentSessionEvent } from "../core/agent-session.ts"; + +type WithoutPartial = T extends { partial: unknown } ? Omit : T; + +type ToJsonEvent = T extends { + type: "message_update"; + assistantMessageEvent: infer TAssistantMessageEvent; +} + ? { + type: "message_update"; + assistantMessageEvent: WithoutPartial; + } + : T; + +/** Session event shape emitted by the JSON and RPC stdout protocols. */ +export type JsonAgentSessionEvent = ToJsonEvent; + +type MessageUpdateEvent = Extract; +type JsonMessageUpdateEvent = Extract; + +/** + * Remove cumulative assistant snapshots from streaming wire events. + * `message_start` provides the initial message, deltas build it, and + * `message_end` provides the final authoritative message. + */ +export function toJsonEvent(event: MessageUpdateEvent): JsonMessageUpdateEvent; +export function toJsonEvent(event: AgentSessionEvent): JsonAgentSessionEvent; +export function toJsonEvent(event: AgentSessionEvent): JsonAgentSessionEvent { + if (event.type !== "message_update") { + return event; + } + + const assistantMessageEvent = event.assistantMessageEvent; + if (!("partial" in assistantMessageEvent)) { + return { type: "message_update", assistantMessageEvent }; + } + + const { partial: _partial, ...deltaEvent } = assistantMessageEvent; + return { type: "message_update", assistantMessageEvent: deltaEvent }; +} diff --git a/packages/coding-agent/src/modes/print-mode.ts b/packages/coding-agent/src/modes/print-mode.ts index ca63157bb..3f677eef6 100644 --- a/packages/coding-agent/src/modes/print-mode.ts +++ b/packages/coding-agent/src/modes/print-mode.ts @@ -8,8 +8,9 @@ import type { AssistantMessage, ImageContent } from "@earendil-works/pi-ai"; import type { AgentSessionRuntime } from "../core/agent-session-runtime.ts"; -import { flushRawStdout, writeRawStdout } from "../core/output-guard.ts"; +import { flushRawStdout, waitForRawStdoutBackpressure, writeRawStdout } from "../core/output-guard.ts"; import { killTrackedDetachedChildren } from "../utils/shell.ts"; +import { toJsonEvent } from "./json-event.ts"; /** * Options for print mode. @@ -34,6 +35,7 @@ export async function runPrintMode(runtimeHost: AgentSessionRuntime, options: Pr let exitCode = 0; let session = runtimeHost.session; let unsubscribe: (() => void) | undefined; + let unsubscribeBackpressure: (() => void) | undefined; let disposed = false; const signalCleanupHandlers: Array<() => void> = []; @@ -41,6 +43,7 @@ export async function runPrintMode(runtimeHost: AgentSessionRuntime, options: Pr if (disposed) return; disposed = true; unsubscribe?.(); + unsubscribeBackpressure?.(); await runtimeHost.dispose(); }; @@ -101,11 +104,18 @@ export async function runPrintMode(runtimeHost: AgentSessionRuntime, options: Pr }); unsubscribe?.(); + unsubscribeBackpressure?.(); unsubscribe = session.subscribe((event) => { if (mode === "json") { - writeRawStdout(`${JSON.stringify(event)}\n`); + writeRawStdout(`${JSON.stringify(toJsonEvent(event))}\n`); } }); + unsubscribeBackpressure = + mode === "json" + ? session.agent.subscribe(async () => { + await waitForRawStdoutBackpressure(); + }) + : undefined; }; try { diff --git a/packages/coding-agent/src/modes/rpc/rpc-client.ts b/packages/coding-agent/src/modes/rpc/rpc-client.ts index 79cd7f2f9..a4cc4b208 100644 --- a/packages/coding-agent/src/modes/rpc/rpc-client.ts +++ b/packages/coding-agent/src/modes/rpc/rpc-client.ts @@ -7,10 +7,11 @@ import { type ChildProcess, spawn } from "node:child_process"; import type { AgentMessage, ThinkingLevel } from "@earendil-works/pi-agent-core"; import type { ImageContent } from "@earendil-works/pi-ai"; -import type { AgentSessionEvent, SessionStats } from "../../core/agent-session.ts"; +import type { SessionStats } from "../../core/agent-session.ts"; import type { BashResult } from "../../core/bash-executor.ts"; import type { CompactionResult } from "../../core/compaction/index.ts"; import type { SessionEntry, SessionTreeNode } from "../../core/session-manager.ts"; +import type { JsonAgentSessionEvent } from "../json-event.ts"; import { attachJsonlLineReader, serializeJsonLine } from "./jsonl.ts"; import type { RpcCommand, RpcResponse, RpcSessionState, RpcSlashCommand } from "./rpc-types.ts"; @@ -46,7 +47,7 @@ export interface ModelInfo { reasoning: boolean; } -export type RpcEventListener = (event: AgentSessionEvent) => void; +export type RpcEventListener = (event: JsonAgentSessionEvent) => void; // ============================================================================ // RPC Client @@ -472,9 +473,9 @@ export class RpcClient { /** * Collect events until agent becomes idle. */ - collectEvents(timeout = 60000): Promise { + collectEvents(timeout = 60000): Promise { return new Promise((resolve, reject) => { - const events: AgentSessionEvent[] = []; + const events: JsonAgentSessionEvent[] = []; const timer = setTimeout(() => { unsubscribe(); reject(new Error(`Timeout collecting events. Stderr: ${this.stderr}`)); @@ -494,7 +495,7 @@ export class RpcClient { /** * Send prompt and wait for completion, returning all events. */ - async promptAndWait(message: string, images?: ImageContent[], timeout = 60000): Promise { + async promptAndWait(message: string, images?: ImageContent[], timeout = 60000): Promise { const eventsPromise = this.collectEvents(timeout); await this.prompt(message, images); return eventsPromise; @@ -518,7 +519,7 @@ export class RpcClient { // Otherwise it's an event for (const listener of this.eventListeners) { - listener(data as AgentSessionEvent); + listener(data as JsonAgentSessionEvent); } } catch { // Ignore non-JSON lines diff --git a/packages/coding-agent/src/modes/rpc/rpc-mode.ts b/packages/coding-agent/src/modes/rpc/rpc-mode.ts index c6db8c3a8..a2e2e65d2 100644 --- a/packages/coding-agent/src/modes/rpc/rpc-mode.ts +++ b/packages/coding-agent/src/modes/rpc/rpc-mode.ts @@ -27,6 +27,7 @@ import { } from "../../core/output-guard.ts"; import { killTrackedDetachedChildren } from "../../utils/shell.ts"; import { type Theme, theme } from "../interactive/theme/theme.ts"; +import { toJsonEvent } from "../json-event.ts"; import { attachJsonlLineReader, serializeJsonLine } from "./jsonl.ts"; import type { RpcCommand, @@ -352,7 +353,7 @@ export async function runRpcMode(runtimeHost: AgentSessionRuntime): Promise { - output(event); + output(toJsonEvent(event)); if (event.type === "agent_settled") { void checkShutdownRequested(); } diff --git a/packages/coding-agent/test/print-mode.test.ts b/packages/coding-agent/test/print-mode.test.ts index 30d7251de..32cec8e87 100644 --- a/packages/coding-agent/test/print-mode.test.ts +++ b/packages/coding-agent/test/print-mode.test.ts @@ -12,7 +12,7 @@ type FakeExtensionRunner = { type FakeSession = { sessionManager: { getHeader: () => object | undefined }; - agent: { waitForIdle: () => Promise }; + agent: { waitForIdle: () => Promise; subscribe: ReturnType }; state: { messages: AssistantMessage[] }; extensionRunner: FakeExtensionRunner; bindExtensions: ReturnType; @@ -65,7 +65,7 @@ function createRuntimeHost(assistantMessage: AssistantMessage): FakeRuntimeHost const session: FakeSession = { sessionManager: { getHeader: () => undefined }, - agent: { waitForIdle: async () => {} }, + agent: { waitForIdle: async () => {}, subscribe: vi.fn(() => () => {}) }, state, extensionRunner, bindExtensions: vi.fn(async () => {}), diff --git a/packages/coding-agent/test/suite/regressions/7290-json-stream-linear.test.ts b/packages/coding-agent/test/suite/regressions/7290-json-stream-linear.test.ts new file mode 100644 index 000000000..cea3b08a3 --- /dev/null +++ b/packages/coding-agent/test/suite/regressions/7290-json-stream-linear.test.ts @@ -0,0 +1,44 @@ +import { fauxAssistantMessage } from "@earendil-works/pi-ai"; +import { afterEach, describe, expect, it } from "vitest"; +import { toJsonEvent } from "../../../src/modes/json-event.ts"; +import { createHarness, type Harness } from "../harness.ts"; + +describe("regression #7290: JSON event streams stay linear", () => { + const harnesses: Harness[] = []; + + afterEach(() => { + while (harnesses.length > 0) { + harnesses.pop()?.cleanup(); + } + }); + + async function measureUpdateBytes(text: string): Promise { + const harness = await createHarness(); + harnesses.push(harness); + harness.setResponses([fauxAssistantMessage(text)]); + + await harness.session.prompt("respond"); + + const sessionUpdates = harness.eventsOfType("message_update"); + for (const update of sessionUpdates) { + expect(update).toHaveProperty("message"); + expect(update.assistantMessageEvent).toHaveProperty("partial"); + } + + const updates = sessionUpdates.map((event) => toJsonEvent(event)); + expect(updates.length).toBeGreaterThan(0); + for (const update of updates) { + expect(update).not.toHaveProperty("message"); + expect(update.assistantMessageEvent).not.toHaveProperty("partial"); + } + return updates.reduce((bytes, event) => bytes + Buffer.byteLength(JSON.stringify(event)), 0); + } + + it("emits delta-only message updates whose size scales linearly", async () => { + const smallBytes = await measureUpdateBytes("x".repeat(2_000)); + const largeBytes = await measureUpdateBytes("x".repeat(4_000)); + + expect(largeBytes).toBeGreaterThan(smallBytes); + expect(largeBytes / smallBytes).toBeLessThan(2.2); + }); +});