diff --git a/packages/agent/src/proxy.ts b/packages/agent/src/proxy.ts index 26192e665fa..96b419f44fe 100644 --- a/packages/agent/src/proxy.ts +++ b/packages/agent/src/proxy.ts @@ -7,29 +7,17 @@ import { type AssistantMessage, type AssistantMessageEvent, + type AssistantMessageEventStream, type Context, - EventStream, + createAssistantMessageEventStream, + createPendingToolCall, type Model, - parseStreamingJson, + type PendingToolCall, type SimpleStreamOptions, type StopReason, type ToolCall, } from "@earendil-works/pi-ai"; -// Create stream class matching ProxyMessageEventStream -class ProxyMessageEventStream extends EventStream { - constructor() { - super( - (event) => event.type === "done" || event.type === "error", - (event) => { - if (event.type === "done") return event.message; - if (event.type === "error") return event.error; - throw new Error("Unexpected event type"); - }, - ); - } -} - /** * Proxy event types - server sends these with partial field stripped to reduce bandwidth. */ @@ -117,8 +105,17 @@ function buildProxyRequestOptions(options: ProxyStreamOptions): ProxySerializabl }; } -export function streamProxy(model: Model, context: Context, options: ProxyStreamOptions): ProxyMessageEventStream { - const stream = new ProxyMessageEventStream(); +export function streamProxy( + model: Model, + context: Context, + options: ProxyStreamOptions, +): AssistantMessageEventStream { + const stream = createAssistantMessageEventStream(); + const pendingCalls = new Map(); + const finishPendingCalls = () => { + for (const pending of pendingCalls.values()) pending.finish(); + pendingCalls.clear(); + }; (async () => { // Initialize the partial message that we'll build up from events @@ -190,9 +187,12 @@ export function streamProxy(model: Model, context: Context, options: ProxyS const data = line.slice(6).trim(); if (!data) return; const proxyEvent = JSON.parse(data) as ProxyAssistantMessageEvent; - const event = processProxyEvent(proxyEvent, partial); + const event = processProxyEvent(proxyEvent, partial, pendingCalls); if (event) { - if (event.type === "done" || event.type === "error") sawTerminalEvent = true; + if (event.type === "done" || event.type === "error") { + finishPendingCalls(); + sawTerminalEvent = true; + } stream.push(event); } }; @@ -231,6 +231,7 @@ export function streamProxy(model: Model, context: Context, options: ProxyS // consumers waiting on a result that never arrives. partial.stopReason = "error"; partial.errorMessage = "Connection closed by proxy server before the response completed"; + finishPendingCalls(); stream.push({ type: "error", reason: "error", @@ -244,6 +245,7 @@ export function streamProxy(model: Model, context: Context, options: ProxyS const reason = options.signal?.aborted ? "aborted" : "error"; partial.stopReason = reason; partial.errorMessage = errorMessage; + finishPendingCalls(); stream.push({ type: "error", reason, @@ -266,6 +268,7 @@ export function streamProxy(model: Model, context: Context, options: ProxyS function processProxyEvent( proxyEvent: ProxyAssistantMessageEvent, partial: AssistantMessage, + pendingCalls: Map, ): AssistantMessageEvent | undefined { switch (proxyEvent.type) { case "start": @@ -335,22 +338,33 @@ function processProxyEvent( throw new Error("Received thinking_end for non-thinking content"); } - case "toolcall_start": - partial.content[proxyEvent.contentIndex] = { - type: "toolCall", - id: proxyEvent.id, - name: proxyEvent.toolName, - arguments: {}, - partialJson: "", - } satisfies ToolCall & { partialJson: string } as ToolCall; + case "toolcall_start": { + const pending = createPendingToolCall( + { + type: "toolCall", + id: proxyEvent.id, + name: proxyEvent.toolName, + arguments: {}, + partialJson: "", + }, + true, + ); + partial.content[proxyEvent.contentIndex] = pending.toolCall; + pendingCalls.set(proxyEvent.contentIndex, pending); return { type: "toolcall_start", contentIndex: proxyEvent.contentIndex, partial }; + } case "toolcall_delta": { const content = partial.content[proxyEvent.contentIndex]; if (content?.type === "toolCall") { - (content as any).partialJson += proxyEvent.delta; - content.arguments = parseStreamingJson((content as any).partialJson) || {}; - partial.content[proxyEvent.contentIndex] = { ...content }; // Trigger reactivity + const block = content as ToolCall & { partialJson: string }; + block.partialJson += proxyEvent.delta; + const pending = pendingCalls.get(proxyEvent.contentIndex)!; + pending.setJson(block.partialJson); + // Trigger reactivity without reading arguments or losing extra metadata. + const copy = pending.copy(); + pendingCalls.set(proxyEvent.contentIndex, copy); + partial.content[proxyEvent.contentIndex] = copy.toolCall; return { type: "toolcall_delta", contentIndex: proxyEvent.contentIndex, @@ -365,6 +379,8 @@ function processProxyEvent( const content = partial.content[proxyEvent.contentIndex]; if (content?.type === "toolCall") { Object.assign(content, proxyEvent.toolCall); + pendingCalls.get(proxyEvent.contentIndex)?.finish(); + pendingCalls.delete(proxyEvent.contentIndex); delete (content as any).partialJson; return { type: "toolcall_end", diff --git a/packages/agent/test/proxy.test.ts b/packages/agent/test/proxy.test.ts index f2ef9825804..56c4491bd10 100644 --- a/packages/agent/test/proxy.test.ts +++ b/packages/agent/test/proxy.test.ts @@ -1,4 +1,4 @@ -import type { AssistantMessage, AssistantMessageEvent, Model } from "@earendil-works/pi-ai"; +import type { AssistantMessage, AssistantMessageEvent, Model, ToolCall } from "@earendil-works/pi-ai"; import { afterEach, describe, expect, it, vi } from "vitest"; import { type ProxyAssistantMessageEvent, streamProxy } from "../src/proxy.ts"; @@ -26,9 +26,150 @@ const usage: AssistantMessage["usage"] = { afterEach(() => { vi.unstubAllGlobals(); + vi.restoreAllMocks(); }); describe("streamProxy", () => { + // #9265: copying blocks for reactivity must not evaluate lazy arguments. + it.each([true, false])("defers argument parsing and settles unread arguments (complete: %s)", async (complete) => { + const content = "a".repeat(32 * 512); + const deltas = ['{"content":"', ...Array(32).fill("a".repeat(512))]; + if (complete) deltas.push('"}'); + const proxyEvents: ProxyAssistantMessageEvent[] = [ + { type: "start" }, + { type: "toolcall_start", contentIndex: 0, id: "call_test", toolName: "write" }, + ...deltas.map((delta): ProxyAssistantMessageEvent => ({ type: "toolcall_delta", contentIndex: 0, delta })), + ...(complete + ? ([ + { + type: "toolcall_end", + contentIndex: 0, + toolCall: { type: "toolCall", id: "call_test", name: "write", arguments: { content } }, + }, + { type: "done", reason: "toolUse", usage }, + ] satisfies ProxyAssistantMessageEvent[]) + : []), + ]; + const body = proxyEvents.map((event) => `data: ${JSON.stringify(event)}\n\n`).join(""); + vi.stubGlobal( + "fetch", + vi.fn(async () => new Response(body)), + ); + const parse = vi.spyOn(JSON, "parse"); + const result = await streamProxy( + model, + { messages: [] }, + { + authToken: "test", + proxyUrl: "https://proxy.example.com", + }, + ).result(); + expect(result.stopReason, result.errorMessage).toBe(complete ? "toolUse" : "error"); + if (complete) { + expect(parse.mock.calls.filter(([text]) => text.startsWith('{"content":'))).toEqual([]); + } + expect(Object.getOwnPropertyDescriptor(result.content[0], "arguments")).toMatchObject({ + value: { content }, + writable: true, + }); + }); + + // #9265: lazy parsing must preserve the proxy's existing `parsed || {}` behavior. + it.each(["null", "false", "0", '""'])("preserves the empty-object fallback for %s arguments", async (json) => { + const proxyEvents: ProxyAssistantMessageEvent[] = [ + { type: "start" }, + { type: "toolcall_start", contentIndex: 0, id: "call_test", toolName: "write" }, + { type: "toolcall_delta", contentIndex: 0, delta: json }, + { type: "error", reason: "error", usage, errorMessage: "Interrupted" }, + ]; + const body = proxyEvents.map((event) => `data: ${JSON.stringify(event)}\n\n`).join(""); + vi.stubGlobal( + "fetch", + vi.fn(async () => new Response(body)), + ); + const result = await streamProxy( + model, + { messages: [] }, + { + authToken: "test", + proxyUrl: "https://proxy.example.com", + }, + ).result(); + expect(result.stopReason).toBe("error"); + const block = result.content[0]; + if (block.type !== "toolCall") throw new Error("Expected tool call"); + expect(block.arguments).toEqual({}); + }); + + // #9265: preserve main's shallow-copy identity and extension metadata during live updates. + it("copies live blocks without losing metadata or changing argument identity", async () => { + let transport!: ReadableStreamDefaultController; + const body = new ReadableStream({ + start: (controller) => { + transport = controller; + }, + }); + const send = (...events: ProxyAssistantMessageEvent[]) => { + transport.enqueue( + new TextEncoder().encode(events.map((event) => `data: ${JSON.stringify(event)}\n\n`).join("")), + ); + }; + send( + { type: "start" }, + { type: "toolcall_start", contentIndex: 0, id: "call_test", toolName: "write" }, + { type: "toolcall_delta", contentIndex: 0, delta: '{"content":"hel' }, + ); + vi.stubGlobal( + "fetch", + vi.fn(async () => new Response(body)), + ); + const stream = streamProxy(model, { messages: [] }, { authToken: "test", proxyUrl: "https://proxy.example.com" }); + const symbol = Symbol("metadata"); + const metadata = { source: "extension" }; + let first: ToolCall | undefined; + let firstArguments: ToolCall["arguments"] | undefined; + let last: ToolCall | undefined; + for await (const event of stream) { + if (event.type !== "toolcall_delta") continue; + const block = event.partial.content[0]; + if (block.type !== "toolCall") throw new Error("Expected tool call"); + if (!first) { + first = block; + firstArguments = block.arguments; + Object.assign(block, { extra: metadata, [symbol]: metadata }); + send({ type: "toolcall_delta", contentIndex: 0, delta: 'lo"}' }); + } else { + last = block; + expect(block).not.toBe(first); + expect(block.arguments).toEqual({ content: "hello" }); + expect(block.arguments).toBe(first.arguments); + expect(firstArguments).toEqual({ content: "hel" }); + expect(Reflect.get(block, "extra")).toBe(metadata); + expect(Reflect.get(block, symbol)).toBe(metadata); + send( + { + type: "toolcall_end", + contentIndex: 0, + toolCall: { + type: "toolCall", + id: "call_test", + name: "write", + arguments: { content: "authoritative" }, + }, + }, + { type: "done", reason: "toolUse", usage }, + ); + transport.close(); + } + } + const result = await stream.result(); + expect(result.content[0]).toBe(last); + expect(Object.getOwnPropertyDescriptor(result.content[0], "arguments")).toMatchObject({ + value: { content: "authoritative" }, + writable: true, + }); + }); + it("preserves tool-call metadata received only on toolcall_end", async () => { const proxyEvents: ProxyAssistantMessageEvent[] = [ { type: "start" }, diff --git a/packages/ai/README.md b/packages/ai/README.md index 97bf64471b5..1bf19c9b2cb 100644 --- a/packages/ai/README.md +++ b/packages/ai/README.md @@ -574,7 +574,9 @@ context.messages.push({ ### Streaming Tool Calls with Partial JSON -During streaming, tool call arguments are progressively parsed as they arrive. This enables real-time UI updates before the complete arguments are available: +During streaming, tool call arguments are parsed on demand when `arguments` is read, and cached until the next argument delta. This enables real-time UI updates before the complete arguments are available without reparsing growing JSON for delta-only consumers. Read partial arguments only when needed; reading them after every delta still reparses each growing prefix. Terminal messages materialize any unread arguments, including on errors and aborts. + +For example: ```typescript const s = models.stream(model, context); diff --git a/packages/ai/src/api/anthropic-messages.ts b/packages/ai/src/api/anthropic-messages.ts index de3a51cf4ff..9c93131ed0f 100644 --- a/packages/ai/src/api/anthropic-messages.ts +++ b/packages/ai/src/api/anthropic-messages.ts @@ -37,6 +37,7 @@ import { appendAssistantMessageDiagnostic } from "../utils/diagnostics.ts"; import { AssistantMessageEventStream } from "../utils/event-stream.ts"; import { headersToRecord } from "../utils/headers.ts"; import { parseJsonWithRepair, parseStreamingJson } from "../utils/json-parse.ts"; +import { createPendingToolCall, type PendingToolCall } from "../utils/pending-tool-call.ts"; import { getPiUserAgent } from "../utils/pi-user-agent.ts"; import { getProviderEnvValue } from "../utils/provider-env.ts"; import { retryProviderRequest } from "../utils/provider-retry.ts"; @@ -530,6 +531,7 @@ export const stream: StreamFunction<"anthropic-messages", AnthropicOptions> = ( timestamp: Date.now(), }; + const pendingCalls = new Map(); try { let client: Anthropic; let isOAuth: boolean; @@ -649,7 +651,7 @@ export const stream: StreamFunction<"anthropic-messages", AnthropicOptions> = ( output.content.push(block); stream.push({ type: "thinking_start", contentIndex: output.content.length - 1, partial: output }); } else if (event.content_block.type === "tool_use") { - const block: Block = { + const pending = createPendingToolCall({ type: "toolCall", id: event.content_block.id, name: isOAuth @@ -658,7 +660,9 @@ export const stream: StreamFunction<"anthropic-messages", AnthropicOptions> = ( arguments: (event.content_block.input as Record) ?? {}, partialJson: "", index: event.index, - }; + }); + const block: Block = pending.toolCall; + pendingCalls.set(block, pending); output.content.push(block); stream.push({ type: "toolcall_start", contentIndex: output.content.length - 1, partial: output }); } @@ -692,7 +696,7 @@ export const stream: StreamFunction<"anthropic-messages", AnthropicOptions> = ( const block = blocks[index]; if (block && block.type === "toolCall") { block.partialJson += event.delta.partial_json; - block.arguments = parseStreamingJson(block.partialJson); + pendingCalls.get(block)!.setJson(block.partialJson); stream.push({ type: "toolcall_delta", contentIndex: index, @@ -729,6 +733,8 @@ export const stream: StreamFunction<"anthropic-messages", AnthropicOptions> = ( }); } else if (block.type === "toolCall") { block.arguments = parseStreamingJson(block.partialJson); + pendingCalls.get(block)!.finish(); + pendingCalls.delete(block); // Finalize in-place and strip the scratch buffer so replay only // carries parsed arguments. delete (block as { partialJson?: string }).partialJson; @@ -803,9 +809,13 @@ export const stream: StreamFunction<"anthropic-messages", AnthropicOptions> = ( }); } + for (const pending of pendingCalls.values()) pending.finish(); + pendingCalls.clear(); stream.push({ type: "done", reason: output.stopReason, message: output }); stream.end(); } catch (error) { + for (const pending of pendingCalls.values()) pending.finish(); + pendingCalls.clear(); for (const block of output.content) { delete (block as { index?: number }).index; // partialJson is only a streaming scratch buffer; never persist it. diff --git a/packages/ai/src/api/bedrock-converse-stream.ts b/packages/ai/src/api/bedrock-converse-stream.ts index 2e944fcd246..3868b2fdc04 100644 --- a/packages/ai/src/api/bedrock-converse-stream.ts +++ b/packages/ai/src/api/bedrock-converse-stream.ts @@ -54,6 +54,7 @@ import { AssistantMessageEventStream } from "../utils/event-stream.ts"; import { providerHeadersToRecord } from "../utils/headers.ts"; import { parseStreamingJson } from "../utils/json-parse.ts"; import { resolveHttpProxyUrlForTarget } from "../utils/node-http-proxy.ts"; +import { createPendingToolCall, type PendingToolCall } from "../utils/pending-tool-call.ts"; import { getProviderEnvValue } from "../utils/provider-env.ts"; import { sanitizeSurrogates } from "../utils/sanitize-unicode.ts"; import { getJsonSchemaToolParameters, resolveJsonSchemaStrictSampling } from "./constrained-sampling.ts"; @@ -140,6 +141,7 @@ export const stream: StreamFunction<"bedrock-converse-stream", BedrockOptions> = }; const blocks = output.content as Block[]; + const pendingCalls = new Map(); // A profile explicitly configured through pi's auth flow (the `profile` // option or scoped `AWS_PROFILE` on the stored credential's env) must win @@ -283,11 +285,11 @@ export const stream: StreamFunction<"bedrock-converse-stream", BedrockOptions> = } stream.push({ type: "start", partial: output }); } else if (item.contentBlockStart) { - handleContentBlockStart(item.contentBlockStart, blocks, output, stream); + handleContentBlockStart(item.contentBlockStart, blocks, output, stream, pendingCalls); } else if (item.contentBlockDelta) { - handleContentBlockDelta(item.contentBlockDelta, blocks, output, stream); + handleContentBlockDelta(item.contentBlockDelta, blocks, output, stream, pendingCalls); } else if (item.contentBlockStop) { - handleContentBlockStop(item.contentBlockStop, blocks, output, stream); + handleContentBlockStop(item.contentBlockStop, blocks, output, stream, pendingCalls); } else if (item.messageStop) { output.rawStopReason = item.messageStop.stopReason; const { stopReason, errorMessage } = mapStopReason(item.messageStop.stopReason); @@ -322,12 +324,12 @@ export const stream: StreamFunction<"bedrock-converse-stream", BedrockOptions> = } // A stream can settle without stopping every block, so finalize here too. - for (const block of output.content) finalizeStreamingBlock(block as Block); + for (const block of output.content) finalizeStreamingBlock(block as Block, pendingCalls); stream.push({ type: "done", reason: output.stopReason, message: output }); stream.end(); } catch (error) { for (const block of output.content) { - finalizeStreamingBlock(block as Block); + finalizeStreamingBlock(block as Block, pendingCalls); } output.stopReason = options.signal?.aborted ? "aborted" : "error"; output.errorMessage = formatBedrockError(error); @@ -564,19 +566,22 @@ function handleContentBlockStart( blocks: Block[], output: AssistantMessage, stream: AssistantMessageEventStream, + pendingCalls: Map, ): void { const index = event.contentBlockIndex!; const start = event.start; if (start?.toolUse) { - const block: Block = { + const pending = createPendingToolCall({ type: "toolCall", id: start.toolUse.toolUseId || "", name: start.toolUse.name || "", arguments: {}, partialJson: "", index, - }; + }); + const block: Block = pending.toolCall; + pendingCalls.set(block, pending); output.content.push(block); stream.push({ type: "toolcall_start", contentIndex: blocks.length - 1, partial: output }); } @@ -587,6 +592,7 @@ function handleContentBlockDelta( blocks: Block[], output: AssistantMessage, stream: AssistantMessageEventStream, + pendingCalls: Map, ): void { const contentBlockIndex = event.contentBlockIndex!; const delta = event.delta; @@ -608,7 +614,7 @@ function handleContentBlockDelta( } } else if (delta?.toolUse && block?.type === "toolCall") { block.partialJson = (block.partialJson || "") + (delta.toolUse.input || ""); - block.arguments = parseStreamingJson(block.partialJson); + pendingCalls.get(block)!.setJson(block.partialJson); stream.push({ type: "toolcall_delta", contentIndex: index, delta: delta.toolUse.input || "", partial: output }); } else if (delta?.reasoningContent) { let thinkingBlock = block; @@ -675,7 +681,11 @@ function flushRedactedContent(block: Block): void { * Strips every streaming scratch field. Runs from the terminal paths as well as * `contentBlockStop`, because a stream can settle without stopping each block. */ -function finalizeStreamingBlock(block: Block): void { +function finalizeStreamingBlock(block: Block, pendingCalls: Map): void { + if (block.type === "toolCall") { + pendingCalls.get(block)?.finish(); + pendingCalls.delete(block); + } delete block.index; // partialJson is only a streaming scratch buffer; never persist it. delete block.partialJson; @@ -702,6 +712,7 @@ function handleContentBlockStop( blocks: Block[], output: AssistantMessage, stream: AssistantMessageEventStream, + pendingCalls: Map, ): void { const index = blocks.findIndex((b) => b.index === event.contentBlockIndex); const block = blocks[index]; @@ -718,6 +729,8 @@ function handleContentBlockStop( break; case "toolCall": block.arguments = parseStreamingJson(block.partialJson); + pendingCalls.get(block)!.finish(); + pendingCalls.delete(block); // Finalize in-place and strip the scratch buffer so replay only // carries parsed arguments. delete (block as Block).partialJson; diff --git a/packages/ai/src/api/mistral-conversations.ts b/packages/ai/src/api/mistral-conversations.ts index 979b467ed35..476789fb725 100644 --- a/packages/ai/src/api/mistral-conversations.ts +++ b/packages/ai/src/api/mistral-conversations.ts @@ -17,6 +17,7 @@ import { AssistantMessageEventStream } from "../utils/event-stream.ts"; import { shortHash } from "../utils/hash.ts"; import { headersToRecord } from "../utils/headers.ts"; import { parseStreamingJson } from "../utils/json-parse.ts"; +import { createPendingToolCall, type PendingToolCall } from "../utils/pending-tool-call.ts"; import { getPiUserAgent } from "../utils/pi-user-agent.ts"; import { sanitizeSurrogates } from "../utils/sanitize-unicode.ts"; import { getJsonSchemaToolParameters, resolveJsonSchemaStrictSampling } from "./constrained-sampling.ts"; @@ -128,6 +129,7 @@ export const stream: StreamFunction<"mistral-conversations", MistralOptions> = ( (async () => { const output = createOutput(model); + const pendingCalls = new Map(); try { const apiKey = options?.apiKey; @@ -145,7 +147,7 @@ export const stream: StreamFunction<"mistral-conversations", MistralOptions> = ( } const mistralStream = await requestMistralStream(model, payload, apiKey, options); stream.push({ type: "start", partial: output }); - await consumeChatStream(model, output, stream, mistralStream); + await consumeChatStream(model, output, stream, mistralStream, pendingCalls); if (options?.signal?.aborted) { throw new Error("Request was aborted"); @@ -161,6 +163,8 @@ export const stream: StreamFunction<"mistral-conversations", MistralOptions> = ( stream.push({ type: "done", reason: output.stopReason, message: output }); stream.end(); } catch (error) { + for (const pending of pendingCalls.values()) pending.finish(); + pendingCalls.clear(); for (const block of output.content) { // partialArgs is only a streaming scratch buffer; never persist it. delete (block as { partialArgs?: string }).partialArgs; @@ -560,6 +564,7 @@ async function consumeChatStream( output: AssistantMessage, stream: AssistantMessageEventStream, mistralStream: AsyncIterable, + pendingCalls: Map, ): Promise { let currentBlock: TextContent | ThinkingContent | null = null; const blocks = output.content; @@ -705,13 +710,15 @@ async function consumeChatStream( } if (!block) { - block = { + const pending = createPendingToolCall({ type: "toolCall", id: callId, name: toolCall.function.name, arguments: {}, partialArgs: "", - }; + }); + block = pending.toolCall; + pendingCalls.set(block, pending); output.content.push(block); toolBlocksByKey.set(key, output.content.length - 1); stream.push({ type: "toolcall_start", contentIndex: output.content.length - 1, partial: output }); @@ -722,7 +729,7 @@ async function consumeChatStream( ? toolCall.function.arguments : JSON.stringify(toolCall.function.arguments || {}); block.partialArgs = (block.partialArgs || "") + argsDelta; - block.arguments = parseStreamingJson>(block.partialArgs); + pendingCalls.get(block)!.setJson(block.partialArgs); stream.push({ type: "toolcall_delta", contentIndex: toolBlocksByKey.get(key)!, @@ -738,6 +745,8 @@ async function consumeChatStream( if (block.type !== "toolCall") continue; const toolBlock = block as ToolCall & { partialArgs?: string }; toolBlock.arguments = parseStreamingJson>(toolBlock.partialArgs); + pendingCalls.get(block)!.finish(); + pendingCalls.delete(block); // Finalize in-place and strip the scratch buffer so replay only // carries parsed arguments. delete toolBlock.partialArgs; diff --git a/packages/ai/src/api/openai-completions.ts b/packages/ai/src/api/openai-completions.ts index 8ac17dbae46..dc1029ee8a4 100644 --- a/packages/ai/src/api/openai-completions.ts +++ b/packages/ai/src/api/openai-completions.ts @@ -41,6 +41,7 @@ import { AssistantMessageEventStream } from "../utils/event-stream.ts"; import { shortHash } from "../utils/hash.ts"; import { headersToRecord } from "../utils/headers.ts"; import { parseStreamingJson } from "../utils/json-parse.ts"; +import { createPendingToolCall, type PendingToolCall } from "../utils/pending-tool-call.ts"; import { getPiUserAgent } from "../utils/pi-user-agent.ts"; import { getProviderEnvValue } from "../utils/provider-env.ts"; import { retryProviderRequest } from "../utils/provider-retry.ts"; @@ -334,6 +335,8 @@ export const stream: StreamFunction<"openai-completions", OpenAICompletionsOptio timestamp: Date.now(), }; + const pendingCalls = new Map(); + // `reasoning_details` are replay metadata, not user-visible stream deltas. // Keep them in memory during streaming and serialize once when the block is finalized. let streamedReasoningDetails: OpenAIReasoningDetail[] | undefined; @@ -454,6 +457,8 @@ export const stream: StreamFunction<"openai-completions", OpenAICompletionsOptio } else { block.arguments = parseStreamingJson(block.partialArgs); } + pendingCalls.get(block)!.finish(); + pendingCalls.delete(block); // Finalize in-place and strip the scratch buffers so replay only // carries parsed arguments. delete block.partialArgs; @@ -500,7 +505,7 @@ export const stream: StreamFunction<"openai-completions", OpenAICompletionsOptio const customInputProperty = toolCall.custom && !toolCall.function ? (grammarToolInputProperties.get(name) ?? "input") : undefined; const hasCustomInput = customInputProperty !== undefined; - block = { + const pending = createPendingToolCall({ type: "toolCall", id: toolCall.id || "", name, @@ -510,7 +515,9 @@ export const stream: StreamFunction<"openai-completions", OpenAICompletionsOptio ? { property: customInputProperty, jsonBuffer: { input: "", started: false, closed: false } } : undefined, streamIndex, - }; + }); + block = pending.toolCall; + pendingCalls.set(block, pending); if (streamIndex !== undefined) { toolCallBlocksByIndex.set(streamIndex, block); } @@ -643,7 +650,7 @@ export const stream: StreamFunction<"openai-completions", OpenAICompletionsOptio if (toolCall.function?.arguments) { delta = toolCall.function.arguments; block.partialArgs = (block.partialArgs ?? "") + toolCall.function.arguments; - block.arguments = parseStreamingJson(block.partialArgs); + pendingCalls.get(block)!.setJson(block.partialArgs); } else if (toolCall.custom?.input) { const nextInput = getCustomToolCallInput(block) + toolCall.custom.input; delta = appendCustomToolCallInput(block, nextInput, false) ?? ""; @@ -695,6 +702,8 @@ export const stream: StreamFunction<"openai-completions", OpenAICompletionsOptio stream.push({ type: "done", reason: output.stopReason, message: output }); stream.end(); } catch (error) { + for (const pending of pendingCalls.values()) pending.finish(); + pendingCalls.clear(); for (const block of output.content) { if (block.type === "thinking") { applyStreamedReasoningDetails(block); diff --git a/packages/ai/src/api/openai-responses-shared.ts b/packages/ai/src/api/openai-responses-shared.ts index 7ac1add0817..3b8584249b0 100644 --- a/packages/ai/src/api/openai-responses-shared.ts +++ b/packages/ai/src/api/openai-responses-shared.ts @@ -31,6 +31,7 @@ import type { import type { AssistantMessageEventStream } from "../utils/event-stream.ts"; import { shortHash } from "../utils/hash.ts"; import { parseStreamingJson } from "../utils/json-parse.ts"; +import { createPendingToolCall, type PendingToolCall } from "../utils/pending-tool-call.ts"; import { sanitizeSurrogates } from "../utils/sanitize-unicode.ts"; import { appendGrammarToolInputJsonDelta, @@ -425,7 +426,7 @@ function appendCustomToolCallInput(block: StreamingToolCall, nextInput: string, type ResponsesOutputSlot = | { type: "thinking"; block: ThinkingContent; contentIndex: number } | { type: "text"; block: TextContent; contentIndex: number } - | { type: "toolCall"; block: StreamingToolCall; contentIndex: number }; + | { type: "toolCall"; block: StreamingToolCall; pending: PendingToolCall; contentIndex: number }; type ToolCallOutputSlot = Extract; @@ -483,18 +484,20 @@ export async function processResponsesStream( return slot; } if (item.type === "function_call") { - const block: StreamingToolCall = { + const pending = createPendingToolCall({ type: "toolCall", id: `${item.call_id}|${item.id}`, name: item.name, arguments: {}, ...(item.namespace !== undefined ? { namespace: item.namespace } : {}), partialJson: item.arguments || "", - }; + }); + const block = pending.toolCall; output.content.push(block); const slot = { type: "toolCall", block, + pending, contentIndex: output.content.length - 1, } satisfies ResponsesOutputSlot; outputSlots.set(outputIndex, slot); @@ -504,7 +507,7 @@ export async function processResponsesStream( if (item.type === "custom_tool_call") { const inputProperty = options?.grammarToolInputProperties?.get(item.name) ?? "input"; const input = item.input || ""; - const block: StreamingToolCall = { + const pending = createPendingToolCall({ type: "toolCall", id: `${item.call_id}|${item.id}`, name: item.name, @@ -514,11 +517,13 @@ export async function processResponsesStream( property: inputProperty, jsonBuffer: { input: "", started: false, closed: false }, }, - }; + }); + const block = pending.toolCall; output.content.push(block); const slot = { type: "toolCall", block, + pending, contentIndex: output.content.length - 1, } satisfies ResponsesOutputSlot; outputSlots.set(outputIndex, slot); @@ -595,168 +600,178 @@ export async function processResponsesStream( } }; - for await (const event of openaiStream) { - if (event.type === "response.created") { - output.responseId = event.response.id; - } else if (event.type === "response.output_item.added") { - createSlot(event.output_index, event.item); - } else if (event.type === "response.reasoning_summary_text.delta") { - const slot = getSlot(event.output_index, "thinking"); - if (!slot) continue; - slot.block.thinking += event.delta; - stream.push({ - type: "thinking_delta", - contentIndex: slot.contentIndex, - delta: event.delta, - partial: output, - }); - } else if (event.type === "response.reasoning_summary_part.done") { - const slot = getSlot(event.output_index, "thinking"); - if (!slot) continue; - slot.block.thinking += "\n\n"; - stream.push({ - type: "thinking_delta", - contentIndex: slot.contentIndex, - delta: "\n\n", - partial: output, - }); - } else if (event.type === "response.reasoning_text.delta") { - const slot = getSlot(event.output_index, "thinking"); - if (!slot) continue; - slot.block.thinking += event.delta; - stream.push({ - type: "thinking_delta", - contentIndex: slot.contentIndex, - delta: event.delta, - partial: output, - }); - } else if (event.type === "response.output_text.delta") { - const slot = getSlot(event.output_index, "text"); - if (!slot) continue; - slot.block.text += event.delta; - stream.push({ - type: "text_delta", - contentIndex: slot.contentIndex, - delta: event.delta, - partial: output, - }); - } else if (event.type === "response.refusal.delta") { - const slot = getSlot(event.output_index, "text"); - if (!slot) continue; - slot.block.text += event.delta; - stream.push({ - type: "text_delta", - contentIndex: slot.contentIndex, - delta: event.delta, - partial: output, - }); - } else if (event.type === "response.function_call_arguments.delta") { - const slot = getSlot(event.output_index, "toolCall"); - if (!slot || slot.block.partialJson === undefined) continue; - slot.block.partialJson += event.delta; - slot.block.arguments = parseStreamingJson(slot.block.partialJson); - pushToolCallDelta(slot, event.delta); - } else if (event.type === "response.function_call_arguments.done") { - const slot = getSlot(event.output_index, "toolCall"); - if (!slot || slot.block.partialJson === undefined) continue; - const previousPartialJson = slot.block.partialJson; - slot.block.partialJson = event.arguments; - slot.block.arguments = parseStreamingJson(slot.block.partialJson); - - if (event.arguments.startsWith(previousPartialJson)) { - const delta = event.arguments.slice(previousPartialJson.length); - if (delta.length > 0) pushToolCallDelta(slot, delta); - } - } else if (event.type === "response.custom_tool_call_input.delta") { - const slot = getSlot(event.output_index, "toolCall"); - if (!slot || !slot.block.customInput) continue; - pushToolCallDelta( - slot, - appendCustomToolCallInput(slot.block, getCustomToolCallInput(slot.block) + event.delta, false), - ); - } else if (event.type === "response.custom_tool_call_input.done") { - const slot = getSlot(event.output_index, "toolCall"); - if (!slot || !slot.block.customInput) continue; - pushToolCallDelta(slot, appendCustomToolCallInput(slot.block, event.input, true)); - } else if (event.type === "response.output_item.done") { - const item = event.item; - applyMessagePhaseStopReason(item); - const slot = getOrCreateSlot(event.output_index, item); - - if (item.type === "reasoning" && slot?.type === "thinking") { - const summaryText = item.summary?.map((s) => s.text).join("\n\n") || ""; - const contentText = item.content?.map((c) => c.text).join("\n\n") || ""; - slot.block.thinking = summaryText || contentText || slot.block.thinking; - slot.block.thinkingSignature = JSON.stringify(item); - reasoningBlocksById.set(item.id, slot.block); + try { + for await (const event of openaiStream) { + if (event.type === "response.created") { + output.responseId = event.response.id; + } else if (event.type === "response.output_item.added") { + createSlot(event.output_index, event.item); + } else if (event.type === "response.reasoning_summary_text.delta") { + const slot = getSlot(event.output_index, "thinking"); + if (!slot) continue; + slot.block.thinking += event.delta; stream.push({ - type: "thinking_end", + type: "thinking_delta", contentIndex: slot.contentIndex, - content: slot.block.thinking, + delta: event.delta, partial: output, }); - outputSlots.delete(event.output_index); - } else if (item.type === "message" && slot?.type === "text") { - slot.block.text = item.content?.map((c) => (c.type === "output_text" ? c.text : c.refusal)).join("") || ""; - slot.block.textSignature = encodeTextSignatureV1(item.id, item.phase ?? undefined); + } else if (event.type === "response.reasoning_summary_part.done") { + const slot = getSlot(event.output_index, "thinking"); + if (!slot) continue; + slot.block.thinking += "\n\n"; stream.push({ - type: "text_end", + type: "thinking_delta", contentIndex: slot.contentIndex, - content: slot.block.text, + delta: "\n\n", partial: output, }); - outputSlots.delete(event.output_index); - } else if ( - item.type === "function_call" && - slot?.type === "toolCall" && - slot.block.partialJson !== undefined - ) { - slot.block.arguments = parseStreamingJson(item.arguments || slot.block.partialJson || "{}"); - if (item.namespace !== undefined) slot.block.namespace = item.namespace; - // Finalize in-place and strip the scratch buffer so replay only - // carries parsed arguments. - delete slot.block.partialJson; + } else if (event.type === "response.reasoning_text.delta") { + const slot = getSlot(event.output_index, "thinking"); + if (!slot) continue; + slot.block.thinking += event.delta; stream.push({ - type: "toolcall_end", + type: "thinking_delta", contentIndex: slot.contentIndex, - toolCall: slot.block, + delta: event.delta, partial: output, }); - outputSlots.delete(event.output_index); - } else if (item.type === "custom_tool_call" && slot?.type === "toolCall" && slot.block.customInput) { - pushToolCallDelta( - slot, - appendCustomToolCallInput(slot.block, item.input ?? getCustomToolCallInput(slot.block), true), - ); - if (item.namespace !== undefined) slot.block.namespace = item.namespace; - delete slot.block.customInput; + } else if (event.type === "response.output_text.delta") { + const slot = getSlot(event.output_index, "text"); + if (!slot) continue; + slot.block.text += event.delta; stream.push({ - type: "toolcall_end", + type: "text_delta", contentIndex: slot.contentIndex, - toolCall: slot.block, + delta: event.delta, partial: output, }); - outputSlots.delete(event.output_index); + } else if (event.type === "response.refusal.delta") { + const slot = getSlot(event.output_index, "text"); + if (!slot) continue; + slot.block.text += event.delta; + stream.push({ + type: "text_delta", + contentIndex: slot.contentIndex, + delta: event.delta, + partial: output, + }); + } else if (event.type === "response.function_call_arguments.delta") { + const slot = getSlot(event.output_index, "toolCall"); + if (!slot || slot.block.partialJson === undefined) continue; + slot.block.partialJson += event.delta; + slot.pending.setJson(slot.block.partialJson); + pushToolCallDelta(slot, event.delta); + } else if (event.type === "response.function_call_arguments.done") { + const slot = getSlot(event.output_index, "toolCall"); + if (!slot || slot.block.partialJson === undefined) continue; + const previousPartialJson = slot.block.partialJson; + slot.block.partialJson = event.arguments; + slot.block.arguments = parseStreamingJson(slot.block.partialJson); + + if (event.arguments.startsWith(previousPartialJson)) { + const delta = event.arguments.slice(previousPartialJson.length); + if (delta.length > 0) pushToolCallDelta(slot, delta); + } + } else if (event.type === "response.custom_tool_call_input.delta") { + const slot = getSlot(event.output_index, "toolCall"); + if (!slot || !slot.block.customInput) continue; + pushToolCallDelta( + slot, + appendCustomToolCallInput(slot.block, getCustomToolCallInput(slot.block) + event.delta, false), + ); + } else if (event.type === "response.custom_tool_call_input.done") { + const slot = getSlot(event.output_index, "toolCall"); + if (!slot || !slot.block.customInput) continue; + pushToolCallDelta(slot, appendCustomToolCallInput(slot.block, event.input, true)); + } else if (event.type === "response.output_item.done") { + const item = event.item; + applyMessagePhaseStopReason(item); + const slot = getOrCreateSlot(event.output_index, item); + + if (item.type === "reasoning" && slot?.type === "thinking") { + const summaryText = item.summary?.map((s) => s.text).join("\n\n") || ""; + const contentText = item.content?.map((c) => c.text).join("\n\n") || ""; + slot.block.thinking = summaryText || contentText || slot.block.thinking; + slot.block.thinkingSignature = JSON.stringify(item); + reasoningBlocksById.set(item.id, slot.block); + stream.push({ + type: "thinking_end", + contentIndex: slot.contentIndex, + content: slot.block.thinking, + partial: output, + }); + outputSlots.delete(event.output_index); + } else if (item.type === "message" && slot?.type === "text") { + slot.block.text = + item.content?.map((c) => (c.type === "output_text" ? c.text : c.refusal)).join("") || ""; + slot.block.textSignature = encodeTextSignatureV1(item.id, item.phase ?? undefined); + stream.push({ + type: "text_end", + contentIndex: slot.contentIndex, + content: slot.block.text, + partial: output, + }); + outputSlots.delete(event.output_index); + } else if ( + item.type === "function_call" && + slot?.type === "toolCall" && + slot.block.partialJson !== undefined + ) { + slot.block.arguments = parseStreamingJson(item.arguments || slot.block.partialJson || "{}"); + if (item.namespace !== undefined) slot.block.namespace = item.namespace; + // Finalize in-place and strip the scratch buffer so replay only + // carries parsed arguments. + delete slot.block.partialJson; + slot.pending.finish(); + stream.push({ + type: "toolcall_end", + contentIndex: slot.contentIndex, + toolCall: slot.block, + partial: output, + }); + outputSlots.delete(event.output_index); + } else if (item.type === "custom_tool_call" && slot?.type === "toolCall" && slot.block.customInput) { + pushToolCallDelta( + slot, + appendCustomToolCallInput(slot.block, item.input ?? getCustomToolCallInput(slot.block), true), + ); + if (item.namespace !== undefined) slot.block.namespace = item.namespace; + delete slot.block.customInput; + slot.pending.finish(); + stream.push({ + type: "toolcall_end", + contentIndex: slot.contentIndex, + toolCall: slot.block, + partial: output, + }); + outputSlots.delete(event.output_index); + } + } else if (event.type === "response.completed" || event.type === "response.incomplete") { + finalizeResponse(event.response); + } else if (event.type === "error") { + throw new Error(`Error Code ${event.code}: ${event.message}` || "Unknown error"); + } else if (event.type === "response.failed") { + sawTerminalResponseEvent = true; + output.rawStopReason = event.response?.status; + const error = event.response?.error; + const details = event.response?.incomplete_details; + const msg = error + ? `${error.code || "unknown"}: ${error.message || "no message"}` + : details?.reason + ? `incomplete: ${details.reason}` + : "Unknown error (no error details in response)"; + throw new Error(msg); } - } else if (event.type === "response.completed" || event.type === "response.incomplete") { - finalizeResponse(event.response); - } else if (event.type === "error") { - throw new Error(`Error Code ${event.code}: ${event.message}` || "Unknown error"); - } else if (event.type === "response.failed") { - sawTerminalResponseEvent = true; - output.rawStopReason = event.response?.status; - const error = event.response?.error; - const details = event.response?.incomplete_details; - const msg = error - ? `${error.code || "unknown"}: ${error.message || "no message"}` - : details?.reason - ? `incomplete: ${details.reason}` - : "Unknown error (no error details in response)"; - throw new Error(msg); } - } - if (!sawTerminalResponseEvent) { - throw new Error("OpenAI Responses stream ended before a terminal response event"); + if (!sawTerminalResponseEvent) { + throw new Error("OpenAI Responses stream ended before a terminal response event"); + } + } finally { + // Covers normal completion and errors from every Responses transport. + for (const slot of outputSlots.values()) { + if (slot.type === "toolCall") slot.pending.finish(); + } } } diff --git a/packages/ai/src/api/pi-messages.ts b/packages/ai/src/api/pi-messages.ts index 253c2a5c3ed..9faeb0b635d 100644 --- a/packages/ai/src/api/pi-messages.ts +++ b/packages/ai/src/api/pi-messages.ts @@ -25,7 +25,7 @@ import type { import { appendAssistantMessageDiagnostic, createAssistantMessageDiagnostic } from "../utils/diagnostics.ts"; import { AssistantMessageEventStream } from "../utils/event-stream.ts"; import { headersToRecord, providerHeadersToRecord } from "../utils/headers.ts"; -import { parseStreamingJson } from "../utils/json-parse.ts"; +import { createPendingToolCall, type PendingToolCall } from "../utils/pending-tool-call.ts"; import { getProviderEnvValue } from "../utils/provider-env.ts"; export interface PiMessagesOptions extends StreamOptions { @@ -175,7 +175,7 @@ function appendRewriteDiagnostic(message: AssistantMessage, rewrite: PiMessagesR }); } -function createEventConverter(model: Model<"pi-messages">) { +function createEventConverter(model: Model<"pi-messages">, pendingCalls: Map) { const partial: AssistantMessage = { role: "assistant", content: [], @@ -240,31 +240,37 @@ function createEventConverter(model: Model<"pi-messages">) { redacted: event.redacted, }); break; - case "toolcall_start": - partial.content[event.contentIndex] = { + case "toolcall_start": { + const pending = createPendingToolCall({ type: "toolCall", id: event.id, name: event.toolName, arguments: {}, - }; + }); + partial.content[event.contentIndex] = pending.toolCall; + pendingCalls.set(pending.toolCall, pending); toolJson.set(event.contentIndex, ""); break; + } case "toolcall_delta": { const json = `${toolJson.get(event.contentIndex) ?? ""}${event.delta}`; toolJson.set(event.contentIndex, json); - (partial.content[event.contentIndex] as ToolCall).arguments = - parseStreamingJson(json); + pendingCalls.get(partial.content[event.contentIndex] as ToolCall)!.setJson(json); break; } - case "toolcall_end": - Object.assign(partial.content[event.contentIndex]!, event.toolCall); + case "toolcall_end": { + const block = partial.content[event.contentIndex] as ToolCall; + Object.assign(block, event.toolCall); + pendingCalls.get(block)!.finish(); + pendingCalls.delete(block); toolJson.delete(event.contentIndex); return { type: "toolcall_end", contentIndex: event.contentIndex, - toolCall: partial.content[event.contentIndex] as ToolCall, + toolCall: block, partial, }; + } } return { ...event, partial } as AssistantMessageEvent; @@ -356,7 +362,12 @@ export const stream: StreamFunction<"pi-messages", PiMessagesOptions> = ( options?: PiMessagesOptions, ): AssistantMessageEventStream => { const eventStream = new AssistantMessageEventStream(); - const convertEvent = createEventConverter(model); + const pendingCalls = new Map(); + const convertEvent = createEventConverter(model, pendingCalls); + const finishPendingCalls = () => { + for (const pending of pendingCalls.values()) pending.finish(); + pendingCalls.clear(); + }; void (async () => { try { @@ -411,6 +422,7 @@ export const stream: StreamFunction<"pi-messages", PiMessagesOptions> = ( for await (const piEvent of readPiMessagesEvents(response.body)) { const event = convertEvent(piEvent); + if (event.type === "done" || event.type === "error") finishPendingCalls(); eventStream.push(event); if (event.type === "done" || event.type === "error") { return; @@ -419,6 +431,7 @@ export const stream: StreamFunction<"pi-messages", PiMessagesOptions> = ( throw new Error(`${model.provider} stream ended without a terminal event`); } catch (error) { + finishPendingCalls(); eventStream.push(createErrorEvent(model, error, options?.signal?.aborted ?? false)); } })(); diff --git a/packages/ai/src/index.ts b/packages/ai/src/index.ts index 8cbeee730d1..057cd9c2dba 100644 --- a/packages/ai/src/index.ts +++ b/packages/ai/src/index.ts @@ -41,6 +41,7 @@ export * from "./utils/diagnostics.ts"; export * from "./utils/event-stream.ts"; export * from "./utils/json-parse.ts"; export * from "./utils/overflow.ts"; +export * from "./utils/pending-tool-call.ts"; export * from "./utils/retry.ts"; export { contentText } from "./utils/text.ts"; export * from "./utils/typebox-helpers.ts"; diff --git a/packages/ai/src/utils/pending-tool-call.ts b/packages/ai/src/utils/pending-tool-call.ts new file mode 100644 index 00000000000..71d074a10b1 --- /dev/null +++ b/packages/ai/src/utils/pending-tool-call.ts @@ -0,0 +1,63 @@ +import type { ToolCall } from "../types.ts"; +import { parseStreamingJson } from "./json-parse.ts"; + +interface ArgumentState { + json?: string; + value: ToolCall["arguments"] | undefined; +} + +/** Provider-owned state for a tool call whose arguments are still streaming. */ +export interface PendingToolCall { + readonly toolCall: T; + setJson(json: string | undefined): void; + finish(): void; + copy(): PendingToolCall; +} + +export function createPendingToolCall(initial: T, fallbackOnFalsy = false): PendingToolCall { + return createPendingView(initial, { value: initial.arguments }, fallbackOnFalsy); +} + +function createPendingView( + initial: T, + initialState: ArgumentState, + fallbackOnFalsy: boolean, +): PendingToolCall { + let state = initialState; + const toolCall: T = { + ...initial, + get arguments() { + if (state.value === undefined) { + state.value = parseStreamingJson(state.json); + if (fallbackOnFalsy) state.value ||= {}; + } + return state.value; + }, + set arguments(value: ToolCall["arguments"]) { + // Assignments replace this view's value without changing earlier copies. + state = { value }; + }, + }; + + return { + toolCall, + setJson(json) { + state = { json, value: undefined }; + }, + finish() { + const value = toolCall.arguments; + Object.defineProperty(toolCall, "arguments", { value, writable: true, enumerable: true, configurable: true }); + state = { value }; + }, + copy() { + // Preserve the proxy's spread semantics for metadata, including symbol keys + // and getter receivers, without evaluating arguments. + const entries = Reflect.ownKeys(toolCall) + .filter((key) => Object.getOwnPropertyDescriptor(toolCall, key)?.enumerable) + .map((key) => [key, key === "arguments" ? {} : Reflect.get(toolCall, key)] as const); + // Reads share a cached parse; later deltas and assignments replace only + // the state of the view receiving them. + return createPendingView(Object.fromEntries(entries) as unknown as T, state, fallbackOnFalsy); + }, + }; +} diff --git a/packages/ai/test/bedrock-redacted-reasoning.test.ts b/packages/ai/test/bedrock-redacted-reasoning.test.ts index 86301e155ae..c64fd6a1f51 100644 --- a/packages/ai/test/bedrock-redacted-reasoning.test.ts +++ b/packages/ai/test/bedrock-redacted-reasoning.test.ts @@ -1,4 +1,4 @@ -import { beforeEach, describe, expect, it, vi } from "vitest"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; /** * OpenAI models served through Bedrock Converse (e.g. `global.openai.gpt-5.6-terra`) @@ -126,6 +126,40 @@ async function capturePayload(context: Context): Promise return capturedPayload; } +// #9265: include the Converse adapter and its interrupted-block cleanup path. +describe("Bedrock lazy tool arguments", () => { + afterEach(() => vi.restoreAllMocks()); + + it.each([true, false])("defers argument parsing and settles unread arguments (complete: %s)", async (complete) => { + const content = "a".repeat(32 * 512); + const deltas = ['{"content":"', ...Array(32).fill("a".repeat(512))]; + if (complete) deltas.push('"}'); + bedrockMock.streamEvents = [ + { messageStart: { role: "assistant" } }, + { contentBlockStart: { contentBlockIndex: 0, start: { toolUse: { toolUseId: "call_test", name: "write" } } } }, + ...deltas.map((delta) => ({ + contentBlockDelta: { contentBlockIndex: 0, delta: { toolUse: { input: delta } } }, + })), + ...(complete + ? [{ contentBlockStop: { contentBlockIndex: 0 } }, { messageStop: { stopReason: "tool_use" } }] + : []), + ]; + const parse = vi.spyOn(JSON, "parse"); + const result = await streamBedrock(gptModel, { messages: [] }).result(); + expect(result.stopReason, result.errorMessage).toBe(complete ? "toolUse" : "error"); + if (complete) { + expect(parse.mock.calls.filter(([text]) => text.startsWith('{"content":')).map(([text]) => text)).toEqual([ + deltas.join(""), + ]); + } + expect(Object.getOwnPropertyDescriptor(result.content[0], "arguments")).toMatchObject({ + value: { content }, + writable: true, + }); + expect(result.content[0]).not.toHaveProperty("partialJson"); + }); +}); + describe("Bedrock redacted reasoning", () => { beforeEach(() => { bedrockMock.streamEvents = undefined; diff --git a/packages/ai/test/streaming-tool-arguments.test.ts b/packages/ai/test/streaming-tool-arguments.test.ts new file mode 100644 index 00000000000..7211789addf --- /dev/null +++ b/packages/ai/test/streaming-tool-arguments.test.ts @@ -0,0 +1,191 @@ +import { afterEach, describe, expect, it, vi } from "vitest"; +import type { AssistantMessage, ToolCall } from "../src/types.ts"; +import { AssistantMessageEventStream } from "../src/utils/event-stream.ts"; +import { parseStreamingJson } from "../src/utils/json-parse.ts"; +import { createPendingToolCall } from "../src/utils/pending-tool-call.ts"; + +function toolCall(): ToolCall { + return { type: "toolCall", id: "call_test", name: "write", arguments: {} }; +} + +afterEach(() => vi.restoreAllMocks()); + +// Regression coverage for #9265: delta-only consumers must not parse growing prefixes. +describe("pending tool calls", () => { + it("parses a large streamed argument only when read and caches the result", () => { + const pending = createPendingToolCall(toolCall()); + const block = pending.toolCall; + const parse = vi.spyOn(JSON, "parse"); + let json = '{"content":"'; + for (let i = 0; i < 128; i++) { + json += "a".repeat(8192); + pending.setJson(json); + } + json += '"}'; + pending.setJson(json); + expect(parse).not.toHaveBeenCalled(); + + const args = block.arguments; + expect(args.content).toHaveLength(1024 * 1024); + expect(block.arguments).toBe(args); + expect(parse).toHaveBeenCalledExactlyOnceWith(json); + pending.finish(); + expect(pending.toolCall).toBe(block); + expect(parse).toHaveBeenCalledExactlyOnceWith(json); + expect(Object.getOwnPropertyDescriptor(block, "arguments")).toEqual({ + value: args, + writable: true, + enumerable: true, + configurable: true, + }); + }); + + it.each([ + undefined, + "", + "not json", + "null", + "false", + "0", + '""', + '{"path":"a.txt","content":"hel', + '{"nested":{"items":[1,true,{"value":"par', + '{"content":"line1\nline2', + String.raw`{"path":"A\H","content":"\uD83D`, + ])("preserves best-effort parsing for %j", (json) => { + const pending = createPendingToolCall(toolCall()); + const expected = parseStreamingJson(json); + pending.setJson(json); + expect(pending.toolCall.arguments).toEqual(expected); + }); + + it("invalidates on the next delta without mutating previously read arguments", () => { + const pending = createPendingToolCall(toolCall()); + pending.setJson('{"content":"hel'); + const first = pending.toolCall.arguments; + expect(first).toEqual({ content: "hel" }); + pending.setJson('{"content":"hello"}'); + expect(pending.toolCall.arguments).toEqual({ content: "hello" }); + expect(first).toEqual({ content: "hel" }); + expect(pending.toolCall.arguments).not.toBe(first); + }); + + it("accepts authoritative assignments without parsing the discarded prefix", () => { + const pending = createPendingToolCall(toolCall()); + const block = pending.toolCall; + const parse = vi.spyOn(JSON, "parse"); + pending.setJson('{"content":"discarded'); + const authoritative = { content: "replacement" }; + Object.assign(block, { arguments: authoritative, namespace: "tools" }); + pending.finish(); + expect(block.arguments).toBe(authoritative); + expect(block.namespace).toBe("tools"); + expect(parse).not.toHaveBeenCalled(); + expect(Object.getOwnPropertyDescriptor(block, "arguments")?.get).toBeUndefined(); + }); + + it("keeps interleaved calls independent", () => { + const first = createPendingToolCall(toolCall()); + const second = createPendingToolCall(toolCall()); + first.setJson('{"content":"one'); + second.setJson('{"content":"two'); + expect(second.toolCall.arguments).toEqual({ content: "two" }); + first.setJson('{"content":"one more"}'); + expect(first.toolCall.arguments).toEqual({ content: "one more" }); + expect(second.toolCall.arguments).toEqual({ content: "two" }); + }); + + it("supports serialization, spreading, and structured cloning during streaming", () => { + const pending = createPendingToolCall(toolCall()); + const block = pending.toolCall; + pending.setJson('{"content":"partial'); + expect(JSON.parse(JSON.stringify(block)).arguments).toEqual({ content: "partial" }); + pending.setJson('{"content":"updated'); + expect({ ...block }.arguments).toEqual({ content: "updated" }); + pending.setJson('{"content":"cloned'); + expect(structuredClone(block).arguments).toEqual({ content: "cloned" }); + }); + + it("shares a cached parse between proxy copies but keeps updates and assignments independent", () => { + const pending = createPendingToolCall(toolCall()); + pending.setJson('{"content":"copied"}'); + const parse = vi.spyOn(JSON, "parse"); + const copy = pending.copy(); + expect(parse).not.toHaveBeenCalled(); + expect(copy.toolCall).not.toBe(pending.toolCall); + expect(copy.toolCall.arguments).toEqual({ content: "copied" }); + expect(pending.toolCall.arguments).toBe(copy.toolCall.arguments); + expect(parse).toHaveBeenCalledTimes(1); + const replacement = { content: "authoritative" }; + copy.toolCall.arguments = replacement; + expect(copy.toolCall.arguments).toBe(replacement); + expect(pending.toolCall.arguments).toEqual({ content: "copied" }); + copy.setJson('{"content":"next"}'); + expect(copy.toolCall.arguments).toEqual({ content: "next" }); + expect(pending.toolCall.arguments).toEqual({ content: "copied" }); + }); + + it("copies enumerable metadata, including symbols, with the proxy's original spread semantics", () => { + const pending = createPendingToolCall(toolCall()); + pending.setJson('{"content":"unread'); + const symbol = Symbol("metadata"); + const metadata = { tag: "extension" }; + Object.assign(pending.toolCall, { extra: metadata, [symbol]: metadata }); + const getter = vi.fn(() => "computed"); + Object.defineProperty(pending.toolCall, "computed", { enumerable: true, get: getter }); + Object.defineProperty(pending.toolCall, "hidden", { value: "private" }); + const parse = vi.spyOn(JSON, "parse"); + const copy = pending.copy().toolCall; + expect(parse).not.toHaveBeenCalled(); + expect(Reflect.get(copy, "extra")).toBe(metadata); + expect(Reflect.get(copy, symbol)).toBe(metadata); + expect(Object.getOwnPropertyDescriptor(copy, "computed")).toMatchObject({ value: "computed", writable: true }); + expect(getter).toHaveBeenCalledTimes(1); + expect(getter.mock.contexts[0]).toBe(pending.toolCall); + expect(copy).not.toHaveProperty("hidden"); + }); + + it.each(["toolUse", "error", "aborted"] as const)( + "lets the provider materialize unread arguments before %s settlement", + async (reason) => { + const pending = createPendingToolCall(toolCall()); + const block = pending.toolCall; + pending.setJson('{"content":"interrupted'); + // #9265: the event stream must not inspect unrelated custom-provider getters. + const foreign = toolCall(); + const foreignGetter = vi.fn(() => { + throw new Error("Custom provider getter must not be read"); + }); + Object.defineProperty(foreign, "arguments", { get: foreignGetter }); + const message: AssistantMessage = { + role: "assistant", + content: [block, foreign], + api: "openai-completions", + provider: "test", + model: "test", + usage: { + input: 0, + output: 0, + cacheRead: 0, + cacheWrite: 0, + totalTokens: 0, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }, + stopReason: reason, + timestamp: 0, + }; + const stream = new AssistantMessageEventStream(); + pending.finish(); + stream.push( + reason === "toolUse" ? { type: "done", reason, message } : { type: "error", reason, error: message }, + ); + expect(Object.getOwnPropertyDescriptor(block, "arguments")).toMatchObject({ + value: { content: "interrupted" }, + writable: true, + }); + expect(await stream.result()).toBe(message); + expect(foreignGetter).not.toHaveBeenCalled(); + expect(Object.getOwnPropertyDescriptor(foreign, "arguments")?.get).toBe(foreignGetter); + }, + ); +}); diff --git a/packages/ai/test/tool-argument-streaming.test.ts b/packages/ai/test/tool-argument-streaming.test.ts new file mode 100644 index 00000000000..1dbe9df806f --- /dev/null +++ b/packages/ai/test/tool-argument-streaming.test.ts @@ -0,0 +1,257 @@ +import { afterEach, describe, expect, it, vi } from "vitest"; +import { stream as streamAnthropic } from "../src/api/anthropic-messages.ts"; +import { stream as streamMistral } from "../src/api/mistral-conversations.ts"; +import { stream as streamCompletions } from "../src/api/openai-completions.ts"; +import { stream as streamResponses } from "../src/api/openai-responses.ts"; +import { stream as streamPiMessages } from "../src/api/pi-messages.ts"; +import type { Api, Model } from "../src/types.ts"; +import type { AssistantMessageEventStream } from "../src/utils/event-stream.ts"; + +function model(api: T): Model { + return { + id: "test", + name: "Test", + api, + provider: "test", + baseUrl: "https://example.com/v1", + reasoning: false, + input: ["text"], + contextWindow: 128000, + maxTokens: 16384, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }, + }; +} + +const content = "a".repeat(32 * 512); +const deltas = ['{"content":"', ...Array(32).fill("a".repeat(512)), '"}']; +const json = deltas.join(""); +const usage = { + input: 0, + output: 0, + cacheRead: 0, + cacheWrite: 0, + totalTokens: 0, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, +}; + +interface ProviderCase { + name: string; + events(deltas: string[], complete: boolean): object[]; + run(response: Response, signal?: AbortSignal): AssistantMessageEventStream; +} + +const providers: ProviderCase[] = [ + { + name: "openai-completions", + events: (deltas, complete) => [ + ...deltas.map((delta) => ({ + choices: [ + { + delta: { tool_calls: [{ index: 0, id: "call_test", function: { name: "write", arguments: delta } }] }, + }, + ], + })), + ...(complete ? [{ choices: [{ delta: {}, finish_reason: "tool_calls" }] }] : []), + ], + run: (response, signal) => + streamCompletions( + model("openai-completions"), + { messages: [] }, + { apiKey: "test", fetch: async () => response, signal }, + ), + }, + { + name: "mistral-conversations", + events: (deltas, complete) => [ + ...deltas.map((delta) => ({ + choices: [ + { + delta: { tool_calls: [{ index: 0, id: "call_test", function: { name: "write", arguments: delta } }] }, + }, + ], + })), + ...(complete ? [{ choices: [{ delta: {}, finish_reason: "tool_calls" }] }] : []), + ], + run: (response, signal) => + streamMistral( + model("mistral-conversations"), + { messages: [] }, + { apiKey: "test", fetch: async () => response, signal }, + ), + }, + { + name: "anthropic-messages", + events: (deltas, complete) => [ + { + type: "message_start", + message: { id: "msg_test", model: "test", usage: { input_tokens: 0, output_tokens: 0 } }, + }, + { + type: "content_block_start", + index: 0, + content_block: { type: "tool_use", id: "call_test", name: "write", input: {} }, + }, + ...deltas.map((delta) => ({ + type: "content_block_delta", + index: 0, + delta: { type: "input_json_delta", partial_json: delta }, + })), + ...(complete + ? [ + { type: "content_block_stop", index: 0 }, + { type: "message_delta", delta: { stop_reason: "tool_use" } }, + { type: "message_stop" }, + ] + : []), + ], + run: (response, signal) => + streamAnthropic( + model("anthropic-messages"), + { messages: [] }, + { apiKey: "test", fetch: async () => response, signal }, + ), + }, + { + name: "openai-responses", + events: (deltas, complete) => [ + { + type: "response.output_item.added", + output_index: 0, + item: { type: "function_call", id: "fc_test", call_id: "call_test", name: "write", arguments: "" }, + }, + ...deltas.map((delta) => ({ type: "response.function_call_arguments.delta", output_index: 0, delta })), + ...(complete + ? [ + { + type: "response.output_item.done", + output_index: 0, + item: { + type: "function_call", + id: "fc_test", + call_id: "call_test", + name: "write", + arguments: json, + }, + }, + { type: "response.completed", response: { id: "resp_test", status: "completed" } }, + ] + : []), + ], + run: (response, signal) => + streamResponses( + model("openai-responses"), + { messages: [] }, + { apiKey: "test", fetch: async () => response, signal }, + ), + }, + { + name: "pi-messages", + events: (deltas, complete) => [ + { type: "start" }, + { type: "toolcall_start", contentIndex: 0, id: "call_test", toolName: "write" }, + ...deltas.map((delta) => ({ type: "toolcall_delta", contentIndex: 0, delta })), + ...(complete + ? [ + { + type: "toolcall_end", + contentIndex: 0, + toolCall: { type: "toolCall", id: "call_test", name: "write", arguments: { content } }, + }, + { type: "done", reason: "toolUse", usage }, + ] + : [{ type: "error", reason: "error", errorMessage: "Interrupted", usage }]), + ], + run: (response) => + streamPiMessages(model("pi-messages"), { messages: [] }, { apiKey: "test", fetch: async () => response }), + }, +]; + +function response(events: object[]): Response { + return new Response( + events + .map((event) => { + const eventName = "type" in event ? `event: ${event.type}\n` : ""; + return `${eventName}data: ${JSON.stringify(event)}\n\n`; + }) + .join(""), + { headers: { "content-type": "text/event-stream" } }, + ); +} + +afterEach(() => vi.restoreAllMocks()); + +// #9265: exercise real adapter loops, not just the lazy argument helper. +describe.each(providers)("$name tool argument parsing", (provider) => { + it("does not parse intermediate prefixes for a delta-only consumer", async () => { + const upstream = response(provider.events(deltas, true)); + const parse = vi.spyOn(JSON, "parse"); + const stream = provider.run(upstream); + let deltaCount = 0; + for await (const event of stream) { + if (event.type === "toolcall_delta") deltaCount++; + } + const result = await stream.result(); + expect(result.stopReason, result.errorMessage).toBe("toolUse"); + expect(deltaCount).toBe(deltas.length); + const argumentParses = parse.mock.calls.filter(([text]) => text.startsWith('{"content":')); + expect(argumentParses.map(([text]) => text)).toEqual(provider.name === "pi-messages" ? [] : [json]); + expect(Object.getOwnPropertyDescriptor(result.content[0], "arguments")?.get).toBeUndefined(); + expect(result.content[0]).toMatchObject({ type: "toolCall", arguments: { content } }); + }); + + it("preserves unread partial arguments when the stream fails", async () => { + const result = await provider.run(response(provider.events(deltas.slice(0, -1), false))).result(); + expect(result.stopReason).toBe("error"); + expect(Object.getOwnPropertyDescriptor(result.content[0], "arguments")).toMatchObject({ + value: { content }, + writable: true, + }); + expect(result.content[0]).not.toHaveProperty("partialJson"); + expect(result.content[0]).not.toHaveProperty("partialArgs"); + }); +}); + +// #9265: unlike clean EOF, a transport error skips Completions/Mistral block-end parsing. +// pi-messages already returns a new, empty message for transport errors; leave that behavior unchanged. +describe.each(providers.filter((provider) => provider.name !== "pi-messages"))( + "$name interrupted tool arguments", + (provider) => { + it.each(["error", "aborted"] as const)( + "materializes unread arguments before toolcall_end (%s)", + async (reason) => { + const partialDeltas = ['{"content":"interrupted']; + const bytes = new Uint8Array(await response(provider.events(partialDeltas, false)).arrayBuffer()); + let transport!: ReadableStreamDefaultController; + const upstream = new Response( + new ReadableStream({ + start(controller) { + transport = controller; + controller.enqueue(bytes); + }, + }), + { headers: { "content-type": "text/event-stream" } }, + ); + const controller = new AbortController(); + const stream = provider.run(upstream, controller.signal); + let deltaCount = 0; + for await (const event of stream) { + expect(event.type).not.toBe("toolcall_end"); + if (event.type === "toolcall_delta" && ++deltaCount === partialDeltas.length) { + if (reason === "aborted") controller.abort(); + transport.error(new Error("Connection interrupted")); + } + } + const result = await stream.result(); + expect(deltaCount).toBe(partialDeltas.length); + expect(result.stopReason).toBe(reason); + if (reason === "error") expect(result.errorMessage).toContain("Connection interrupted"); + expect(Object.getOwnPropertyDescriptor(result.content[0], "arguments")).toMatchObject({ + value: { content: "interrupted" }, + writable: true, + }); + expect(result.content[0]).not.toHaveProperty("partialJson"); + expect(result.content[0]).not.toHaveProperty("partialArgs"); + }, + ); + }, +);