From b403a7a4580d92d0ef79fc84eae8ef6eb24ec6ba Mon Sep 17 00:00:00 2001 From: Alex Castillo Date: Wed, 28 Oct 2026 11:25:06 -0400 Subject: [PATCH] fix(server-utils): Keep Workers AI streamed tool calls that have no index Co-Authored-By: Claude Sonnet 5.5 --- .../src/ai/workers-ai/streaming.ts | 53 +++++++++++++----- .../lib/tracing/workers-ai-streaming.test.ts | 55 +++++++++++++++++++ 2 files changed, 95 insertions(+), 13 deletions(-) diff --git a/packages/server-utils/src/ai/workers-ai/streaming.ts b/packages/server-utils/src/ai/workers-ai/streaming.ts index 37aefd9a6046..1d4493c78c25 100644 --- a/packages/server-utils/src/ai/workers-ai/streaming.ts +++ b/packages/server-utils/src/ai/workers-ai/streaming.ts @@ -5,7 +5,7 @@ import type { WorkersAiUsage } from './types'; import { setOutputMessagesAttribute } from './utils'; interface WorkersAiStreamingToolCall { - index?: number; + index?: number | null; id?: string; type?: string; function?: { name?: string; arguments?: string }; @@ -28,14 +28,36 @@ interface WorkersAiStreamChunk { usage?: WorkersAiUsage & { prompt_tokens?: number; completion_tokens?: number; total_tokens?: number }; } +/** + * Tool calls accumulated from a stream: those that carry an `index` are keyed by it, those that don't + * are kept in arrival order so they can never collide with a real index. + */ +interface ToolCallAccumulator { + byIndex: Record; + unindexed: WorkersAiStreamingToolCall[]; +} + +/** + * Find the call a tool call delta without an `index` belongs to. Some providers omit the index and send each + * parallel call whole in its own chunk: a delta with a known id continues that call, a new id (or a name without + * an id) starts a new call, and anything else continues the latest one. + */ +function getUnindexedToolCall( + toolCall: WorkersAiStreamingToolCall, + name: string | undefined, + accumulator: ToolCallAccumulator, +): WorkersAiStreamingToolCall | undefined { + if (toolCall.id) { + return accumulator.unindexed.find(call => call.id === toolCall.id); + } + return name ? undefined : accumulator.unindexed[accumulator.unindexed.length - 1]; +} + /** * Accumulate a fragmented OpenAI-compatible tool call (delivered across multiple - * `choices[].delta.tool_calls` chunks) into the index-keyed accumulator. + * `choices[].delta.tool_calls` chunks) into the accumulator. */ -function accumulateStreamingToolCalls( - toolCalls: WorkersAiStreamingToolCall[], - accumulator: Record, -): void { +function accumulateStreamingToolCalls(toolCalls: WorkersAiStreamingToolCall[], accumulator: ToolCallAccumulator): void { for (const toolCall of toolCalls) { // Normalize both shapes: name/arguments nested under `function`, or at the top level. const name = toolCall.function?.name ?? toolCall.name; @@ -46,12 +68,12 @@ function accumulateStreamingToolCalls( continue; } - const index = toolCall.index ?? 0; - const existing = accumulator[index]; + const existing = + toolCall.index == null ? getUnindexedToolCall(toolCall, name, accumulator) : accumulator.byIndex[toolCall.index]; if (!existing) { - accumulator[index] = { - index, + const created: WorkersAiStreamingToolCall = { + index: toolCall.index ?? accumulator.unindexed.length, id: toolCall.id, type: toolCall.type, function: { @@ -59,6 +81,11 @@ function accumulateStreamingToolCalls( arguments: args ?? '', }, }; + if (toolCall.index == null) { + accumulator.unindexed.push(created); + } else { + accumulator.byIndex[toolCall.index] = created; + } } else if (existing.function) { if (name && !existing.function.name) { existing.function.name = name; @@ -81,7 +108,7 @@ function processLine( line: string, state: StreamResponseState, recordOutputs: boolean, - toolCallAccumulator: Record, + toolCallAccumulator: ToolCallAccumulator, ): void { const trimmed = line.trim(); if (!trimmed.startsWith('data:')) { @@ -162,7 +189,7 @@ export function instrumentWorkersAiStream( // OpenAI-compatible tool calls arrive fragmented across chunks and are keyed by index; // accumulate them here and flatten into `state.toolCalls` once the stream ends. - const toolCallAccumulator: Record = {}; + const toolCallAccumulator: ToolCallAccumulator = { byIndex: {}, unindexed: [] }; let buffer = ''; let spanEnded = false; @@ -174,7 +201,7 @@ export function instrumentWorkersAiStream( spanEnded = true; if (recordOutputs) { - const accumulatedToolCalls = Object.values(toolCallAccumulator); + const accumulatedToolCalls = [...Object.values(toolCallAccumulator.byIndex), ...toolCallAccumulator.unindexed]; if (accumulatedToolCalls.length > 0) { state.toolCalls.push(...accumulatedToolCalls); } diff --git a/packages/server-utils/test/ai/lib/tracing/workers-ai-streaming.test.ts b/packages/server-utils/test/ai/lib/tracing/workers-ai-streaming.test.ts index 433a0774cd20..56e4e0a415de 100644 --- a/packages/server-utils/test/ai/lib/tracing/workers-ai-streaming.test.ts +++ b/packages/server-utils/test/ai/lib/tracing/workers-ai-streaming.test.ts @@ -162,6 +162,61 @@ describe('instrumentWorkersAiStream', () => { ]); }); + it('keeps parallel tool calls that stream without an index as separate calls', async () => { + const { span, attributes } = createMockSpan(); + const chunks = [ + 'data: {"choices":[{"delta":{"tool_calls":[{"id":"call_1","type":"function","function":{"name":"getWeather","arguments":"{\\"city\\":\\"Paris\\"}"}}]}}]}\n\n', + 'data: {"choices":[{"delta":{"tool_calls":[{"id":"call_2","type":"function","function":{"name":"getTime","arguments":"{\\"tz\\":\\"CET\\"}"}}]}}]}\n\n', + 'data: {"choices":[{"delta":{},"finish_reason":"tool_calls"}]}\n\ndata: [DONE]\n\n', + ]; + + const instrumented = instrumentWorkersAiStream(streamFromChunks(chunks), span, true); + await new Response(instrumented).text(); + + const toolCalls = JSON.parse(attributes[GEN_AI_RESPONSE_TOOL_CALLS] as string); + expect(toolCalls).toEqual([ + { index: 0, id: 'call_1', type: 'function', function: { name: 'getWeather', arguments: '{"city":"Paris"}' } }, + { index: 1, id: 'call_2', type: 'function', function: { name: 'getTime', arguments: '{"tz":"CET"}' } }, + ]); + }); + + it('merges argument fragments of a tool call that streams without an index', async () => { + const { span, attributes } = createMockSpan(); + const chunks = [ + 'data: {"choices":[{"delta":{"tool_calls":[{"id":"call_1","type":"function","function":{"name":"getRepoInfo","arguments":"{\\"owner\\":"}}]}}]}\n\n', + 'data: {"choices":[{"delta":{"tool_calls":[{"function":{"arguments":"\\"cloudflare\\"}"}}]}}]}\n\n', + 'data: {"choices":[{"delta":{},"finish_reason":"tool_calls"}]}\n\ndata: [DONE]\n\n', + ]; + + const instrumented = instrumentWorkersAiStream(streamFromChunks(chunks), span, true); + await new Response(instrumented).text(); + + const toolCalls = JSON.parse(attributes[GEN_AI_RESPONSE_TOOL_CALLS] as string); + expect(toolCalls).toEqual([ + { + index: 0, + id: 'call_1', + type: 'function', + function: { name: 'getRepoInfo', arguments: '{"owner":"cloudflare"}' }, + }, + ]); + }); + + it('does not merge a tool call without an index into one that has the same index number', async () => { + const { span, attributes } = createMockSpan(); + const chunks = [ + 'data: {"choices":[{"delta":{"tool_calls":[{"index":0,"id":"call_1","type":"function","function":{"name":"getWeather","arguments":"{}"}}]}}]}\n\n', + 'data: {"choices":[{"delta":{"tool_calls":[{"id":"call_2","type":"function","function":{"name":"getTime","arguments":"{}"}}]}}]}\n\n', + 'data: {"choices":[{"delta":{},"finish_reason":"tool_calls"}]}\n\ndata: [DONE]\n\n', + ]; + + const instrumented = instrumentWorkersAiStream(streamFromChunks(chunks), span, true); + await new Response(instrumented).text(); + + const toolCalls = JSON.parse(attributes[GEN_AI_RESPONSE_TOOL_CALLS] as string); + expect(toolCalls.map((call: { id: string }) => call.id)).toEqual(['call_1', 'call_2']); + }); + // Some Workers AI models (e.g. `@cf/moonshotai/kimi-k2.6`) stream tool calls with the // name/arguments at the top level of the tool-call object rather than nested under `function`. it('assembles tool calls whose name/arguments are at the top level of the delta', async () => {