From f715a15c101ca8a44d759bceeb7ec2a51024c81c Mon Sep 17 00:00:00 2001 From: Mario Zechner Date: Sat, 24 Oct 2026 08:00:07 +0000 Subject: [PATCH] fix(ai,chord,durable,coding-agent): parse streamed tool-call arguments incrementally (fixes #9265) - ai: StreamingJsonParser / ToolCallArgumentsStream read each argument delta once (tolerant, immutable snapshots, linear cost); every provider, pi-messages and the faux provider use it and keep stream state (partialJson, partialArgs, index, customInput, streamIndex, Bedrock reasoning bytes) off the content blocks; grammar tool input appends per delta - agent: streamProxy uses the same parser - chord: Delta append ops own their text (ownedSlice), so ops and replicas no longer keep the whole growing string alive - durable, ai, mcp: bounded error excerpts and summaries own their characters - coding-agent: tool calls render streamed arguments once per frame, not once per delta --- packages/agent/CHANGELOG.md | 4 + packages/agent/src/proxy.ts | 22 +- packages/agent/test/proxy.test.ts | 53 ++ packages/ai/CHANGELOG.md | 12 + packages/ai/src/api/anthropic-messages.ts | 100 +-- packages/ai/src/api/azure-openai-responses.ts | 6 - .../ai/src/api/bedrock-converse-stream.ts | 163 ++-- packages/ai/src/api/constrained-sampling.ts | 55 +- packages/ai/src/api/google-generative-ai.ts | 6 - packages/ai/src/api/google-vertex.ts | 6 - packages/ai/src/api/mistral-conversations.ts | 62 +- packages/ai/src/api/openai-codex-responses.ts | 5 - packages/ai/src/api/openai-completions.ts | 114 +-- .../ai/src/api/openai-responses-shared.ts | 154 ++-- packages/ai/src/api/openai-responses.ts | 6 - packages/ai/src/api/pi-messages.ts | 18 +- packages/ai/src/index.ts | 1 + packages/ai/src/providers/faux.ts | 4 + packages/ai/src/providers/radius-config.ts | 3 +- packages/ai/src/telemetry.ts | 3 +- packages/ai/src/utils/error-body.ts | 4 +- packages/ai/src/utils/streaming-json.ts | 559 ++++++++++++ packages/ai/src/utils/text.ts | 10 + packages/ai/test/bedrock-stream-state.test.ts | 158 ++++ packages/ai/test/constrained-sampling.test.ts | 77 +- packages/ai/test/data/json-test-suite.json | 322 +++++++ packages/ai/test/error-body.test.ts | 26 +- packages/ai/test/memory-helpers.ts | 39 + packages/ai/test/streaming-json.test.ts | 848 ++++++++++++++++++ .../ai/test/streaming-tool-arguments.test.ts | 420 +++++++++ packages/ai/test/telemetry.test.ts | 23 + packages/chord/CHANGELOG.md | 5 + packages/chord/src/delta/diff.ts | 7 +- packages/chord/src/delta/index.ts | 10 + packages/chord/src/delta/tracker.ts | 19 +- .../test/delta-op-text-retention.test.ts | 80 ++ packages/coding-agent/CHANGELOG.md | 1 + .../custom-provider-anthropic/index.ts | 58 +- .../interactive/components/tool-execution.ts | 11 +- .../test/tool-execution-component.test.ts | 53 ++ packages/durable/CHANGELOG.md | 6 + packages/durable/src/codemode/codemode.ts | 4 +- packages/durable/src/harness/scheduler.ts | 3 +- packages/durable/src/harness/tool.ts | 9 +- packages/durable/test/harness-events.test.ts | 89 ++ packages/mcp/CHANGELOG.md | 1 + .../mcp/src/transports/streamable-http.ts | 3 +- 47 files changed, 3187 insertions(+), 455 deletions(-) create mode 100644 packages/ai/src/utils/streaming-json.ts create mode 100644 packages/ai/test/bedrock-stream-state.test.ts create mode 100644 packages/ai/test/data/json-test-suite.json create mode 100644 packages/ai/test/memory-helpers.ts create mode 100644 packages/ai/test/streaming-json.test.ts create mode 100644 packages/ai/test/streaming-tool-arguments.test.ts create mode 100644 packages/chord/test/delta-op-text-retention.test.ts diff --git a/packages/agent/CHANGELOG.md b/packages/agent/CHANGELOG.md index b5367030f..2ae23784d 100644 --- a/packages/agent/CHANGELOG.md +++ b/packages/agent/CHANGELOG.md @@ -2,6 +2,10 @@ ## [Unreleased] +### Fixed + +- Fixed `streamProxy()` re-parsing the accumulated tool-call arguments on every delta and keeping their raw text on the partial content block ([#9265](https://github.com/earendil-works/pi/issues/9265)) + ## [1.4.0] - 2026-10-21 ### Added diff --git a/packages/agent/src/proxy.ts b/packages/agent/src/proxy.ts index c5f6388ed..8aae7b72c 100644 --- a/packages/agent/src/proxy.ts +++ b/packages/agent/src/proxy.ts @@ -3,17 +3,16 @@ * The server manages auth and proxies requests to LLM providers. */ -// Internal import for JSON parsing utility import { type AssistantMessage, type AssistantMessageEvent, type AssistantMessageEventStream, createAssistantMessageEventStream, type Model, - parseStreamingJson, type SimpleStreamOptions, type StopReason, type ToolCall, + ToolCallArgumentsStream, type TranscriptContext, } from "@earendil-works/pi-ai"; @@ -175,13 +174,15 @@ export function streamProxy( const decoder = new TextDecoder(); let buffer = ""; let sawTerminalEvent = false; + // Argument streams of open tool calls by content index, kept off the content blocks. + const toolArguments = new Map(); const processLine = (line: string): void => { if (!line.startsWith("data: ")) return; 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, toolArguments); if (event) { if (event.type === "done" || event.type === "error") sawTerminalEvent = true; stream.push(event); @@ -257,6 +258,7 @@ export function streamProxy( function processProxyEvent( proxyEvent: ProxyAssistantMessageEvent, partial: AssistantMessage, + toolArguments: Map, ): AssistantMessageEvent | undefined { switch (proxyEvent.type) { case "start": @@ -332,16 +334,16 @@ function processProxyEvent( id: proxyEvent.id, name: proxyEvent.toolName, arguments: {}, - partialJson: "", - } satisfies ToolCall & { partialJson: string } as ToolCall; + }; + toolArguments.set(proxyEvent.contentIndex, new ToolCallArgumentsStream()); 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 args = toolArguments.get(proxyEvent.contentIndex); + if (content?.type === "toolCall" && args !== undefined) { + // A new block object, so reactive consumers see the change. + partial.content[proxyEvent.contentIndex] = { ...content, arguments: args.append(proxyEvent.delta) }; return { type: "toolcall_delta", contentIndex: proxyEvent.contentIndex, @@ -356,7 +358,7 @@ function processProxyEvent( const content = partial.content[proxyEvent.contentIndex]; if (content?.type === "toolCall") { Object.assign(content, proxyEvent.toolCall); - delete (content as any).partialJson; + toolArguments.delete(proxyEvent.contentIndex); return { type: "toolcall_end", contentIndex: proxyEvent.contentIndex, diff --git a/packages/agent/test/proxy.test.ts b/packages/agent/test/proxy.test.ts index 43e5174c3..822cc5391 100644 --- a/packages/agent/test/proxy.test.ts +++ b/packages/agent/test/proxy.test.ts @@ -115,4 +115,57 @@ describe("streamProxy", () => { expect(result.stopReason).toBe("error"); expect(result.errorMessage).toContain("Connection closed by proxy server"); }); + + // #9265 + it("shows growing partial arguments without keeping the raw text on the block, in linear time", async () => { + const content = 'line "quoted" \\ end\n'.repeat(20_000); + const text = JSON.stringify({ path: "a.ts", content }); + const proxyEvents: ProxyAssistantMessageEvent[] = [ + { type: "start" }, + { type: "toolcall_start", contentIndex: 0, id: "call_1", toolName: "write" }, + ]; + for (let index = 0; index < text.length; index += 32) { + proxyEvents.push({ type: "toolcall_delta", contentIndex: 0, delta: text.slice(index, index + 32) }); + } + proxyEvents.push( + { + type: "toolcall_end", + contentIndex: 0, + toolCall: { type: "toolCall", id: "call_1", name: "write", arguments: { path: "a.ts", content } }, + }, + { type: "done", reason: "toolUse", usage }, + ); + const body = proxyEvents.map((event) => `data: ${JSON.stringify(event)}\n\n`).join(""); + vi.stubGlobal( + "fetch", + vi.fn(async () => new Response(body, { status: 200 })), + ); + + const started = performance.now(); + const stream = streamProxy(model, normalizeContext({ systemPrompt: "", messages: [] }), { + authToken: "test-token", + proxyUrl: "https://proxy.example.com", + }); + let shown = 0; + let deltas = 0; + for await (const event of stream) { + if (event.type !== "toolcall_delta") continue; + const block = event.partial.content[0]; + expect(Object.keys(block!).sort()).toEqual(["arguments", "id", "name", "type"]); + const partial = block?.type === "toolCall" ? block.arguments.content : undefined; + if (typeof partial !== "string") continue; + // Comparing the whole text on every delta would make this test itself quadratic. + if (deltas++ % 500 === 0) expect(content.startsWith(partial)).toBe(true); + expect(partial.length).toBeGreaterThanOrEqual(shown); + shown = partial.length; + } + const result = await stream.result(); + + expect(shown).toBeGreaterThan(content.length - 64); + expect(result.content).toEqual([ + { type: "toolCall", id: "call_1", name: "write", arguments: { path: "a.ts", content } }, + ]); + // Re-parsing the accumulated 520 KB of arguments on each of its 16,000 deltas took minutes. + expect(performance.now() - started).toBeLessThan(10_000); + }); }); diff --git a/packages/ai/CHANGELOG.md b/packages/ai/CHANGELOG.md index 63efd0e37..9ea77c63b 100644 --- a/packages/ai/CHANGELOG.md +++ b/packages/ai/CHANGELOG.md @@ -2,14 +2,26 @@ ## [Unreleased] +### Breaking Changes + +- `appendGrammarToolInputJsonDelta()` (`api/constrained-sampling`) takes the input delta instead of the whole input so far and no longer closes the input; close it with the new `closeGrammarToolInputJson()` ([#9265](https://github.com/earendil-works/pi/issues/9265)) + ### Added - Added optional `openBrowserOnConfirm` and `copyCodeOnConfirm` to `device_code` auth events and `OAuthDeviceCodeInfo`. Both default to `true`; set them to `false` to stop clients from offering to open the verification page or copy the code ([#9282](https://github.com/earendil-works/pi/issues/9282), [#9301](https://github.com/earendil-works/pi/pull/9301)) +- Added `StreamingJsonParser`, an incremental, tolerant parser for JSON text that arrives in pieces (each piece is read once; snapshots of the partial value never change once returned), and `ToolCallArgumentsStream`, which every provider now uses for streamed tool-call arguments; its `append()` also returns the changes each piece made (`JsonChange`: a value set at a path, or text appended to a string), for writers that keep their own copy ([#9265](https://github.com/earendil-works/pi/issues/9265)) + +### Changed + +- The faux provider streams partial tool-call arguments, so `toolcall_delta` partials show growing arguments as real providers do ([#9265](https://github.com/earendil-works/pi/issues/9265)) ### Fixed - Fixed Claude models that accept mid-conversation system messages rejecting a session with a 400 (`role 'system' must precede an 'assistant' message`) when a system update was followed by an assistant reply with nothing to send to Anthropic, such as an empty reply or another model's reasoning-only reply ([#11091](https://github.com/earendil-works/pi/issues/11091)) - Fixed Claude models with mid-conversation effort failing with `Invalid effort level` on accounts whose organization caps effort below `high`, even at `low` or `medium` thinking, by sending a fixed `low` top-level effort ([#10063](https://github.com/earendil-works/pi/issues/10063)) +- Fixed streamed tool-call arguments being re-parsed in full on every delta, which made CPU time and memory grow quadratically with the size of a tool call (a 90 KB `write` on Anthropic peaked at about 140 MB of heap; a 500 KB one took minutes). Every provider and the pi-messages client now read each delta once, and keep the raw argument text and their other stream state (`partialJson`, `partialArgs`, `index`, `customInput`, `streamIndex`, Bedrock's encrypted reasoning bytes) off the content blocks, so it no longer reaches partial messages or transcripts of failed streams ([#9265](https://github.com/earendil-works/pi/issues/9265)) +- Fixed grammar (custom) tool input streaming copying the whole input so far on every delta ([#9265](https://github.com/earendil-works/pi/issues/9265)) +- Fixed bounded error texts (telemetry span messages, truncated provider error bodies, pi-messages diagnostics, Radius gateway errors) keeping the whole text they were cut from alive ([#9265](https://github.com/earendil-works/pi/issues/9265)) ## [1.4.0] - 2026-10-21 diff --git a/packages/ai/src/api/anthropic-messages.ts b/packages/ai/src/api/anthropic-messages.ts index 00d820341..53764b406 100644 --- a/packages/ai/src/api/anthropic-messages.ts +++ b/packages/ai/src/api/anthropic-messages.ts @@ -32,16 +32,14 @@ import type { StreamFunction, StreamOptions, TextContent, - ThinkingContent, Tool, - ToolCall, ToolResultMessage, } from "../types.ts"; import { appendAssistantMessageDiagnostic } from "../utils/diagnostics.ts"; import { AssistantMessageEventStream } from "../utils/event-stream.ts"; import { headersToRecord } from "../utils/headers.ts"; import { type ClientIdentity, getIdentityHeaders } from "../utils/identity.ts"; -import { parseJsonWithRepair, parseStreamingJson, parseToolCallArguments } from "../utils/json-parse.ts"; +import { parseJsonWithRepair } from "../utils/json-parse.ts"; import { applyOpenRouterReportedCost, isOpenRouterModel, @@ -50,6 +48,7 @@ import { import { getProviderEnvValue } from "../utils/provider-env.ts"; import { retryProviderRequest } from "../utils/provider-retry.ts"; import { sanitizeSurrogates } from "../utils/sanitize-unicode.ts"; +import { ToolCallArgumentsStream } from "../utils/streaming-json.ts"; import { getSystemMessageText, renderSystemMessageUpdate } from "../utils/text.ts"; import { getCurrentTools, @@ -695,8 +694,12 @@ export const stream: StreamFunction<"anthropic-messages", AnthropicOptions> = ( await options?.onResponse?.({ status: response.status, headers: headersToRecord(response.headers) }, model); stream.push({ type: "start", partial: output }); - type Block = (ThinkingContent | TextContent | (ToolCall & { partialJson: string })) & { index: number }; - const blocks = output.content as Block[]; + // Stream state per Anthropic block index, kept off the content blocks so it never reaches partials. + const openBlocks = new Map(); + const openBlock = (anthropicIndex: number) => { + const open = openBlocks.get(anthropicIndex); + return open === undefined ? undefined : { ...open, block: output.content[open.contentIndex]! }; + }; for await (const event of iterateAnthropicEvents(response, options?.signal)) { await options?.onProviderStreamEvent?.(event, model); @@ -732,97 +735,84 @@ export const stream: StreamFunction<"anthropic-messages", AnthropicOptions> = ( continue; } if (event.content_block.type === "text") { - const block: Block = { - type: "text", - text: event.content_block.text ?? "", - index: event.index, - }; - output.content.push(block); + output.content.push({ type: "text", text: event.content_block.text ?? "" }); + openBlocks.set(event.index, { contentIndex: output.content.length - 1 }); stream.push({ type: "text_start", contentIndex: output.content.length - 1, partial: output }); } else if (event.content_block.type === "thinking") { - const block: Block = { + output.content.push({ type: "thinking", thinking: event.content_block.thinking ?? "", thinkingSignature: event.content_block.signature ?? "", - index: event.index, - }; - output.content.push(block); + }); + openBlocks.set(event.index, { contentIndex: output.content.length - 1 }); stream.push({ type: "thinking_start", contentIndex: output.content.length - 1, partial: output }); } else if (event.content_block.type === "redacted_thinking") { - const block: Block = { + output.content.push({ type: "thinking", thinking: "[Reasoning redacted]", thinkingSignature: event.content_block.data, redacted: true, - index: event.index, - }; - output.content.push(block); + }); + openBlocks.set(event.index, { contentIndex: output.content.length - 1 }); stream.push({ type: "thinking_start", contentIndex: output.content.length - 1, partial: output }); } else if (event.content_block.type === "tool_use") { - const block: Block = { + output.content.push({ type: "toolCall", id: event.content_block.id, name: isOAuth ? fromClaudeCodeName(event.content_block.name, currentTools) : event.content_block.name, arguments: (event.content_block.input as Record) ?? {}, - partialJson: "", - index: event.index, - }; - output.content.push(block); + }); + openBlocks.set(event.index, { + contentIndex: output.content.length - 1, + arguments: new ToolCallArgumentsStream(), + }); stream.push({ type: "toolcall_start", contentIndex: output.content.length - 1, partial: output }); } } else if (event.type === "content_block_delta") { + const open = openBlock(event.index); if (event.delta.type === "text_delta") { - const index = blocks.findIndex((b) => b.index === event.index); - const block = blocks[index]; - if (block && block.type === "text") { - block.text += event.delta.text; + if (open?.block.type === "text") { + open.block.text += event.delta.text; stream.push({ type: "text_delta", - contentIndex: index, + contentIndex: open.contentIndex, delta: event.delta.text, partial: output, }); } } else if (event.delta.type === "thinking_delta") { - const index = blocks.findIndex((b) => b.index === event.index); - const block = blocks[index]; - if (block && block.type === "thinking") { - block.thinking += event.delta.thinking; + if (open?.block.type === "thinking") { + open.block.thinking += event.delta.thinking; stream.push({ type: "thinking_delta", - contentIndex: index, + contentIndex: open.contentIndex, delta: event.delta.thinking, partial: output, }); } } else if (event.delta.type === "input_json_delta") { - const index = blocks.findIndex((b) => b.index === event.index); - const block = blocks[index]; - if (block && block.type === "toolCall") { - block.partialJson += event.delta.partial_json; - block.arguments = parseStreamingJson(block.partialJson); + if (open?.block.type === "toolCall" && open.arguments !== undefined) { + open.block.arguments = open.arguments.append(event.delta.partial_json); stream.push({ type: "toolcall_delta", - contentIndex: index, + contentIndex: open.contentIndex, delta: event.delta.partial_json, partial: output, }); } } else if (event.delta.type === "signature_delta") { - const index = blocks.findIndex((b) => b.index === event.index); - const block = blocks[index]; - if (block && block.type === "thinking") { - block.thinkingSignature = block.thinkingSignature || ""; - block.thinkingSignature += event.delta.signature; + if (open?.block.type === "thinking") { + open.block.thinkingSignature = open.block.thinkingSignature || ""; + open.block.thinkingSignature += event.delta.signature; } } } else if (event.type === "content_block_stop") { - const index = blocks.findIndex((b) => b.index === event.index); - const block = blocks[index]; - if (block) { - delete (block as any).index; + const open = openBlock(event.index); + openBlocks.delete(event.index); + if (open) { + const { block, contentIndex: index } = open; if (block.type === "text") { stream.push({ type: "text_end", @@ -837,11 +827,8 @@ export const stream: StreamFunction<"anthropic-messages", AnthropicOptions> = ( content: block.thinking, partial: output, }); - } else if (block.type === "toolCall") { - Object.assign(block, parseToolCallArguments(block.partialJson)); - // Finalize in-place and strip the scratch buffer so replay only - // carries parsed arguments. - delete (block as { partialJson?: string }).partialJson; + } else if (block.type === "toolCall" && open.arguments !== undefined) { + Object.assign(block, open.arguments.finish()); stream.push({ type: "toolcall_end", contentIndex: index, @@ -931,11 +918,6 @@ export const stream: StreamFunction<"anthropic-messages", AnthropicOptions> = ( stream.push({ type: "done", reason: output.stopReason, message: output }); stream.end(); } catch (error) { - for (const block of output.content) { - delete (block as { index?: number }).index; - // partialJson is only a streaming scratch buffer; never persist it. - delete (block as { partialJson?: string }).partialJson; - } output.stopReason = options?.signal?.aborted ? "aborted" : "error"; output.errorMessage = error instanceof Error ? error.message : JSON.stringify(error); stream.push({ type: "error", reason: output.stopReason, error: output }); diff --git a/packages/ai/src/api/azure-openai-responses.ts b/packages/ai/src/api/azure-openai-responses.ts index 07ce1a289..24e73fd05 100644 --- a/packages/ai/src/api/azure-openai-responses.ts +++ b/packages/ai/src/api/azure-openai-responses.ts @@ -118,12 +118,6 @@ export const stream: StreamFunction<"azure-openai-responses", AzureOpenAIRespons stream.push({ type: "done", reason: output.stopReason, message: output }); stream.end(); } catch (error) { - for (const block of output.content) { - delete (block as { index?: number }).index; - // Streaming scratch buffers are only used during parsing; never persist them. - delete (block as { partialJson?: string }).partialJson; - delete (block as { customInput?: unknown }).customInput; - } output.stopReason = options?.signal?.aborted ? "aborted" : "error"; output.errorMessage = formatAzureOpenAIError(error); stream.push({ type: "error", reason: output.stopReason, error: output }); diff --git a/packages/ai/src/api/bedrock-converse-stream.ts b/packages/ai/src/api/bedrock-converse-stream.ts index 8c6df2531..7a7740f06 100644 --- a/packages/ai/src/api/bedrock-converse-stream.ts +++ b/packages/ai/src/api/bedrock-converse-stream.ts @@ -43,10 +43,8 @@ import type { StreamOptions, TextContent, ThinkingBudgets, - ThinkingContent, ThinkingLevel, Tool, - ToolCall, ToolResultMessage, } from "../types.ts"; import { appendAssistantMessageDiagnostic } from "../utils/diagnostics.ts"; @@ -54,10 +52,10 @@ import { normalizeProviderError } from "../utils/error-body.ts"; import { AssistantMessageEventStream } from "../utils/event-stream.ts"; import { providerHeadersToRecord } from "../utils/headers.ts"; import { type ClientIdentity, getClientVersion } from "../utils/identity.ts"; -import { parseStreamingJson, parseToolCallArguments } from "../utils/json-parse.ts"; import { resolveHttpProxyUrlForTarget } from "../utils/node-http-proxy.ts"; import { getProviderEnvValue } from "../utils/provider-env.ts"; import { sanitizeSurrogates } from "../utils/sanitize-unicode.ts"; +import { ToolCallArgumentsStream } from "../utils/streaming-json.ts"; import { getSystemMessageText } from "../utils/text.ts"; import { collapseSystemMessages, @@ -111,13 +109,19 @@ export interface BedrockOptions extends StreamOptions { bearerToken?: string; } -type Block = (TextContent | ThinkingContent | ToolCall) & { - index?: number; - partialJson?: string; - /** Scratch buffer for encrypted reasoning deltas, joined into `thinkingSignature`. */ +/** + * Stream state of one open Bedrock content block, keyed by `contentBlockIndex` and kept off the content block so it + * never reaches partial messages. + */ +type OpenBlock = { + readonly contentIndex: number; + readonly arguments?: ToolCallArgumentsStream; + /** Encrypted reasoning deltas, joined into `thinkingSignature` when the block ends. */ redactedChunks?: Uint8Array[]; }; +type OpenBlocks = Map; + const EMPTY_TEXT_PLACEHOLDER = ""; const THINKING_BINDING_CONTROLS_BETA = "thinking-binding-controls-2026-08-01"; @@ -153,7 +157,7 @@ export const stream: StreamFunction<"bedrock-converse-stream", BedrockOptions> = timestamp: Date.now(), }; - const blocks = output.content as Block[]; + const openBlocks: OpenBlocks = 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 @@ -305,11 +309,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, openBlocks, output, stream); } else if (item.contentBlockDelta) { - handleContentBlockDelta(item.contentBlockDelta, blocks, output, stream); + handleContentBlockDelta(item.contentBlockDelta, openBlocks, output, stream); } else if (item.contentBlockStop) { - handleContentBlockStop(item.contentBlockStop, blocks, output, stream); + handleContentBlockStop(item.contentBlockStop, openBlocks, output, stream); } else if (item.messageStop) { output.rawStopReason = item.messageStop.stopReason; const { stopReason, errorMessage } = mapStopReason(item.messageStop.stopReason); @@ -344,13 +348,11 @@ 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); + flushOpenBlocks(openBlocks, output); stream.push({ type: "done", reason: output.stopReason, message: output }); stream.end(); } catch (error) { - for (const block of output.content) { - finalizeStreamingBlock(block as Block); - } + flushOpenBlocks(openBlocks, output); output.stopReason = options.signal?.aborted ? "aborted" : "error"; output.errorMessage = formatBedrockError(error); if (output.stopReason === "error") { @@ -593,70 +595,70 @@ export const streamSimple: StreamFunction<"bedrock-converse-stream", SimpleStrea function handleContentBlockStart( event: ContentBlockStartEvent, - blocks: Block[], + openBlocks: OpenBlocks, output: AssistantMessage, stream: AssistantMessageEventStream, ): void { - const index = event.contentBlockIndex!; const start = event.start; if (start?.toolUse) { - const block: Block = { + output.content.push({ type: "toolCall", id: start.toolUse.toolUseId || "", name: start.toolUse.name || "", arguments: {}, - partialJson: "", - index, - }; - output.content.push(block); - stream.push({ type: "toolcall_start", contentIndex: blocks.length - 1, partial: output }); + }); + const contentIndex = output.content.length - 1; + openBlocks.set(event.contentBlockIndex!, { contentIndex, arguments: new ToolCallArgumentsStream() }); + stream.push({ type: "toolcall_start", contentIndex, partial: output }); } } function handleContentBlockDelta( event: ContentBlockDeltaEvent, - blocks: Block[], + openBlocks: OpenBlocks, output: AssistantMessage, stream: AssistantMessageEventStream, ): void { const contentBlockIndex = event.contentBlockIndex!; const delta = event.delta; - let index = blocks.findIndex((b) => b.index === contentBlockIndex); - let block = blocks[index]; + let open = openBlocks.get(contentBlockIndex); + let block = open === undefined ? undefined : output.content[open.contentIndex]; if (delta?.text !== undefined) { // If no text block exists yet, create one, as `handleContentBlockStart` is not sent for text blocks - if (!block) { - const newBlock: Block = { type: "text", text: "", index: contentBlockIndex }; - output.content.push(newBlock); - index = blocks.length - 1; - block = blocks[index]; - stream.push({ type: "text_start", contentIndex: index, partial: output }); + if (!open) { + output.content.push({ type: "text", text: "" }); + open = { contentIndex: output.content.length - 1 }; + openBlocks.set(contentBlockIndex, open); + block = output.content[open.contentIndex]; + stream.push({ type: "text_start", contentIndex: open.contentIndex, partial: output }); } - if (block.type === "text") { + if (block?.type === "text") { block.text += delta.text; - stream.push({ type: "text_delta", contentIndex: index, delta: delta.text, partial: output }); + stream.push({ type: "text_delta", contentIndex: open.contentIndex, delta: delta.text, partial: output }); } - } else if (delta?.toolUse && block?.type === "toolCall") { - block.partialJson = (block.partialJson || "") + (delta.toolUse.input || ""); - block.arguments = parseStreamingJson(block.partialJson); - stream.push({ type: "toolcall_delta", contentIndex: index, delta: delta.toolUse.input || "", partial: output }); + } else if (delta?.toolUse && block?.type === "toolCall" && open?.arguments !== undefined) { + block.arguments = open.arguments.append(delta.toolUse.input || ""); + stream.push({ + type: "toolcall_delta", + contentIndex: open.contentIndex, + delta: delta.toolUse.input || "", + partial: output, + }); } else if (delta?.reasoningContent) { - let thinkingBlock = block; - let thinkingIndex = index; - - if (!thinkingBlock) { - const newBlock: Block = { type: "thinking", thinking: "", thinkingSignature: "", index: contentBlockIndex }; - output.content.push(newBlock); - thinkingIndex = blocks.length - 1; - thinkingBlock = blocks[thinkingIndex]; - stream.push({ type: "thinking_start", contentIndex: thinkingIndex, partial: output }); + if (!open) { + output.content.push({ type: "thinking", thinking: "", thinkingSignature: "" }); + open = { contentIndex: output.content.length - 1 }; + openBlocks.set(contentBlockIndex, open); + block = output.content[open.contentIndex]; + stream.push({ type: "thinking_start", contentIndex: open.contentIndex, partial: output }); } - if (thinkingBlock?.type === "thinking") { + if (block?.type === "thinking") { + const thinkingIndex = open.contentIndex; if (delta.reasoningContent.text) { - thinkingBlock.thinking += delta.reasoningContent.text; + block.thinking += delta.reasoningContent.text; stream.push({ type: "thinking_delta", contentIndex: thinkingIndex, @@ -666,18 +668,17 @@ function handleContentBlockDelta( } // `thinkingSignature` holds either an Anthropic signature or an opaque redacted // payload, never both: mixing them would corrupt whichever arrived first. - if (delta.reasoningContent.signature && !thinkingBlock.redacted) { - thinkingBlock.thinkingSignature = - (thinkingBlock.thinkingSignature || "") + delta.reasoningContent.signature; + if (delta.reasoningContent.signature && !block.redacted) { + block.thinkingSignature = (block.thinkingSignature || "") + delta.reasoningContent.signature; } if (delta.reasoningContent.redactedContent?.length) { // Encrypted reasoning from non-Anthropic models on Bedrock (e.g. OpenAI GPT-5.6). // The payload is opaque, so keep it verbatim in `thinkingSignature` the way the // Anthropic path stores redacted thinking, and replay it on the next turn. - if (!thinkingBlock.redacted) { - thinkingBlock.redacted = true; - thinkingBlock.thinkingSignature = ""; - thinkingBlock.thinking += REDACTED_THINKING_PLACEHOLDER; + if (!block.redacted) { + block.redacted = true; + block.thinkingSignature = ""; + block.thinking += REDACTED_THINKING_PLACEHOLDER; stream.push({ type: "thinking_delta", contentIndex: thinkingIndex, @@ -685,33 +686,29 @@ function handleContentBlockDelta( partial: output, }); } - thinkingBlock.redactedChunks ??= []; - thinkingBlock.redactedChunks.push(delta.reasoningContent.redactedContent); + open.redactedChunks ??= []; + open.redactedChunks.push(delta.reasoningContent.redactedContent); } } } } /** - * Encodes buffered encrypted reasoning into `thinkingSignature` and drops the scratch - * buffer, which must never reach a persisted message: `Uint8Array` serializes to an - * index-keyed object roughly ten times the size of the base64 payload. + * Encodes buffered encrypted reasoning into `thinkingSignature`. The chunks stay in the stream state, never on the + * block: `Uint8Array` is not JSON and serializes to an index-keyed object roughly ten times the size of the base64 + * payload. */ -function flushRedactedContent(block: Block): void { - if (block.type !== "thinking" || !block.redactedChunks) return; - block.thinkingSignature = bytesToBase64(block.redactedChunks); - delete block.redactedChunks; +function flushRedactedContent(open: OpenBlock, output: AssistantMessage): void { + const block = output.content[open.contentIndex]; + if (block?.type !== "thinking" || !open.redactedChunks) return; + block.thinkingSignature = bytesToBase64(open.redactedChunks); + open.redactedChunks = undefined; } -/** - * 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 { - delete block.index; - // partialJson is only a streaming scratch buffer; never persist it. - delete block.partialJson; - flushRedactedContent(block); +/** Ends the stream state of blocks a stream settled without stopping. */ +function flushOpenBlocks(openBlocks: OpenBlocks, output: AssistantMessage): void { + for (const open of openBlocks.values()) flushRedactedContent(open, output); + openBlocks.clear(); } function handleMetadata( @@ -735,28 +732,26 @@ function handleMetadata( function handleContentBlockStop( event: ContentBlockStopEvent, - blocks: Block[], + openBlocks: OpenBlocks, output: AssistantMessage, stream: AssistantMessageEventStream, ): void { - const index = blocks.findIndex((b) => b.index === event.contentBlockIndex); - const block = blocks[index]; - if (!block) return; - delete (block as Block).index; + const open = openBlocks.get(event.contentBlockIndex!); + if (!open) return; + openBlocks.delete(event.contentBlockIndex!); + const index = open.contentIndex; + const block = output.content[index]!; switch (block.type) { case "text": stream.push({ type: "text_end", contentIndex: index, content: block.text, partial: output }); break; case "thinking": - flushRedactedContent(block); + flushRedactedContent(open, output); stream.push({ type: "thinking_end", contentIndex: index, content: block.thinking, partial: output }); break; case "toolCall": - Object.assign(block, parseToolCallArguments(block.partialJson)); - // Finalize in-place and strip the scratch buffer so replay only - // carries parsed arguments. - delete (block as Block).partialJson; + if (open.arguments !== undefined) Object.assign(block, open.arguments.finish()); stream.push({ type: "toolcall_end", contentIndex: index, toolCall: block, partial: output }); break; } diff --git a/packages/ai/src/api/constrained-sampling.ts b/packages/ai/src/api/constrained-sampling.ts index 3f1cb262e..62e03a105 100644 --- a/packages/ai/src/api/constrained-sampling.ts +++ b/packages/ai/src/api/constrained-sampling.ts @@ -167,36 +167,51 @@ export function getGrammarToolInput( return input; } +/** + * Append `inputDelta` to a streaming grammar tool input and return the matching piece of its JSON arguments text, or + * `undefined` for an empty delta. Costs time linear in the delta, never in the input so far. + */ export function appendGrammarToolInputJsonDelta( buffer: GrammarToolInputJsonBuffer, inputProperty: string, - nextInput: string, - close: boolean, + inputDelta: string, ): string | undefined { if (buffer.closed) { - if (close && nextInput === buffer.input) return undefined; throw new Error(`grammar tool input for property "${inputProperty}" changed after it was closed`); } - if (!nextInput.startsWith(buffer.input)) { + if (inputDelta.length === 0) return undefined; + buffer.input += inputDelta; + return grammarToolInputJsonStart(buffer, inputProperty) + JSON.stringify(inputDelta).slice(1, -1); +} + +/** + * Close a streaming grammar tool input and return the rest of its JSON arguments text. `finalInput` is the complete + * input when the provider sends one; it must continue the input streamed so far. Closing again with the same input + * returns `undefined`. + */ +export function closeGrammarToolInputJson( + buffer: GrammarToolInputJsonBuffer, + inputProperty: string, + finalInput: string | undefined, +): string | undefined { + if (buffer.closed) { + if (finalInput === undefined || finalInput === buffer.input) return undefined; + throw new Error(`grammar tool input for property "${inputProperty}" changed after it was closed`); + } + if (finalInput !== undefined && !finalInput.startsWith(buffer.input)) { throw new Error(`grammar tool input for property "${inputProperty}" changed non-monotonically`); } + const rest = finalInput === undefined ? "" : finalInput.slice(buffer.input.length); + if (finalInput !== undefined) buffer.input = finalInput; + buffer.closed = true; + return `${grammarToolInputJsonStart(buffer, inputProperty)}${JSON.stringify(rest).slice(1, -1)}"}`; +} - const inputDelta = nextInput.slice(buffer.input.length); - if (!close && inputDelta.length === 0) return undefined; - - let delta = ""; - if (!buffer.started) { - delta += `{${JSON.stringify(inputProperty)}:"`; - buffer.started = true; - } - delta += JSON.stringify(inputDelta).slice(1, -1); - buffer.input = nextInput; - - if (close) { - delta += '"}'; - buffer.closed = true; - } - return delta; +/** The opening of the JSON arguments text, once. */ +function grammarToolInputJsonStart(buffer: GrammarToolInputJsonBuffer, inputProperty: string): string { + if (buffer.started) return ""; + buffer.started = true; + return `{${JSON.stringify(inputProperty)}:"`; } function inferGrammarInputProperty(tool: Tool): string { diff --git a/packages/ai/src/api/google-generative-ai.ts b/packages/ai/src/api/google-generative-ai.ts index 5e44ce2a4..213eff5e8 100644 --- a/packages/ai/src/api/google-generative-ai.ts +++ b/packages/ai/src/api/google-generative-ai.ts @@ -286,12 +286,6 @@ export const stream: StreamFunction<"google-generative-ai", GoogleOptions> = ( stream.push({ type: "done", reason: output.stopReason, message: output }); stream.end(); } catch (error) { - // Remove internal index property used during streaming - for (const block of output.content) { - if ("index" in block) { - delete (block as { index?: number }).index; - } - } output.stopReason = options?.signal?.aborted ? "aborted" : "error"; output.errorMessage = formatProviderError(normalizeProviderError(error)); stream.push({ type: "error", reason: output.stopReason, error: output }); diff --git a/packages/ai/src/api/google-vertex.ts b/packages/ai/src/api/google-vertex.ts index 1865cbceb..07f78afe8 100644 --- a/packages/ai/src/api/google-vertex.ts +++ b/packages/ai/src/api/google-vertex.ts @@ -301,12 +301,6 @@ export const stream: StreamFunction<"google-vertex", GoogleVertexOptions> = ( stream.push({ type: "done", reason: output.stopReason, message: output }); stream.end(); } catch (error) { - // Remove internal index property used during streaming - for (const block of output.content) { - if ("index" in block) { - delete (block as { index?: number }).index; - } - } output.stopReason = options?.signal?.aborted ? "aborted" : "error"; output.errorMessage = formatProviderError(normalizeProviderError(error)); stream.push({ type: "error", reason: output.stopReason, error: output }); diff --git a/packages/ai/src/api/mistral-conversations.ts b/packages/ai/src/api/mistral-conversations.ts index 9ac710399..69e3f1dd1 100644 --- a/packages/ai/src/api/mistral-conversations.ts +++ b/packages/ai/src/api/mistral-conversations.ts @@ -14,12 +14,13 @@ import type { TranscriptContext, } from "../types.ts"; import { combineAbortSignals } from "../utils/abort-signals.ts"; +import { truncateErrorText } from "../utils/error-body.ts"; import { AssistantMessageEventStream } from "../utils/event-stream.ts"; import { shortHash } from "../utils/hash.ts"; import { headersToRecord } from "../utils/headers.ts"; import { getIdentityHeaders } from "../utils/identity.ts"; -import { parseStreamingJson, parseToolCallArguments } from "../utils/json-parse.ts"; import { sanitizeSurrogates } from "../utils/sanitize-unicode.ts"; +import { ToolCallArgumentsStream } from "../utils/streaming-json.ts"; import { getSystemMessageText, renderSystemMessageUpdate } from "../utils/text.ts"; import { getCurrentTools, resolveTranscript } from "../utils/transcript.ts"; import { getJsonSchemaToolParameters, resolveJsonSchemaStrictSampling } from "./constrained-sampling.ts"; @@ -167,10 +168,6 @@ export const stream: StreamFunction<"mistral-conversations", MistralOptions> = ( stream.push({ type: "done", reason: output.stopReason, message: output }); stream.end(); } catch (error) { - for (const block of output.content) { - // partialArgs is only a streaming scratch buffer; never persist it. - delete (block as { partialArgs?: string }).partialArgs; - } output.stopReason = options?.signal?.aborted ? "aborted" : "error"; output.errorMessage = formatMistralError(error); stream.push({ type: "error", reason: output.stopReason, error: output }); @@ -281,11 +278,6 @@ function formatMistralError(error: unknown): string { return safeJsonStringify(error); } -function truncateErrorText(text: string, maxChars: number): string { - if (text.length <= maxChars) return text; - return `${text.slice(0, maxChars)}... [truncated ${text.length - maxChars} chars]`; -} - function safeJsonStringify(value: unknown): string { try { const serialized = JSON.stringify(value); @@ -583,7 +575,8 @@ async function consumeChatStream( let currentBlock: TextContent | ThinkingContent | null = null; const blocks = output.content; const blockIndex = () => blocks.length - 1; - const toolBlocksByKey = new Map(); + // Stream state per Mistral tool call, kept off the content blocks so it never reaches partials. + const toolCallsByKey = new Map(); const finishCurrentBlock = (block?: typeof currentBlock) => { if (!block) return; @@ -718,38 +711,22 @@ async function consumeChatStream( ? toolCall.id : deriveMistralToolCallId(`toolcall:${toolCall.index ?? 0}`, 0); const key = toolCall.index ?? callId; - const existingIndex = toolBlocksByKey.get(key); - let block: (ToolCall & { partialArgs?: string }) | undefined; - - if (existingIndex !== undefined) { - const existing = output.content[existingIndex]; - if (existing?.type === "toolCall") { - block = existing as ToolCall & { partialArgs?: string }; - } - } - - if (!block) { - block = { - type: "toolCall", - id: callId, - name: toolCall.function.name, - arguments: {}, - partialArgs: "", - }; - output.content.push(block); - toolBlocksByKey.set(key, output.content.length - 1); - stream.push({ type: "toolcall_start", contentIndex: output.content.length - 1, partial: output }); + let open = toolCallsByKey.get(key); + if (!open) { + output.content.push({ type: "toolCall", id: callId, name: toolCall.function.name, arguments: {} }); + open = { contentIndex: output.content.length - 1, arguments: new ToolCallArgumentsStream() }; + toolCallsByKey.set(key, open); + stream.push({ type: "toolcall_start", contentIndex: open.contentIndex, partial: output }); } const argsDelta = typeof toolCall.function.arguments === "string" ? toolCall.function.arguments : JSON.stringify(toolCall.function.arguments || {}); - block.partialArgs = (block.partialArgs || "") + argsDelta; - block.arguments = parseStreamingJson(block.partialArgs); + (output.content[open.contentIndex] as ToolCall).arguments = open.arguments.append(argsDelta); stream.push({ type: "toolcall_delta", - contentIndex: toolBlocksByKey.get(key)!, + contentIndex: open.contentIndex, delta: argsDelta, partial: output, }); @@ -757,18 +734,13 @@ async function consumeChatStream( } finishCurrentBlock(currentBlock); - for (const index of toolBlocksByKey.values()) { - const block = output.content[index]; - if (block.type !== "toolCall") continue; - const toolBlock = block as ToolCall & { partialArgs?: string }; - Object.assign(toolBlock, parseToolCallArguments(toolBlock.partialArgs)); - // Finalize in-place and strip the scratch buffer so replay only - // carries parsed arguments. - delete toolBlock.partialArgs; + for (const open of toolCallsByKey.values()) { + const block = output.content[open.contentIndex] as ToolCall; + Object.assign(block, open.arguments.finish()); stream.push({ type: "toolcall_end", - contentIndex: index, - toolCall: toolBlock, + contentIndex: open.contentIndex, + toolCall: block, partial: output, }); } diff --git a/packages/ai/src/api/openai-codex-responses.ts b/packages/ai/src/api/openai-codex-responses.ts index 90a599e66..5165e1ccf 100644 --- a/packages/ai/src/api/openai-codex-responses.ts +++ b/packages/ai/src/api/openai-codex-responses.ts @@ -497,11 +497,6 @@ export const stream: StreamFunction<"openai-codex-responses", OpenAICodexRespons stream.push({ type: "done", reason: output.stopReason, message: output }); stream.end(); } catch (error) { - for (const block of output.content) { - // Streaming scratch buffers are only used during parsing; never persist them. - delete (block as { partialJson?: string }).partialJson; - delete (block as { customInput?: unknown }).customInput; - } output.stopReason = options?.signal?.aborted ? "aborted" : "error"; output.errorMessage = formatProviderError(normalizeProviderError(error)); stream.push({ type: "error", reason: output.stopReason, error: output }); diff --git a/packages/ai/src/api/openai-completions.ts b/packages/ai/src/api/openai-completions.ts index a752d304e..0e5b606e6 100644 --- a/packages/ai/src/api/openai-completions.ts +++ b/packages/ai/src/api/openai-completions.ts @@ -40,7 +40,6 @@ import { AssistantMessageEventStream } from "../utils/event-stream.ts"; import { shortHash } from "../utils/hash.ts"; import { headersToRecord } from "../utils/headers.ts"; import { type ClientIdentity, getIdentityHeaders } from "../utils/identity.ts"; -import { parseStreamingJson, parseToolCallArguments } from "../utils/json-parse.ts"; import { applyOpenRouterReportedCost, isOpenRouterModel, @@ -49,6 +48,7 @@ import { import { getProviderEnvValue } from "../utils/provider-env.ts"; import { retryProviderRequest } from "../utils/provider-retry.ts"; import { sanitizeSurrogates } from "../utils/sanitize-unicode.ts"; +import { ToolCallArgumentsStream } from "../utils/streaming-json.ts"; import { getSystemMessageText, renderSystemMessageUpdate } from "../utils/text.ts"; import { getDeclaredTools, @@ -58,6 +58,7 @@ import { } from "../utils/transcript.ts"; import { appendGrammarToolInputJsonDelta, + closeGrammarToolInputJson, createGrammarToolInputProperties, type GrammarToolInputJsonBuffer, getGrammarToolInput, @@ -389,15 +390,13 @@ export const stream: StreamFunction<"openai-completions", OpenAICompletionsOptio await options?.onResponse?.({ status: response.status, headers: headersToRecord(response.headers) }, model); stream.push({ type: "start", partial: output }); - interface StreamingToolCallBlock extends ToolCall { - partialArgs?: string; - customInput?: { - property: string; - jsonBuffer: GrammarToolInputJsonBuffer; - }; - streamIndex?: number; - } - type StreamingBlock = TextContent | ThinkingContent | StreamingToolCallBlock; + // Stream state of a tool call, kept off the content block so it never reaches partials: JSON arguments, or + // the input of a grammar (custom) tool. + type ToolCallInput = + | { readonly kind: "json"; readonly arguments: ToolCallArgumentsStream } + | { readonly kind: "custom"; readonly property: string; readonly buffer: GrammarToolInputJsonBuffer }; + type ToolCallState = { streamIndex: number | undefined; input: ToolCallInput }; + type StreamingBlock = TextContent | ThinkingContent | ToolCall; type StreamingToolCallDelta = { index?: number; id?: string; @@ -409,32 +408,16 @@ export const stream: StreamFunction<"openai-completions", OpenAICompletionsOptio let textBlock: TextContent | null = null; let thinkingBlock: ThinkingContent | null = null; let hasFinishReason = false; - const toolCallBlocksByIndex = new Map(); - const toolCallBlocksById = new Map(); + const toolCallBlocksByIndex = new Map(); + const toolCallBlocksById = new Map(); + const toolCallStates = new Map(); const blocks = output.content as StreamingBlock[]; const getContentIndex = (block: StreamingBlock) => blocks.indexOf(block); - const getCustomToolCallInput = (block: StreamingToolCallBlock): string => { - const property = block.customInput?.property; - if (property === undefined) return ""; - const value = block.arguments[property]; - return typeof value === "string" ? value : ""; - }; - const appendCustomToolCallInput = ( - block: StreamingToolCallBlock, - nextInput: string, - close: boolean, - ): string | undefined => { - const customInput = block.customInput; - if (!customInput) return undefined; - const delta = appendGrammarToolInputJsonDelta( - customInput.jsonBuffer, - customInput.property, - nextInput, - close, - ); - block.arguments = { [customInput.property]: nextInput }; - return delta; - }; + const customInput = (property: string): ToolCallInput => ({ + kind: "custom", + property, + buffer: { input: "", started: false, closed: false }, + }); const finishBlock = (block: StreamingBlock) => { const contentIndex = getContentIndex(block); if (contentIndex === -1) { @@ -456,8 +439,10 @@ export const stream: StreamFunction<"openai-completions", OpenAICompletionsOptio partial: output, }); } else if (block.type === "toolCall") { - if (block.customInput) { - const delta = appendCustomToolCallInput(block, getCustomToolCallInput(block), true); + const input = toolCallStates.get(block)!.input; + toolCallStates.delete(block); + if (input.kind === "custom") { + const delta = closeGrammarToolInputJson(input.buffer, input.property, undefined); if (delta !== undefined) { stream.push({ type: "toolcall_delta", @@ -467,13 +452,8 @@ export const stream: StreamFunction<"openai-completions", OpenAICompletionsOptio }); } } else { - Object.assign(block, parseToolCallArguments(block.partialArgs)); + Object.assign(block, input.arguments.finish()); } - // Finalize in-place and strip the scratch buffers so replay only - // carries parsed arguments. - delete block.partialArgs; - delete block.customInput; - delete block.streamIndex; stream.push({ type: "toolcall_end", contentIndex, @@ -514,18 +494,19 @@ export const stream: StreamFunction<"openai-completions", OpenAICompletionsOptio // a tool we don't knwo about, we at least have a place to stash our stuff. const customInputProperty = toolCall.custom && !toolCall.function ? (grammarToolInputProperties.get(name) ?? "input") : undefined; - const hasCustomInput = customInputProperty !== undefined; block = { type: "toolCall", id: toolCall.id || "", name, - arguments: hasCustomInput ? { [customInputProperty]: "" } : {}, - partialArgs: hasCustomInput ? undefined : "", - customInput: hasCustomInput - ? { property: customInputProperty, jsonBuffer: { input: "", started: false, closed: false } } - : undefined, - streamIndex, + arguments: customInputProperty !== undefined ? { [customInputProperty]: "" } : {}, }; + toolCallStates.set(block, { + streamIndex, + input: + customInputProperty !== undefined + ? customInput(customInputProperty) + : { kind: "json", arguments: new ToolCallArgumentsStream() }, + }); if (streamIndex !== undefined) { toolCallBlocksByIndex.set(streamIndex, block); } @@ -539,8 +520,9 @@ export const stream: StreamFunction<"openai-completions", OpenAICompletionsOptio partial: output, }); } - if (streamIndex !== undefined && block.streamIndex === undefined) { - block.streamIndex = streamIndex; + const state = toolCallStates.get(block)!; + if (streamIndex !== undefined && state.streamIndex === undefined) { + state.streamIndex = streamIndex; toolCallBlocksByIndex.set(streamIndex, block); } if (toolCall.id) { @@ -549,16 +531,12 @@ export const stream: StreamFunction<"openai-completions", OpenAICompletionsOptio if (!block.name && name) { block.name = name; } - if (toolCall.custom && !toolCall.function && !block.customInput) { + if (toolCall.custom && !toolCall.function && state.input.kind !== "custom") { const customInputProperty = grammarToolInputProperties.get(block.name) ?? "input"; block.arguments = { [customInputProperty]: "" }; - block.customInput = { - property: customInputProperty, - jsonBuffer: { input: "", started: false, closed: false }, - }; - delete block.partialArgs; + state.input = customInput(customInputProperty); } - return block; + return { block, state }; }; for await (const chunk of openaiStream) { @@ -660,7 +638,7 @@ export const stream: StreamFunction<"openai-completions", OpenAICompletionsOptio if (choice?.delta?.tool_calls) { for (const toolCall of choice.delta.tool_calls as StreamingToolCallDelta[]) { - const block = ensureToolCallBlock(toolCall); + const { block, state } = ensureToolCallBlock(toolCall); if (!block.id && toolCall.id) { block.id = toolCall.id; toolCallBlocksById.set(toolCall.id, block); @@ -671,13 +649,14 @@ export const stream: StreamFunction<"openai-completions", OpenAICompletionsOptio } let delta = ""; - if (toolCall.function?.arguments) { + const input = state.input; + if (toolCall.function?.arguments && input.kind === "json") { delta = toolCall.function.arguments; - block.partialArgs = (block.partialArgs ?? "") + toolCall.function.arguments; - block.arguments = parseStreamingJson(block.partialArgs); - } else if (toolCall.custom?.input) { - const nextInput = getCustomToolCallInput(block) + toolCall.custom.input; - delta = appendCustomToolCallInput(block, nextInput, false) ?? ""; + block.arguments = input.arguments.append(delta); + } else if (toolCall.custom?.input && input.kind === "custom") { + delta = + appendGrammarToolInputJsonDelta(input.buffer, input.property, toolCall.custom.input) ?? ""; + block.arguments = { [input.property]: input.buffer.input }; } stream.push({ type: "toolcall_delta", @@ -717,11 +696,6 @@ export const stream: StreamFunction<"openai-completions", OpenAICompletionsOptio if (block.type === "thinking") { applyStreamedReasoningDetails(block); } - delete (block as { index?: number }).index; - // Streaming scratch buffers are only used during parsing; never persist them. - delete (block as { partialArgs?: string }).partialArgs; - delete (block as { customInput?: unknown }).customInput; - delete (block as { streamIndex?: number }).streamIndex; } output.stopReason = options?.signal?.aborted ? "aborted" : "error"; output.errorMessage = formatProviderError(normalizeProviderError(error)); diff --git a/packages/ai/src/api/openai-responses-shared.ts b/packages/ai/src/api/openai-responses-shared.ts index 3f6371053..ecc104043 100644 --- a/packages/ai/src/api/openai-responses-shared.ts +++ b/packages/ai/src/api/openai-responses-shared.ts @@ -33,12 +33,13 @@ import type { } from "../types.ts"; import type { AssistantMessageEventStream } from "../utils/event-stream.ts"; import { shortHash } from "../utils/hash.ts"; -import { parseStreamingJson, parseToolCallArguments } from "../utils/json-parse.ts"; import { sanitizeSurrogates } from "../utils/sanitize-unicode.ts"; +import { ToolCallArgumentsStream } from "../utils/streaming-json.ts"; import { getSystemMessageText, renderSystemMessageUpdate } from "../utils/text.ts"; import { resolveTranscript, resolveTranscriptTools, type TranscriptToolSupport } from "../utils/transcript.ts"; import { appendGrammarToolInputJsonDelta, + closeGrammarToolInputJson, type GrammarToolInputJsonBuffer, getGrammarToolInput, getJsonSchemaToolParameters, @@ -478,33 +479,24 @@ export function convertResponsesTools(tools: readonly Tool[], options?: ConvertR // Stream processing // ============================================================================= -type StreamingToolCall = ToolCall & { - partialJson?: string; - customInput?: { - property: string; - jsonBuffer: GrammarToolInputJsonBuffer; - }; -}; - -function getCustomToolCallInput(block: StreamingToolCall): string { - const property = block.customInput?.property; - if (property === undefined) return ""; - const value = block.arguments[property]; - return typeof value === "string" ? value : ""; -} - -function appendCustomToolCallInput(block: StreamingToolCall, nextInput: string, close: boolean): string | undefined { - const customInput = block.customInput; - if (!customInput) return undefined; - const delta = appendGrammarToolInputJsonDelta(customInput.jsonBuffer, customInput.property, nextInput, close); - block.arguments = { [customInput.property]: nextInput }; - return delta; -} +/** + * Stream state of a tool call, kept in its output slot and never on the content block, so it never reaches partials: + * the JSON arguments of a function call, or the input of a custom (grammar) tool call. + */ +type ToolCallInput = + | { readonly kind: "json"; readonly arguments: ToolCallArgumentsStream } + | { + readonly kind: "custom"; + readonly property: string; + readonly buffer: GrammarToolInputJsonBuffer; + /** JSON text of the item's initial input, sent with the first delta. */ + unsent: string; + }; type ResponsesOutputSlot = | { type: "thinking"; block: ThinkingContent; contentIndex: number } | { type: "text"; block: TextContent; contentIndex: number } - | { type: "toolCall"; block: StreamingToolCall; contentIndex: number }; + | { type: "toolCall"; block: ToolCall; contentIndex: number; input: ToolCallInput }; type ToolCallOutputSlot = Extract; @@ -517,6 +509,8 @@ export async function processResponsesStream( ): Promise { let sawTerminalResponseEvent = false; const outputSlots = new Map(); + // Tool calls whose output_item.done has not arrived yet. + const unfinishedToolCalls = new Set(); const reasoningBlocksById = new Map(); const applyMessagePhaseStopReason = (item: ResponseOutputItem): void => { if (item.type === "message" && item.phase === "final_answer") { @@ -562,19 +556,21 @@ export async function processResponsesStream( return slot; } if (item.type === "function_call") { - const block: StreamingToolCall = { + const args = new ToolCallArgumentsStream(); + const block: ToolCall = { type: "toolCall", id: `${item.call_id}|${item.id}`, name: item.name, - arguments: {}, + arguments: item.arguments ? args.append(item.arguments) : {}, ...(item.namespace !== undefined ? { namespace: item.namespace } : {}), - partialJson: item.arguments || "", }; output.content.push(block); + unfinishedToolCalls.add(block); const slot = { type: "toolCall", block, contentIndex: output.content.length - 1, + input: { kind: "json", arguments: args }, } satisfies ResponsesOutputSlot; outputSlots.set(outputIndex, slot); stream.push({ type: "toolcall_start", contentIndex: slot.contentIndex, partial: output }); @@ -583,22 +579,20 @@ 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 block: ToolCall = { type: "toolCall", id: `${item.call_id}|${item.id}`, name: item.name, arguments: { [inputProperty]: input }, ...(item.namespace !== undefined ? { namespace: item.namespace } : {}), - customInput: { - property: inputProperty, - jsonBuffer: { input: "", started: false, closed: false }, - }, }; output.content.push(block); + unfinishedToolCalls.add(block); const slot = { type: "toolCall", block, contentIndex: output.content.length - 1, + input: customInput(inputProperty, input), } satisfies ResponsesOutputSlot; outputSlots.set(outputIndex, slot); stream.push({ type: "toolcall_start", contentIndex: slot.contentIndex, partial: output }); @@ -732,32 +726,23 @@ export async function processResponsesStream( }); } 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); + if (slot?.input.kind !== "json") continue; + slot.block.arguments = slot.input.arguments.append(event.delta); 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); - } + if (slot?.input.kind !== "json") continue; + const replaced = slot.input.arguments.replace(event.arguments); + slot.block.arguments = replaced.arguments; + if (replaced.delta) pushToolCallDelta(slot, replaced.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), - ); + if (slot?.input.kind !== "custom") continue; + pushToolCallDelta(slot, appendCustomInput(slot.block, slot.input, event.delta)); } 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)); + if (slot?.input.kind !== "custom") continue; + pushToolCallDelta(slot, closeCustomInput(slot.block, slot.input, event.input)); } else if (event.type === "response.output_item.done") { const item = event.item; applyMessagePhaseStopReason(item); @@ -786,16 +771,11 @@ export async function processResponsesStream( partial: output, }); outputSlots.delete(event.output_index); - } else if ( - item.type === "function_call" && - slot?.type === "toolCall" && - slot.block.partialJson !== undefined - ) { - Object.assign(slot.block, parseToolCallArguments(item.arguments || slot.block.partialJson)); + } else if (item.type === "function_call" && slot?.type === "toolCall" && slot.input.kind === "json") { + if (item.arguments) slot.input.arguments.replace(item.arguments); + Object.assign(slot.block, slot.input.arguments.finish()); 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; + unfinishedToolCalls.delete(slot.block); stream.push({ type: "toolcall_end", contentIndex: slot.contentIndex, @@ -803,13 +783,10 @@ export async function processResponsesStream( 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), - ); + } else if (item.type === "custom_tool_call" && slot?.type === "toolCall" && slot.input.kind === "custom") { + pushToolCallDelta(slot, closeCustomInput(slot.block, slot.input, item.input)); if (item.namespace !== undefined) slot.block.namespace = item.namespace; - delete slot.block.customInput; + unfinishedToolCalls.delete(slot.block); stream.push({ type: "toolcall_end", contentIndex: slot.contentIndex, @@ -840,20 +817,47 @@ export async function processResponsesStream( } // The agent runs every tool call in the final message. Refuse to hand over calls whose // output_item.done never arrived: their arguments may be cut off or mixed up, e.g. when a - // non-compliant server omits output_index. Finished calls have their scratch buffers removed. + // non-compliant server omits output_index. if (output.stopReason === "toolUse") { - for (const block of output.content) { - if (block.type !== "toolCall") continue; - const toolCall = block as StreamingToolCall; - if (toolCall.partialJson !== undefined || toolCall.customInput !== undefined) { - throw new Error( - `OpenAI Responses stream completed with an unfinished tool call: ${toolCall.name} (${toolCall.id})`, - ); - } + for (const toolCall of unfinishedToolCalls) { + throw new Error( + `OpenAI Responses stream completed with an unfinished tool call: ${toolCall.name} (${toolCall.id})`, + ); } } } +type CustomInput = Extract; + +/** Stream state of a custom tool call that starts with `initialInput`. */ +function customInput(property: string, initialInput: string): CustomInput { + const buffer = { input: "", started: false, closed: false }; + return { + kind: "custom", + property, + buffer, + unsent: appendGrammarToolInputJsonDelta(buffer, property, initialInput) ?? "", + }; +} + +/** Append to a streaming custom tool input; returns the matching piece of its JSON arguments text. */ +function appendCustomInput(block: ToolCall, input: CustomInput, delta: string): string | undefined { + return sendCustomInput(block, input, appendGrammarToolInputJsonDelta(input.buffer, input.property, delta)); +} + +/** Close a custom tool input with the provider's complete input, if it sent one. */ +function closeCustomInput(block: ToolCall, input: CustomInput, finalInput: string | undefined): string | undefined { + return sendCustomInput(block, input, closeGrammarToolInputJson(input.buffer, input.property, finalInput)); +} + +function sendCustomInput(block: ToolCall, input: CustomInput, jsonDelta: string | undefined): string | undefined { + block.arguments = { [input.property]: input.buffer.input }; + if (input.unsent === "") return jsonDelta; + const sent = input.unsent + (jsonDelta ?? ""); + input.unsent = ""; + return sent; +} + function mapStopReason( status: OpenAI.Responses.ResponseStatus | undefined, incompleteReason?: string, diff --git a/packages/ai/src/api/openai-responses.ts b/packages/ai/src/api/openai-responses.ts index 8ddaa6944..c29e6adfb 100644 --- a/packages/ai/src/api/openai-responses.ts +++ b/packages/ai/src/api/openai-responses.ts @@ -249,12 +249,6 @@ export const stream: StreamFunction<"openai-responses", OpenAIResponsesOptions> stream.push({ type: "done", reason: output.stopReason, message: output }); stream.end(); } catch (error) { - for (const block of output.content) { - delete (block as { index?: number }).index; - // Streaming scratch buffers are only used during parsing; never persist them. - delete (block as { partialJson?: string }).partialJson; - delete (block as { customInput?: unknown }).customInput; - } output.stopReason = options?.signal?.aborted ? "aborted" : "error"; const errorMessage = formatProviderError( normalizeProviderError(error), diff --git a/packages/ai/src/api/pi-messages.ts b/packages/ai/src/api/pi-messages.ts index 83ddc06c7..59b54567e 100644 --- a/packages/ai/src/api/pi-messages.ts +++ b/packages/ai/src/api/pi-messages.ts @@ -28,8 +28,9 @@ import { appendAssistantMessageDiagnostic, createAssistantMessageDiagnostic } fr import { AssistantMessageEventStream } from "../utils/event-stream.ts"; import { headersToRecord, providerHeadersToRecord } from "../utils/headers.ts"; import { getIdentityHeaders } from "../utils/identity.ts"; -import { parseStreamingJson } from "../utils/json-parse.ts"; import { getProviderEnvValue } from "../utils/provider-env.ts"; +import { ToolCallArgumentsStream } from "../utils/streaming-json.ts"; +import { ownedSlice } from "../utils/text.ts"; export interface PiMessagesOptions extends StreamOptions { reasoning?: ThinkingLevel; @@ -122,7 +123,7 @@ function parsePiMessagesErrorBody(body: string): PiMessagesErrorBody | undefined function truncateDiagnosticString(value: string): string { const maxLength = 8192; - return value.length > maxLength ? `${value.slice(0, maxLength)}…` : value; + return value.length > maxLength ? `${ownedSlice(value, 0, maxLength)}…` : value; } function formatPiMessagesResponseError( @@ -191,7 +192,7 @@ function createEventConverter(model: Model<"pi-messages">) { stopReason: "pending", timestamp: Date.now(), }; - const toolJson = new Map(); + const toolArguments = new Map(); return (event: PiMessagesEvent): AssistantMessageEvent => { switch (event.type) { @@ -258,18 +259,17 @@ function createEventConverter(model: Model<"pi-messages">) { name: event.toolName, arguments: {}, }; - toolJson.set(event.contentIndex, ""); + toolArguments.set(event.contentIndex, new ToolCallArgumentsStream()); 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); + const args = toolArguments.get(event.contentIndex); + if (args !== undefined) + (partial.content[event.contentIndex] as ToolCall).arguments = args.append(event.delta); break; } case "toolcall_end": Object.assign(partial.content[event.contentIndex]!, event.toolCall); - toolJson.delete(event.contentIndex); + toolArguments.delete(event.contentIndex); return { type: "toolcall_end", contentIndex: event.contentIndex, diff --git a/packages/ai/src/index.ts b/packages/ai/src/index.ts index d58b2a259..54a0345f1 100644 --- a/packages/ai/src/index.ts +++ b/packages/ai/src/index.ts @@ -41,6 +41,7 @@ export * from "./utils/identity.ts"; export * from "./utils/json-parse.ts"; export * from "./utils/overflow.ts"; export * from "./utils/retry.ts"; +export * from "./utils/streaming-json.ts"; export { contentText, getSystemMessageText, renderSystemMessageUpdate } from "./utils/text.ts"; export * from "./utils/transcript.ts"; export * from "./utils/typebox-helpers.ts"; diff --git a/packages/ai/src/providers/faux.ts b/packages/ai/src/providers/faux.ts index 0f4188e60..87797dbfa 100644 --- a/packages/ai/src/providers/faux.ts +++ b/packages/ai/src/providers/faux.ts @@ -19,6 +19,7 @@ import type { Usage, } from "../types.ts"; import { createAssistantMessageEventStream } from "../utils/event-stream.ts"; +import { ToolCallArgumentsStream } from "../utils/streaming-json.ts"; import { getSystemMessageText } from "../utils/text.ts"; const DEFAULT_API = "faux"; @@ -425,6 +426,8 @@ async function streamWithDeltas( partial.content = [...partial.content, { type: "toolCall", id: block.id, name: block.name, arguments: {} }]; stream.push({ type: "toolcall_start", contentIndex: index, partial: { ...partial } }); + // Partial arguments come from the same parser real providers use, so consumers see them grow. + const toolArguments = new ToolCallArgumentsStream(); for (const chunk of splitStringByTokenSize(JSON.stringify(block.arguments), minTokenSize, maxTokenSize)) { await scheduleChunk(chunk, tokensPerSecond); if (signal?.aborted) { @@ -433,6 +436,7 @@ async function streamWithDeltas( stream.end(aborted); return; } + (partial.content[index] as ToolCall).arguments = toolArguments.append(chunk); stream.push({ type: "toolcall_delta", contentIndex: index, delta: chunk, partial: { ...partial } }); } (partial.content[index] as ToolCall).arguments = block.arguments; diff --git a/packages/ai/src/providers/radius-config.ts b/packages/ai/src/providers/radius-config.ts index a2715aaec..4bf82cf9f 100644 --- a/packages/ai/src/providers/radius-config.ts +++ b/packages/ai/src/providers/radius-config.ts @@ -1,5 +1,6 @@ import type { OAuthCredential } from "../auth/types.ts"; import type { Model, ThinkingLevelMap } from "../types.ts"; +import { ownedSlice } from "../utils/text.ts"; export const DEFAULT_RADIUS_GATEWAY = "https://radius.pi.dev"; @@ -74,7 +75,7 @@ export function getRadiusModels(providerId: string, credential: OAuthCredential function truncateHttpBody(body: string): string { const trimmed = body.trim(); - return trimmed.length > 512 ? `${trimmed.slice(0, 512)}…` : trimmed; + return trimmed.length > 512 ? `${ownedSlice(trimmed, 0, 512)}…` : trimmed; } export async function loadRadiusGatewayConfig( diff --git a/packages/ai/src/telemetry.ts b/packages/ai/src/telemetry.ts index 683339939..7969bfc23 100644 --- a/packages/ai/src/telemetry.ts +++ b/packages/ai/src/telemetry.ts @@ -10,6 +10,7 @@ import { } from "@earendil-works/pi-telemetry"; import type { LazyStreamObserver } from "./api/lazy.ts"; import type { Api, Model } from "./types.ts"; +import { ownedSlice } from "./utils/text.ts"; const provider = { type: "string", description: "Provider identifier.", required: true } as const; const model = { type: "string", description: "Requested model identifier.", required: true } as const; @@ -191,7 +192,7 @@ export function observeModelRequest( if (message.stopReason === "error") { span.setStatus({ status: "error", - error: { name: "provider_error", message: message.errorMessage?.slice(0, 500) ?? "" }, + error: { name: "provider_error", message: ownedSlice(message.errorMessage ?? "", 0, 500) }, }); } requestDuration.record(performance.now() - started, { ...attributes, stop: message.stopReason }, span); diff --git a/packages/ai/src/utils/error-body.ts b/packages/ai/src/utils/error-body.ts index cb8c1859a..d35f51a0e 100644 --- a/packages/ai/src/utils/error-body.ts +++ b/packages/ai/src/utils/error-body.ts @@ -13,6 +13,8 @@ // Anthropic / `@google/genai` happy path where the SDK already folded the body // into the message, so providers can preserve it without double-printing. +import { ownedSlice } from "./text.ts"; + export const MAX_PROVIDER_ERROR_BODY_CHARS = 4000; export interface NormalizedProviderError { @@ -136,7 +138,7 @@ export function formatProviderError(norm: NormalizedProviderError, prefix?: stri export function truncateErrorText(text: string, maxChars: number): string { if (text.length <= maxChars) return text; - return `${text.slice(0, maxChars)}... [truncated ${text.length - maxChars} chars]`; + return `${ownedSlice(text, 0, maxChars)}... [truncated ${text.length - maxChars} chars]`; } export function safeJsonStringify(value: unknown): string { diff --git a/packages/ai/src/utils/streaming-json.ts b/packages/ai/src/utils/streaming-json.ts new file mode 100644 index 000000000..9f71ae7ec --- /dev/null +++ b/packages/ai/src/utils/streaming-json.ts @@ -0,0 +1,559 @@ +/** + * Incremental, tolerant parsing of JSON text that arrives in pieces, such as streamed tool-call arguments. + * + * `StreamingJsonParser.append()` reads each piece once, so a stream costs time and memory linear in its length, and + * `snapshot()` returns the value of the text read so far. Snapshots are plain JSON values and never change once + * returned: the parser copies the containers that are still open (the "open path") into every new snapshot and shares + * completed containers and strings. The rules for incomplete text follow `parseStreamingJson()`: + * + * - An unterminated string shows the characters read so far; an unterminated escape is held back until complete. + * - An unterminated container shows the members and elements read so far. + * - A member whose key or value has not started is omitted. + * - An unterminated literal shows its value once it is a prefix of `true`, `false` or `null`; an unterminated number + * shows its value when the digits read so far form a number (an exponent without digits is ignored). + * + * Malformed text is read leniently, the way `repairJson()` and `parseStreamingJson()` treat it: raw control characters + * in strings are kept, an invalid escape such as `\d` stays literal text, missing commas and colons are tolerated, + * values in invalid positions and unknown words are dropped, and text after the root value is ignored. The parser never + * throws. Callers that need the exact value of the finished text parse it once, e.g. with `parseToolCallArguments()`. + * + * Differences from `parseStreamingJson()` on incomplete text: trailing whitespace inside an unterminated string is + * kept (it was trimmed), a string that contains a raw control character or an invalid escape stays visible while it + * streams (it was dropped until the document was complete), and `NaN` and `Infinity` words are dropped. + * + * `append()` also returns the changes the piece made, for a writer that keeps its own copy of the value and applies + * them instead of re-assigning snapshots: `set` when a key, an array item or a finished value appears (a new string is + * set once with its first text, a container empty), `append` when the piece extended an open string. Unterminated + * numbers and literals are not reported until they end. A container nested more than 32 levels deep is reported as one + * `set` of its whole value when it closes, so paths stay short. + * + * Cost: copying the open path costs its member and element count. Each `append()` earns copy credit proportional to its + * length plus a constant; `snapshot()` spends it and returns the previous snapshot while the open path costs more than + * the credit. So a document that keeps a very large container open (thousands of elements) shows a preview that lags + * a little behind the text, and no input makes the total work grow faster than the text. + */ + +import type { JsonValue, ToolCall } from "../types.ts"; +import { parseToolCallArguments } from "./json-parse.ts"; + +type JsonContainer = { [key: string]: JsonValue } | JsonValue[]; + +/** Where a value sits in its parent; absent for a value in an invalid position, which is read but dropped. */ +type Slot = + | { readonly kind: "root" } + | { readonly kind: "member"; readonly key: string } + | { readonly kind: "element"; readonly index: number }; + +/** A path from the root value: member keys and array indices. */ +export type JsonPath = readonly (string | number)[]; + +/** One change a piece of text made to the value; see the module comment. */ +export type JsonChange = + | { readonly kind: "set"; readonly path: JsonPath; readonly value: JsonValue } + | { readonly kind: "append"; readonly path: JsonPath; readonly text: string }; + +/** A change still being extended by the piece being read. */ +type OpenChange = + | { kind: "set"; readonly path: JsonPath; value: JsonValue } + | { kind: "append"; readonly path: JsonPath; text: string }; + +/** An open object or array. Its `container` is never handed out while open; snapshots get copies. */ +type Frame = { + readonly container: JsonContainer; + readonly slot: Slot | undefined; + /** Where the container sits, when its changes are reported. */ + readonly path: JsonPath | undefined; + /** Whether changes inside it are reported; when not, its whole value is reported as it closes. */ + readonly reportsInside: boolean; + /** Members or elements; what a copy of the container costs. */ + size: number; + /** Objects: the key read last, where the member value goes. */ + key: string | undefined; + phase: "key" | "colon" | "value" | "comma"; +}; + +/** An unterminated string. */ +type OpenString = { + readonly target: "key" | "value"; + readonly slot: Slot | undefined; + text: string; + /** Escape read so far: "" for none, otherwise a backslash and up to five more characters. */ + escape: string; + /** Where a value string sits, when its changes are reported. */ + readonly path: JsonPath | undefined; + /** The change this string extends, while it is the newest change of the piece. */ + change: OpenChange | undefined; +}; + +// Changes are reported at most this deep, so a path costs O(1) to build. +const MAX_REPORTED_DEPTH = 32; + +const ROOT_SLOT: Slot = { kind: "root" }; +// Copy credit per append and per appended character. The floor keeps small documents fresh at every snapshot. +const CREDIT_PER_APPEND = 256; +const CREDIT_PER_CHARACTER = 8; +// An unterminated number longer than this is not shown until it ends, so `snapshot()` stays cheap. +const MAX_PENDING_SCALAR_LENGTH = 64; +const NUMBER = /^-?(?:0|[1-9]\d*)(?:\.\d+)?(?:[eE][+-]?\d+)?$/; +const NUMBER_WITHOUT_EXPONENT_DIGITS = /^(-?(?:0|[1-9]\d*)(?:\.\d+)?)[eE][+-]?$/; +const LITERALS: readonly [string, JsonValue][] = [ + ["true", true], + ["false", false], + ["null", null], +]; + +const QUOTE = 0x22; +const BACKSLASH = 0x5c; + +/** Incremental, tolerant JSON parser; see the module comment for its rules and costs. */ +export class StreamingJsonParser { + #root: JsonValue | undefined = undefined; + #done = false; + readonly #stack: Frame[] = []; + #string: OpenString | undefined = undefined; + #scalar = ""; + #scalarSlot: Slot | undefined = undefined; + /** Sum of `size + 1` over the open path: what a fresh snapshot copies. */ + #pathCost = 0; + #credit = 0; + #dirty = false; + #snapshot: JsonValue | undefined = undefined; + /** Changes of the piece being read. */ + #changes: OpenChange[] = []; + + /** + * Read the next piece of text; returns the changes it made, in order. Costs time linear in `text.length`. + */ + append(text: string): readonly JsonChange[] { + this.#credit += CREDIT_PER_APPEND + CREDIT_PER_CHARACTER * text.length; + this.#changes = []; + if (this.#string !== undefined) this.#string.change = undefined; + let index = 0; + while (index < text.length && !this.#done) { + const open = this.#string; + if (open !== undefined) index = this.#readString(open, text, index); + else if (this.#scalar !== "") index = this.#readScalar(text, index); + else this.#structural(text[index++]!); + } + return this.#changes; + } + + /** + * The value of the text read so far: `undefined` before a value starts, otherwise a plain JSON value that never + * changes afterwards. Returns the previous snapshot while the text read since did not change the value, or while a + * fresh copy of the open path is not yet paid for (see the module comment). + */ + snapshot(): JsonValue | undefined { + if (!this.#dirty) return this.#snapshot; + const pending = this.#pendingScalar(); + const cost = this.#pathCost + (pending === undefined ? 0 : 1); + if (cost > this.#credit) return this.#snapshot; + this.#credit -= cost; + this.#dirty = false; + this.#snapshot = this.#copyOpenPath(pending); + return this.#snapshot; + } + + /** Read string characters from `start` until the string ends or the text does; returns the next index. */ + #readString(open: OpenString, text: string, start: number): number { + let runStart = start; + for (let index = start; index < text.length; index++) { + const code = text.charCodeAt(index); + if (open.escape !== "") { + if (this.#continueEscape(open, text[index]!)) { + runStart = index + 1; + continue; + } + // Not part of the escape: read it as a normal string character below. + runStart = index; + } + if (code === QUOTE) { + if (index > runStart) this.#appendString(open, text.slice(runStart, index)); + this.#endString(open); + return index + 1; + } + if (code === BACKSLASH) { + if (index > runStart) this.#appendString(open, text.slice(runStart, index)); + open.escape = "\\"; + runStart = index + 1; + } + } + if (text.length > runStart) this.#appendString(open, text.slice(runStart)); + return text.length; + } + + /** Read number or literal characters from `start` up to the next delimiter; returns the delimiter's index. */ + #readScalar(text: string, start: number): number { + let end = start; + while (end < text.length && !isDelimiter(text[end]!)) end++; + if (end > start) { + this.#scalar += text.slice(start, end); + if (this.#scalarSlot !== undefined) this.#dirty = true; + } + if (end < text.length) this.#endScalar(); + return end; + } + + /** Handle one character outside strings. */ + #structural(char: string): void { + switch (char) { + case " ": + case "\t": + case "\n": + case "\r": + return; + case '"': + this.#startString(); + return; + case "{": + case "[": + this.#openContainer(char === "{" ? {} : []); + return; + case "}": + case "]": + this.#closeContainer(); + return; + case ":": { + const frame = this.#stack.at(-1); + if (frame !== undefined && !Array.isArray(frame.container) && frame.phase === "colon") + frame.phase = "value"; + return; + } + case ",": { + const frame = this.#stack.at(-1); + if (frame !== undefined) frame.phase = Array.isArray(frame.container) ? "value" : "key"; + return; + } + default: + this.#scalar = char; + this.#scalarSlot = this.#valueSlot(); + if (this.#scalarSlot !== undefined) this.#dirty = true; + } + } + + /** Where a value starting now goes, or `undefined` when this position takes no value. */ + #valueSlot(): Slot | undefined { + const frame = this.#stack.at(-1); + if (frame === undefined) return this.#root === undefined ? ROOT_SLOT : undefined; + if (Array.isArray(frame.container)) return { kind: "element", index: frame.container.length }; + // A missing colon is tolerated: `{"a" 1}` reads as `{"a": 1}`. + if ((frame.phase === "value" || frame.phase === "colon") && frame.key !== undefined) { + return { kind: "member", key: frame.key }; + } + return undefined; + } + + #startString(): void { + const frame = this.#stack.at(-1); + // A missing comma is tolerated: `{"a": 1 "b": 2}` reads the second member. + if ( + frame !== undefined && + !Array.isArray(frame.container) && + (frame.phase === "key" || frame.phase === "comma") + ) { + this.#string = { target: "key", slot: undefined, text: "", escape: "", path: undefined, change: undefined }; + return; + } + const slot = this.#valueSlot(); + const path = this.#reportedPath(slot); + this.#string = { target: "value", slot, text: "", escape: "", path, change: undefined }; + this.#place(slot, ""); + if (path !== undefined) this.#string.change = this.#report({ kind: "set", path, value: "" }); + } + + /** The path of a value at `slot` in the innermost open container, when its changes are reported. */ + #reportedPath(slot: Slot | undefined): JsonPath | undefined { + if (slot === undefined) return undefined; + if (slot.kind === "root") return []; + const frame = this.#stack.at(-1)!; + if (!frame.reportsInside) return undefined; + return [...frame.path!, slot.kind === "member" ? slot.key : slot.index]; + } + + #report(change: OpenChange): OpenChange { + this.#changes.push(change); + return change; + } + + /** Continue an escape with `char`; returns false when `char` is not part of it and must be read normally. */ + #continueEscape(open: OpenString, char: string): boolean { + if (open.escape === "\\") { + const decoded = SIMPLE_ESCAPES[char]; + if (decoded !== undefined) { + open.escape = ""; + this.#appendString(open, decoded); + return true; + } + if (char === "u") { + open.escape = "\\u"; + return true; + } + // Invalid escape: keep the backslash as text and read the character normally, as repairJson() does. + open.escape = ""; + this.#appendString(open, "\\"); + return false; + } + if (isHexDigit(char)) { + open.escape += char; + if (open.escape.length < 6) return true; + const decoded = String.fromCharCode(Number.parseInt(open.escape.slice(2), 16)); + open.escape = ""; + this.#appendString(open, decoded); + return true; + } + // `\u` without four hex digits: keep it as text and read the character normally, as repairJson() does. + const literal = open.escape; + open.escape = ""; + this.#appendString(open, literal); + return false; + } + + #appendString(open: OpenString, text: string): void { + open.text += text; + if (open.target !== "value") return; + this.#place(open.slot, open.text); + if (open.path === undefined) return; + // Extend this string's newest change while nothing came after it; otherwise report an append. + const change = open.change; + if (change !== undefined && this.#changes.at(-1) === change) { + if (change.kind === "set") change.value = (change.value as string) + text; + else change.text += text; + return; + } + open.change = this.#report({ kind: "append", path: open.path, text }); + } + + #endString(open: OpenString): void { + this.#string = undefined; + const frame = this.#stack.at(-1); + if (open.target === "key") { + if (frame === undefined) return; + frame.key = open.text; + frame.phase = "colon"; + return; + } + this.#valueEnded(frame); + } + + #endScalar(): void { + const token = this.#scalar; + const slot = this.#scalarSlot; + this.#scalar = ""; + this.#scalarSlot = undefined; + const value = scalarValue(token); + if (value !== undefined) { + this.#place(slot, value); + const path = this.#reportedPath(slot); + if (path !== undefined) this.#report({ kind: "set", path, value }); + } + this.#valueEnded(this.#stack.at(-1)); + } + + /** The value of the unterminated scalar at the end of the text, if it shows. */ + #pendingScalar(): { readonly slot: Slot; readonly value: JsonValue } | undefined { + const slot = this.#scalarSlot; + const token = this.#scalar; + if (slot === undefined || token === "" || token.length > MAX_PENDING_SCALAR_LENGTH) return undefined; + const value = pendingScalarValue(token); + return value === undefined ? undefined : { slot, value }; + } + + #openContainer(container: JsonContainer): void { + const slot = this.#valueSlot(); + const path = this.#reportedPath(slot); + this.#place(slot, container); + this.#stack.push({ + container, + slot, + path, + reportsInside: path !== undefined && path.length < MAX_REPORTED_DEPTH, + size: 0, + key: undefined, + phase: Array.isArray(container) ? "value" : "key", + }); + this.#pathCost += 1; + this.#dirty = true; + // A fresh container: the working one keeps changing, and a change never does. + if (path !== undefined) this.#report({ kind: "set", path, value: Array.isArray(container) ? [] : {} }); + } + + #closeContainer(): void { + const frame = this.#stack.pop(); + // A stray closer outside any container is ignored. + if (frame === undefined) return; + this.#pathCost -= frame.size + 1; + this.#dirty = true; + // Too deep to report inside: the whole value, final now that it is closed. + if (frame.path !== undefined && !frame.reportsInside) { + this.#report({ kind: "set", path: frame.path, value: frame.container }); + } + this.#valueEnded(this.#stack.at(-1)); + } + + /** A value in `frame`, or the root value when there is no frame, has ended. */ + #valueEnded(frame: Frame | undefined): void { + if (frame === undefined) { + // Text after the root value is ignored, but only once a root value exists. + if (this.#root !== undefined) this.#done = true; + return; + } + frame.phase = "comma"; + } + + /** Set the value at `slot` in the open path; nothing for a value in an invalid position. */ + #place(slot: Slot | undefined, value: JsonValue): void { + if (slot === undefined) return; + this.#dirty = true; + if (slot.kind === "root") { + this.#root = value; + return; + } + const frame = this.#stack.at(-1)!; + if (slot.kind === "element") { + const elements = frame.container as JsonValue[]; + if (slot.index === elements.length) { + frame.size++; + this.#pathCost++; + } + elements[slot.index] = value; + return; + } + const members = frame.container as { [key: string]: JsonValue }; + if (!Object.hasOwn(members, slot.key)) { + frame.size++; + this.#pathCost++; + } + setMember(members, slot.key, value); + } + + /** A fresh snapshot: copies of the open containers, innermost first, each holding the copy of its open child. */ + #copyOpenPath(pending: { readonly slot: Slot; readonly value: JsonValue } | undefined): JsonValue | undefined { + if (this.#stack.length === 0) return pending?.slot.kind === "root" ? pending.value : this.#root; + let child: JsonValue | undefined; + let childSlot: Slot | undefined; + for (let depth = this.#stack.length - 1; depth >= 0; depth--) { + const frame = this.#stack[depth]!; + const copy: JsonContainer = Array.isArray(frame.container) ? frame.container.slice() : { ...frame.container }; + if (depth === this.#stack.length - 1 && pending !== undefined) setSlot(copy, pending.slot, pending.value); + if (child !== undefined && childSlot !== undefined) setSlot(copy, childSlot, child); + child = copy; + childSlot = frame.slot; + } + // The outermost open container is the root unless it sits in an invalid position. + return childSlot?.kind === "root" ? child : this.#root; + } +} + +/** + * The arguments of one streaming tool call: the raw text, for the one full parse when the call ends, and a + * `StreamingJsonParser` for the partial arguments in between. Providers keep one per call in their own stream state, + * never on the content block, so the raw text never reaches partial messages or transcripts. + */ +export class ToolCallArgumentsStream { + #text = ""; + #parser = new StreamingJsonParser(); + #snapshot: JsonValue | undefined = undefined; + #arguments: ToolCall["arguments"] = {}; + + /** The raw argument text received so far. */ + get text(): string { + return this.#text; + } + + /** Append a piece of argument text; returns the partial arguments, a new object whenever they changed. */ + append(delta: string): ToolCall["arguments"] { + this.#text += delta; + this.#parser.append(delta); + return this.#partial(); + } + + /** + * Take the provider's complete argument text. When it continues the text so far, returns the missing piece as + * `delta`; otherwise starts over with `text`. Either way returns the partial arguments of `text`. + */ + replace(text: string): { readonly delta: string | undefined; readonly arguments: ToolCall["arguments"] } { + if (text.startsWith(this.#text)) { + const delta = text.slice(this.#text.length); + return { delta, arguments: delta === "" ? this.#partial() : this.append(delta) }; + } + this.#text = ""; + this.#parser = new StreamingJsonParser(); + return { delta: undefined, arguments: this.append(text) }; + } + + /** The final arguments: one full parse of the complete text, which reports malformed arguments. */ + finish(): Pick { + return parseToolCallArguments(this.#text); + } + + #partial(): ToolCall["arguments"] { + const snapshot = this.#parser.snapshot(); + if (snapshot === this.#snapshot) return this.#arguments; + this.#snapshot = snapshot; + // Tool arguments are an object; any other partial value shows as no arguments yet. + this.#arguments = + typeof snapshot === "object" && snapshot !== null && !Array.isArray(snapshot) + ? (snapshot as ToolCall["arguments"]) + : {}; + return this.#arguments; + } +} + +const SIMPLE_ESCAPES: Readonly> = { + '"': '"', + "\\": "\\", + "/": "/", + b: "\b", + f: "\f", + n: "\n", + r: "\r", + t: "\t", +}; + +function isDelimiter(char: string): boolean { + return ( + char === " " || + char === "\t" || + char === "\n" || + char === "\r" || + char === "," || + char === ":" || + char === "{" || + char === "}" || + char === "[" || + char === "]" || + char === '"' + ); +} + +function isHexDigit(char: string): boolean { + return (char >= "0" && char <= "9") || (char >= "a" && char <= "f") || (char >= "A" && char <= "F"); +} + +/** The value of a complete number or literal token; `undefined` for anything else, which is dropped. */ +function scalarValue(token: string): JsonValue | undefined { + for (const [literal, value] of LITERALS) if (token === literal) return value; + return NUMBER.test(token) ? Number(token) : undefined; +} + +/** The value an unterminated number or literal token shows. */ +function pendingScalarValue(token: string): JsonValue | undefined { + for (const [literal, value] of LITERALS) if (literal.startsWith(token)) return value; + if (NUMBER.test(token)) return Number(token); + const mantissa = NUMBER_WITHOUT_EXPONENT_DIGITS.exec(token); + return mantissa === null ? undefined : Number(mantissa[1]); +} + +function setSlot(container: JsonContainer, slot: Slot, value: JsonValue): void { + if (slot.kind === "element") (container as JsonValue[])[slot.index] = value; + else if (slot.kind === "member") setMember(container as { [key: string]: JsonValue }, slot.key, value); +} + +function setMember(members: { [key: string]: JsonValue }, key: string, value: JsonValue): void { + // `__proto__` must become an own member, as in JSON.parse(), not the prototype. + if (key === "__proto__") { + Object.defineProperty(members, key, { value, writable: true, enumerable: true, configurable: true }); + } else { + members[key] = value; + } +} diff --git a/packages/ai/src/utils/text.ts b/packages/ai/src/utils/text.ts index 5798ed346..8f90a1e9a 100644 --- a/packages/ai/src/utils/text.ts +++ b/packages/ai/src/utils/text.ts @@ -38,3 +38,13 @@ export function renderSystemMessageUpdate(message: SystemMessage): string { } return parts.join("\n\n"); } + +/** + * `text.slice(start, end)` as a string of its own. V8 represents a slice of 13 or more characters as a view that keeps + * all of `text` alive, so a short excerpt kept from a large text (a 500-character error message cut from a megabyte + * response body) would keep the whole text alive as long as the excerpt lives. Appending a character and slicing it off + * again makes V8 copy just the slice. + */ +export function ownedSlice(text: string, start: number, end = text.length): string { + return `${text.slice(start, end)} `.slice(0, -1); +} diff --git a/packages/ai/test/bedrock-stream-state.test.ts b/packages/ai/test/bedrock-stream-state.test.ts new file mode 100644 index 000000000..190a8389b --- /dev/null +++ b/packages/ai/test/bedrock-stream-state.test.ts @@ -0,0 +1,158 @@ +import { beforeEach, describe, expect, it, vi } from "vitest"; +import { allocatedBytes } from "./memory-helpers.ts"; + +// #9265: Bedrock kept `index`, `partialJson` and the encrypted reasoning bytes on the content blocks while streaming, +// and re-parsed the accumulated tool-call arguments on every delta. + +const bedrockMock = vi.hoisted(() => ({ streamEvents: undefined as (() => Iterable) | undefined })); + +vi.mock("@aws-sdk/client-bedrock-runtime", () => { + class BedrockRuntimeServiceException extends Error {} + + class BedrockRuntimeClient { + send(): Promise { + const events = bedrockMock.streamEvents; + if (!events) return Promise.reject(new Error("mock send")); + return Promise.resolve({ + $metadata: { httpStatusCode: 200 }, + stream: (async function* () { + yield* events(); + })(), + }); + } + } + + class ConverseStreamCommand { + readonly input: unknown; + + constructor(input: unknown) { + this.input = input; + } + } + + return { + BedrockRuntimeClient, + BedrockRuntimeServiceException, + ConverseStreamCommand, + StopReason: { + END_TURN: "end_turn", + STOP_SEQUENCE: "stop_sequence", + MAX_TOKENS: "max_tokens", + TOOL_USE: "tool_use", + }, + CachePointType: { DEFAULT: "default" }, + CacheTTL: { ONE_HOUR: "ONE_HOUR" }, + ConversationRole: { ASSISTANT: "assistant", USER: "user" }, + ImageFormat: { JPEG: "jpeg", PNG: "png", GIF: "gif", WEBP: "webp" }, + ToolResultStatus: { ERROR: "error", SUCCESS: "success" }, + }; +}); + +import { stream as streamBedrock } from "../src/api/bedrock-converse-stream.ts"; +import type { AssistantMessage, Model } from "../src/types.ts"; +import { normalizeContext } from "../src/utils/transcript.ts"; + +const model: Model<"bedrock-converse-stream"> = { + id: "global.anthropic.claude-test", + name: "Claude Test", + api: "bedrock-converse-stream", + provider: "amazon-bedrock", + baseUrl: "https://bedrock-runtime.us-east-1.amazonaws.com", + reasoning: false, + input: ["text"], + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }, + contextWindow: 200_000, + maxTokens: 64_000, +}; + +const context = normalizeContext({ messages: [{ role: "user", content: "write", timestamp: 0 }] }); + +const content = Array.from({ length: 20_000 }, (_, line) => `const line${line} = "value ${line}"; // \\ done\n`).join( + "", +); + +function writeCall(size: number) { + const args = { path: "/src/file.ts", content: content.slice(0, size) }; + return { args, text: JSON.stringify(args) }; +} + +/** Encrypted reasoning, then a `write` tool call whose arguments arrive in 32-character deltas. */ +function* events(text: string): Generator { + yield { messageStart: { role: "assistant" } }; + yield { + contentBlockDelta: { + contentBlockIndex: 0, + delta: { reasoningContent: { redactedContent: new Uint8Array([1, 2]) } }, + }, + }; + yield { + contentBlockDelta: { + contentBlockIndex: 0, + delta: { reasoningContent: { redactedContent: new Uint8Array([3]) } }, + }, + }; + yield { contentBlockStop: { contentBlockIndex: 0 } }; + yield { contentBlockStart: { contentBlockIndex: 1, start: { toolUse: { toolUseId: "tool_1", name: "write" } } } }; + for (let index = 0; index < text.length; index += 32) { + yield { + contentBlockDelta: { contentBlockIndex: 1, delta: { toolUse: { input: text.slice(index, index + 32) } } }, + }; + } + yield { contentBlockStop: { contentBlockIndex: 1 } }; + yield { messageStop: { stopReason: "tool_use" } }; +} + +const BLOCK_KEYS: Record = { + text: ["type", "text", "textSignature"], + thinking: ["type", "thinking", "thinkingSignature", "redacted"], + toolCall: ["type", "id", "name", "arguments", "argumentsError"], +}; + +function expectNoStreamState(message: AssistantMessage): void { + for (const block of message.content) { + for (const key of Object.keys(block)) expect(BLOCK_KEYS[block.type], `${block.type}.${key}`).toContain(key); + } +} + +describe("Bedrock stream state", () => { + beforeEach(() => { + bedrockMock.streamEvents = undefined; + }); + + it("keeps stream state off partial blocks and finishes the tool call with its exact arguments", async () => { + const { args, text } = writeCall(5_000); + bedrockMock.streamEvents = () => events(text); + let shown = -1; + const stream = streamBedrock(model, context); + for await (const event of stream) { + if (event.type === "done" || event.type === "error") continue; + expectNoStreamState(event.partial); + if (event.type !== "toolcall_delta") continue; + const call = event.partial.content[event.contentIndex]; + const partial = call?.type === "toolCall" ? call.arguments.content : undefined; + if (typeof partial !== "string") continue; + expect(args.content.startsWith(partial)).toBe(true); + expect(partial.length).toBeGreaterThanOrEqual(shown); + shown = partial.length; + } + const message = await stream.result(); + expect(message.stopReason, message.errorMessage).toBe("toolUse"); + expectNoStreamState(message); + expect(message.content).toEqual([ + { type: "thinking", thinking: "[Reasoning redacted]", thinkingSignature: "AQID", redacted: true }, + { type: "toolCall", id: "tool_1", name: "write", arguments: args }, + ]); + }); + + it("allocates linearly in the argument text", async () => { + const run = async (size: number) => { + const { text } = writeCall(size); + bedrockMock.streamEvents = () => events(text); + return allocatedBytes(() => streamBedrock(model, context).result()); + }; + const small = await run(25_000); + const large = await run(250_000); + expect(large / small).toBeLessThan(25); + expect(large / writeCall(250_000).text.length).toBeLessThan(5_000); + }); +}); diff --git a/packages/ai/test/constrained-sampling.test.ts b/packages/ai/test/constrained-sampling.test.ts index 8eb9ea8da..d31807413 100644 --- a/packages/ai/test/constrained-sampling.test.ts +++ b/packages/ai/test/constrained-sampling.test.ts @@ -3,6 +3,7 @@ import { Type } from "typebox"; import { describe, expect, it } from "vitest"; import { appendGrammarToolInputJsonDelta, + closeGrammarToolInputJson, makeStrictJsonSchema, resolveJsonSchemaStrictSampling, } from "../src/api/constrained-sampling.ts"; @@ -14,6 +15,7 @@ import { import type { AssistantMessage, Model, Tool, ToolCall } from "../src/types.ts"; import { AssistantMessageEventStream } from "../src/utils/event-stream.ts"; import { normalizeContext } from "../src/utils/transcript.ts"; +import { allocatedBytes } from "./memory-helpers.ts"; function makeModel(): Model<"openai-responses"> { return { @@ -292,14 +294,81 @@ describe("constrained tool sampling", () => { it("keeps grammar input JSON deltas append-only", () => { const buffer = { input: "", started: false, closed: false }; - const first = appendGrammarToolInputJsonDelta(buffer, "payload", 'a"', false); - const second = appendGrammarToolInputJsonDelta(buffer, "payload", 'a"\nb', true); + const first = appendGrammarToolInputJsonDelta(buffer, "payload", 'a"'); + const second = closeGrammarToolInputJson(buffer, "payload", 'a"\nb'); expect(JSON.parse(`${first}${second}`)).toEqual({ payload: 'a"\nb' }); - expect(appendGrammarToolInputJsonDelta(buffer, "payload", 'a"\nb', true)).toBeUndefined(); - expect(() => appendGrammarToolInputJsonDelta(buffer, "payload", "changed", true)).toThrow( + expect(closeGrammarToolInputJson(buffer, "payload", 'a"\nb')).toBeUndefined(); + expect(() => closeGrammarToolInputJson(buffer, "payload", "changed")).toThrow( 'grammar tool input for property "payload" changed after it was closed', ); + expect(() => appendGrammarToolInputJsonDelta(buffer, "payload", "x")).toThrow( + 'grammar tool input for property "payload" changed after it was closed', + ); + }); + + it("rejects a final grammar input that does not continue the streamed input", () => { + const buffer = { input: "", started: false, closed: false }; + appendGrammarToolInputJsonDelta(buffer, "payload", "abc"); + expect(() => closeGrammarToolInputJson(buffer, "payload", "abd")).toThrow( + 'grammar tool input for property "payload" changed non-monotonically', + ); + }); + + // #9265 + it("streams custom tool input in time and memory linear in its length", async () => { + const run = async (size: number) => { + const deltas = Array.from({ length: size / 32 }, (_, index) => + `chunk ${index} "quoted" \\ \n`.padEnd(32, "x"), + ); + const events = [ + { + type: "response.output_item.added", + output_index: 0, + item: { type: "custom_tool_call", call_id: "call_1", id: "ctc_1", name: "sample_tool", input: "" }, + }, + ...deltas.map((delta) => ({ + type: "response.custom_tool_call_input.delta", + output_index: 0, + item_id: "ctc_1", + delta, + })), + { type: "response.custom_tool_call_input.done", output_index: 0, item_id: "ctc_1", input: deltas.join("") }, + { + type: "response.output_item.done", + output_index: 0, + item: { + type: "custom_tool_call", + call_id: "call_1", + id: "ctc_1", + name: "sample_tool", + input: deltas.join(""), + }, + }, + { + type: "response.completed", + response: { status: "completed", usage: { input_tokens: 1, output_tokens: 1, total_tokens: 2 } }, + }, + ] as ResponseStreamEvent[]; + const output = makeOutput(); + const stream = new AssistantMessageEventStream(); + const jsonDeltas: string[] = []; + const bytes = await allocatedBytes(async () => { + const consumed = (async () => { + for await (const event of stream) if (event.type === "toolcall_delta") jsonDeltas.push(event.delta); + })(); + await processResponsesStream(iterateEvents(events), output, stream, makeModel(), { + grammarToolInputProperties: new Map([["sample_tool", "payload"]]), + }); + stream.end(); + await consumed; + }); + expect(JSON.parse(jsonDeltas.join(""))).toEqual({ payload: deltas.join("") }); + return bytes; + }; + const small = await run(25_600); + const large = await run(256_000); + expect(large / small).toBeLessThan(25); }); it("starts custom Responses tool calls with their initial input", async () => { diff --git a/packages/ai/test/data/json-test-suite.json b/packages/ai/test/data/json-test-suite.json new file mode 100644 index 000000000..4dbe7d575 --- /dev/null +++ b/packages/ai/test/data/json-test-suite.json @@ -0,0 +1,322 @@ +{ + "source": "https://github.com/nst/JSONTestSuite (test_parsing, files up to 2 KB)", + "license": "MIT License, Copyright (c) 2016 Nicolas Seriot", + "files": { + "i_number_double_huge_neg_exp.json": "WzEyMy40NTZlLTc4OV0=", + "i_number_huge_exp.json": "WzAuNGUwMDY2OTk5OTk5OTk5OTk5OTk5OTk5OTk5OTk5OTk5OTk5OTk5OTk5OTk5OTk5OTk5OTk5OTk5OTk5OTk5OTk5OTk5OTk5OTk5OTk5OTk5OTk5OTk5OTk5OTk5OTk5OTk5OTk5OTk5OTk5OTk5OTk5OTk5OTk5OTk2OTk5OTk5OTAwNl0=", + "i_number_neg_int_huge_exp.json": "Wy0xZSs5OTk5XQ==", + "i_number_pos_double_huge_exp.json": "WzEuNWUrOTk5OV0=", + "i_number_real_neg_overflow.json": "Wy0xMjMxMjNlMTAwMDAwXQ==", + "i_number_real_pos_overflow.json": "WzEyMzEyM2UxMDAwMDBd", + "i_number_real_underflow.json": "WzEyM2UtMTAwMDAwMDBd", + "i_number_too_big_neg_int.json": "Wy0xMjMxMjMxMjMxMjMxMjMxMjMxMjMxMjMxMjMxMjNd", + "i_number_too_big_pos_int.json": "WzEwMDAwMDAwMDAwMDAwMDAwMDAwMF0=", + "i_number_very_big_negative_int.json": "Wy0yMzc0NjIzNzQ2NzMyNzY4OTQyNzk4MzI3NDk4MzI0MjM0Nzk4MjMyNDYzMjc4NDZd", + "i_object_key_lone_2nd_surrogate.json": "eyJcdURGQUEiOjB9", + "i_string_1st_surrogate_but_2nd_missing.json": "WyJcdURBREEiXQ==", + "i_string_1st_valid_surrogate_2nd_invalid.json": "WyJcdUQ4ODhcdTEyMzQiXQ==", + "i_string_UTF-16LE_with_BOM.json": "//5bACIA6QAiAF0A", + "i_string_UTF-8_invalid_sequence.json": "WyLml6XRiPoiXQ==", + "i_string_UTF8_surrogate_U+D800.json": "WyLtoIAiXQ==", + "i_string_incomplete_surrogate_and_escape_valid.json": "WyJcdUQ4MDBcbiJd", + "i_string_incomplete_surrogate_pair.json": "WyJcdURkMWVhIl0=", + "i_string_incomplete_surrogates_escape_valid.json": "WyJcdUQ4MDBcdUQ4MDBcbiJd", + "i_string_invalid_lonely_surrogate.json": "WyJcdWQ4MDAiXQ==", + "i_string_invalid_surrogate.json": "WyJcdWQ4MDBhYmMiXQ==", + "i_string_invalid_utf-8.json": "WyL/Il0=", + "i_string_inverted_surrogates_U+1D11E.json": "WyJcdURkMWVcdUQ4MzQiXQ==", + "i_string_iso_latin_1.json": "WyLpIl0=", + "i_string_lone_second_surrogate.json": "WyJcdURGQUEiXQ==", + "i_string_lone_utf8_continuation_byte.json": "WyKBIl0=", + "i_string_not_in_unicode_range.json": "WyL0v7+/Il0=", + "i_string_overlong_sequence_2_bytes.json": "WyLAryJd", + "i_string_overlong_sequence_6_bytes.json": "WyL8g7+/v78iXQ==", + "i_string_overlong_sequence_6_bytes_null.json": "WyL8gICAgIAiXQ==", + "i_string_truncated-utf-8.json": "WyLg/yJd", + "i_string_utf16BE_no_BOM.json": "AFsAIgDpACIAXQ==", + "i_string_utf16LE_no_BOM.json": "WwAiAOkAIgBdAA==", + "i_structure_500_nested_arrays.json": "W1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tbW1tdXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXV1dXQ==", + "i_structure_UTF-8_BOM_empty_object.json": "77u/e30=", + "n_array_1_true_without_comma.json": "WzEgdHJ1ZV0=", + "n_array_a_invalid_utf8.json": "W2HlXQ==", + "n_array_colon_instead_of_comma.json": "WyIiOiAxXQ==", + "n_array_comma_after_close.json": "WyIiXSw=", + "n_array_comma_and_number.json": "WywxXQ==", + "n_array_double_comma.json": "WzEsLDJd", + "n_array_double_extra_comma.json": "WyJ4IiwsXQ==", + "n_array_extra_close.json": "WyJ4Il1d", + "n_array_extra_comma.json": "WyIiLF0=", + "n_array_incomplete.json": "WyJ4Ig==", + "n_array_incomplete_invalid_value.json": "W3g=", + "n_array_inner_array_no_comma.json": "WzNbNF1d", + "n_array_invalid_utf8.json": "W/9d", + "n_array_items_separated_by_semicolon.json": "WzE6Ml0=", + "n_array_just_comma.json": "Wyxd", + "n_array_just_minus.json": "Wy1d", + "n_array_missing_value.json": "WyAgICwgIiJd", + "n_array_newlines_unclosed.json": "WyJhIiwKNAosMSw=", + "n_array_number_and_comma.json": "WzEsXQ==", + "n_array_number_and_several_commas.json": "WzEsLF0=", + "n_array_spaces_vertical_tab_formfeed.json": "WyILYSJcZl0=", + "n_array_star_inside.json": "Wypd", + "n_array_unclosed.json": "WyIi", + "n_array_unclosed_trailing_comma.json": "WzEs", + "n_array_unclosed_with_new_lines.json": "WzEsCjEKLDE=", + "n_array_unclosed_with_object_inside.json": "W3t9", + "n_incomplete_false.json": "W2ZhbHNd", + "n_incomplete_null.json": "W251bF0=", + "n_incomplete_true.json": "W3RydV0=", + "n_multidigit_number_then_00.json": "MTIzAA==", + "n_number_++.json": "WysrMTIzNF0=", + "n_number_+1.json": "WysxXQ==", + "n_number_+Inf.json": "WytJbmZd", + "n_number_-01.json": "Wy0wMV0=", + "n_number_-1.0..json": "Wy0xLjAuXQ==", + "n_number_-2..json": "Wy0yLl0=", + "n_number_-NaN.json": "Wy1OYU5d", + "n_number_.-1.json": "Wy4tMV0=", + "n_number_.2e-3.json": "Wy4yZS0zXQ==", + "n_number_0.1.2.json": "WzAuMS4yXQ==", + "n_number_0.3e+.json": "WzAuM2UrXQ==", + "n_number_0.3e.json": "WzAuM2Vd", + "n_number_0.e1.json": "WzAuZTFd", + "n_number_0_capital_E+.json": "WzBFK10=", + "n_number_0_capital_E.json": "WzBFXQ==", + "n_number_0e+.json": "WzBlK10=", + "n_number_0e.json": "WzBlXQ==", + "n_number_1.0e+.json": "WzEuMGUrXQ==", + "n_number_1.0e-.json": "WzEuMGUtXQ==", + "n_number_1.0e.json": "WzEuMGVd", + "n_number_1_000.json": "WzEgMDAwLjBd", + "n_number_1eE2.json": "WzFlRTJd", + "n_number_2.e+3.json": "WzIuZSszXQ==", + "n_number_2.e-3.json": "WzIuZS0zXQ==", + "n_number_2.e3.json": "WzIuZTNd", + "n_number_9.e+.json": "WzkuZStd", + "n_number_Inf.json": "W0luZl0=", + "n_number_NaN.json": "W05hTl0=", + "n_number_U+FF11_fullwidth_digit_one.json": "W++8kV0=", + "n_number_expression.json": "WzErMl0=", + "n_number_hex_1_digit.json": "WzB4MV0=", + "n_number_hex_2_digits.json": "WzB4NDJd", + "n_number_infinity.json": "W0luZmluaXR5XQ==", + "n_number_invalid+-.json": "WzBlKy0xXQ==", + "n_number_invalid-negative-real.json": "Wy0xMjMuMTIzZm9vXQ==", + "n_number_invalid-utf-8-in-bigger-int.json": "WzEyM+Vd", + "n_number_invalid-utf-8-in-exponent.json": "WzFlMeVd", + "n_number_invalid-utf-8-in-int.json": "WzDlXQo=", + "n_number_minus_infinity.json": "Wy1JbmZpbml0eV0=", + "n_number_minus_sign_with_trailing_garbage.json": "Wy1mb29d", + "n_number_minus_space_1.json": "Wy0gMV0=", + "n_number_neg_int_starting_with_zero.json": "Wy0wMTJd", + "n_number_neg_real_without_int_part.json": "Wy0uMTIzXQ==", + "n_number_neg_with_garbage_at_end.json": "Wy0xeF0=", + "n_number_real_garbage_after_e.json": "WzFlYV0=", + "n_number_real_with_invalid_utf8_after_e.json": "WzFl5V0=", + "n_number_real_without_fractional_part.json": "WzEuXQ==", + "n_number_starting_with_dot.json": "Wy4xMjNd", + "n_number_with_alpha.json": "WzEuMmEtM10=", + "n_number_with_alpha_char.json": "WzEuODAxMTY3MDAzMzM3NjUxNEgtMzA4XQ==", + "n_number_with_leading_zero.json": "WzAxMl0=", + "n_object_bad_value.json": "WyJ4IiwgdHJ1dGhd", + "n_object_bracket_key.json": "e1s6ICJ4In0K", + "n_object_comma_instead_of_colon.json": "eyJ4IiwgbnVsbH0=", + "n_object_double_colon.json": "eyJ4Ijo6ImIifQ==", + "n_object_emoji.json": "e/Cfh6jwn4etfQ==", + "n_object_garbage_at_end.json": "eyJhIjoiYSIgMTIzfQ==", + "n_object_key_with_single_quotes.json": "e2tleTogJ3ZhbHVlJ30=", + "n_object_lone_continuation_byte_in_key_and_trailing_comma.json": "eyK5IjoiMCIsfQ==", + "n_object_missing_colon.json": "eyJhIiBifQ==", + "n_object_missing_key.json": "ezoiYiJ9", + "n_object_missing_semicolon.json": "eyJhIiAiYiJ9", + "n_object_missing_value.json": "eyJhIjo=", + "n_object_no-colon.json": "eyJhIg==", + "n_object_non_string_key.json": "ezE6MX0=", + "n_object_non_string_key_but_huge_number_instead.json": "ezk5OTlFOTk5OToxfQ==", + "n_object_repeated_null_null.json": "e251bGw6bnVsbCxudWxsOm51bGx9", + "n_object_several_trailing_commas.json": "eyJpZCI6MCwsLCwsfQ==", + "n_object_single_quote.json": "eydhJzowfQ==", + "n_object_trailing_comma.json": "eyJpZCI6MCx9", + "n_object_trailing_comment.json": "eyJhIjoiYiJ9LyoqLw==", + "n_object_trailing_comment_open.json": "eyJhIjoiYiJ9LyoqLy8=", + "n_object_trailing_comment_slash_open.json": "eyJhIjoiYiJ9Ly8=", + "n_object_trailing_comment_slash_open_incomplete.json": "eyJhIjoiYiJ9Lw==", + "n_object_two_commas_in_a_row.json": "eyJhIjoiYiIsLCJjIjoiZCJ9", + "n_object_unquoted_key.json": "e2E6ICJiIn0=", + "n_object_unterminated-value.json": "eyJhIjoiYQ==", + "n_object_with_single_string.json": "eyAiZm9vIiA6ICJiYXIiLCAiYSIgfQ==", + "n_object_with_trailing_garbage.json": "eyJhIjoiYiJ9Iw==", + "n_single_space.json": "IA==", + "n_string_1_surrogate_then_escape.json": "WyJcdUQ4MDBcIl0=", + "n_string_1_surrogate_then_escape_u.json": "WyJcdUQ4MDBcdSJd", + "n_string_1_surrogate_then_escape_u1.json": "WyJcdUQ4MDBcdTEiXQ==", + "n_string_1_surrogate_then_escape_u1x.json": "WyJcdUQ4MDBcdTF4Il0=", + "n_string_accentuated_char_no_quotes.json": "W8OpXQ==", + "n_string_backslash_00.json": "WyJcACJd", + "n_string_escape_x.json": "WyJceDAwIl0=", + "n_string_escaped_backslash_bad.json": "WyJcXFwiXQ==", + "n_string_escaped_ctrl_char_tab.json": "WyJcCSJd", + "n_string_escaped_emoji.json": "WyJc8J+MgCJd", + "n_string_incomplete_escape.json": "WyJcIl0=", + "n_string_incomplete_escaped_character.json": "WyJcdTAwQSJd", + "n_string_incomplete_surrogate.json": "WyJcdUQ4MzRcdURkIl0=", + "n_string_incomplete_surrogate_escape_invalid.json": "WyJcdUQ4MDBcdUQ4MDBceCJd", + "n_string_invalid-utf-8-in-escape.json": "WyJcdeUiXQ==", + "n_string_invalid_backslash_esc.json": "WyJcYSJd", + "n_string_invalid_unicode_escape.json": "WyJcdXFxcXEiXQ==", + "n_string_invalid_utf8_after_escape.json": "WyJc5SJd", + "n_string_leading_uescaped_thinspace.json": "W1x1MDAyMCJhc2QiXQ==", + "n_string_no_quotes_with_bad_escape.json": "W1xuXQ==", + "n_string_single_doublequote.json": "Ig==", + "n_string_single_quote.json": "WydzaW5nbGUgcXVvdGUnXQ==", + "n_string_single_string_no_double_quotes.json": "YWJj", + "n_string_start_escape_unclosed.json": "WyJc", + "n_string_unescaped_ctrl_char.json": "WyJhAGEiXQ==", + "n_string_unescaped_newline.json": "WyJuZXcKbGluZSJd", + "n_string_unescaped_tab.json": "WyIJIl0=", + "n_string_unicode_CapitalU.json": "IlxVQTY2RCI=", + "n_string_with_trailing_garbage.json": "IiJ4", + "n_structure_U+2060_word_joined.json": "W+KBoF0=", + "n_structure_UTF8_BOM_no_data.json": "77u/", + "n_structure_angle_bracket_..json": "PC4+", + "n_structure_angle_bracket_null.json": "WzxudWxsPl0=", + "n_structure_array_trailing_garbage.json": "WzFdeA==", + "n_structure_array_with_extra_array_close.json": "WzFdXQ==", + "n_structure_array_with_unclosed_string.json": "WyJhc2Rd", + "n_structure_ascii-unicode-identifier.json": "YcOl", + "n_structure_capitalized_True.json": "W1RydWVd", + "n_structure_close_unopened_array.json": "MV0=", + "n_structure_comma_instead_of_closing_brace.json": "eyJ4IjogdHJ1ZSw=", + "n_structure_double_array.json": "W11bXQ==", + "n_structure_end_array.json": "XQ==", + "n_structure_incomplete_UTF8_BOM.json": "77t7fQ==", + "n_structure_lone-invalid-utf-8.json": "5Q==", + "n_structure_lone-open-bracket.json": "Ww==", + "n_structure_no_data.json": "", + "n_structure_null-byte-outside-string.json": "WwBd", + "n_structure_number_with_trailing_garbage.json": "MkA=", + "n_structure_object_followed_by_closing_object.json": "e319", + "n_structure_object_unclosed_no_value.json": "eyIiOg==", + "n_structure_object_with_comment.json": "eyJhIjovKmNvbW1lbnQqLyJiIn0=", + "n_structure_object_with_trailing_garbage.json": "eyJhIjogdHJ1ZX0gIngi", + "n_structure_open_array_apostrophe.json": "Wyc=", + "n_structure_open_array_comma.json": "Wyw=", + "n_structure_open_array_open_object.json": "W3s=", + "n_structure_open_array_open_string.json": "WyJh", + "n_structure_open_array_string.json": "WyJhIg==", + "n_structure_open_object.json": "ew==", + "n_structure_open_object_close_array.json": "e10=", + "n_structure_open_object_comma.json": "eyw=", + "n_structure_open_object_open_array.json": "e1s=", + "n_structure_open_object_open_string.json": "eyJh", + "n_structure_open_object_string_with_apostrophes.json": "eydhJw==", + "n_structure_open_open.json": "WyJce1siXHtbIlx7WyJcew==", + "n_structure_single_eacute.json": "6Q==", + "n_structure_single_star.json": "Kg==", + "n_structure_trailing_#.json": "eyJhIjoiYiJ9I3t9", + "n_structure_uescaped_LF_before_string.json": "W1x1MDAwQSIiXQ==", + "n_structure_unclosed_array.json": "WzE=", + "n_structure_unclosed_array_partial_null.json": "WyBmYWxzZSwgbnVs", + "n_structure_unclosed_array_unfinished_false.json": "WyB0cnVlLCBmYWxz", + "n_structure_unclosed_array_unfinished_true.json": "WyBmYWxzZSwgdHJ1", + "n_structure_unclosed_object.json": "eyJhc2QiOiJhc2Qi", + "n_structure_unicode-identifier.json": "w6U=", + "n_structure_whitespace_U+2060_word_joiner.json": "W+KBoF0=", + "n_structure_whitespace_formfeed.json": "Wwxd", + "y_array_arraysWithSpaces.json": "W1tdICAgXQ==", + "y_array_empty-string.json": "WyIiXQ==", + "y_array_empty.json": "W10=", + "y_array_ending_with_newline.json": "WyJhIl0=", + "y_array_false.json": "W2ZhbHNlXQ==", + "y_array_heterogeneous.json": "W251bGwsIDEsICIxIiwge31d", + "y_array_null.json": "W251bGxd", + "y_array_with_1_and_newline.json": "WzEKXQ==", + "y_array_with_leading_space.json": "IFsxXQ==", + "y_array_with_several_null.json": "WzEsbnVsbCxudWxsLG51bGwsMl0=", + "y_array_with_trailing_space.json": "WzJdIA==", + "y_number.json": "WzEyM2U2NV0=", + "y_number_0e+1.json": "WzBlKzFd", + "y_number_0e1.json": "WzBlMV0=", + "y_number_after_space.json": "WyA0XQ==", + "y_number_double_close_to_zero.json": "Wy0wLjAwMDAwMDAwMDAwMDAwMDAwMDAwMDAwMDAwMDAwMDAwMDAwMDAwMDAwMDAwMDAwMDAwMDAwMDAwMDAwMDAwMDAwMDAwMDAwMDAwMDAwMV0K", + "y_number_int_with_exp.json": "WzIwZTFd", + "y_number_minus_zero.json": "Wy0wXQ==", + "y_number_negative_int.json": "Wy0xMjNd", + "y_number_negative_one.json": "Wy0xXQ==", + "y_number_negative_zero.json": "Wy0wXQ==", + "y_number_real_capital_e.json": "WzFFMjJd", + "y_number_real_capital_e_neg_exp.json": "WzFFLTJd", + "y_number_real_capital_e_pos_exp.json": "WzFFKzJd", + "y_number_real_exponent.json": "WzEyM2U0NV0=", + "y_number_real_fraction_exponent.json": "WzEyMy40NTZlNzhd", + "y_number_real_neg_exp.json": "WzFlLTJd", + "y_number_real_pos_exponent.json": "WzFlKzJd", + "y_number_simple_int.json": "WzEyM10=", + "y_number_simple_real.json": "WzEyMy40NTY3ODld", + "y_object.json": "eyJhc2QiOiJzZGYiLCAiZGZnIjoiZmdoIn0=", + "y_object_basic.json": "eyJhc2QiOiJzZGYifQ==", + "y_object_duplicated_key.json": "eyJhIjoiYiIsImEiOiJjIn0=", + "y_object_duplicated_key_and_value.json": "eyJhIjoiYiIsImEiOiJiIn0=", + "y_object_empty.json": "e30=", + "y_object_empty_key.json": "eyIiOjB9", + "y_object_escaped_null_in_key.json": "eyJmb29cdTAwMDBiYXIiOiA0Mn0=", + "y_object_extreme_numbers.json": "eyAibWluIjogLTEuMGUrMjgsICJtYXgiOiAxLjBlKzI4IH0=", + "y_object_long_strings.json": "eyJ4IjpbeyJpZCI6ICJ4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4In1dLCAiaWQiOiAieHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eCJ9", + "y_object_simple.json": "eyJhIjpbXX0=", + "y_object_string_unicode.json": "eyJ0aXRsZSI6Ilx1MDQxZlx1MDQzZVx1MDQzYlx1MDQ0Mlx1MDQzZVx1MDQ0MFx1MDQzMCBcdTA0MTdcdTA0MzVcdTA0M2NcdTA0M2JcdTA0MzVcdTA0M2FcdTA0M2VcdTA0M2ZcdTA0MzAiIH0=", + "y_object_with_newlines.json": "ewoiYSI6ICJiIgp9", + "y_string_1_2_3_bytes_UTF-8_sequences.json": "WyJcdTAwNjBcdTAxMmFcdTEyQUIiXQ==", + "y_string_accepted_surrogate_pair.json": "WyJcdUQ4MDFcdWRjMzciXQ==", + "y_string_accepted_surrogate_pairs.json": "WyJcdWQ4M2RcdWRlMzlcdWQ4M2RcdWRjOGQiXQ==", + "y_string_allowed_escapes.json": "WyJcIlxcXC9cYlxmXG5cclx0Il0=", + "y_string_backslash_and_u_escaped_zero.json": "WyJcXHUwMDAwIl0=", + "y_string_backslash_doublequotes.json": "WyJcIiJd", + "y_string_comments.json": "WyJhLypiKi9jLypkLy9lIl0=", + "y_string_double_escape_a.json": "WyJcXGEiXQ==", + "y_string_double_escape_n.json": "WyJcXG4iXQ==", + "y_string_escaped_control_character.json": "WyJcdTAwMTIiXQ==", + "y_string_escaped_noncharacter.json": "WyJcdUZGRkYiXQ==", + "y_string_in_array.json": "WyJhc2QiXQ==", + "y_string_in_array_with_leading_space.json": "WyAiYXNkIl0=", + "y_string_last_surrogates_1_and_2.json": "WyJcdURCRkZcdURGRkYiXQ==", + "y_string_nbsp_uescaped.json": "WyJuZXdcdTAwQTBsaW5lIl0=", + "y_string_nonCharacterInUTF-8_U+10FFFF.json": "WyL0j7+/Il0=", + "y_string_nonCharacterInUTF-8_U+FFFF.json": "WyLvv78iXQ==", + "y_string_null_escape.json": "WyJcdTAwMDAiXQ==", + "y_string_one-byte-utf-8.json": "WyJcdTAwMmMiXQ==", + "y_string_pi.json": "WyLPgCJd", + "y_string_reservedCharacterInUTF-8_U+1BFFF.json": "WyLwm7+/Il0=", + "y_string_simple_ascii.json": "WyJhc2QgIl0=", + "y_string_space.json": "IiAi", + "y_string_surrogates_U+1D11E_MUSICAL_SYMBOL_G_CLEF.json": "WyJcdUQ4MzRcdURkMWUiXQ==", + "y_string_three-byte-utf-8.json": "WyJcdTA4MjEiXQ==", + "y_string_two-byte-utf-8.json": "WyJcdTAxMjMiXQ==", + "y_string_u+2028_line_sep.json": "WyLigKgiXQ==", + "y_string_u+2029_par_sep.json": "WyLigKkiXQ==", + "y_string_uEscape.json": "WyJcdTAwNjFcdTMwYWZcdTMwRUFcdTMwYjkiXQ==", + "y_string_uescaped_newline.json": "WyJuZXdcdTAwMEFsaW5lIl0=", + "y_string_unescaped_char_delete.json": "WyJ/Il0=", + "y_string_unicode.json": "WyJcdUE2NkQiXQ==", + "y_string_unicodeEscapedBackslash.json": "WyJcdTAwNUMiXQ==", + "y_string_unicode_2.json": "WyLijYLjiLTijYIiXQ==", + "y_string_unicode_U+10FFFE_nonchar.json": "WyJcdURCRkZcdURGRkUiXQ==", + "y_string_unicode_U+1FFFE_nonchar.json": "WyJcdUQ4M0ZcdURGRkUiXQ==", + "y_string_unicode_U+200B_ZERO_WIDTH_SPACE.json": "WyJcdTIwMEIiXQ==", + "y_string_unicode_U+2064_invisible_plus.json": "WyJcdTIwNjQiXQ==", + "y_string_unicode_U+FDD0_nonchar.json": "WyJcdUZERDAiXQ==", + "y_string_unicode_U+FFFE_nonchar.json": "WyJcdUZGRkUiXQ==", + "y_string_unicode_escaped_double_quote.json": "WyJcdTAwMjIiXQ==", + "y_string_utf8.json": "WyLigqzwnYSeIl0=", + "y_string_with_del_character.json": "WyJhf2EiXQ==", + "y_structure_lonely_false.json": "ZmFsc2U=", + "y_structure_lonely_int.json": "NDI=", + "y_structure_lonely_negative_real.json": "LTAuMQ==", + "y_structure_lonely_null.json": "bnVsbA==", + "y_structure_lonely_string.json": "ImFzZCI=", + "y_structure_lonely_true.json": "dHJ1ZQ==", + "y_structure_string_empty.json": "IiI=", + "y_structure_trailing_newline.json": "WyJhIl0K", + "y_structure_true_in_array.json": "W3RydWVd", + "y_structure_whitespace_array.json": "IFtdIA==" + } +} \ No newline at end of file diff --git a/packages/ai/test/error-body.test.ts b/packages/ai/test/error-body.test.ts index 15cf5e11c..c2824e8be 100644 --- a/packages/ai/test/error-body.test.ts +++ b/packages/ai/test/error-body.test.ts @@ -5,8 +5,15 @@ // Bedrock ServiceException), plus the non-Error fallback, truncation, the empty // parsed-body edge case, and the formatProviderError compose helper. +import { getHeapStatistics } from "node:v8"; import { describe, expect, it } from "vitest"; -import { formatProviderError, MAX_PROVIDER_ERROR_BODY_CHARS, normalizeProviderError } from "../src/utils/error-body.ts"; +import { + formatProviderError, + MAX_PROVIDER_ERROR_BODY_CHARS, + normalizeProviderError, + truncateErrorText, +} from "../src/utils/error-body.ts"; +import { collectGarbage } from "./memory-helpers.ts"; describe("normalizeProviderError", () => { it("extracts status and body from a Mistral-shaped error", () => { @@ -224,3 +231,20 @@ describe("formatProviderError", () => { expect(formatProviderError(norm)).toBe('{"reason":"boom"}'); }); }); + +// #9265 +describe("truncateErrorText", () => { + it("keeps an excerpt that does not hold the whole body alive", async () => { + await collectGarbage(); + const before = getHeapStatistics().used_heap_size; + const kept = (() => { + const body = "".repeat(1_000_000); + return truncateErrorText(body, MAX_PROVIDER_ERROR_BODY_CHARS); + })(); + await collectGarbage(); + expect(kept.startsWith("")).toBe(true); + expect(kept.endsWith(`[truncated ${6_000_000 - MAX_PROVIDER_ERROR_BODY_CHARS} chars]`)).toBe(true); + // A slice would keep the 6 MB body alive with the excerpt. + expect(getHeapStatistics().used_heap_size - before).toBeLessThan(1_000_000); + }); +}); diff --git a/packages/ai/test/memory-helpers.ts b/packages/ai/test/memory-helpers.ts new file mode 100644 index 000000000..cf025a15a --- /dev/null +++ b/packages/ai/test/memory-helpers.ts @@ -0,0 +1,39 @@ +// GC-based measurements for cost guards and leak tests: bytes a function allocates, and collecting garbage on demand. +import { GCProfiler, getHeapStatistics, setFlagsFromString } from "node:v8"; +import { runInNewContext } from "node:vm"; + +setFlagsFromString("--expose-gc"); +const gc = runInNewContext("gc") as () => void; + +/** Run a full garbage collection a few times, yielding in between so finalizers and weak references settle. */ +export async function collectGarbage(): Promise { + for (let round = 0; round < 4; round++) { + await new Promise((resolve) => setImmediate(resolve)); + gc(); + } +} + +/** + * Bytes `run` allocates on the JavaScript heap, including garbage collected while it ran: the heap growth plus what + * every collection during the run freed. + */ +export async function allocatedBytes(run: () => unknown): Promise { + gc(); + const profiler = new GCProfiler(); + profiler.start(); + const before = getHeapStatistics().used_heap_size; + await run(); + const after = getHeapStatistics().used_heap_size; + const { statistics } = profiler.stop(); + let freed = 0; + for (const collection of statistics) { + freed += collection.beforeGC.heapStatistics.usedHeapSize - collection.afterGC.heapStatistics.usedHeapSize; + } + return after - before + freed; +} + +/** Whether `reference` is collected once nothing else holds its target. */ +export async function isCollected(reference: WeakRef): Promise { + await collectGarbage(); + return reference.deref() === undefined; +} diff --git a/packages/ai/test/streaming-json.test.ts b/packages/ai/test/streaming-json.test.ts new file mode 100644 index 000000000..e2be9bce0 --- /dev/null +++ b/packages/ai/test/streaming-json.test.ts @@ -0,0 +1,848 @@ +import { readFileSync } from "node:fs"; +import { describe, expect, it } from "vitest"; +import type { JsonValue } from "../src/types.ts"; +import { parseStreamingJson, parseToolCallArguments } from "../src/utils/json-parse.ts"; +import { type JsonChange, StreamingJsonParser, ToolCallArgumentsStream } from "../src/utils/streaming-json.ts"; +import { allocatedBytes, collectGarbage } from "./memory-helpers.ts"; + +// #9265: streamed tool-call arguments were re-parsed in full on every delta. + +/** Seeded random numbers, so every run checks the same documents. */ +function random(seed: number): () => number { + let state = seed; + return () => { + state = (state * 1103515245 + 12345) & 0x7fffffff; + return state / 0x7fffffff; + }; +} + +type Random = () => number; + +const pick = (rand: Random, items: readonly T[]): T => items[Math.floor(rand() * items.length)]!; + +const STRING_PIECES = [ + "a", + "b", + "x y", + " ", + "\n", + "\t", + "\\", + '"', + "/", + "é", + "😀", + "\u0001", + "\u001f", + "\u2028", + "{", + "]", + ",", +]; + +function randomString(rand: Random): string { + let text = ""; + const length = Math.floor(rand() * 6); + for (let index = 0; index < length; index++) text += pick(rand, STRING_PIECES); + return text; +} + +function randomScalar(rand: Random): JsonValue { + return pick(rand, [ + () => randomString(rand), + () => Math.floor(rand() * 2000) - 1000, + () => rand() * 100, + () => 1e21, + () => -5e-8, + () => -0, + () => 0, + () => true, + () => false, + () => null, + ])(); +} + +function randomValue(rand: Random, depth: number): JsonValue { + const roll = rand(); + if (depth > 3 || roll < 0.3) return randomScalar(rand); + if (roll < 0.65) { + const members: { [key: string]: JsonValue } = {}; + const count = Math.floor(rand() * 4); + for (let index = 0; index < count; index++) members[randomString(rand)] = randomValue(rand, depth + 1); + return members; + } + return Array.from({ length: Math.floor(rand() * 4) }, () => randomValue(rand, depth + 1)); +} + +/** JSON text of `value`, with random insignificant whitespace when `spaced`. */ +function render(value: JsonValue, rand: Random, spaced: boolean): string { + const space = () => (spaced ? pick(rand, ["", "", " ", "\n ", "\t", "\r\n"]) : ""); + if (Array.isArray(value)) { + return `[${space()}${value.map((item) => render(item, rand, spaced)).join(`${space()},${space()}`)}${space()}]`; + } + if (value !== null && typeof value === "object") { + const members = Object.entries(value).map( + ([key, item]) => `${JSON.stringify(key)}${space()}:${space()}${render(item, rand, spaced)}`, + ); + return `{${space()}${members.join(`${space()},${space()}`)}${space()}}`; + } + // JSON.stringify(-0) is "0"; keep the sign so the text parses back to the same value. + return Object.is(value, -0) ? "-0" : JSON.stringify(value); +} + +function randomDocuments(seed: number, count: number, spaced: boolean): string[] { + const rand = random(seed); + return Array.from({ length: count }, () => + render( + rand() < 0.8 ? { path: randomString(rand), args: randomValue(rand, 0) } : randomValue(rand, 0), + rand, + spaced, + ), + ); +} + +/** Whether the reference shows unterminated numbers and literals, as snapshots do; changes leave them out. */ +let pendingScalars = true; + +/** + * Reference for `StreamingJsonParser` on a prefix of valid JSON (strings may also hold raw control characters and + * invalid escapes): a plain recursive-descent reading of the whole prefix that follows the documented rules. It is + * written independently of the incremental parser so the two check each other. + */ +function referenceSnapshot(text: string, options: { pendingScalars?: boolean } = {}): JsonValue | undefined { + pendingScalars = options.pendingScalars ?? true; + const read = readValue(text, skipSpace(text, 0)); + return read === undefined ? undefined : read.value; +} + +type Read = { readonly value: JsonValue; readonly complete: boolean; readonly end: number }; + +const SPACE = " \t\n\r"; + +function skipSpace(text: string, start: number): number { + let index = start; + while (index < text.length && SPACE.includes(text[index]!)) index++; + return index; +} + +function readValue(text: string, start: number): Read | undefined { + if (start >= text.length) return undefined; + const char = text[start]; + if (char === '"') return readString(text, start); + if (char === "{") return readObject(text, start); + if (char === "[") return readArray(text, start); + return readScalar(text, start); +} + +const SIMPLE: Record = { '"': '"', "\\": "\\", "/": "/", b: "\b", f: "\f", n: "\n", r: "\r", t: "\t" }; + +function readString(text: string, start: number): Read & { readonly value: string } { + let value = ""; + let index = start + 1; + while (index < text.length) { + const char = text[index]!; + if (char === '"') return { value, complete: true, end: index + 1 }; + if (char !== "\\") { + value += char; + index++; + continue; + } + const next = text[index + 1]; + if (next === undefined) return { value, complete: false, end: text.length }; + if (SIMPLE[next] !== undefined) { + value += SIMPLE[next]; + index += 2; + continue; + } + if (next === "u") { + let end = index + 2; + while (end < index + 6 && end < text.length && /[0-9a-fA-F]/.test(text[end]!)) end++; + if (end === index + 6) { + value += String.fromCharCode(Number.parseInt(text.slice(index + 2, end), 16)); + index = end; + continue; + } + if (end === text.length) return { value, complete: false, end }; + // Not four hex digits: the escape stays literal and the next character is read normally. + value += text.slice(index, end); + index = end; + continue; + } + // Invalid escape: the backslash stays literal and the next character is read normally. + value += "\\"; + index++; + } + return { value, complete: false, end: text.length }; +} + +function setOwn(members: { [key: string]: JsonValue }, key: string, value: JsonValue): void { + Object.defineProperty(members, key, { value, writable: true, enumerable: true, configurable: true }); +} + +function readObject(text: string, start: number): Read { + const members: { [key: string]: JsonValue } = {}; + const partial = { value: members, complete: false, end: text.length }; + let index = skipSpace(text, start + 1); + if (index >= text.length) return partial; + if (text[index] === "}") return { value: members, complete: true, end: index + 1 }; + for (;;) { + if (text[index] !== '"') throw new Error(`reference reads valid JSON only: ${JSON.stringify(text)}`); + const key = readString(text, index); + if (!key.complete) return partial; + index = skipSpace(text, key.end); + if (index >= text.length) return partial; + index = skipSpace(text, index + 1); + const member = readValue(text, index); + if (member === undefined) return partial; + setOwn(members, key.value, member.value); + if (!member.complete) return partial; + index = skipSpace(text, member.end); + if (index >= text.length) return partial; + if (text[index] === "}") return { value: members, complete: true, end: index + 1 }; + index = skipSpace(text, index + 1); + if (index >= text.length) return partial; + } +} + +function readArray(text: string, start: number): Read { + const elements: JsonValue[] = []; + const partial = { value: elements, complete: false, end: text.length }; + let index = skipSpace(text, start + 1); + if (index >= text.length) return partial; + if (text[index] === "]") return { value: elements, complete: true, end: index + 1 }; + for (;;) { + const element = readValue(text, index); + if (element === undefined) return partial; + elements.push(element.value); + if (!element.complete) return partial; + index = skipSpace(text, element.end); + if (index >= text.length) return partial; + if (text[index] === "]") return { value: elements, complete: true, end: index + 1 }; + index = skipSpace(text, index + 1); + if (index >= text.length) return partial; + } +} + +function readScalar(text: string, start: number): Read | undefined { + let end = start; + while (end < text.length && !`${SPACE},:{}[]"`.includes(text[end]!)) end++; + const token = text.slice(start, end); + if (end < text.length) return { value: JSON.parse(token) as JsonValue, complete: true, end }; + if (!pendingScalars) return undefined; + // At the end of the text: a literal prefix, a number, or a number whose exponent has no digits yet. + for (const [literal, value] of [ + ["true", true], + ["false", false], + ["null", null], + ] as const) { + if (literal.startsWith(token)) return { value, complete: false, end }; + } + if (/^-?(?:0|[1-9]\d*)(?:\.\d+)?(?:[eE][+-]?\d+)?$/.test(token)) + return { value: Number(token), complete: false, end }; + const mantissa = /^(-?(?:0|[1-9]\d*)(?:\.\d+)?)[eE][+-]?$/.exec(token); + return mantissa === null ? undefined : { value: Number(mantissa[1]), complete: false, end }; +} + +type Step = { readonly read: string; readonly snapshot: JsonValue | undefined; readonly copy: JsonValue | undefined }; + +/** Feed `text` in pieces; returns the snapshot after each piece, a copy taken right then, and the text read by then. */ +function stream(text: string, chunkSize: number | Random): Step[] { + const parser = new StreamingJsonParser(); + const steps: Step[] = []; + for (let index = 0; index < text.length; ) { + const size = typeof chunkSize === "number" ? chunkSize : 1 + Math.floor(chunkSize() * 12); + parser.append(text.slice(index, index + size)); + index += size; + const snapshot = parser.snapshot(); + steps.push({ read: text.slice(0, index), snapshot, copy: structuredClone(snapshot) }); + } + return steps; +} + +function finalSnapshot(text: string, chunkSize = text.length || 1): JsonValue | undefined { + return stream(text, chunkSize).at(-1)?.snapshot; +} + +const CHUNK_SIZES: readonly (number | "random")[] = [1, 2, 3, 7, "random"]; + +/** A stable seed from the text, so random chunk sizes differ per document. */ +function read32(text: string): number { + let hash = 7; + for (let index = 0; index < text.length; index++) hash = (hash * 31 + text.charCodeAt(index)) | 0; + return Math.abs(hash) + 1; +} + +describe("StreamingJsonParser differential fuzzing", () => { + it("matches the reference after every piece, at every chunk size, with insignificant whitespace", () => { + for (const text of randomDocuments(1, 250, true)) { + for (const size of CHUNK_SIZES) { + for (const { read, snapshot } of stream(text, size === "random" ? random(read32(text)) : size)) { + expect(snapshot, `${JSON.stringify(read)} (chunk size ${size})`).toEqual(referenceSnapshot(read)); + } + } + } + }); + + it("matches today's parseStreamingJson at every prefix of compact documents", () => { + for (const text of randomDocuments(2, 400, false)) { + for (const { read, snapshot } of stream(text, 1)) { + // Documented differences: parseStreamingJson trims trailing whitespace, and shows an unfinished root + // `null` as {}. + if (/\s$/.test(read)) continue; + const expected = parseStreamingJson(read); + if (snapshot === null && "null".startsWith(read)) continue; + expect(snapshot ?? {}, JSON.stringify(read)).toEqual(expected); + } + } + }); + + it("reads complete documents exactly like JSON.parse", () => { + for (const text of [...randomDocuments(3, 300, true), ...randomDocuments(4, 300, false)]) { + const expected = JSON.parse(text) as JsonValue; + for (const size of CHUNK_SIZES) { + expect(finalSnapshot(text, size === "random" ? 5 : size), text).toEqual(expected); + } + } + }); + + it("never changes a snapshot once it is returned", () => { + for (const text of randomDocuments(5, 150, true)) { + for (const { snapshot, copy } of stream(text, random(read32(text)))) expect(snapshot).toEqual(copy); + } + }); + + it("repairs raw control characters and invalid escapes like parseStreamingJson, and shows them while streaming", () => { + const rand = random(6); + const malformed = ["\n", "\t", "\u0000", "\\d", "\\x41", "\\u12", "\\uZZZZ", "\\u00e9", "\\\\", "\\q"]; + for (let round = 0; round < 300; round++) { + let content = ""; + for (let index = 0; index < 6; index++) content += pick(rand, [...malformed, "a", " ", "é"]); + const text = `{"path":"p","content":"${content}","n":${round}}`; + // Complete text: the same repaired value as today's parser. + expect(finalSnapshot(text, 1), JSON.stringify(text)).toEqual(parseStreamingJson(text)); + for (const { read, snapshot } of stream(text, 1)) { + expect(snapshot, JSON.stringify(read)).toEqual(referenceSnapshot(read)); + } + } + }); +}); + +describe("StreamingJsonParser strings", () => { + it("holds an escape back until it is complete, at every split position", () => { + for (const [text, value] of [ + [String.raw`{"a":"x\u00e9y"}`, "xéy"], + [String.raw`{"a":"x\ny"}`, "x\ny"], + [String.raw`{"a":"x\\y"}`, "x\\y"], + [String.raw`{"a":"x\"y"}`, 'x"y'], + [String.raw`{"a":"x\/y"}`, "x/y"], + ] as const) { + for (let split = 0; split <= text.length; split++) { + const parser = new StreamingJsonParser(); + parser.append(text.slice(0, split)); + expect(parser.snapshot(), text.slice(0, split)).toEqual(referenceSnapshot(text.slice(0, split))); + parser.append(text.slice(split)); + expect(parser.snapshot()).toEqual({ a: value }); + } + } + }); + + it("joins surrogate pairs split across pieces, escaped or raw", () => { + const escaped = String.raw`{"a":"\ud83d\ude00!"}`; + const raw = '{"a":"😀!"}'; + for (const text of [escaped, raw]) { + for (let split = 0; split <= text.length; split++) { + const parser = new StreamingJsonParser(); + parser.append(text.slice(0, split)); + parser.snapshot(); + parser.append(text.slice(split)); + expect(parser.snapshot()).toEqual({ a: "😀!" }); + } + } + // A high surrogate whose low half has not arrived shows alone, like JSON.parse() of the escape. + const parser = new StreamingJsonParser(); + parser.append(String.raw`{"a":"\ud83d`); + expect(parser.snapshot()).toEqual({ a: "\ud83d" }); + }); + + it("keeps trailing whitespace of an unterminated string", () => { + const parser = new StreamingJsonParser(); + parser.append('{"code":"x\t\t '); + expect(parser.snapshot()).toEqual({ code: "x\t\t " }); + }); + + it("keeps raw control characters and invalid escapes visible while streaming", () => { + const parser = new StreamingJsonParser(); + parser.append('{"a":"x\ny \\d \\u12"'); + expect(parser.snapshot()).toEqual({ a: "x\ny \\d \\u12" }); + }); + + it("keeps invalid \\u escapes literal without dropping fields (#10868)", () => { + expect(finalSnapshot(String.raw`{"path":"a.tex","content":"\usepackage{x}","mode":"w"}`, 1)).toEqual({ + path: "a.tex", + content: "\\usepackage{x}", + mode: "w", + }); + expect(finalSnapshot(String.raw`{"text":"caf\u00e9 \url"}`, 3)).toEqual({ text: "café \\url" }); + const parser = new StreamingJsonParser(); + parser.append(String.raw`{"text":"ab\u00`); + expect(parser.snapshot()).toEqual({ text: "ab" }); + }); + + it("makes `__proto__` an own member, as JSON.parse does", () => { + const text = '{"__proto__":{"polluted":true},"b":{"__proto__":"x"}}'; + for (let split = 0; split <= text.length; split++) { + const parser = new StreamingJsonParser(); + parser.append(text.slice(0, split)); + const partial = parser.snapshot() as { [key: string]: JsonValue } | undefined; + if (partial !== undefined) expect(Object.getPrototypeOf(partial)).toBe(Object.prototype); + parser.append(text.slice(split)); + const value = parser.snapshot() as { [key: string]: JsonValue }; + expect(value).toEqual(JSON.parse(text)); + expect(Object.getPrototypeOf(value)).toBe(Object.prototype); + expect(Object.hasOwn(value, "__proto__")).toBe(true); + } + expect(({} as { polluted?: boolean }).polluted).toBeUndefined(); + }); + + it("keeps the last value of a duplicate key in the first key's position, as JSON.parse does", () => { + const text = '{"a":1,"b":2,"a":"three"}'; + expect(Object.keys(finalSnapshot(text, 1) as object)).toEqual(Object.keys(JSON.parse(text))); + expect(finalSnapshot(text, 1)).toEqual(JSON.parse(text)); + }); +}); + +/** Expected values of the inputs, from other partial-JSON parsers' public test suites (see each table's source). */ +type Vector = readonly [input: string, expected: JsonValue | undefined]; + +// partial-json (MIT, https://github.com/promplate/partial-json-parser-js/blob/main/tests/examples.test.js), with every +// partial type allowed, as parseStreamingJson uses it. +const PARTIAL_JSON_VECTORS: readonly Vector[] = [ + ['"', ""], + ['["', [""]], + ['{"": "', { "": "" }], + ["n", null], + ["t", true], + ["f", false], + ["0", 0], + ["-1.25e+4", -1.25e4], + ["-1.25e+", -1.25], + ["-1.25e", -1.25], + // Differences, documented in streaming-json.ts: an invalid escape stays visible (partial-json cuts the string at + // the backslash), and `Infinity`/`NaN` words are not JSON and are dropped. + ['" \\x12', " \\x12"], + ["I", undefined], + ["-I", undefined], + ["N", undefined], +]; + +// Vercel AI SDK fixJson (Apache-2.0, https://github.com/vercel/ai/blob/main/packages/ai/src/util/fix-json.test.ts): +// the expected values are JSON.parse() of its repaired texts. +const FIX_JSON_VECTORS: readonly Vector[] = [ + ["", undefined], + ["nul", null], + ["t", true], + ["fals", false], + ["12.2", 12.2], + ["-12", -12], + ["-", undefined], + ["2.5e", 2.5], + ["2.5e-", 2.5], + ["2.5e3", 2.5e3], + ["-2.5e3", -2.5e3], + ["2.5E", 2.5], + ["2.5E-", 2.5], + ["2.5E3", 2.5e3], + ["-2.5E3", -2.5e3], + ["12.34e", 12.34], + ["5e", 5], + ['"abc', "abc"], + ['"value with \\"quoted\\" text and \\\\ escape', 'value with "quoted" text and \\ escape'], + ['"value with \\', "value with "], + ['"\\u', ""], + ['"\\u12', ""], + ['"text \\u00', "text "], + ['{"a":"\\u12', { a: "" }], + ['"value with unicode \u003C"', "value with unicode <"], + ["[", []], + ["[[1], [2", [[1], [2]]], + ['[["1"], ["2', [["1"], ["2"]]], + ["[[false], [nu", [[false], [null]]], + ["[[[]], [[]", [[[]], [[]]]], + ["[[{}], [{", [[{}], [{}]]], + ["[1, ", [1]], + ["[[], 123", [[], 123]], + ['{"key":', {}], + ['{"a": {"b": 1}, "c": {"d": 2', { a: { b: 1 }, c: { d: 2 } }], + ['{"a": {"b": "1"}, "c": {"d": 2', { a: { b: "1" }, c: { d: 2 } }], + ['{"a": {"b": false}, "c": {"d": 2', { a: { b: false }, c: { d: 2 } }], + ['{"a": {"b": []}, "c": {"d": 2', { a: { b: [] }, c: { d: 2 } }], + ['{"a": {"b": {}}, "c": {"d": 2', { a: { b: {} }, c: { d: 2 } }], + ['{"ke', {}], + ['{"k1": 1, "k2', { k1: 1 }], + ['{"k1": 1, "k2":', { k1: 1 }], + ['{"key": "value" ', { key: "value" }], + ['{"a": {"b": {}', { a: { b: {} } }], + ["[1, [2, 3, [", [1, [2, 3, []]]], + ["[false, [true, [", [false, [true, []]]], + ['{"key": {"subKey":', { key: {} }], + ['{"key": 123, "key2": {"subKey":', { key: 123, key2: {} }], + ['{"key": null, "key2": {"subKey":', { key: null, key2: {} }], + ['{"key": [1, 2, {', { key: [1, 2, {}] }], + ['[1, 2, {"key": "value",', [1, 2, { key: "value" }]], + ['{"a": {"b": ["c", {"d": "e",', { a: { b: ["c", { d: "e" }] } }], + ['{"a": {"b": {"c": {"d":', { a: { b: { c: {} } } }], + ['{"a": 1, "b": [', { a: 1, b: [] }], + ['{"a": 1, "b": {', { a: 1, b: {} }], + ['{"a": 1, "b": "', { a: 1, b: "" }], + [ + '{\n "a": [\n {\n "a1": "v1",\n "a2": "v2",\n "a3": "v3"\n }\n ],\n "b": [\n {\n "b1": "n', + { a: [{ a1: "v1", a2: "v2", a3: "v3" }], b: [{ b1: "n" }] }, + ], + ['{"type":"div","children":[{"type":"Card","props":{}', { type: "div", children: [{ type: "Card", props: {} }] }], + // Difference: fixJson completes `12.` to 12; like parseStreamingJson, a number cut after its dot is omitted. + ["12.", undefined], + ["12.e", undefined], +]; + +describe("StreamingJsonParser ported test vectors", () => { + it.each([...PARTIAL_JSON_VECTORS, ...FIX_JSON_VECTORS])("reads %j", (input, expected) => { + for (const size of [input.length || 1, 1, 2]) { + expect(finalSnapshot(input, size)).toEqual(expected); + } + }); + + it("reads every JSONTestSuite document JSON.parse accepts exactly like it, and never throws on the rest", () => { + const suite = JSON.parse(readFileSync(new URL("./data/json-test-suite.json", import.meta.url), "utf8")) as { + files: Record; + }; + for (const [name, base64] of Object.entries(suite.files)) { + const text = Buffer.from(base64, "base64").toString("utf8"); + let expected: JsonValue | undefined; + try { + expected = JSON.parse(text) as JsonValue; + } catch { + expected = undefined; + } + for (const size of [text.length || 1, 1, 3]) { + const snapshot = finalSnapshot(text, size); + if (expected !== undefined) expect(snapshot, name).toEqual(expected); + else expectJson(snapshot, name); + } + } + }); +}); + +/** Asserts `value` is plain, finite JSON data: what consumers such as Durable's strict JSON copy accept. */ +function expectJson(value: JsonValue | undefined, label: string): void { + if (value === undefined || value === null || typeof value === "string" || typeof value === "boolean") return; + if (typeof value === "number") { + expect(Number.isFinite(value), label).toBe(true); + return; + } + expect([Object.prototype, Array.prototype], label).toContain(Object.getPrototypeOf(value)); + for (const item of Object.values(value)) expectJson(item, label); +} + +describe("StreamingJsonParser malformed text", () => { + const MALFORMED = [ + '{"a":1 "b":2}', + '{"a" 1, "b": 2}', + "{a:1, 'b': 2}", + '{"a":1,,"b":2,}', + "[1 2, , 3,]", + '{"a":[1,2}, "b":3]', + "]]]}}}{", + '{"a":NaN,"b":Infinity,"c":-Infinity,"d":undefined,"e":1}', + '{"a":01,"b":1.,"c":.5,"d":+1,"e":0x10,"f":1e}', + '{"a":"x"} trailing garbage {"b":2}', + '{"a": tru, "b": nulll, "c": falsey}', + '{"a":{"b":{"c":', + '{"a\\": 1}', + '"unterminated \\', + '// comment\n{"a":1}', + '\u0000{\u0001"a"\u0002:\u00031}', + '{"a":"x"}}}]]]', + ]; + + it("never throws, keeps snapshots plain JSON and unchanged, at every chunk size", () => { + for (const text of MALFORMED) { + for (const size of CHUNK_SIZES) { + const steps = stream(text, size === "random" ? random(read32(text)) : size); + for (const { snapshot, copy } of steps) { + expectJson(snapshot, text); + expect(snapshot).toEqual(copy); + } + // Every chunk size reads malformed text the same way. + expect(steps.at(-1)?.snapshot, text).toEqual(finalSnapshot(text)); + } + } + }); + + it("tolerates missing commas and colons and drops values in invalid positions", () => { + expect(finalSnapshot('{"a":1 "b":2}')).toEqual({ a: 1, b: 2 }); + expect(finalSnapshot('{"a" 1, "b": 2}')).toEqual({ a: 1, b: 2 }); + expect(finalSnapshot("[1 2, , 3,]")).toEqual([1, 2, 3]); + expect(finalSnapshot('{"a":NaN,"b":Infinity,"e":1}')).toEqual({ e: 1 }); + expect(finalSnapshot('{"a":"x"} trailing garbage {"b":2}')).toEqual({ a: "x" }); + }); + + it("survives random corruption of valid documents", () => { + const rand = random(7); + const noise = ['"', "\\", "{", "}", "[", "]", ",", ":", "x", "1", "-", "\n", "t"]; + for (const text of randomDocuments(8, 300, true)) { + const characters = [...text]; + for (let edit = 0; edit < 3; edit++) { + const at = Math.floor(rand() * (characters.length + 1)); + if (rand() < 0.5) characters.splice(at, 1); + else characters.splice(at, 0, pick(rand, noise)); + } + const corrupted = characters.join(""); + for (const { snapshot } of stream(corrupted, random(read32(corrupted)))) expectJson(snapshot, corrupted); + // The changes keep a copy equal to the parser's own value, malformed or not. + const parser = new StreamingJsonParser(); + const mirror = new Mirror(); + for (let index = 0; index < corrupted.length; index += 4) + mirror.apply(parser.append(corrupted.slice(index, index + 4))); + // A space ends a number or literal at the end, which snapshots show and changes do not report yet. + mirror.apply(parser.append(" ")); + expect(mirror.value, corrupted).toEqual(parser.snapshot()); + } + }); +}); + +/** A copy of a value kept current by applying the parser's changes, as a writer such as Durable keeps one. */ +class Mirror { + value: JsonValue | undefined = undefined; + + apply(changes: readonly JsonChange[]): void { + for (const change of changes) { + const key = change.path.at(-1); + let parent = this.value as Record; + for (const segment of change.path.slice(0, -1)) parent = parent[segment] as Record; + if (change.kind === "set") { + const value = structuredClone(change.value); + if (key === undefined) this.value = value; + else parent[key] = value; + } else if (key === undefined) { + this.value = `${this.value as string}${change.text}`; + } else { + parent[key] = `${parent[key] as string}${change.text}`; + } + } + } +} + +/** Feed `text` in pieces into a parser and a mirror of its changes; checks the mirror after each piece. */ +function mirrorStream(text: string, size: number | Random): Mirror { + const parser = new StreamingJsonParser(); + const mirror = new Mirror(); + for (let index = 0; index < text.length; ) { + const length = typeof size === "number" ? size : 1 + Math.floor(size() * 12); + mirror.apply(parser.append(text.slice(index, index + length))); + index += length; + // Changes leave out unterminated numbers and literals, which snapshots show. + expect(mirror.value, JSON.stringify(text.slice(0, index))).toEqual( + referenceSnapshot(text.slice(0, index), { pendingScalars: false }), + ); + } + return mirror; +} + +describe("StreamingJsonParser changes", () => { + it("keep a copy equal to the value read so far, at every chunk size, and to JSON.parse at the end", () => { + for (const text of randomDocuments(9, 200, true)) { + for (const size of CHUNK_SIZES) { + const mirror = mirrorStream(text, size === "random" ? random(read32(text)) : size); + expect(mirror.value, text).toEqual(JSON.parse(text)); + } + } + }); + + it("report a number that ends and a string that starts in one piece, in order", () => { + const parser = new StreamingJsonParser(); + expect(parser.append('{"n":1')).toEqual([{ kind: "set", path: [], value: {} }]); + expect(parser.append('2,"s":"ab')).toEqual([ + { kind: "set", path: ["n"], value: 12 }, + { kind: "set", path: ["s"], value: "ab" }, + ]); + expect(parser.append('c","t":tru')).toEqual([{ kind: "append", path: ["s"], text: "c" }]); + expect(parser.append("e}")).toEqual([{ kind: "set", path: ["t"], value: true }]); + expect(parser.append(' {"ignored":1}')).toEqual([]); + }); + + it("report nested objects and arrays as empty containers, then their members and items", () => { + const parser = new StreamingJsonParser(); + expect(parser.append('{"a":[1,{"b":["x"')).toEqual([ + { kind: "set", path: [], value: {} }, + { kind: "set", path: ["a"], value: [] }, + { kind: "set", path: ["a", 0], value: 1 }, + { kind: "set", path: ["a", 1], value: {} }, + { kind: "set", path: ["a", 1, "b"], value: [] }, + { kind: "set", path: ["a", 1, "b", 0], value: "x" }, + ]); + expect(parser.append(',"y')).toEqual([{ kind: "set", path: ["a", 1, "b", 1], value: "y" }]); + expect(parser.append('z"]}],"a":"again"}')).toEqual([ + { kind: "append", path: ["a", 1, "b", 1], text: "z" }, + { kind: "set", path: ["a"], value: "again" }, + ]); + }); + + it("hold split escapes back and report their decoded text once complete", () => { + const text = String.raw`{"a":"x\u00e9\n\\\"\ud83d\ude00\q"}`; + const expected = (JSON.parse(String.raw`{"a":"x\u00e9\n\\\"\ud83d\ude00\\q"}`) as { a: string }).a; + for (let first = 1; first < text.length; first++) { + for (let second = first; second < text.length; second++) { + const parser = new StreamingJsonParser(); + const mirror = new Mirror(); + for (const piece of [text.slice(0, first), text.slice(first, second), text.slice(second)]) { + mirror.apply(parser.append(piece)); + } + expect(mirror.value).toEqual({ a: expected }); + } + } + const parser = new StreamingJsonParser(); + parser.append('{"a":"x'); + expect(parser.append("\\u00")).toEqual([]); + expect(parser.append("e9")).toEqual([{ kind: "append", path: ["a"], text: "é" }]); + }); + + it("report a container nested deeper than 32 levels once, whole, when it closes", () => { + const depth = 40; + const text = `${'{"a":'.repeat(depth)}"deep"${"}".repeat(depth)}`; + const parser = new StreamingJsonParser(); + const mirror = new Mirror(); + for (let index = 0; index < text.length; index += 3) mirror.apply(parser.append(text.slice(index, index + 3))); + expect(mirror.value).toEqual(JSON.parse(text)); + const changes = new StreamingJsonParser().append(text); + // The containers down to path length 32, then the one at length 32 again, whole, as it closes. + expect(changes).toHaveLength(34); + const last = changes.at(-1)!; + expect(last.kind === "set" && last.path.length).toBe(32); + expect(last.kind === "set" && last.value).toEqual(JSON.parse(text.slice(32 * 5, -32))); + }); + + it("report nothing for values in invalid positions and text after the root", () => { + const parser = new StreamingJsonParser(); + expect(parser.append('{"a":1 2, "b" "x"}')).toEqual([ + { kind: "set", path: [], value: {} }, + { kind: "set", path: ["a"], value: 1 }, + { kind: "set", path: ["b"], value: "x" }, + ]); + expect(parser.append('{"c":3}')).toEqual([]); + }); +}); + +describe("StreamingJsonParser cost", () => { + const content = Array.from( + { length: 20_000 }, + (_, line) => `const line${line} = "value ${line}"; // \\ done\n`, + ).join(""); + + /** `write`-style arguments with `size` characters of file content, streamed in 32-character deltas. */ + const writeArguments = (size: number) => JSON.stringify({ path: "/src/file.ts", content: content.slice(0, size) }); + + async function streamAllocation(text: string, delta: number): Promise { + return allocatedBytes(() => { + const args = new ToolCallArgumentsStream(); + for (let index = 0; index < text.length; index += delta) args.append(text.slice(index, index + delta)); + return args.finish(); + }); + } + + it("allocates linearly in the streamed text, never per accumulated text", async () => { + const small = writeArguments(50_000); + const large = writeArguments(500_000); + const smallBytes = await streamAllocation(small, 32); + const largeBytes = await streamAllocation(large, 32); + // Measured about 15 bytes per character; re-parsing the accumulated text cost about 1 MB per character. + expect(largeBytes / large.length).toBeLessThan(200); + // Ten times the text: linear growth is about 10x, quadratic about 100x. + expect(largeBytes / smallBytes).toBeLessThan(20); + }); + + it("keeps very large open arrays linear by letting their preview lag", async () => { + const elements = (count: number) => JSON.stringify({ items: Array.from({ length: count }, (_, index) => index) }); + const small = elements(10_000); + const large = elements(100_000); + const smallBytes = await streamAllocation(small, 16); + const largeBytes = await streamAllocation(large, 16); + expect(largeBytes / large.length).toBeLessThan(2000); + expect(largeBytes / smallBytes).toBeLessThan(20); + + const parser = new StreamingJsonParser(); + let shown = 0; + for (let index = 0; index < large.length; index += 16) { + parser.append(large.slice(index, index + 16)); + const items = (parser.snapshot() as { items?: number[] } | undefined)?.items ?? []; + // The preview only moves forward. + expect(items.length).toBeGreaterThanOrEqual(shown); + shown = items.length; + } + expect(parser.snapshot()).toEqual(JSON.parse(large)); + }); + + it("reads deep nesting without recursion and in linear time", () => { + const depth = 200_000; + const text = `${"[".repeat(depth)}1${"]".repeat(depth)}`; + const started = performance.now(); + const parser = new StreamingJsonParser(); + for (let index = 0; index < text.length; index += 7) { + parser.append(text.slice(index, index + 7)); + parser.snapshot(); + } + let value = parser.snapshot(); + for (let level = 0; level < depth; level++) value = (value as JsonValue[])[0]!; + expect(value).toBe(1); + expect(performance.now() - started).toBeLessThan(5_000); + }); + + it("reads one huge piece, such as a provider's complete arguments, in linear time", async () => { + const text = writeArguments(500_000); + const bytes = await allocatedBytes(() => { + const parser = new StreamingJsonParser(); + parser.append(text); + return parser.snapshot(); + }); + expect(bytes / text.length).toBeLessThan(100); + }); +}); + +describe("ToolCallArgumentsStream", () => { + it("shows partial arguments as an object and finishes with one full parse", () => { + const args = new ToolCallArgumentsStream(); + expect(args.append(" ")).toEqual({}); + expect(args.append('{"path":"a.ts","content":"x')).toEqual({ path: "a.ts", content: "x" }); + const before = args.append("y"); + expect(args.append("")).toBe(before); + expect(args.text).toBe(' {"path":"a.ts","content":"xy'); + expect(args.finish()).toEqual(parseToolCallArguments(' {"path":"a.ts","content":"xy')); + args.append('"}'); + expect(args.finish()).toEqual({ arguments: { path: "a.ts", content: "xy" } }); + }); + + it("shows non-object roots as no arguments yet", () => { + const args = new ToolCallArgumentsStream(); + expect(args.append("[1,2")).toEqual({}); + args.append("]"); + expect(args.finish().argumentsError).toMatch(/^Tool arguments must be a JSON object/); + }); + + it("replaces the text with the provider's complete text, returning the missing piece", () => { + const args = new ToolCallArgumentsStream(); + args.append('{"a":"x'); + expect(args.replace('{"a":"xyz"}')).toEqual({ delta: 'yz"}', arguments: { a: "xyz" } }); + expect(args.replace('{"b":2}')).toEqual({ delta: undefined, arguments: { b: 2 } }); + expect(args.text).toBe('{"b":2}'); + }); + + it("is collected with its snapshots once the stream is dropped", async () => { + const references: WeakRef[] = []; + (() => { + const args = new ToolCallArgumentsStream(); + for (let index = 0; index < 100; index++) references.push(new WeakRef(args.append(`{"a${index}":"${index}`))); + references.push(new WeakRef(args)); + })(); + await collectGarbage(); + expect(references.filter((reference) => reference.deref() !== undefined)).toEqual([]); + }); +}); 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 000000000..5b685e577 --- /dev/null +++ b/packages/ai/test/streaming-tool-arguments.test.ts @@ -0,0 +1,420 @@ +import type { ResponseStreamEvent } from "openai/resources/responses/responses.js"; +import { describe, expect, it } 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 streamOpenAICompletions } from "../src/api/openai-completions.ts"; +import { processResponsesStream } from "../src/api/openai-responses-shared.ts"; +import { stream as streamPiMessages } from "../src/api/pi-messages.ts"; +import { fauxAssistantMessage, fauxProvider, fauxToolCall } from "../src/providers/faux.ts"; +import type { AssistantMessage, AssistantMessageEvent, Model, TranscriptContext } from "../src/types.ts"; +import { AssistantMessageEventStream } from "../src/utils/event-stream.ts"; +import { normalizeContext } from "../src/utils/transcript.ts"; +import { allocatedBytes, collectGarbage } from "./memory-helpers.ts"; + +// #9265: every provider re-parsed the accumulated tool-call arguments on each delta and kept the raw text on the +// content block, where it reached partial messages, Durable commits and transcripts. + +const content = Array.from({ length: 20_000 }, (_, line) => `const line${line} = "value ${line}"; // \\ done\n`).join( + "", +); + +/** `write` arguments with `size` characters of file content, and their JSON text. */ +function writeCall(size: number) { + const args = { path: "/src/file.ts", content: content.slice(0, size) }; + return { args, text: JSON.stringify(args) }; +} + +function deltasOf(text: string, size: number): string[] { + const deltas: string[] = []; + for (let index = 0; index < text.length; index += size) deltas.push(text.slice(index, index + size)); + return deltas; +} + +/** A fetch whose response streams `events` as server-sent events, produced lazily. */ +function sseFetch(events: () => Iterable): typeof fetch { + return async () => { + const iterator = events()[Symbol.iterator](); + const encoder = new TextEncoder(); + const body = new ReadableStream({ + pull(controller) { + let chunk = ""; + for (let count = 0; count < 16; count++) { + const next = iterator.next(); + if (next.done) { + if (chunk) controller.enqueue(encoder.encode(chunk)); + controller.close(); + return; + } + chunk += next.value; + } + controller.enqueue(encoder.encode(chunk)); + }, + }); + return new Response(body, { status: 200, headers: { "content-type": "text/event-stream" } }); + }; +} + +const sse = (data: unknown, event?: string) => `${event ? `event: ${event}\n` : ""}data: ${JSON.stringify(data)}\n\n`; + +const baseModel = { + reasoning: false, + input: ["text" as const], + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }, + contextWindow: 200_000, + maxTokens: 64_000, +}; + +const context: TranscriptContext = normalizeContext({ messages: [{ role: "user", content: "write", timestamp: 0 }] }); + +function* anthropicEvents(deltas: readonly string[]): Generator { + yield sse( + { + type: "message_start", + message: { id: "msg_1", model: "claude-test", usage: { input_tokens: 1, output_tokens: 1 } }, + }, + "message_start", + ); + yield sse( + { + type: "content_block_start", + index: 0, + content_block: { type: "tool_use", id: "toolu_1", name: "write", input: {} }, + }, + "content_block_start", + ); + for (const delta of deltas) { + yield sse( + { type: "content_block_delta", index: 0, delta: { type: "input_json_delta", partial_json: delta } }, + "content_block_delta", + ); + } + yield sse({ type: "content_block_stop", index: 0 }, "content_block_stop"); + yield sse( + { type: "message_delta", delta: { stop_reason: "tool_use" }, usage: { output_tokens: 10 } }, + "message_delta", + ); + yield sse({ type: "message_stop" }, "message_stop"); +} + +function* chatEvents(deltas: readonly string[], withIdEveryChunk: boolean): Generator { + const base = { id: "chat_1", object: "chat.completion.chunk", created: 1, model: "test" }; + const call = (fields: object) => ({ + ...base, + choices: [{ index: 0, delta: { tool_calls: [{ index: 0, ...fields }] } }], + }); + yield sse(call({ id: "call_1", type: "function", function: { name: "write", arguments: "" } })); + for (const delta of deltas) { + yield sse( + call( + withIdEveryChunk + ? { id: "call_1", function: { name: "write", arguments: delta } } + : { function: { arguments: delta } }, + ), + ); + } + yield sse({ + ...base, + choices: [{ index: 0, delta: {}, finish_reason: "tool_calls" }], + usage: { prompt_tokens: 1, completion_tokens: 1, total_tokens: 2 }, + }); + yield "data: [DONE]\n\n"; +} + +function* piMessagesEvents(deltas: readonly string[], args: object): Generator { + yield sse({ type: "start" }); + yield sse({ type: "toolcall_start", contentIndex: 0, id: "call_1", toolName: "write" }); + for (const delta of deltas) yield sse({ type: "toolcall_delta", contentIndex: 0, delta }); + yield sse({ + type: "toolcall_end", + contentIndex: 0, + toolCall: { type: "toolCall", id: "call_1", name: "write", arguments: args }, + }); + yield sse({ + type: "done", + reason: "toolUse", + usage: { + input: 1, + output: 1, + cacheRead: 0, + cacheWrite: 0, + totalTokens: 2, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }, + }); +} + +function* responsesEvents(deltas: readonly string[], text: string): Generator { + const item = { + type: "function_call", + id: "fc_1", + call_id: "call_1", + name: "write", + arguments: "", + status: "in_progress", + }; + yield { type: "response.output_item.added", output_index: 0, sequence_number: 0, item } as ResponseStreamEvent; + for (const delta of deltas) { + yield { + type: "response.function_call_arguments.delta", + output_index: 0, + item_id: "fc_1", + sequence_number: 0, + delta, + } as ResponseStreamEvent; + } + yield { + type: "response.function_call_arguments.done", + output_index: 0, + item_id: "fc_1", + sequence_number: 0, + arguments: text, + } as ResponseStreamEvent; + yield { + type: "response.output_item.done", + output_index: 0, + sequence_number: 0, + item: { ...item, arguments: text, status: "completed" }, + } as ResponseStreamEvent; + yield { + type: "response.completed", + sequence_number: 0, + response: { + id: "resp_1", + status: "completed", + output: [], + usage: { input_tokens: 1, output_tokens: 1, total_tokens: 2 }, + }, + } as unknown as ResponseStreamEvent; +} + +type Path = { + readonly name: string; + /** Stream one `write` call whose argument text arrives as `deltas`. */ + readonly stream: (deltas: readonly string[], args: object, text: string) => AssistantMessageEventStream; +}; + +const PATHS: readonly Path[] = [ + { + name: "anthropic-messages", + stream: (deltas) => + streamAnthropic( + { + ...baseModel, + id: "claude-test", + name: "c", + api: "anthropic-messages", + provider: "anthropic", + baseUrl: "https://api.anthropic.com", + }, + context, + { apiKey: "sk-ant-test", fetch: sseFetch(() => anthropicEvents(deltas)) }, + ), + }, + { + name: "openai-completions", + stream: (deltas) => + streamOpenAICompletions( + { + ...baseModel, + id: "gpt-test", + name: "g", + api: "openai-completions", + provider: "openai", + baseUrl: "https://api.openai.com/v1", + }, + context, + { apiKey: "sk-test", fetch: sseFetch(() => chatEvents(deltas, false)) }, + ), + }, + { + name: "openai-responses", + stream: (deltas, _args, text) => { + const output: AssistantMessage = { + role: "assistant", + content: [], + api: "openai-responses", + provider: "openai", + model: "gpt-test", + usage: { + input: 0, + output: 0, + cacheRead: 0, + cacheWrite: 0, + totalTokens: 0, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }, + stopReason: "pending", + timestamp: 0, + }; + const events = new AssistantMessageEventStream(); + const model = { + ...baseModel, + id: "gpt-test", + name: "g", + api: "openai-responses", + provider: "openai", + baseUrl: "https://api.openai.com/v1", + } satisfies Model<"openai-responses">; + void (async () => { + events.push({ type: "start", partial: output }); + await processResponsesStream( + (async function* () { + yield* responsesEvents(deltas, text); + })(), + output, + events, + model, + ); + events.push({ type: "done", reason: "toolUse", message: output }); + events.end(); + })(); + return events; + }, + }, + { + name: "mistral-conversations", + stream: (deltas) => + streamMistral( + { + ...baseModel, + id: "devstral", + name: "d", + api: "mistral-conversations", + provider: "mistral", + baseUrl: "https://api.mistral.ai", + }, + context, + { apiKey: "test", fetch: sseFetch(() => chatEvents(deltas, true)) }, + ), + }, + { + name: "pi-messages", + stream: (deltas, args) => + streamPiMessages( + { + ...baseModel, + id: "auto", + name: "a", + api: "pi-messages", + provider: "radius", + baseUrl: "https://example.test/v1", + }, + context, + { apiKey: "test", fetch: sseFetch(() => piMessagesEvents(deltas, args)) }, + ), + }, + { + name: "faux", + stream: (_deltas, args) => { + const faux = fauxProvider({ tokenSize: { min: 8, max: 8 } }); + faux.setResponses([ + fauxAssistantMessage([fauxToolCall("write", args as never, { id: "call_1" })], { stopReason: "toolUse" }), + ]); + return faux.provider.streamSimple!(faux.getModel(), context, {}) as AssistantMessageEventStream; + }, + }, +]; + +/** Own keys a content block may carry; anything else is stream state leaking into partial messages. */ +const BLOCK_KEYS: Record = { + text: ["type", "text", "textSignature"], + thinking: ["type", "thinking", "thinkingSignature", "redacted"], + toolCall: ["type", "id", "name", "arguments", "argumentsError", "thoughtSignature", "namespace"], +}; + +function expectNoStreamState(message: AssistantMessage, label: string): void { + for (const block of message.content) { + for (const key of Object.keys(block)) + expect(BLOCK_KEYS[block.type], `${label}: ${block.type}.${key}`).toContain(key); + } +} + +async function consume( + events: AssistantMessageEventStream, + onEvent: (event: AssistantMessageEvent) => void = () => {}, +) { + for await (const event of events) onEvent(event); + return events.result(); +} + +describe.each(PATHS)("$name tool-call argument streaming", (path) => { + it("shows growing partial arguments, keeps stream state off every partial, and finishes with the exact arguments", async () => { + const { args, text } = writeCall(5_000); + let shown = -1; + let partials = 0; + const message = await consume(path.stream(deltasOf(text, 24), args, text), (event) => { + const partial = event.type === "done" ? event.message : event.type === "error" ? event.error : event.partial; + expectNoStreamState(partial, event.type); + if (event.type !== "toolcall_delta") return; + const call = partial.content[event.contentIndex]; + const shownContent = call?.type === "toolCall" ? call.arguments.content : undefined; + if (typeof shownContent !== "string") return; + partials++; + expect(args.content.startsWith(shownContent)).toBe(true); + expect(shownContent.length).toBeGreaterThanOrEqual(shown); + shown = shownContent.length; + }); + expect(message.stopReason).toBe("toolUse"); + expect(partials).toBeGreaterThan(100); + expect(message.content).toEqual([{ type: "toolCall", id: expect.any(String), name: "write", arguments: args }]); + }); + + it("allocates linearly in the argument text", async () => { + const run = (size: number) => { + const { args, text } = writeCall(size); + return allocatedBytes(() => consume(path.stream(deltasOf(text, 32), args, text))); + }; + const small = await run(25_000); + const large = await run(250_000); + // Ten times the text: linear growth is about 10x, re-parsing the accumulated text about 100x. + expect(large / small).toBeLessThan(25); + expect(large / writeCall(250_000).text.length).toBeLessThan(5_000); + }); + + it("releases the stream, its partial messages and argument snapshots once the caller drops them", async () => { + const references: WeakRef[] = []; + await (async () => { + const { args, text } = writeCall(2_000); + const events = path.stream(deltasOf(text, 16), args, text); + references.push(new WeakRef(events)); + await consume(events, (event) => { + if (event.type !== "toolcall_delta") return; + const call = event.partial.content[event.contentIndex]; + if (call?.type === "toolCall") references.push(new WeakRef(call.arguments)); + }); + })(); + await collectGarbage(); + expect(references.filter((reference) => reference.deref() !== undefined)).toHaveLength(0); + }); +}); + +describe("tool-call stream state on errors", () => { + it("leaves no stream state on the error message of a stream cut off mid-call", async () => { + const { text } = writeCall(2_000); + const cut = deltasOf(text, 16).slice(0, 20); + const message = await consume( + streamAnthropic( + { + ...baseModel, + id: "claude-test", + name: "c", + api: "anthropic-messages", + provider: "anthropic", + baseUrl: "https://api.anthropic.com", + }, + context, + { + apiKey: "sk-ant-test", + fetch: sseFetch(function* () { + const events = [...anthropicEvents(cut)]; + // Drop the block stop and the message end, as a dropped connection does. + yield* events.slice(0, -3); + }), + }, + ), + ); + expect(message.stopReason).toBe("error"); + expectNoStreamState(message, "error"); + expect(message.content[0]).toMatchObject({ type: "toolCall", arguments: { path: "/src/file.ts" } }); + }); +}); diff --git a/packages/ai/test/telemetry.test.ts b/packages/ai/test/telemetry.test.ts index 41b9ba457..6412ffe22 100644 --- a/packages/ai/test/telemetry.test.ts +++ b/packages/ai/test/telemetry.test.ts @@ -1,3 +1,4 @@ +import { getHeapStatistics } from "node:v8"; import { InMemoryTelemetryContext, type TelemetrySpan } from "@earendil-works/pi-telemetry"; import { describe, expect, it, vi } from "vitest"; import { lazyStream } from "../src/api/lazy.ts"; @@ -5,6 +6,7 @@ import { AI_TELEMETRY_SCHEMA, createModels, createProvider } from "../src/models import { observeModelRequest } from "../src/telemetry.ts"; import type { AssistantMessage, AssistantMessageEvent, DeferredHandle, Model } from "../src/types.ts"; import { AssistantMessageEventStream } from "../src/utils/event-stream.ts"; +import { collectGarbage } from "./memory-helpers.ts"; const model: Model<"telemetry-test"> = { id: "model", @@ -195,6 +197,27 @@ describe("AI telemetry", () => { expect(JSON.stringify(span.events)).not.toContain("xxx"); }); + // #9265 + it("keeps a bounded error message that does not hold the provider's whole message alive", async () => { + const telemetry = new InMemoryTelemetryContext(); + await collectGarbage(); + const before = getHeapStatistics().used_heap_size; + await (async () => { + const result = message("error"); + result.errorMessage = `${"x".repeat(4_000_000)}tail`; + const { models } = setup(result); + await telemetry.startSpan({ name: "caller" }, (span) => + models.fetchDeferred(model, handle, { telemetry: span }), + ); + await settled(telemetry); + })(); + await collectGarbage(); + const [, span] = telemetry.getSpans(); + expect(span?.status).toEqual({ status: "error", error: { name: "provider_error", message: "x".repeat(500) } }); + // A slice of the 4 MB message would keep all of it alive with the span. + expect(getHeapStatistics().used_heap_size - before).toBeLessThan(2_000_000); + }); + it("includes setup failures and omits non-scalar native reasoning", async () => { const telemetry = new InMemoryTelemetryContext(); const models = createModels(); diff --git a/packages/chord/CHANGELOG.md b/packages/chord/CHANGELOG.md index 92189ad9d..aa6634e23 100644 --- a/packages/chord/CHANGELOG.md +++ b/packages/chord/CHANGELOG.md @@ -13,6 +13,11 @@ - `Watch`, the one watch handle implementation, moved from Durable: a source creates one per watch and feeds it with `advance(value, ops, context)`, `retire(context)` (last frame `[["r", null]]`, then `retired`), `end()` or `end({ error })`, and `observeCancellation(signal)`. More than 100 undelivered frames collapse to one root replacement `[["r", newest]]`; termination drops the listener and every reference into the source. - `Watchable`: a value that changes and can be watched (`value`, `watch(context)` returning a `WatchHandle`), the common interface of documents, derives and other observable values. - `NativeWatchHandle` and its implementation `NativeWatch`: the same lifecycle for in-process values, with listener `(value, context)` and no ops; an overflow collapses to the newest value. +- `ownedSlice(text, start, end?)` in `@earendil-works/chord/delta`: a slice of a string that does not keep the whole string alive ([#9265](https://github.com/earendil-works/pi/issues/9265)) + +### Fixed + +- Fixed append ops from the delta tracker and `diffRevisions()` keeping the whole string they were cut from alive (V8 represents a slice as a view of its source), so a delta chain, a watch backlog or a replica built from appends to a growing text held one copy of the text per commit ([#9265](https://github.com/earendil-works/pi/issues/9265)) ## [1.4.0] - 2026-10-21 diff --git a/packages/chord/src/delta/diff.ts b/packages/chord/src/delta/diff.ts index 774ad1b59..ff0bf2121 100644 --- a/packages/chord/src/delta/diff.ts +++ b/packages/chord/src/delta/diff.ts @@ -1,5 +1,5 @@ import type { JsonValue } from "../json.ts"; -import { type NonEmptyPath, type Op, overlap, type Path, RESERVED_SEGMENTS, type Seg } from "./index.ts"; +import { type NonEmptyPath, type Op, overlap, ownedSlice, type Path, RESERVED_SEGMENTS, type Seg } from "./index.ts"; const DEFAULT_OVERLAP_SCAN = 65_536; const MAX_DELTA_OPERATIONS = 4_096; @@ -60,10 +60,11 @@ const permutation = (before: readonly JsonValue[], after: readonly JsonValue[]): return result; }; +// O(length): comparing two revisions reads both strings, and slicing a rope makes V8 flatten all of it. const emitString = (before: string, after: string, path: NonEmptyPath, operations: Op[]): void => { if (before === after) return; if (after.length > before.length && after.slice(0, before.length) === before) { - emitOperation(operations, ["a", path, after.slice(before.length)]); + emitOperation(operations, ["a", path, ownedSlice(after, before.length)]); return; } const shared = overlap(before, after, DEFAULT_OVERLAP_SCAN); @@ -72,7 +73,7 @@ const emitString = (before: string, after: string, path: NonEmptyPath, operation return; } emitOperation(operations, ["t", path, before.length - shared]); - if (after.length > shared) emitOperation(operations, ["a", path, after.slice(shared)]); + if (after.length > shared) emitOperation(operations, ["a", path, ownedSlice(after, shared)]); }; type ArrayMatch = readonly [before: number, after: number]; diff --git a/packages/chord/src/delta/index.ts b/packages/chord/src/delta/index.ts index ad4ed57d6..11d94b47c 100644 --- a/packages/chord/src/delta/index.ts +++ b/packages/chord/src/delta/index.ts @@ -109,6 +109,16 @@ export function overlap(a: string, b: string, scan: number, probe = 64, maxCandi return 0; } +/** + * `text.slice(start, end)` as a string of its own. V8 represents a slice of 13 or more characters as a view that keeps + * all of `text` alive, so an append op cut from a long, growing string would retain that whole string for as long as + * the op lives: in a delta chain, a watch backlog, or a value built by applying the op. Appending a character and + * slicing it off again makes V8 copy just the slice. + */ +export function ownedSlice(text: string, start: number, end = text.length): string { + return `${text.slice(start, end)} `.slice(0, -1); +} + // ─── Immutable revision tracking ───────────────────────────────────────────── export { diffRevisions } from "./diff.ts"; diff --git a/packages/chord/src/delta/tracker.ts b/packages/chord/src/delta/tracker.ts index 8e01705fa..bb267bd78 100644 --- a/packages/chord/src/delta/tracker.ts +++ b/packages/chord/src/delta/tracker.ts @@ -2,7 +2,15 @@ import type { JsonValue } from "../json.ts"; import { copyJson } from "../json.ts"; import { applyImmutableTrusted } from "./apply-immutable-trusted.ts"; import type { Draft } from "./draft.ts"; -import { applyImmutable, type NonEmptyPath, type Op, overlap, type Path, RESERVED_SEGMENTS } from "./index.ts"; +import { + applyImmutable, + type NonEmptyPath, + type Op, + overlap, + ownedSlice, + type Path, + RESERVED_SEGMENTS, +} from "./index.ts"; export type { Draft, JsonValue, Op }; @@ -2045,16 +2053,17 @@ function emitChangedValue( if (!isContainer(after) && before === after) return false; if (isContainer(before) && isContainer(after) && equalTrustedJson(before, after)) return false; if (typeof before === "string" && typeof after === "string") { - // Slice/equality stays on V8's rope fast path; startsWith scans the entire - // accumulated prefix on every append. + // O(length): any slice of a rope makes V8 flatten all of it, so this prefix check copies the whole string once per + // change, even when the writer only appended. The op text gets a copy of its own, so it does not keep the whole + // of `after` alive. if (after.length > before.length && after.slice(0, before.length) === before) { - operations.push(["a", path, after.slice(before.length)]); + operations.push(["a", path, ownedSlice(after, before.length)]); return true; } const shared = overlap(before, after, 65_536); if (shared > 0) { operations.push(["t", path, before.length - shared]); - if (after.length > shared) operations.push(["a", path, after.slice(shared)]); + if (after.length > shared) operations.push(["a", path, ownedSlice(after, shared)]); return true; } } diff --git a/packages/chord/test/delta-op-text-retention.test.ts b/packages/chord/test/delta-op-text-retention.test.ts new file mode 100644 index 000000000..906feb56b --- /dev/null +++ b/packages/chord/test/delta-op-text-retention.test.ts @@ -0,0 +1,80 @@ +import { getHeapStatistics, setFlagsFromString } from "node:v8"; +import { runInNewContext } from "node:vm"; +import { describe, expect, it } from "vitest"; +import { applyImmutable, diffRevisions, type Op, track } from "../src/delta/index.ts"; + +// #9265: an append op's text was a V8 sliced string of the whole new string, so every op kept a full copy of the +// string it was cut from alive: a delta chain or a replica built from N appends to a growing text held N copies. + +const gc = (() => { + setFlagsFromString("--expose-gc"); + return runInNewContext("gc") as () => void; +})(); + +/** Heap bytes `build` leaves reachable through its result, and the result. */ +function retained(build: () => T): { bytes: number; kept: T } { + gc(); + const before = getHeapStatistics().used_heap_size; + const kept = build(); + gc(); + return { bytes: getHeapStatistics().used_heap_size - before, kept }; +} + +const BASE = "x".repeat(100_000); +const COMMITS = 300; +const piece = (index: number) => `line ${index} of a long streamed text\n`; + +describe("append op text", () => { + it("tracker append ops do not retain the string they were cut from", () => { + const { bytes } = retained(() => { + const tracker = track({ text: BASE }); + const ops: Op[] = []; + for (let index = 0; index < COMMITS; index++) { + const change = tracker.beginChange(); + // A rope, as a streaming writer builds it. + change.state.text = tracker.value.text + piece(index); + const prepared = change.prepare(); + ops.push(...prepared.ops); + tracker.adopt(prepared); + } + expect(ops.every((op) => op[0] === "a")).toBe(true); + return { ops, value: tracker.value }; + }); + // One copy of the text plus the op texts: well under a megabyte. Sliced op texts held one copy per op (30 MB). + expect(bytes).toBeLessThan(2_000_000); + }); + + it("diffRevisions append ops do not retain the string they were cut from", () => { + const { bytes } = retained(() => { + const ops: Op[] = []; + let text = BASE; + for (let index = 0; index < COMMITS; index++) { + const next = text + piece(index); + ops.push(...diffRevisions({ text }, { text: next })); + text = next; + } + expect(ops.every((op) => op[0] === "a")).toBe(true); + return { ops, text }; + }); + expect(bytes).toBeLessThan(2_000_000); + }); + + it("a replica that applies a long stream of appends retains the text once, not once per commit", () => { + const { bytes, kept } = retained(() => { + const tracker = track({ text: BASE }); + // Like a watch, a remote replica or Durable's conversation view: a value kept current by applying ops. + let replica: { text: string } = { text: BASE }; + for (let index = 0; index < COMMITS; index++) { + const change = tracker.beginChange(); + change.state.text = tracker.value.text + piece(index); + const prepared = change.prepare(); + tracker.adopt(prepared); + replica = applyImmutable(replica, prepared.ops); + } + // Comparing here would flatten the replica's text and drop what it retains; compare after measuring. + return { replica, tracker }; + }); + expect(bytes).toBeLessThan(2_000_000); + expect(kept.replica).toEqual(kept.tracker.value); + }); +}); diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index c1b5fee7a..246c1225e 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -20,6 +20,7 @@ - Fixed HTML export showing a blank page for sessions with thousands of entries ([#11225](https://github.com/earendil-works/pi/issues/11225), [#11293](https://github.com/earendil-works/pi/pull/11293)) - Fixed HTML export not showing system messages: the system prompt and tools, and later section updates such as after `/reload`. `pi --export` now includes the system prompt and tools too ([#11261](https://github.com/earendil-works/pi/issues/11261), [#3464](https://github.com/earendil-works/pi/issues/3464), [#11290](https://github.com/earendil-works/pi/pull/11290)) - Fixed aborting a bash command on macOS sometimes running the next command of a list such as `long_step; next_step`. pi now stops the process group before killing it ([#11252](https://github.com/earendil-works/pi/issues/11252), [#11289](https://github.com/earendil-works/pi/pull/11289)) +- Fixed streamed tool calls re-running their call renderer on every argument delta; the interactive mode now renders the latest arguments once per frame, so a large streamed `write` no longer re-highlights its whole content per delta ([#9265](https://github.com/earendil-works/pi/issues/9265)) ## [1.4.0] - 2026-10-21 diff --git a/packages/coding-agent/examples/extensions/custom-provider-anthropic/index.ts b/packages/coding-agent/examples/extensions/custom-provider-anthropic/index.ts index 5f116651f..c0c18f505 100644 --- a/packages/coding-agent/examples/extensions/custom-provider-anthropic/index.ts +++ b/packages/coding-agent/examples/extensions/custom-provider-anthropic/index.ts @@ -42,7 +42,7 @@ import { type TextContent, type ThinkingContent, type Tool, - type ToolCall, + ToolCallArgumentsStream, type ToolResultMessage, type TranscriptContext, } from "@earendil-works/pi-ai"; @@ -454,8 +454,9 @@ function streamCustomAnthropic( const anthropicStream = client.messages.stream({ ...params }, { signal: options?.signal }); stream.push({ type: "start", partial: output }); - type Block = (ThinkingContent | TextContent | (ToolCall & { partialJson: string })) & { index: number }; - const blocks = output.content as Block[]; + // Stream state per Anthropic block index, kept off the content blocks so it never reaches partial messages. + // ToolCallArgumentsStream reads each argument delta once and shows the arguments parsed so far. + const openBlocks = new Map(); for await (const event of anthropicStream) { if (event.type === "message_start") { @@ -468,15 +469,12 @@ function streamCustomAnthropic( calculateCost(model, output.usage); } else if (event.type === "content_block_start") { if (event.content_block.type === "text") { - output.content.push({ type: "text", text: "", index: event.index } as any); + output.content.push({ type: "text", text: "" }); + openBlocks.set(event.index, { contentIndex: output.content.length - 1 }); stream.push({ type: "text_start", contentIndex: output.content.length - 1, partial: output }); } else if (event.content_block.type === "thinking") { - output.content.push({ - type: "thinking", - thinking: "", - thinkingSignature: "", - index: event.index, - } as any); + output.content.push({ type: "thinking", thinking: "", thinkingSignature: "" }); + openBlocks.set(event.index, { contentIndex: output.content.length - 1 }); stream.push({ type: "thinking_start", contentIndex: output.content.length - 1, partial: output }); } else if (event.content_block.type === "tool_use") { output.content.push({ @@ -484,15 +482,18 @@ function streamCustomAnthropic( id: event.content_block.id, name: isOAuth ? fromClaudeCodeName(event.content_block.name, tools) : event.content_block.name, arguments: {}, - partialJson: "", - index: event.index, - } as any); + }); + openBlocks.set(event.index, { + contentIndex: output.content.length - 1, + arguments: new ToolCallArgumentsStream(), + }); stream.push({ type: "toolcall_start", contentIndex: output.content.length - 1, partial: output }); } } else if (event.type === "content_block_delta") { - const index = blocks.findIndex((b) => b.index === event.index); - const block = blocks[index]; - if (!block) continue; + const open = openBlocks.get(event.index); + if (!open) continue; + const index = open.contentIndex; + const block = output.content[index]!; if (event.delta.type === "text_delta" && block.type === "text") { block.text += event.delta.text; @@ -505,11 +506,8 @@ function streamCustomAnthropic( delta: event.delta.thinking, partial: output, }); - } else if (event.delta.type === "input_json_delta" && block.type === "toolCall") { - (block as any).partialJson += event.delta.partial_json; - try { - block.arguments = JSON.parse((block as any).partialJson); - } catch {} + } else if (event.delta.type === "input_json_delta" && block.type === "toolCall" && open.arguments) { + block.arguments = open.arguments.append(event.delta.partial_json); stream.push({ type: "toolcall_delta", contentIndex: index, @@ -520,20 +518,19 @@ function streamCustomAnthropic( block.thinkingSignature = (block.thinkingSignature || "") + (event.delta as any).signature; } } else if (event.type === "content_block_stop") { - const index = blocks.findIndex((b) => b.index === event.index); - const block = blocks[index]; - if (!block) continue; + const open = openBlocks.get(event.index); + if (!open) continue; + openBlocks.delete(event.index); + const index = open.contentIndex; + const block = output.content[index]!; - delete (block as any).index; if (block.type === "text") { stream.push({ type: "text_end", contentIndex: index, content: block.text, partial: output }); } else if (block.type === "thinking") { stream.push({ type: "thinking_end", contentIndex: index, content: block.thinking, partial: output }); - } else if (block.type === "toolCall") { - try { - block.arguments = JSON.parse((block as any).partialJson); - } catch {} - delete (block as any).partialJson; + } else if (block.type === "toolCall" && open.arguments) { + // One full parse of the complete text; malformed arguments are reported in `argumentsError`. + Object.assign(block, open.arguments.finish()); stream.push({ type: "toolcall_end", contentIndex: index, toolCall: block, partial: output }); } } else if (event.type === "message_delta") { @@ -563,7 +560,6 @@ function streamCustomAnthropic( stream.push({ type: "done", reason: output.stopReason, message: output }); stream.end(); } catch (error) { - for (const block of output.content) delete (block as any).index; output.stopReason = options?.signal?.aborted ? "aborted" : "error"; output.errorMessage = error instanceof Error ? error.message : JSON.stringify(error); stream.push({ type: "error", reason: output.stopReason, error: output }); diff --git a/packages/coding-agent/src/modes/interactive/components/tool-execution.ts b/packages/coding-agent/src/modes/interactive/components/tool-execution.ts index bf1cd27f9..e739fc503 100644 --- a/packages/coding-agent/src/modes/interactive/components/tool-execution.ts +++ b/packages/coding-agent/src/modes/interactive/components/tool-execution.ts @@ -54,6 +54,8 @@ export class ToolExecutionComponent extends Container { private cwd: string; private executionStarted = false; private argsComplete = false; + /** `updateArgs()` ran since the last display update; the next render applies it. */ + private argsPending = false; private result?: { content: Array<{ type: string; text?: string; data?: string; mimeType?: string }>; isError: boolean; @@ -167,9 +169,14 @@ export class ToolExecutionComponent extends Container { }); } + /** + * Show new partial arguments. Applied at the next render, so a call streaming many argument deltas between two + * frames runs its renderer once per frame, not once per delta (#9265). + */ updateArgs(args: any): void { this.args = args; - this.updateDisplay(); + this.argsPending = true; + this.ui.requestRender(); } markExecutionStarted(): void { @@ -225,6 +232,7 @@ export class ToolExecutionComponent extends Container { } override render(width: number): string[] { + if (this.argsPending) this.updateDisplay(); if (this.hideComponent) { return []; } @@ -268,6 +276,7 @@ export class ToolExecutionComponent extends Container { } private updateDisplay(): void { + this.argsPending = false; const bgFn = this.isPartial ? (text: string) => theme.bg("toolPendingBg", text) : this.result?.isError diff --git a/packages/coding-agent/test/tool-execution-component.test.ts b/packages/coding-agent/test/tool-execution-component.test.ts index f334f7d6b..f88091049 100644 --- a/packages/coding-agent/test/tool-execution-component.test.ts +++ b/packages/coding-agent/test/tool-execution-component.test.ts @@ -531,6 +531,59 @@ describe("ToolExecutionComponent parity", () => { expect(rendered).not.toContain("two\n\n"); }); + // #9265: every argument delta re-ran the call renderer; a streamed `write` re-highlighted its whole content each time. + test("renders streamed partial arguments once per frame, with the latest arguments", () => { + const write = createWriteToolDefinition(process.cwd()); + let renderCalls = 0; + const counted: typeof write = { + ...write, + renderCall: (...args) => { + renderCalls++; + return write.renderCall!(...args); + }, + }; + let renderRequests = 0; + const tui = { requestRender: () => renderRequests++ } as unknown as TUI; + const component = new ToolExecutionComponent("write", "tool-9", {}, {}, counted, tui, "."); + const lines = Array.from({ length: 200 }, (_, index) => `const line${index} = ${index};`); + const callsBefore = renderCalls; + let content = ""; + for (const line of lines) { + content += `${line}\n`; + component.updateArgs({ path: "src/file.ts", content }); + } + expect(renderCalls).toBe(callsBefore); + expect(renderRequests).toBeGreaterThan(0); + + const streamed = stripAnsi(component.render(120).join("\n")); + expect(renderCalls - callsBefore).toBe(1); + // The same preview as rendering the final arguments directly. + const direct = new ToolExecutionComponent( + "write", + "tool-10", + { path: "src/file.ts", content }, + {}, + write, + tui, + ".", + ); + expect(streamed).toBe(stripAnsi(direct.render(120).join("\n"))); + + // More deltas before the next frame: one more renderer run, which shows them all. + component.updateArgs({ path: "src/file.ts", content: `${content}const tail = 1;\n` }); + const final = { path: "src/file.ts", content: `${content}const tail = 2;\n` }; + component.updateArgs(final); + const next = stripAnsi(component.render(120).join("\n")); + expect(renderCalls - callsBefore).toBe(2); + expect(next).toContain("201 total"); + expect(next).toBe( + stripAnsi(new ToolExecutionComponent("write", "t", final, {}, write, tui, ".").render(120).join("\n")), + ); + // Rendering again without new arguments does not run the renderer. + component.render(120); + expect(renderCalls - callsBefore).toBe(2); + }); + test("trims trailing blank display lines from read results", () => { const component = new ToolExecutionComponent( "read", diff --git a/packages/durable/CHANGELOG.md b/packages/durable/CHANGELOG.md index 76511515a..41cf5091a 100644 --- a/packages/durable/CHANGELOG.md +++ b/packages/durable/CHANGELOG.md @@ -17,6 +17,12 @@ - Document, conversation view, task graph, and event watches use Chord's `Watch` implementation; behaviour is unchanged. +### Fixed + +- Fixed a conversation watcher (`watchEvents()`, conversation views) holding a copy of a streamed text or tool call per partial commit, through Chord append ops that kept the whole string alive; a 140 KB answer followed by a 140 KB tool call held well over 100 MB ([#9265](https://github.com/earendil-works/pi/issues/9265)) +- Fixed a nested call's error summary keeping the call's whole output alive through a slice of it ([#9265](https://github.com/earendil-works/pi/issues/9265)) +- Fixed task span error messages and codemode call previews keeping the whole text they were cut from alive ([#9265](https://github.com/earendil-works/pi/issues/9265)) + ## [1.4.0] - 2026-10-21 ### Breaking Changes diff --git a/packages/durable/src/codemode/codemode.ts b/packages/durable/src/codemode/codemode.ts index d1f04d223..dba1d1e57 100644 --- a/packages/durable/src/codemode/codemode.ts +++ b/packages/durable/src/codemode/codemode.ts @@ -35,6 +35,7 @@ * call with an error result. */ import { type Context, copyJson, type JsonValue } from "@earendil-works/chord"; +import { ownedSlice } from "@earendil-works/chord/delta"; import type { TextContent } from "@earendil-works/pi-ai"; import { type CodemodeHostCall, @@ -198,7 +199,8 @@ const ARGS_PREVIEW_CHARS = 200; const ERROR_PREVIEW_CHARS = 500; function cut(text: string, max: number): string { - return text.length > max ? `${text.slice(0, max - 3)}...` : text; + // A copy: the row outlives the arguments or error text it previews. + return text.length > max ? `${ownedSlice(text, 0, max - 3)}...` : text; } /** A function scripts call as a top-level global; its result is journaled. */ diff --git a/packages/durable/src/harness/scheduler.ts b/packages/durable/src/harness/scheduler.ts index 0ce1b3de3..8fde30b63 100644 --- a/packages/durable/src/harness/scheduler.ts +++ b/packages/durable/src/harness/scheduler.ts @@ -1,5 +1,6 @@ import { type Context, copyJson, type JsonValue } from "@earendil-works/chord"; import { awaitWithContext, withAbortSignal, withoutAbortSignal, withTelemetry } from "@earendil-works/chord/context"; +import { ownedSlice } from "@earendil-works/chord/delta"; import type { Models } from "@earendil-works/pi-ai"; import { createTypedMetrics, @@ -1801,7 +1802,7 @@ export class TaskScheduler { const message = terminal.status === "orphaned" ? terminal.reason : terminal.error.message; span.setStatus({ status: "error", - error: { name: `task_${terminal.status}`, message: message.slice(0, 500) }, + error: { name: `task_${terminal.status}`, message: ownedSlice(message, 0, 500) }, }); } else if (outcome === "aborted") { span.setStatus({ status: "ok" }); diff --git a/packages/durable/src/harness/tool.ts b/packages/durable/src/harness/tool.ts index 52bc4a3ad..215e8f063 100644 --- a/packages/durable/src/harness/tool.ts +++ b/packages/durable/src/harness/tool.ts @@ -1,6 +1,6 @@ import { type Context, copyJson, type Draft, type JsonRepresentation, type JsonValue } from "@earendil-works/chord"; import { awaitWithContext } from "@earendil-works/chord/context"; -import { overlap } from "@earendil-works/chord/delta"; +import { overlap, ownedSlice } from "@earendil-works/chord/delta"; import { boundOutput, OutputBuffer, @@ -864,9 +864,12 @@ function truncated( /** Error text a nested call's summary keeps. */ const MAX_SUMMARY_ERROR_CHARS = 500; -/** How a nested call ended, for its slot: error text, from its diagnostics or else its `output`, bounded. */ +/** + * How a nested call ended, for its slot: error text, from its diagnostics or else its `output`, bounded. The bound is a + * copy, so the slot never keeps a large output alive through a slice of it. + */ function summaryOf(result: NestedToolExecutionResult, output: Content): NestedToolSummary { - const error = result.isError ? errorTextOf(result, output).slice(0, MAX_SUMMARY_ERROR_CHARS) : ""; + const error = result.isError ? ownedSlice(errorTextOf(result, output), 0, MAX_SUMMARY_ERROR_CHARS) : ""; return { isError: result.isError, ...(result.durationMs === undefined ? {} : { durationMs: result.durationMs }), diff --git a/packages/durable/test/harness-events.test.ts b/packages/durable/test/harness-events.test.ts index 7171f7a86..ce136a674 100644 --- a/packages/durable/test/harness-events.test.ts +++ b/packages/durable/test/harness-events.test.ts @@ -1,3 +1,5 @@ +import { getHeapStatistics, setFlagsFromString } from "node:v8"; +import { runInNewContext } from "node:vm"; import type { Draft } from "@earendil-works/chord"; import { type AssistantMessage, @@ -620,3 +622,90 @@ describe("agent events", () => { await harness.close(context); }); }); + +const gc = (() => { + setFlagsFromString("--expose-gc"); + return runInNewContext("gc") as () => void; +})(); + +/** Heap bytes in use after a full collection. */ +async function heapInUse(): Promise { + for (let round = 0; round < 3; round++) { + await drained(); + gc(); + } + return getHeapStatistics().used_heap_size; +} + +// #9265 +describe("streamed partials", () => { + it("streams a tool-call partial as toolcall_delta appends at its argument paths", async () => { + const setup = chatSetup({ tokensPerSecond: 500, tokenSize: { min: 1, max: 1 } }); + setup.settings.progress = { partialIntervalMs: 20, outputIntervalMs: 100 }; + const args = { path: "src/file.ts", content: "const a = 1;\n".repeat(150) }; + setup.faux.setResponses([ + fauxAssistantMessage([fauxToolCall("missing", args, { id: "c1" })], { stopReason: "toolUse" }), + fauxAssistantMessage([fauxText("done")]), + ]); + const { harness, root } = await openChat(new MemoryStorage(), setup); + const partials = partialsOf(harness); + const { stream, events } = await listen(harness, root); + await (await root.submit({ type: "input", content: "go" }, context)).wait(context); + await drained(); + const changes = events().flatMap((event) => (event.type === "message_update" ? event.changes : [])); + const deltas = changes.filter((change) => change.type === "toolcall_delta"); + expect(deltas.length).toBeGreaterThan(10); + expect(deltas.every((change) => ["path", "content"].includes(String(change.path)))).toBe(true); + // The call is resent whole only when a member appears (at most once per argument); growing strings are appends. + // With the raw argument text on the block, every commit resent the whole call. + expect(changes.filter((change) => change.type === "message")).toEqual([]); + expect(changes.filter((change) => change.type === "block").length).toBeLessThanOrEqual(2); + // The committed partials carry only the parsed arguments, never the raw argument text. + for (const partial of partials) { + for (const block of partial.content) { + expect(Object.keys(block).sort()).toEqual(["arguments", "id", "name", "type"]); + } + } + await stream.stop(); + await harness.close(context); + }); + + it("keeps a watcher's memory proportional to a long streamed text and tool call, not to their commits", async () => { + const setup = chatSetup({ tokensPerSecond: 16_000, tokenSize: { min: 16, max: 16 } }); + setup.settings.progress = { partialIntervalMs: 1, outputIntervalMs: 100 }; + const line = (index: number) => `line ${index}: some streamed source text with "quotes" and \\ escapes\n`; + const text = Array.from({ length: 2_000 }, (_, index) => line(index)).join(""); + const content = Array.from({ length: 2_000 }, (_, index) => line(index)).join(""); + setup.faux.setResponses([ + fauxAssistantMessage([fauxText(text), fauxToolCall("missing", { path: "a.ts", content }, { id: "c1" })], { + stopReason: "toolUse", + }), + fauxAssistantMessage([fauxText("done")]), + ]); + const { harness, root } = await openChat(new MemoryStorage(), setup); + const { stream, batches } = await listen(harness, root); + const before = await heapInUse(); + const submission = await root.submit({ type: "input", content: "go" }, context); + const partialContent = async () => + (await harness.snapshot(LiveDoc, root.id, context))?.generation?.message?.content; + const shown = async (index: number): Promise => { + const block = (await partialContent())?.[index]; + if (block?.type === "text") return block.text.length; + return block?.type === "toolCall" ? String(block.arguments.content ?? "").length : 0; + }; + const retained: number[] = []; + // Measure while the text, then the tool call, is almost complete: the watcher's view holds the most then. + await waitFor(async () => (await shown(0)) > text.length * 0.95, 30_000); + retained.push((await heapInUse()) - before); + await waitFor(async () => (await shown(1)) > content.length * 0.95, 30_000); + retained.push((await heapInUse()) - before); + await submission.wait(context); + const commits = batches.length; + expect(commits).toBeGreaterThan(500); + // The two 140 KB strings, the watcher's recorded batches and the rest of the stream: about 3 MB. Sliced append ops + // made the watcher's view hold a copy of the text per commit: well over 100 MB. + for (const bytes of retained) expect(bytes).toBeLessThan(20_000_000); + await stream.stop(); + await harness.close(context); + }, 60_000); +}); diff --git a/packages/mcp/CHANGELOG.md b/packages/mcp/CHANGELOG.md index 412df6f3f..74908a649 100644 --- a/packages/mcp/CHANGELOG.md +++ b/packages/mcp/CHANGELOG.md @@ -5,6 +5,7 @@ ### Fixed - OAuth metadata parsing and `startAuthorization()` now accept only `http:` and `https:` URLs. Before, a server could supply a `file:` or custom protocol authorization endpoint (for example `search-ms:`), which the platform URL handler would launch on sign-in. +- Fixed the excerpt of an HTTP error body in Streamable HTTP errors keeping the whole response body alive ([#9265](https://github.com/earendil-works/pi/issues/9265)) ## [1.4.0] - 2026-10-21 diff --git a/packages/mcp/src/transports/streamable-http.ts b/packages/mcp/src/transports/streamable-http.ts index af9c922d2..5ea76682e 100644 --- a/packages/mcp/src/transports/streamable-http.ts +++ b/packages/mcp/src/transports/streamable-http.ts @@ -338,7 +338,8 @@ export class StreamableHttpTransport extends TransportEvents implements McpTrans private async checkResponse(response: Response): Promise { if (response.ok) return; - const body = (await response.text().catch(() => "")).slice(0, MAX_ERROR_BODY_BYTES); + // A copy, not a slice: V8 slices keep the whole body alive, and the error keeps this excerpt. + const body = `${(await response.text().catch(() => "")).slice(0, MAX_ERROR_BODY_BYTES)} `.slice(0, -1); if (response.status === 401) throw new McpAuthRequiredError(response, body); if (response.status === 404 && this.sessionIdValue) throw new McpSessionExpiredError(body); throw new McpHttpError(response.status, describeHttpFailure(response.status, body), body);