diff --git a/packages/ai/src/internal/runtime.ts b/packages/ai/src/internal/runtime.ts index 537d2f146ff2..211c5148e930 100644 --- a/packages/ai/src/internal/runtime.ts +++ b/packages/ai/src/internal/runtime.ts @@ -13,3 +13,4 @@ export * from "../utils/reasoning-tag-text-partitioner.js"; export * from "../utils/sanitize-unicode.js"; export * from "../utils/stream-first-event-timeout.js"; export * from "../utils/streaming-byte-guard.js"; +export { parseTerminalToolCallArguments } from "../transports/transport-stream-shared.js"; diff --git a/packages/ai/src/providers/anthropic.ts b/packages/ai/src/providers/anthropic.ts index acc658d32ee5..4c84029327e4 100644 --- a/packages/ai/src/providers/anthropic.ts +++ b/packages/ai/src/providers/anthropic.ts @@ -53,7 +53,12 @@ import type { } from "../types.js"; import { createDeferredEventBuffer } from "../utils/deferred-event-buffer.js"; import { AssistantMessageEventStream } from "../utils/event-stream.js"; -import { parseJsonWithRepair, parseStreamingJson } from "../utils/json-parse.js"; +import { + createToolArgumentPreviewSchedule, + parseJsonWithRepair, + parseStreamingJson, + type ToolArgumentPreviewSchedule, +} from "../utils/json-parse.js"; import { notifyLlmRequestActivity } from "../utils/llm-request-activity.js"; import { projectProviderError } from "../utils/provider-error.js"; import { sanitizeSurrogates } from "../utils/sanitize-unicode.js"; @@ -432,6 +437,11 @@ export const streamAnthropic: StreamFunction<"anthropic-messages", AnthropicComp }; const blocks = output.content as Block[]; const blockIndexes = new Map(); + // Preview schedules are per active tool call; WeakMap keys die with the block. + const toolArgumentPreviewSchedules = new WeakMap< + Extract, + ToolArgumentPreviewSchedule + >(); const sealedToolCalls: Array<{ block: Extract; contentIndex: number; @@ -565,6 +575,7 @@ export const streamAnthropic: StreamFunction<"anthropic-messages", AnthropicComp }; output.content.push(block); blockIndexes.set(event.index, output.content.length - 1); + toolArgumentPreviewSchedules.set(block, createToolArgumentPreviewSchedule()); eventSink.push({ type: "toolcall_start", contentIndex: output.content.length - 1, @@ -604,7 +615,11 @@ export const streamAnthropic: StreamFunction<"anthropic-messages", AnthropicComp const block = index === undefined ? undefined : blocks[index]; if (index !== undefined && block?.type === "toolCall") { block.partialJson = (block.partialJson ?? "") + event.delta.partial_json; - block.arguments = parseStreamingJson(block.partialJson); + // Preview refresh is scheduled geometrically; the terminal + // finalize re-parses the full buffer authoritatively either way. + if (toolArgumentPreviewSchedules.get(block)?.(block.partialJson.length)) { + block.arguments = parseStreamingJson(block.partialJson); + } eventSink.push({ type: "toolcall_delta", contentIndex: index, diff --git a/packages/ai/src/providers/mistral.ts b/packages/ai/src/providers/mistral.ts index 97c1ca0572f2..791b784c6ea9 100644 --- a/packages/ai/src/providers/mistral.ts +++ b/packages/ai/src/providers/mistral.ts @@ -33,7 +33,11 @@ import type { } from "../types.js"; import { AssistantMessageEventStream } from "../utils/event-stream.js"; import { shortHash } from "../utils/hash.js"; -import { parseStreamingJson } from "../utils/json-parse.js"; +import { + createToolArgumentPreviewSchedule, + parseStreamingJson, + type ToolArgumentPreviewSchedule, +} from "../utils/json-parse.js"; import { sortPromptCacheToolsByName } from "../utils/prompt-cache-stability.js"; import { projectProviderError } from "../utils/provider-error.js"; import { sanitizeSurrogates } from "../utils/sanitize-unicode.js"; @@ -408,6 +412,11 @@ async function consumeChatStream( // Persist every identity fact across chunks. The SDK defaults omitted indexes // to zero, so only a unique compatible candidate may receive later arguments. const toolBlockIdentities = new Map(); + // Preview schedules are per active tool call; WeakMap keys die with the block. + const toolArgumentPreviewSchedules = new WeakMap< + ToolCall & { partialArgs?: string }, + ToolArgumentPreviewSchedule + >(); const normalizeMissingToolCallId = createMistralToolCallIdNormalizer(); // Some Mistral-compatible endpoints omit tool-call ids. Their streamed index // is only response-local, so namespace the fallback before strict-9 hashing. @@ -714,6 +723,7 @@ async function consumeChatStream( partialArgs: "", }; output.content.push(block); + toolArgumentPreviewSchedules.set(block, createToolArgumentPreviewSchedule()); toolBlockIdentities.set(contentIndex, { explicitIds: new Set(providedCallId ? [providedCallId] : []), functionNames: new Set(functionName ? [functionName] : []), @@ -753,7 +763,11 @@ async function consumeChatStream( ? toolCall.function.arguments : JSON.stringify(toolCall.function.arguments || {}); block.partialArgs = (block.partialArgs || "") + argsDelta; - block.arguments = parseStreamingJson(block.partialArgs); + // Preview refresh is scheduled geometrically; the terminal strict parse + // below re-reads the full buffer authoritatively either way. + if (toolArgumentPreviewSchedules.get(block)?.(block.partialArgs.length)) { + block.arguments = parseStreamingJson(block.partialArgs); + } stream.push({ type: "toolcall_delta", contentIndex, diff --git a/packages/ai/src/providers/openai-completions.ts b/packages/ai/src/providers/openai-completions.ts index 20cab4f377ce..14183a8e8aca 100644 --- a/packages/ai/src/providers/openai-completions.ts +++ b/packages/ai/src/providers/openai-completions.ts @@ -52,7 +52,11 @@ import { type PendingCommentaryTags, } from "../utils/assistant-text-phase.js"; import { AssistantMessageEventStream } from "../utils/event-stream.js"; -import { parseStreamingJson } from "../utils/json-parse.js"; +import { + createToolArgumentPreviewSchedule, + parseStreamingJson, + type ToolArgumentPreviewSchedule, +} from "../utils/json-parse.js"; import { notifyLlmRequestActivity } from "../utils/llm-request-activity.js"; import { sortPromptCacheToolsByName } from "../utils/prompt-cache-stability.js"; import { projectProviderError } from "../utils/provider-error.js"; @@ -206,6 +210,11 @@ export const streamOpenAICompletions: StreamFunction< const toolCallBlocksByIndex = new Map(); const toolCallBlocksById = new Map(); const toolCallBlocksByFirstId = new Map(); + // Preview schedules are per active tool call; WeakMap keys die with the block. + const toolArgumentPreviewSchedules = new WeakMap< + StreamingToolCallBlock, + ToolArgumentPreviewSchedule + >(); const normalizeToolCallDeltas = createOpenAICompletionsToolCallDeltaNormalizer(); const pendingReasoningDetailsByToolCallId = new Map(); const blocks = output.content as StreamingBlock[]; @@ -381,6 +390,7 @@ export const streamOpenAICompletions: StreamFunction< toolCallBlocksById.set(toolCall.id, block); rememberFirstToolCallById(toolCall.id, block); } + toolArgumentPreviewSchedules.set(block, createToolArgumentPreviewSchedule()); appendBlock(block); stream.push({ type: "toolcall_start", @@ -576,7 +586,11 @@ export const streamOpenAICompletions: StreamFunction< if (toolCall.function?.arguments) { delta = toolCall.function.arguments; block.partialArgs = (block.partialArgs ?? "") + toolCall.function.arguments; - block.arguments = parseStreamingJson(block.partialArgs); + // Preview refresh is scheduled geometrically; the terminal + // finalize re-parses the full buffer authoritatively either way. + if (toolArgumentPreviewSchedules.get(block)?.(block.partialArgs.length)) { + block.arguments = parseStreamingJson(block.partialArgs); + } } stream.push({ type: "toolcall_delta", diff --git a/packages/ai/src/transports/anthropic-transport-stream.test.ts b/packages/ai/src/transports/anthropic-transport-stream.test.ts index d18ca6fb8f18..633fae7f3104 100644 --- a/packages/ai/src/transports/anthropic-transport-stream.test.ts +++ b/packages/ai/src/transports/anthropic-transport-stream.test.ts @@ -2138,6 +2138,125 @@ describe("anthropic transport stream", () => { expect(result.content.some((block) => block.type === "toolCall")).toBe(false); }); + it("refreshes streamed tool argument previews on geometric checkpoints instead of every delta", async () => { + // ~1165 chars of argument JSON split into twelve ~100-char deltas: below + // the first preview checkpoint for five deltas, crossing it on the sixth. + const argsJson = JSON.stringify({ content: "x".repeat(1150) }); + const deltaSize = 100; + const chunks: string[] = []; + for (let offset = 0; offset < argsJson.length; offset += deltaSize) { + chunks.push(argsJson.slice(offset, offset + deltaSize)); + } + // Frame delivery waits for consumer observation: partial output blocks are + // mutated in place by the handler, so a buffered producer would let later + // deltas overwrite the state an earlier snapshot wants to capture. + const gates: Array<() => void> = []; + const deltaObservedGates = chunks.map( + () => + new Promise((resolve) => { + gates.push(resolve); + }), + ); + const encoder = new TextEncoder(); + const frames: Array<{ encoded: Uint8Array; waitFor?: Promise }> = [ + { + encoded: encoder.encode( + `data: ${JSON.stringify({ + type: "message_start", + message: { id: "msg_preview_schedule", usage: { input_tokens: 2, output_tokens: 0 } }, + })}\n\n`, + ), + }, + { + encoded: encoder.encode( + `data: ${JSON.stringify({ + type: "content_block_start", + index: 0, + content_block: { type: "tool_use", id: "call_preview", name: "write", input: {} }, + })}\n\n`, + ), + }, + ...chunks.map((chunk, index) => ({ + encoded: encoder.encode( + `data: ${JSON.stringify({ + type: "content_block_delta", + index: 0, + delta: { type: "input_json_delta", partial_json: chunk }, + })}\n\n`, + ), + ...(index > 0 ? { waitFor: deltaObservedGates[index - 1] } : {}), + })), + ...[ + `data: ${JSON.stringify({ type: "content_block_stop", index: 0 })}\n\n`, + `data: ${JSON.stringify({ + type: "message_delta", + delta: { stop_reason: "tool_use" }, + usage: { input_tokens: 2, output_tokens: 2 }, + })}\n\n`, + `data: ${JSON.stringify({ type: "message_stop" })}\n\n`, + ].map((payload) => ({ + encoded: encoder.encode(payload), + waitFor: deltaObservedGates[chunks.length - 1], + })), + ]; + let nextFrame = 0; + const body = new ReadableStream({ + async pull(controller) { + const frame = frames[nextFrame]; + if (!frame) { + controller.close(); + return; + } + await frame.waitFor; + controller.enqueue(frame.encoded); + nextFrame += 1; + }, + }); + guardedFetchMock.mockResolvedValueOnce( + new Response(body, { + status: 200, + headers: { "content-type": "text/event-stream" }, + }), + ); + const streamFn = createAnthropicMessagesTransportStreamFn(); + const stream = await Promise.resolve( + streamFn( + makeAnthropicTransportModel(), + { messages: [{ role: "user", content: "write" }] } as AnthropicStreamContext, + { apiKey: "sk-ant-api" } as AnthropicStreamOptions, + ), + ); + const previews: Array> = []; + for await (const event of stream as AsyncIterable<{ + type: string; + partial?: { content?: Array<{ type: string; arguments?: Record }> }; + }>) { + if (event.type !== "toolcall_delta") { + continue; + } + const block = event.partial?.content?.find((entry) => entry.type === "toolCall"); + previews.push(structuredClone(block?.arguments ?? {})); + gates.shift()?.(); + } + const result = await stream.result(); + + // Preview parses are preview-only: below the first checkpoint no parse has + // run yet, and the checkpoint refresh keeps later deltas stale until the + // next doubling. The terminal parse stays authoritative. + expect(previews.length).toBe(chunks.length); + for (const snapshot of previews.slice(0, 5)) { + expect(snapshot).toEqual({}); + } + expect(previews[5]).not.toEqual({}); + expect(result.stopReason).toBe("toolUse"); + const toolCall = result.content.find((block) => block.type === "toolCall"); + expect(toolCall).toMatchObject({ + type: "toolCall", + name: "write", + arguments: { content: "x".repeat(1150) }, + }); + }); + it("uses seeded Anthropic tool input when no argument deltas arrive", async () => { guardedFetchMock.mockResolvedValueOnce( createSseResponse([ diff --git a/packages/ai/src/transports/anthropic-transport-stream.ts b/packages/ai/src/transports/anthropic-transport-stream.ts index 8aa3387aa246..656e015d3b15 100644 --- a/packages/ai/src/transports/anthropic-transport-stream.ts +++ b/packages/ai/src/transports/anthropic-transport-stream.ts @@ -81,7 +81,11 @@ import { } from "../providers/tool-result-text.js"; import { tagPendingCommentaryText } from "../utils/assistant-text-phase.js"; import { createDeferredEventBuffer } from "../utils/deferred-event-buffer.js"; -import { parseStreamingJson } from "../utils/json-parse.js"; +import { + createToolArgumentPreviewSchedule, + parseStreamingJson, + type ToolArgumentPreviewSchedule, +} from "../utils/json-parse.js"; import { notifyLlmRequestActivity } from "../utils/llm-request-activity.js"; import { buildAnthropicReplayPlan, @@ -1205,6 +1209,11 @@ export function createAnthropicMessagesTransportStreamFn(): StreamFn { } const blocks = output.content; const blockIndexes = new Map(); + // Preview schedules are per active tool call; WeakMap keys die with the block. + const toolArgumentPreviewSchedules = new WeakMap< + Extract, + ToolArgumentPreviewSchedule + >(); const sealedToolCalls: Array<{ block: Extract; contentIndex: number; @@ -1521,6 +1530,7 @@ export function createAnthropicMessagesTransportStreamFn(): StreamFn { }; output.content.push(block); blockIndexes.set(index, output.content.length - 1); + toolArgumentPreviewSchedules.set(block, createToolArgumentPreviewSchedule()); eventSink.push({ type: "toolcall_start", contentIndex: output.content.length - 1, @@ -1625,7 +1635,11 @@ export function createAnthropicMessagesTransportStreamFn(): StreamFn { ) { const partialJson = `${block.partialJson ?? ""}${delta.partial_json}`; block.partialJson = partialJson; - block.arguments = parseAnthropicToolCallArguments(partialJson); + // Preview refresh is scheduled geometrically; content_block_stop + // re-parses the full buffer authoritatively either way. + if (toolArgumentPreviewSchedules.get(block)?.(partialJson.length)) { + block.arguments = parseAnthropicToolCallArguments(partialJson); + } eventSink.push({ type: "toolcall_delta", contentIndex: index, diff --git a/packages/ai/src/transports/openai-completions-stream.ts b/packages/ai/src/transports/openai-completions-stream.ts index 8ab70d6cbc72..8b0017030d43 100644 --- a/packages/ai/src/transports/openai-completions-stream.ts +++ b/packages/ai/src/transports/openai-completions-stream.ts @@ -15,7 +15,11 @@ import { tagUnresolvedTextAsCommentary, type PendingCommentaryTags, } from "../utils/assistant-text-phase.js"; -import { parseStreamingJson } from "../utils/json-parse.js"; +import { + createToolArgumentPreviewSchedule, + parseStreamingJson, + type ToolArgumentPreviewSchedule, +} from "../utils/json-parse.js"; import { notifyLlmRequestActivity } from "../utils/llm-request-activity.js"; import { createReasoningTagTextPartitioner } from "../utils/reasoning-tag-text-partitioner.js"; import { withFirstStreamEventTimeout } from "../utils/stream-first-event-timeout.js"; @@ -114,6 +118,8 @@ export async function processCompletionsStream( let isFlushingPendingPostToolCallDeltas = false; const toolCallBlocksByIndex = new Map(); const toolCallBlocksById = new Map(); + // Preview schedules are per active tool call; WeakMap keys die with the block. + const toolArgumentPreviewSchedules = new WeakMap(); const provisionalCommentaryTags: PendingCommentaryTags = new Map(); const toolCallBlockIndices = new WeakMap(); let explicitVisibleTextBlocks: Set | undefined; @@ -268,6 +274,7 @@ export async function processCompletionsStream( arguments: toolCall.arguments, partialArgs: toolCall.partialArgs, }; + toolArgumentPreviewSchedules.set(block, createToolArgumentPreviewSchedule()); currentBlock = block; output.content.push(block); toolCallBlockIndices.set(block, output.content.length - 1); @@ -515,6 +522,7 @@ export async function processCompletionsStream( partialArgs: "", ...(initialSig ? { thoughtSignature: initialSig } : {}), }; + toolArgumentPreviewSchedules.set(block, createToolArgumentPreviewSchedule()); output.content.push(block); toolCallBlockIndices.set(block, output.content.length - 1); pushStreamEvent({ @@ -540,7 +548,11 @@ export async function processCompletionsStream( } if (toolCall.function?.arguments) { block.partialArgs += toolCall.function.arguments; - block.arguments = parseStreamingJson(block.partialArgs); + // Preview refresh is scheduled geometrically; the terminal + // finalize re-parses the full buffer authoritatively either way. + if (toolArgumentPreviewSchedules.get(block)?.(block.partialArgs.length)) { + block.arguments = parseStreamingJson(block.partialArgs); + } pushStreamEvent({ type: "toolcall_delta", contentIndex: toolCallBlockIndices.get(block) ?? -1, diff --git a/packages/ai/src/transports/openai-responses-stream-internal.ts b/packages/ai/src/transports/openai-responses-stream-internal.ts index 9d61d6f7f2d0..5349c43fffe8 100644 --- a/packages/ai/src/transports/openai-responses-stream-internal.ts +++ b/packages/ai/src/transports/openai-responses-stream-internal.ts @@ -1,15 +1,8 @@ import { isRecord } from "@openclaw/normalization-core/record-coerce"; -import type { - ResponseCreateParamsStreaming, - ResponseOutputItem, - ResponseOutputMessage, - ResponseStreamEvent, -} from "openai/resources/responses/responses.js"; +import type { ResponseOutputItem } from "openai/resources/responses/responses.js"; import { AZURE_RESPONSES_TEXT_CONTENT_PART_TYPE, OPENAI_RESPONSES_OUTPUT_TEXT_CONTENT_PART_TYPE, - type AzureResponsesTextContentPart, - type AzureResponsesTextDeltaEvent, isAzureResponsesTextDeltaEvent, isResponsesTextContentPartType, resolveResponsesMessageSnapshotCollapse, @@ -19,18 +12,16 @@ import { readResponsesToolCallItemIdentity, type ResponsesToolCallState, } from "../providers/openai-responses-tool-call-tracker.js"; -import type { Api, AssistantMessage, Model, TextContent, ToolCall, Usage } from "../types.js"; -import { parseStreamingJson } from "../utils/json-parse.js"; +import type { Api, AssistantMessage, Model, TextContent, ToolCall } from "../types.js"; +import { + createToolArgumentPreviewSchedule, + parseStreamingJson, + type ToolArgumentPreviewSchedule, +} from "../utils/json-parse.js"; import { notifyLlmRequestActivity } from "../utils/llm-request-activity.js"; -import { - type FirstStreamEventInternalOptions, - withFirstStreamEventTimeout, -} from "../utils/stream-first-event-timeout.js"; +import { withFirstStreamEventTimeout } from "../utils/stream-first-event-timeout.js"; import { createCompactionTracker } from "./openai-responses-compaction-replay.js"; -import { - OPENAI_RESPONSES_REASONING_REPLAY_BLOCK_META_KEY, - type OpenAIResponsesReasoningReplayMetadata, -} from "./openai-responses-contracts.js"; +import { OPENAI_RESPONSES_REASONING_REPLAY_BLOCK_META_KEY } from "./openai-responses-contracts.js"; import { normalizeResponsesFailedEvent, ResponsesStreamFailure } from "./openai-responses-debug.js"; import { encodeTextSignatureV1 } from "./openai-responses-replay-internal.js"; import { adaptResponsesStream } from "./openai-responses-stream-observer-internal.js"; @@ -49,71 +40,14 @@ import { type ResponsesThinkingBlock, type TextBlockReference, } from "./openai-responses-stream-terminal-internal.js"; +import type { + CompletedResponse, + ResponsesStreamOptions, + ResponsesStreamOutputMessage, +} from "./openai-responses-stream-types-internal.js"; import { transportAbortError } from "./transport-stream-shared.js"; -type ResponsesConsumedEventType = - | "error" - | "response.completed" - | "response.content_part.added" - | "response.created" - | "response.failed" - | "response.function_call_arguments.delta" - | "response.function_call_arguments.done" - | "response.incomplete" - | "response.output_item.added" - | "response.output_item.done" - | "response.output_text.delta" - | "response.reasoning_summary_part.added" - | "response.reasoning_summary_part.done" - | "response.reasoning_summary_text.delta" - | "response.reasoning_text.delta" - | "response.refusal.delta"; - -type OpenAIResponsesConsumedEvent = Extract< - ResponseStreamEvent, - { type: ResponsesConsumedEventType } ->; -type CompletedResponse = Extract["response"]; -type OpenAIResponsesIgnoredSdkEvent = Exclude; -type ResponsesTextContentPart = - | ResponseOutputMessage["content"][number] - | AzureResponsesTextContentPart; -type ResponsesStreamOutputMessage = Omit & { - content: ResponsesTextContentPart[] | null; -}; -type ResponsesContentPartAddedEvent = Extract< - ResponseStreamEvent, - { type: "response.content_part.added" } ->; -type ResponsesOutputItemDoneEvent = Extract< - ResponseStreamEvent, - { type: "response.output_item.done" } ->; - -export type OpenAIResponsesStreamEvent = - | OpenAIResponsesConsumedEvent - | OpenAIResponsesIgnoredSdkEvent - | (Omit & { - part: Extract; - }) - | (Omit & { - item: ResponsesStreamOutputMessage; - }) - | AzureResponsesTextDeltaEvent; - -type ResponsesStreamOptions = FirstStreamEventInternalOptions & { - serviceTier?: ResponseCreateParamsStreaming["service_tier"]; - resolveServiceTier?: ( - responseServiceTier: ResponseCreateParamsStreaming["service_tier"] | undefined, - requestServiceTier: ResponseCreateParamsStreaming["service_tier"] | undefined, - ) => ResponseCreateParamsStreaming["service_tier"] | undefined; - applyServiceTierPricing?: ( - usage: Usage, - serviceTier: ResponseCreateParamsStreaming["service_tier"] | undefined, - ) => void; - signal?: AbortSignal; - reasoningReplayMetadata?: OpenAIResponsesReasoningReplayMetadata; -}; +export type { OpenAIResponsesStreamEvent } from "./openai-responses-stream-types-internal.js"; export async function processResponsesStream( openaiStream: AsyncIterable, @@ -126,6 +60,8 @@ export async function processResponsesStream( type StreamingToolCallState = ResponsesToolCallState & { block: StreamingToolCallBlock; contentIndex: number; + // Preview refresh schedule for streamed arguments; done/terminal parses stay authoritative. + previewSchedule: ToolArgumentPreviewSchedule; }; type ResponsesOutputSlot = ResponsesStreamOutputSlot< ResponsesStreamOutputMessage, @@ -318,6 +254,7 @@ export async function processResponsesStream( block: toolCallBlock, contentIndex, argumentStreamReliable: true, + previewSchedule: createToolArgumentPreviewSchedule(), ...readResponsesToolCallItemIdentity(item), }; streamingToolCalls.register(event, toolCallState); @@ -467,7 +404,11 @@ export async function processResponsesStream( const toolCall = streamingToolCalls.resolve(event); if (toolCall) { toolCall.block.partialJson += event.delta; - toolCall.block.arguments = parseStreamingJson(toolCall.block.partialJson); + // Preview refresh is geometric; the done event and terminal finalize + // re-parse the full buffer authoritatively either way. + if (toolCall.previewSchedule(toolCall.block.partialJson.length)) { + toolCall.block.arguments = parseStreamingJson(toolCall.block.partialJson); + } stream.push({ type: "toolcall_delta", contentIndex: toolCall.contentIndex, diff --git a/packages/ai/src/transports/openai-responses-stream-types-internal.ts b/packages/ai/src/transports/openai-responses-stream-types-internal.ts new file mode 100644 index 000000000000..e9fd4aebfd8e --- /dev/null +++ b/packages/ai/src/transports/openai-responses-stream-types-internal.ts @@ -0,0 +1,81 @@ +import type { + ResponseCreateParamsStreaming, + ResponseOutputMessage, + ResponseStreamEvent, +} from "openai/resources/responses/responses.js"; +import type { + AzureResponsesTextContentPart, + AzureResponsesTextDeltaEvent, +} from "../providers/openai-responses-stream-compat.js"; +import type { Usage } from "../types.js"; +import type { FirstStreamEventInternalOptions } from "../utils/stream-first-event-timeout.js"; +import type { OpenAIResponsesReasoningReplayMetadata } from "./openai-responses-contracts.js"; + +type ResponsesConsumedEventType = + | "error" + | "response.completed" + | "response.content_part.added" + | "response.created" + | "response.failed" + | "response.function_call_arguments.delta" + | "response.function_call_arguments.done" + | "response.incomplete" + | "response.output_item.added" + | "response.output_item.done" + | "response.output_text.delta" + | "response.reasoning_summary_part.added" + | "response.reasoning_summary_part.done" + | "response.reasoning_summary_text.delta" + | "response.reasoning_text.delta" + | "response.refusal.delta"; + +type OpenAIResponsesConsumedEvent = Extract< + ResponseStreamEvent, + { type: ResponsesConsumedEventType } +>; +type OpenAIResponsesIgnoredSdkEvent = Exclude; +type ResponsesTextContentPart = + | ResponseOutputMessage["content"][number] + | AzureResponsesTextContentPart; +type ResponsesContentPartAddedEvent = Extract< + ResponseStreamEvent, + { type: "response.content_part.added" } +>; +type ResponsesOutputItemDoneEvent = Extract< + ResponseStreamEvent, + { type: "response.output_item.done" } +>; + +export type CompletedResponse = Extract< + ResponseStreamEvent, + { type: "response.completed" } +>["response"]; + +export type ResponsesStreamOutputMessage = Omit & { + content: ResponsesTextContentPart[] | null; +}; + +export type OpenAIResponsesStreamEvent = + | OpenAIResponsesConsumedEvent + | OpenAIResponsesIgnoredSdkEvent + | (Omit & { + part: Extract; + }) + | (Omit & { + item: ResponsesStreamOutputMessage; + }) + | AzureResponsesTextDeltaEvent; + +export type ResponsesStreamOptions = FirstStreamEventInternalOptions & { + serviceTier?: ResponseCreateParamsStreaming["service_tier"]; + resolveServiceTier?: ( + responseServiceTier: ResponseCreateParamsStreaming["service_tier"] | undefined, + requestServiceTier: ResponseCreateParamsStreaming["service_tier"] | undefined, + ) => ResponseCreateParamsStreaming["service_tier"] | undefined; + applyServiceTierPricing?: ( + usage: Usage, + serviceTier: ResponseCreateParamsStreaming["service_tier"] | undefined, + ) => void; + signal?: AbortSignal; + reasoningReplayMetadata?: OpenAIResponsesReasoningReplayMetadata; +}; diff --git a/packages/ai/src/utils/json-parse.ts b/packages/ai/src/utils/json-parse.ts index 2839cd8a439f..d72373c9602b 100644 --- a/packages/ai/src/utils/json-parse.ts +++ b/packages/ai/src/utils/json-parse.ts @@ -140,3 +140,26 @@ export function parseStreamingJson(partialJson: string | undefined): Record boolean; + +/** + * Streamed tool-call arguments are preview-only; the terminal parse re-reads + * the full buffer authoritatively at content_block_stop. Reparsing every delta + * scans an ever-growing buffer and makes assembly quadratic in the argument + * size, so refresh previews on a geometric length schedule instead — bounded + * staleness, linear total parse work. + */ +export function createToolArgumentPreviewSchedule(): ToolArgumentPreviewSchedule { + let nextCheckpointChars = TOOL_ARGUMENT_PREVIEW_FIRST_CHECKPOINT_CHARS; + return (accumulatedChars: number): boolean => { + if (accumulatedChars < nextCheckpointChars) { + return false; + } + nextCheckpointChars = accumulatedChars * 2; + return true; + }; +} diff --git a/src/agents/runtime/proxy.test.ts b/src/agents/runtime/proxy.test.ts index 7544ea47559f..981f27bc0f4d 100644 --- a/src/agents/runtime/proxy.test.ts +++ b/src/agents/runtime/proxy.test.ts @@ -44,6 +44,28 @@ function responseFromText(text: string): Response { ); } +function responseFromSseFrames(frames: unknown[]): Response { + const encoder = new TextEncoder(); + const chunks = frames.map((frame) => encoder.encode(`data: ${JSON.stringify(frame)}\n\n`)); + const reader = { + read: vi.fn(async () => { + await new Promise((resolve) => { + setImmediate(resolve); + }); + const value = chunks.shift(); + return value ? { done: false, value } : { done: true, value: undefined }; + }), + cancel: vi.fn(async () => undefined), + releaseLock: vi.fn(), + } as unknown as ReadableStreamDefaultReader; + + return { + ok: true, + status: 200, + body: { getReader: () => reader }, + } as Response; +} + function responseFromReaderText(text: string, releaseLock: () => void): Response { const chunks: Array> = [ { done: false, value: new TextEncoder().encode(text) }, @@ -160,6 +182,74 @@ describe("streamProxy", () => { }); }); + it("delays tool argument previews while preserving exact terminal arguments", async () => { + const initialContent = "a".repeat(128); + const checkpointContent = "b".repeat(400); + const deltas = [`{"content":"${initialContent}`, checkpointContent, `","terminal":"exact"}`]; + vi.stubGlobal( + "fetch", + vi.fn(async () => + responseFromSseFrames([ + { type: "toolcall_start", contentIndex: 0, id: "call-1", toolName: "write" }, + ...deltas.map((delta) => ({ type: "toolcall_delta", contentIndex: 0, delta })), + { type: "toolcall_end", contentIndex: 0 }, + { type: "done", reason: "toolUse", usage }, + ]), + ), + ); + + const stream = streamProxy(model, context, { + authToken: "token", + proxyUrl: "https://proxy.example", + }); + const argumentSnapshots: Array> = []; + let terminalArguments: Record | undefined; + for await (const event of stream) { + if (event.type === "toolcall_delta") { + const content = event.partial.content[event.contentIndex]; + if (content?.type === "toolCall") { + argumentSnapshots.push(structuredClone(content.arguments)); + } + } else if (event.type === "toolcall_end") { + terminalArguments = structuredClone(event.toolCall.arguments); + } + } + + const checkpointPreview = { content: initialContent + checkpointContent }; + expect(argumentSnapshots).toEqual([{}, checkpointPreview, checkpointPreview]); + const exactArguments = { + content: initialContent + checkpointContent, + terminal: "exact", + }; + expect(terminalArguments).toEqual(exactArguments); + await expect(stream.result()).resolves.toMatchObject({ + content: [{ type: "toolCall", arguments: exactArguments }], + }); + }); + + it("preserves empty arguments for terminal-only tool calls", async () => { + vi.stubGlobal( + "fetch", + vi.fn(async () => + responseFromSseFrames([ + { type: "toolcall_start", contentIndex: 0, id: "call-1", toolName: "list" }, + { type: "toolcall_end", contentIndex: 0 }, + { type: "done", reason: "toolUse", usage }, + ]), + ), + ); + + const stream = streamProxy(model, context, { + authToken: "token", + proxyUrl: "https://proxy.example", + }); + + await expect(stream.result()).resolves.toMatchObject({ + stopReason: "toolUse", + content: [{ type: "toolCall", id: "call-1", name: "list", arguments: {} }], + }); + }); + it("flushes a final SSE frame without a trailing newline", async () => { // Provider proxies can close immediately after the last SSE frame; the // parser still has to emit the terminal done event. diff --git a/src/agents/runtime/proxy.ts b/src/agents/runtime/proxy.ts index b927a359a461..1aa9bd4a20e7 100644 --- a/src/agents/runtime/proxy.ts +++ b/src/agents/runtime/proxy.ts @@ -4,9 +4,12 @@ */ import { + createToolArgumentPreviewSchedule, createSseByteGuard, parseStreamingJson, + parseTerminalToolCallArguments, type SseByteGuard, + type ToolArgumentPreviewSchedule, } from "@openclaw/ai/internal/runtime"; import { resolvePositiveTimerTimeoutMs } from "@openclaw/normalization-core/number-coercion"; import { readResponseWithLimit } from "../../infra/http-body.js"; @@ -27,7 +30,9 @@ const PROXY_SSE_STREAM_MAX_BYTES = 16 * 1024 * 1024; const PROXY_SSE_PENDING_BUFFER_MAX_BYTES = PROXY_SSE_STREAM_MAX_BYTES; const PROXY_SSE_READ_IDLE_TIMEOUT_MS = 120_000; -type StreamingToolCall = ToolCall & { partialJson?: string }; +type StreamingToolCall = ToolCall & { + partialJson: string; +}; // Create stream class matching ProxyMessageEventStream class ProxyMessageEventStream extends EventStream { @@ -348,6 +353,7 @@ export function streamProxy( const decoder = new TextDecoder(); let buffer = ""; let terminalEventSeen = false; + const toolArgumentPreviewSchedules = new Map(); const processSseLine = (line: string) => { if (!line.startsWith("data: ")) { @@ -358,7 +364,7 @@ export function streamProxy( return; } const proxyEvent = JSON.parse(data) as ProxyAssistantMessageEvent; - const event = processProxyEvent(proxyEvent, partial); + const event = processProxyEvent(proxyEvent, partial, toolArgumentPreviewSchedules); if (!event) { return; } @@ -431,6 +437,7 @@ export function streamProxy( function processProxyEvent( proxyEvent: ProxyAssistantMessageEvent, partial: AssistantMessage, + toolArgumentPreviewSchedules: Map, ): AssistantMessageEvent | undefined { switch (proxyEvent.type) { case "start": @@ -508,22 +515,34 @@ function processProxyEvent( throw new Error("Received thinking_end for non-thinking content"); } - case "toolcall_start": - partial.content[proxyEvent.contentIndex] = { + case "toolcall_start": { + const content = { type: "toolCall", id: proxyEvent.id, name: proxyEvent.toolName, arguments: {}, partialJson: "", - } satisfies ToolCall & { partialJson: string } as ToolCall; + } satisfies StreamingToolCall; + partial.content[proxyEvent.contentIndex] = content; + toolArgumentPreviewSchedules.set( + proxyEvent.contentIndex, + createToolArgumentPreviewSchedule(), + ); return { type: "toolcall_start", contentIndex: proxyEvent.contentIndex, partial }; + } case "toolcall_delta": { const content = partial.content[proxyEvent.contentIndex]; if (content?.type === "toolCall") { const streamingContent = content as StreamingToolCall; - streamingContent.partialJson = `${streamingContent.partialJson ?? ""}${proxyEvent.delta}`; - content.arguments = parseStreamingJson(streamingContent.partialJson) || {}; + streamingContent.partialJson += proxyEvent.delta; + const previewSchedule = toolArgumentPreviewSchedules.get(proxyEvent.contentIndex); + if (!previewSchedule) { + throw new Error("Received toolcall_delta without a preview schedule"); + } + if (previewSchedule(streamingContent.partialJson.length)) { + content.arguments = parseStreamingJson(streamingContent.partialJson); + } partial.content[proxyEvent.contentIndex] = { ...content }; // Trigger reactivity return { type: "toolcall_delta", @@ -538,7 +557,12 @@ function processProxyEvent( case "toolcall_end": { const content = partial.content[proxyEvent.contentIndex]; if (content?.type === "toolCall") { - delete (content as StreamingToolCall).partialJson; + const streamingContent = content as StreamingToolCall; + content.arguments = streamingContent.partialJson + ? parseTerminalToolCallArguments(streamingContent.partialJson) + : {}; + toolArgumentPreviewSchedules.delete(proxyEvent.contentIndex); + delete (content as Partial).partialJson; return { type: "toolcall_end", contentIndex: proxyEvent.contentIndex, diff --git a/src/worker/inference-stream.runtime.test.ts b/src/worker/inference-stream.runtime.test.ts new file mode 100644 index 000000000000..736a8c956021 --- /dev/null +++ b/src/worker/inference-stream.runtime.test.ts @@ -0,0 +1,91 @@ +import { expect, it, vi } from "vitest"; +import type { + WorkerInferenceEventParams, + WorkerInferenceModelRef, + WorkerInferenceTerminalOutcome, +} from "../../packages/gateway-protocol/src/schema/worker-inference.js"; +import type { Usage } from "../llm/types.js"; +import { createWorkerInferenceStreamAdapter } from "./inference-stream.runtime.js"; +import type { WorkerInferenceProxyClient } from "./worker-rpc-clients.js"; + +const modelRef: WorkerInferenceModelRef = { provider: "test", model: "test-model" }; +const usage: Usage = { + input: 1, + output: 2, + cacheRead: 0, + cacheWrite: 0, + totalTokens: 3, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, +}; + +it("delays worker tool argument previews while preserving exact terminal arguments", async () => { + const initialContent = "a".repeat(128); + const checkpointContent = "b".repeat(400); + const deltas = [`{"content":"${initialContent}`, checkpointContent, `","terminal":"exact"}`]; + const terminalArguments = { + content: initialContent + checkpointContent, + terminal: "exact", + }; + const start: WorkerInferenceProxyClient["start"] = async (request, handlers) => { + const identity = { + runEpoch: request.runEpoch, + sessionId: request.sessionId, + runId: request.runId, + turnId: request.turnId, + }; + const streamEvents: WorkerInferenceEventParams["event"][] = [ + { type: "toolcall_start", contentIndex: 0, id: "call-1", toolName: "write" }, + ...deltas.map((delta) => ({ type: "toolcall_delta" as const, contentIndex: 0, delta })), + { type: "toolcall_end", contentIndex: 0 }, + ]; + for (const [index, event] of streamEvents.entries()) { + handlers?.onEvent?.({ ...identity, seq: index + 1, event }); + await new Promise((resolve) => { + setImmediate(resolve); + }); + } + return { + type: "done", + message: { + role: "assistant", + content: [{ type: "toolCall", id: "call-1", name: "write", arguments: terminalArguments }], + api: "openai-responses", + provider: modelRef.provider, + model: modelRef.model, + stopReason: "toolUse", + usage, + timestamp: 1, + }, + } satisfies WorkerInferenceTerminalOutcome; + }; + const client = { start, cancel: vi.fn() } as unknown as WorkerInferenceProxyClient; + const streamFn = createWorkerInferenceStreamAdapter({ + client, + sessionId: "session-1", + runEpoch: 1, + runId: "run-1", + turnId: "turn-1", + modelRef, + }); + + const stream = streamFn({ modelRef, context: { messages: [] }, options: {} }); + const argumentSnapshots: Array> = []; + let endArguments: Record | undefined; + for await (const event of stream) { + if (event.type === "toolcall_delta") { + const content = event.partial.content[event.contentIndex]; + if (content?.type === "toolCall") { + argumentSnapshots.push(structuredClone(content.arguments)); + } + } else if (event.type === "toolcall_end") { + endArguments = structuredClone(event.toolCall.arguments); + } + } + + const checkpointPreview = { content: initialContent + checkpointContent }; + expect(argumentSnapshots).toEqual([{}, checkpointPreview, checkpointPreview]); + expect(endArguments).toEqual(terminalArguments); + await expect(stream.result()).resolves.toMatchObject({ + content: [{ type: "toolCall", arguments: terminalArguments }], + }); +}); diff --git a/src/worker/inference-stream.runtime.ts b/src/worker/inference-stream.runtime.ts index 8f2c3801dd84..477332defc69 100644 --- a/src/worker/inference-stream.runtime.ts +++ b/src/worker/inference-stream.runtime.ts @@ -1,4 +1,9 @@ -import { parseStreamingJson } from "@openclaw/ai/internal/runtime"; +import { + createToolArgumentPreviewSchedule, + parseStreamingJson, + parseTerminalToolCallArguments, + type ToolArgumentPreviewSchedule, +} from "@openclaw/ai/internal/runtime"; import { WORKER_PROTOCOL_MAX_IDENTIFIER_LENGTH } from "../../packages/gateway-protocol/src/schema/worker-admission.js"; import type { WorkerInferenceContext, @@ -18,7 +23,9 @@ import { createAssistantMessageEventStream } from "../llm/utils/event-stream.js" import { isWorkerTranscriptMessageFrameSafe } from "./transcript-message.js"; import type { WorkerInferenceProxyClient } from "./worker-rpc-clients.js"; -type StreamingToolCall = ToolCall & { partialJson?: string }; +type StreamingToolCall = ToolCall & { + partialJson: string; +}; type WorkerInferenceStreamAdapterOptions = { client: WorkerInferenceProxyClient; @@ -59,6 +66,7 @@ function emptyAssistantMessage(modelRef: WorkerInferenceModelRef): AssistantMess function processInferenceEvent( payload: WorkerInferenceEventParams, partial: AssistantMessage, + toolArgumentPreviewSchedules: Map, tolerateMissingState: boolean, ): AssistantMessageEvent | undefined { const event = payload.event; @@ -146,13 +154,15 @@ function processInferenceEvent( }; } case "toolcall_start": { - partial.content[event.contentIndex] = { + const content = { type: "toolCall", id: event.id, name: event.toolName, arguments: {}, partialJson: "", - } satisfies StreamingToolCall as ToolCall; + } satisfies StreamingToolCall; + partial.content[event.contentIndex] = content; + toolArgumentPreviewSchedules.set(event.contentIndex, createToolArgumentPreviewSchedule()); return { type: "toolcall_start", contentIndex: event.contentIndex, partial }; } case "toolcall_delta": { @@ -164,8 +174,14 @@ function processInferenceEvent( throw new Error("worker inference tool delta has no active tool call"); } const streaming = content as StreamingToolCall; - streaming.partialJson = `${streaming.partialJson ?? ""}${event.delta}`; - content.arguments = parseStreamingJson(streaming.partialJson); + streaming.partialJson += event.delta; + const previewSchedule = toolArgumentPreviewSchedules.get(event.contentIndex); + if (!previewSchedule) { + throw new Error("worker inference tool delta has no preview schedule"); + } + if (previewSchedule(streaming.partialJson.length)) { + content.arguments = parseStreamingJson(streaming.partialJson); + } return { type: "toolcall_delta", contentIndex: event.contentIndex, @@ -181,7 +197,10 @@ function processInferenceEvent( } throw new Error("worker inference tool end has no active tool call"); } - delete (content as StreamingToolCall).partialJson; + const streaming = content as StreamingToolCall; + content.arguments = parseTerminalToolCallArguments(streaming.partialJson); + toolArgumentPreviewSchedules.delete(event.contentIndex); + delete (content as Partial).partialJson; return { type: "toolcall_end", contentIndex: event.contentIndex, toolCall: content, partial }; } } @@ -220,6 +239,7 @@ export function createWorkerInferenceStreamAdapter( return (inferenceRequest) => { const stream = createAssistantMessageEventStream(); const partial = emptyAssistantMessage(adapter.modelRef); + const toolArgumentPreviewSchedules = new Map(); let streamHasGap = false; let settled = false; modelCallSeq += 1; @@ -277,7 +297,12 @@ export function createWorkerInferenceStreamAdapter( streamHasGap = true; }, onEvent: (event) => { - const projected = processInferenceEvent(event, partial, streamHasGap); + const projected = processInferenceEvent( + event, + partial, + toolArgumentPreviewSchedules, + streamHasGap, + ); if (projected) { stream.push(projected); }