diff --git a/packages/aikit/src/llm/stream.ts b/packages/aikit/src/llm/stream.ts index 38d68c1f..4221e99f 100644 --- a/packages/aikit/src/llm/stream.ts +++ b/packages/aikit/src/llm/stream.ts @@ -4,6 +4,7 @@ import * as Message from "../message/message.ts"; import * as Model from "../model/model.ts"; import { AssistantMessageEventStream } from "../utils/eventstream.ts"; import { compact } from "../utils/helpers.ts"; +import { ToolCallArgumentsStream } from "../utils/streamjson.ts"; import * as Failure from "./failure.ts"; import { Options, type ProviderOptionBag, type RuntimeOptions } from "./options.ts"; import * as Pricing from "./pricing.ts"; @@ -27,14 +28,12 @@ import { mapFinishReason, mapUsage, toolCallFromPart, - updateToolCallFromInput, - type StreamingToolCallBlock, } from "./transform.ts"; type TextBlock = Message.TextContent & { streamId?: string }; type ThinkingBlock = Message.ThinkingContent & { streamId?: string }; -type StreamBlock = TextBlock | ThinkingBlock | StreamingToolCallBlock; +type StreamBlock = TextBlock | ThinkingBlock | Message.ToolCallPendingPart; function mergeProviderOptions(...sources: Array): ProviderOptionBag { const result: ProviderOptionBag = {}; @@ -148,16 +147,16 @@ function ensureToolCallBlock( name: string, stream: AssistantMessageEventStream, emitStart: boolean, -): StreamingToolCallBlock { +): Message.ToolCallPendingPart { const existing = output.parts.find((part) => part.type === "toolCall" && part.callID === id) as - | StreamingToolCallBlock + | Message.ToolCallPendingPart | undefined; if (existing) { if (!existing.name && name) existing.name = name; return existing; } - const block: StreamingToolCallBlock = { + const block: Message.ToolCallPendingPart = { type: "toolCall", callID: id, name, @@ -209,11 +208,10 @@ function finishThinkingBlock( function finishToolInput(output: Message.AssistantMessage, id: string, stream: AssistantMessageEventStream): void { const block = output.parts.find((part) => part.type === "toolCall" && part.callID === id) as - | StreamingToolCallBlock + | Message.ToolCallPendingPart | undefined; if (!block) return; const index = partIndex(output, block); - delete block.partialJson; stream.push({ type: "toolcall.end", partIndex: index, toolCall: block, partial: output }); } @@ -224,7 +222,7 @@ function finalizeToolCall( stream: AssistantMessageEventStream, ): void { const existing = output.parts.find((item) => item.type === "toolCall" && item.callID === part.toolCallId) as - | StreamingToolCallBlock + | Message.ToolCallPendingPart | undefined; const block = existing ?? toolCallFromPart(part); block.callID = part.toolCallId; @@ -236,7 +234,6 @@ function finalizeToolCall( if (thoughtSignature) block.thoughtSignature = thoughtSignature; if (typeof namespace === "string") block.namespace = namespace; block.time.end = Date.now(); - delete block.partialJson; if (!existing) { output.parts.push(block); @@ -263,6 +260,7 @@ function handlePart( model: Model.Info, stream: AssistantMessageEventStream, pricing: Pricing, + toolArgs: Map, ): void { switch (part.type) { case "text-start": @@ -315,11 +313,14 @@ function handlePart( } case "tool-input-start": ensureToolCallBlock(output, part.id, part.toolName, stream, true); + toolArgs.set(part.id, new ToolCallArgumentsStream()); break; case "tool-input-delta": { const block = ensureToolCallBlock(output, part.id, "", stream, true); - const partialJson = (block.partialJson ?? "") + part.delta; - updateToolCallFromInput(block, partialJson); + let args = toolArgs.get(part.id); + if (!args) toolArgs.set(part.id, (args = new ToolCallArgumentsStream())); + block.arguments = args.append(part.delta); + block.time.end = Date.now(); stream.push({ type: "toolcall.delta", partIndex: partIndex(output, block), @@ -329,9 +330,11 @@ function handlePart( break; } case "tool-input-end": + toolArgs.delete(part.id); finishToolInput(output, part.id, stream); break; case "tool-call": + toolArgs.delete(part.toolCallId); finalizeToolCall(output, part, model, stream); break; case "start-step": @@ -485,11 +488,12 @@ export const stream: Protocol.StreamFunction(); stream.push({ type: "start", partial: output }); for (;;) { try { for await (const part of streamText(built.params).fullStream) { - handlePart(part, output, model, stream, pricing); + handlePart(part, output, model, stream, pricing, toolArgs); } break; } catch (error) { @@ -531,7 +535,6 @@ export const stream: Protocol.StreamFunction, { type: "tool-call" }>, -): StreamingToolCallBlock { +): Message.ToolCallPendingPart { return { type: "toolCall", callID: part.toolCallId, @@ -808,9 +803,3 @@ export function toolCallFromPart( }, }; } - -export function updateToolCallFromInput(block: StreamingToolCallBlock, partialJson: string): void { - block.partialJson = partialJson; - block.arguments = parseStreamingJson>(partialJson); - block.time.end = Date.now(); -} diff --git a/packages/aikit/src/utils/streamjson.ts b/packages/aikit/src/utils/streamjson.ts new file mode 100644 index 00000000..8fe5240a --- /dev/null +++ b/packages/aikit/src/utils/streamjson.ts @@ -0,0 +1,443 @@ +/** + * 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. Incomplete text: + * + * - 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: 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. + * + * 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 shows a preview that lags a little behind the + * text, and no input makes the total work grow faster than the text. + */ + +export type JsonValue = string | number | boolean | null | JsonValue[] | { [key: string]: JsonValue }; + +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 }; + +/** 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; + /** 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; +}; + +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; + + /** Read the next piece of text. Costs time linear in `text.length`. */ + append(text: string): void { + this.#credit += CREDIT_PER_APPEND + CREDIT_PER_CHARACTER * text.length; + 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++]!); + } + } + + /** + * 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: "" }; + return; + } + const slot = this.#valueSlot(); + this.#string = { target: "value", slot, text: "", escape: "" }; + this.#place(slot, ""); + } + + /** 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. + 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. + 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") this.#place(open.slot, open.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); + 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(); + this.#place(slot, container); + this.#stack.push({ + container, + slot, + size: 0, + key: undefined, + phase: Array.isArray(container) ? "value" : "key", + }); + this.#pathCost += 1; + this.#dirty = true; + } + + #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; + 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; + } +} + +/** + * Partial arguments of one streaming tool call. Kept in the stream's own state, never on the content block, so the + * raw text never reaches partial messages; the final arguments come parsed with the provider's `tool-call` event. + */ +export class ToolCallArgumentsStream { + readonly #parser = new StreamingJsonParser(); + #snapshot: JsonValue | undefined = undefined; + #arguments: Record = {}; + + /** Append a piece of argument text; returns the partial arguments, a new object whenever they changed. */ + append(delta: string): Record { + this.#parser.append(delta); + 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 : {}; + 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/aikit/test/e2e/__artifacts__/streaming-tool-args.json b/packages/aikit/test/e2e/__artifacts__/streaming-tool-args.json new file mode 100644 index 00000000..6773b282 --- /dev/null +++ b/packages/aikit/test/e2e/__artifacts__/streaming-tool-args.json @@ -0,0 +1,33 @@ +{ + "argumentBytes": 618778, + "contentLength": 512064, + "deltas": 9220, + "snapshots": 9220, + "argumentsSha256": "26c7e051c674cbd4f8b49011779e969ccc238c61b075e7099bfc684549a9bfb3", + "events": [ + { + "type": "start", + "count": 1 + }, + { + "type": "toolcall.start", + "count": 1 + }, + { + "type": "toolcall.delta", + "count": 9220 + }, + { + "type": "toolcall.end", + "count": 1 + }, + { + "type": "toolcall.final", + "count": 1 + }, + { + "type": "done", + "count": 1 + } + ] +} diff --git a/packages/aikit/test/e2e/streaming.tool.args.e2e.test.ts b/packages/aikit/test/e2e/streaming.tool.args.e2e.test.ts new file mode 100644 index 00000000..1ae2a541 --- /dev/null +++ b/packages/aikit/test/e2e/streaming.tool.args.e2e.test.ts @@ -0,0 +1,169 @@ +import { createHash } from "node:crypto"; +import { createServer, type Server } from "node:http"; +import type { AddressInfo } from "node:net"; +import Type from "typebox"; +import { afterAll, beforeAll, describe, expect, it } from "vite-plus/test"; +import * as Message from "../../src/message/message.ts"; +import { stream } from "../../src/stream.ts"; +import { getAnthropicModel } from "../utils/llm.ts"; + +const writeTool = Message.defineTool({ + name: "write", + description: "Write a file", + parameters: Type.Object({ + path: Type.String(), + content: Type.String(), + }), +}); + +const LINE = 'const s = "say \\"hi\\"";\n\tif (a < b) return \'ok\'; // ünïcode 🚀 C:\\path\r\n'; +const content = LINE.repeat(Math.ceil((500 * 1024) / LINE.length)); +const expected = { path: "src/big.ts", content }; +const argumentText = JSON.stringify(expected); + +/** ~64-character pieces of the argument text, never splitting a surrogate pair. */ +function splitDeltas(text: string, size: number): string[] { + const pieces: string[] = []; + let start = 0; + while (start < text.length) { + let end = Math.min(start + size, text.length); + const last = text.charCodeAt(end - 1); + if (end < text.length && last >= 0xd800 && last <= 0xdbff) end++; + pieces.push(text.slice(start, end)); + start = end; + } + return pieces; +} + +const deltas = splitDeltas(argumentText, 64); + +function contentLength(args: Record): number { + return typeof args.content === "string" ? args.content.length : -1; +} + +function sse(type: string, data: Record): string { + return `event: ${type}\ndata: ${JSON.stringify({ type, ...data })}\n\n`; +} + +function anthropicStream(): string { + const events = [ + sse("message_start", { + message: { + id: "msg_local", + type: "message", + role: "assistant", + model: "claude-sonnet-4-6", + content: [], + stop_reason: null, + usage: { input_tokens: 10, output_tokens: 1 }, + }, + }), + sse("content_block_start", { + index: 0, + content_block: { type: "tool_use", id: "toolu_local", name: "write", input: {} }, + }), + ...deltas.map((partial_json) => + sse("content_block_delta", { index: 0, delta: { type: "input_json_delta", partial_json } }), + ), + sse("content_block_stop", { index: 0 }), + sse("message_delta", { + delta: { stop_reason: "tool_use", stop_sequence: null }, + usage: { output_tokens: deltas.length }, + }), + sse("message_stop", {}), + ]; + return events.join(""); +} + +describe("streaming tool arguments", () => { + let server: Server; + let baseURL: string; + + beforeAll(async () => { + const body = anthropicStream(); + server = createServer((request, response) => { + request.resume(); + request.on("end", () => { + response.writeHead(200, { "content-type": "text/event-stream", "cache-control": "no-cache" }); + response.end(body); + }); + }); + await new Promise((resolve) => server.listen(0, "127.0.0.1", resolve)); + baseURL = `http://127.0.0.1:${(server.address() as AddressInfo).port}/v1`; + }); + + afterAll(async () => { + await new Promise((resolve) => server.close(() => resolve())); + }); + + it("parses ~500 KB of streamed arguments in linear time", async () => { + const model = await getAnthropicModel("claude-sonnet-4-6"); + const context: Message.Context = { + messages: [ + Message.createUserMessage({ + role: "user", + parts: [{ type: "text", text: "Write src/big.ts" }], + time: { created: Date.now() }, + }), + ], + tools: [writeTool], + }; + + const types: { type: string; count: number }[] = []; + // Only the latest snapshot is kept: holding every one would retain every intermediate string. + let snapshots = 0; + let lastPartial: Record | undefined; + let lastLength = -1; + let mutated = false; + let deltaCount = 0; + let partialJsonSeen = false; + + const started = performance.now(); + const response = stream(model, context, { apiKey: "test-key", baseURL, maxRetries: 0 }); + for await (const event of response) { + const previous = types.at(-1); + if (previous?.type === event.type) previous.count++; + else types.push({ type: event.type, count: 1 }); + + if (event.type === "toolcall.delta") { + deltaCount++; + const part = event.partial.parts[event.partIndex]; + if (part?.type !== "toolCall") throw new Error("toolcall.delta without a toolCall part"); + if ("partialJson" in part) partialJsonSeen = true; + if (part.arguments !== lastPartial) { + // Snapshots are never mutated after they are handed out. + if (lastPartial && contentLength(lastPartial) !== lastLength) mutated = true; + snapshots++; + lastPartial = part.arguments; + lastLength = contentLength(lastPartial); + } + } + } + const message = await response.result(); + const elapsed = performance.now() - started; + + expect(message.stopReason).toBe("toolUse"); + const toolCall = message.parts.find((part) => part.type === "toolCall"); + expect(toolCall?.arguments).toEqual(expected); + expect(toolCall && "partialJson" in toolCall).toBe(false); + expect(partialJsonSeen).toBe(false); + expect(deltaCount).toBe(deltas.length); + expect(snapshots).toBeGreaterThan(1); + expect(mutated).toBe(false); + expect(typeof lastPartial?.content).toBe("string"); + expect(content.startsWith(lastPartial?.content as string)).toBe(true); + expect(elapsed).toBeLessThan(2000); + + const artifact = { + argumentBytes: Buffer.byteLength(argumentText, "utf8"), + contentLength: content.length, + deltas: deltas.length, + snapshots, + argumentsSha256: createHash("sha256").update(JSON.stringify(toolCall?.arguments)).digest("hex"), + events: types, + }; + await expect(JSON.stringify(artifact, null, 2) + "\n").toMatchFileSnapshot( + "./__artifacts__/streaming-tool-args.json", + ); + }); +});