From 2a2aa9fe59525f4bd84a457e89eb7d7bea39d409 Mon Sep 17 00:00:00 2001 From: oldschoola Date: Tue, 23 Jun 2026 09:40:32 -0700 Subject: [PATCH] fix(proxy): preserve partialJson on content during streaming, scrub at terminal events Address review feedback: downstream renderers (event-controller.ts:535) read content.partialJson during toolcall_delta to pace streaming previews (bash env assignments, write/edit smooth streaming). Revised approach: - toolcall_start: initialize partialJson on content via typed ToolCall & { partialJson: string } intersection (not as any) - toolcall_delta: accumulate in side-channel Map, write onto content via typed intersection cast - toolcall_end: delete partialJson from content + side-channel map - done/error: scrubPartialJson() cleans any remaining blocks that never got toolcall_end (the original leak bug, now fixed for all terminal paths) Added test verifying partialJson IS present during streaming and IS absent after completion. --- packages/agent/CHANGELOG.md | 2 +- packages/agent/src/proxy.ts | 31 +++++++--- .../test/proxy-toolcall-partial-json.test.ts | 60 +++++++++++++++++++ 3 files changed, 85 insertions(+), 8 deletions(-) diff --git a/packages/agent/CHANGELOG.md b/packages/agent/CHANGELOG.md index df9d7f462..fffb47ac6 100644 --- a/packages/agent/CHANGELOG.md +++ b/packages/agent/CHANGELOG.md @@ -4,7 +4,7 @@ ### Fixed -- Fixed `streamProxy` leaking internal `partialJson` streaming state onto typed `ToolCall` objects via `as any` casts. Streaming JSON accumulation now uses a side-channel `Map` keyed by `contentIndex`, eliminating all `as any` casts and guaranteeing `partialJson` never appears on the final `AssistantMessage` content — even when the stream ends without a `toolcall_end` event. +- Fixed `streamProxy` leaking internal `partialJson` streaming state onto the final `AssistantMessage` when the stream ended without a `toolcall_end` event. The field is now accumulated in a side-channel `Map` (eliminating `as any` casts on the accumulation path), written onto the content object via a typed `ToolCall & { partialJson: string }` intersection so downstream renderers can still read it during streaming, and scrubbed from all content blocks at `toolcall_end`, `done`, and `error` — guaranteeing it never appears on the final message. ## [16.1.16] - 2026-06-23 diff --git a/packages/agent/src/proxy.ts b/packages/agent/src/proxy.ts index 68a81d800..4b0a2638a 100644 --- a/packages/agent/src/proxy.ts +++ b/packages/agent/src/proxy.ts @@ -201,14 +201,27 @@ export function streamProxy(model: Model, context: Context, options: ProxyStream return stream; } +/** + * Remove the `partialJson` streaming field from any tool-call content blocks + * that still carry it (e.g. when the stream ended without a `toolcall_end`). + */ +function scrubPartialJson(partial: AssistantMessage): void { + for (const block of partial.content) { + if (block?.type === "toolCall") { + delete (block as ToolCall & { partialJson?: string }).partialJson; + } + } +} + /** * Process a proxy event and update the partial message. * - * Streaming `partialJson` for in-progress tool calls is kept in a side-channel - * map keyed by `contentIndex` rather than stored on the `ToolCall` object - * itself. This avoids `as any` casts to smuggle non-spec fields through the - * typed `ToolCall` interface and guarantees the field never leaks into the - * final message if `toolcall_end` is skipped (e.g. on stream error). + * Streaming `partialJson` for in-progress tool calls is accumulated in a + * side-channel map keyed by `contentIndex` and also written onto the content + * object (as a typed intersection field) so downstream renderers can read it + * during streaming. The field is deleted at `toolcall_end` and scrubbed from + * any remaining blocks at `done`/`error` to guarantee it never leaks into the + * final `AssistantMessage`. */ function processProxyEvent( model: Model, @@ -302,16 +315,17 @@ function processProxyEvent( id: proxyEvent.id, name: proxyEvent.toolName, arguments: {}, - } satisfies ToolCall; + partialJson: "", + } as ToolCall & { partialJson: string }; partialJsonByIndex.set(proxyEvent.contentIndex, ""); return { type: "toolcall_start", contentIndex: proxyEvent.contentIndex, partial }; - case "toolcall_delta": { const content = partial.content[proxyEvent.contentIndex]; if (content?.type === "toolCall") { const acc = (partialJsonByIndex.get(proxyEvent.contentIndex) ?? "") + proxyEvent.delta; partialJsonByIndex.set(proxyEvent.contentIndex, acc); content.arguments = parseStreamingJson(acc) || {}; + (content as ToolCall & { partialJson: string }).partialJson = acc; partial.content[proxyEvent.contentIndex] = { ...content }; // Trigger reactivity return { type: "toolcall_delta", @@ -327,6 +341,7 @@ function processProxyEvent( const content = partial.content[proxyEvent.contentIndex]; if (content?.type === "toolCall") { partialJsonByIndex.delete(proxyEvent.contentIndex); + delete (content as ToolCall & { partialJson?: string }).partialJson; return { type: "toolcall_end", contentIndex: proxyEvent.contentIndex, @@ -341,6 +356,7 @@ function processProxyEvent( partial.stopReason = proxyEvent.reason; partial.usage = proxyEvent.usage; calculateCost(model, partial.usage); + scrubPartialJson(partial); return { type: "done", reason: proxyEvent.reason, message: partial }; case "error": @@ -348,6 +364,7 @@ function processProxyEvent( partial.errorMessage = proxyEvent.errorMessage; partial.usage = proxyEvent.usage; calculateCost(model, partial.usage); + scrubPartialJson(partial); return { type: "error", reason: proxyEvent.reason, error: partial }; } } diff --git a/packages/agent/test/proxy-toolcall-partial-json.test.ts b/packages/agent/test/proxy-toolcall-partial-json.test.ts index 6e491c329..316e4d5bf 100644 --- a/packages/agent/test/proxy-toolcall-partial-json.test.ts +++ b/packages/agent/test/proxy-toolcall-partial-json.test.ts @@ -105,6 +105,66 @@ describe("streamProxy — tool-call streaming and partialJson isolation", () => expect(toolCall.arguments).toEqual({ command: "ls" }); }); + it("exposes partialJson on content during streaming for renderers", async () => { + // Downstream renderers (event-controller.ts) read content.partialJson + // during toolcall_delta to pace streaming previews. The field must be + // present on the partial snapshot while streaming is in progress. + // Note: partial is a shared mutable reference, so we snapshot the + // partialJson value during iteration — by the time the stream completes, + // scrubPartialJson will have deleted it. + const events: ProxyAssistantMessageEvent[] = [ + { type: "start" }, + { type: "toolcall_start", contentIndex: 0, id: "call_1", toolName: "bash" }, + { type: "toolcall_delta", contentIndex: 0, delta: '{"comm' }, + { type: "toolcall_delta", contentIndex: 0, delta: 'and":"ls"}' }, + { type: "toolcall_end", contentIndex: 0 }, + { type: "done", reason: "toolUse", usage: { ...baseUsage } }, + ]; + const body = buildSseBody(events); + const fetchMock: FetchImpl = () => Promise.resolve(new Response(body, { status: 200 })); + + const stream = streamProxy(mockModel, mockContext, { + proxyUrl: "http://localhost:0", + authToken: "test", + fetch: fetchMock, + }); + + // Collect delta events and snapshot partialJson during iteration, + // before the done event scrubs it from the shared partial reference. + const deltaSnapshots: Array<{ hasPartialJson: boolean; value: string | undefined }> = []; + const iterator = stream[Symbol.asyncIterator](); + const deadline = Date.now() + 2000; + while (Date.now() < deadline) { + const { promise: timeoutPromise, resolve: timeoutResolve } = + Promise.withResolvers>(); + const timer = setTimeout( + () => timeoutResolve({ value: undefined, done: true } as IteratorResult), + 2000, + ); + const result = await Promise.race([iterator.next(), timeoutPromise]); + clearTimeout(timer); + if (result.done) break; + if (result.value.type === "toolcall_delta") { + const content = result.value.partial.content[0]; + deltaSnapshots.push({ + hasPartialJson: "partialJson" in (content ?? {}), + value: (content as (ToolCall & { partialJson?: string }) | undefined)?.partialJson, + }); + } + } + + expect(deltaSnapshots.length).toBe(2); + for (const snap of deltaSnapshots) { + expect(snap.hasPartialJson).toBe(true); + expect(snap.value).toBeTruthy(); + } + + // After completion, partialJson must be gone + const result = await stream.result(); + const toolCall = extractToolCall(result); + expect("partialJson" in toolCall).toBe(false); + }); + it("does not leak partialJson field into the final ToolCall object", async () => { const events: ProxyAssistantMessageEvent[] = [ { type: "start" },