From b403a7a4580d92d0ef79fc84eae8ef6eb24ec6ba Mon Sep 17 00:00:00 2001 From: Alex Castillo Date: Wed, 28 Oct 2026 11:25:06 -0400 Subject: [PATCH 1/3] 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 () => { From 2725da766acb0a5747fc1edd1dee6dac8365295f Mon Sep 17 00:00:00 2001 From: Alex Castillo Date: Sat, 7 Nov 2026 10:01:59 -0400 Subject: [PATCH 2/3] fix(server-utils): Match Workers AI streamed tool calls to OpenAI handling Reuse the id/index grouping of the OpenAI streaming accumulator: a new id on a taken index opens a new call, argument fragments without index or id continue the latest call, and calls keep their arrival order. Co-Authored-By: Claude Sonnet 5.5 --- .../src/ai/workers-ai/streaming.ts | 69 ++++++++++++------- .../lib/tracing/workers-ai-streaming.test.ts | 48 +++++++++++++ 2 files changed, 94 insertions(+), 23 deletions(-) diff --git a/packages/server-utils/src/ai/workers-ai/streaming.ts b/packages/server-utils/src/ai/workers-ai/streaming.ts index 1d4493c78c25..0f1636f3a1bd 100644 --- a/packages/server-utils/src/ai/workers-ai/streaming.ts +++ b/packages/server-utils/src/ai/workers-ai/streaming.ts @@ -29,28 +29,55 @@ interface WorkersAiStreamChunk { } /** - * 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. + * Tool calls accumulated from a stream, keyed in the order they were opened (so output order is arrival order). */ interface ToolCallAccumulator { - byIndex: Record; - unindexed: WorkersAiStreamingToolCall[]; + calls: Map; + /** Key of the latest call, which argument fragments with neither `index`, `id` nor name continue. */ + lastKey?: string; + /** Key of the call currently open for a reported `index`, when a new `id` reused that index. */ + keyByIndex: Map; } /** - * 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. + * Find the key of the call a tool call delta belongs to. A numeric `index` groups deltas, except that a different + * non-empty `id` arriving on an index that is already taken opens a new call (a known `id` continues its own call). + * Some providers omit the index and send each parallel call whole: 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( +function getToolCallKey( toolCall: WorkersAiStreamingToolCall, name: string | undefined, accumulator: ToolCallAccumulator, -): WorkersAiStreamingToolCall | undefined { +): string { + const { calls } = accumulator; + const findById = (id: string): string | undefined => { + for (const [key, call] of calls) { + if (call.id === id) return key; + } + return undefined; + }; + + if (typeof toolCall.index === 'number') { + const key = accumulator.keyByIndex.get(toolCall.index) ?? `index-${toolCall.index}`; + const existingId = calls.get(key)?.id; + if (existingId && toolCall.id && existingId !== toolCall.id) { + const next = findById(toolCall.id) ?? `index-${toolCall.index}-${calls.size}`; + accumulator.keyByIndex.set(toolCall.index, next); + return next; + } + return key; + } if (toolCall.id) { - return accumulator.unindexed.find(call => call.id === toolCall.id); + const known = findById(toolCall.id); + if (known !== undefined) { + return known; + } + } + if (toolCall.id || name || accumulator.lastKey === undefined) { + return `no-index-${calls.size}`; } - return name ? undefined : accumulator.unindexed[accumulator.unindexed.length - 1]; + return accumulator.lastKey; } /** @@ -68,24 +95,19 @@ function accumulateStreamingToolCalls(toolCalls: WorkersAiStreamingToolCall[], a continue; } - const existing = - toolCall.index == null ? getUnindexedToolCall(toolCall, name, accumulator) : accumulator.byIndex[toolCall.index]; + const key = getToolCallKey(toolCall, name, accumulator); + const existing = accumulator.calls.get(key); if (!existing) { - const created: WorkersAiStreamingToolCall = { - index: toolCall.index ?? accumulator.unindexed.length, + accumulator.calls.set(key, { + index: toolCall.index ?? accumulator.calls.size, id: toolCall.id, type: toolCall.type, function: { name, 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; @@ -94,6 +116,7 @@ function accumulateStreamingToolCalls(toolCalls: WorkersAiStreamingToolCall[], a existing.function.arguments = `${existing.function.arguments ?? ''}${args}`; } } + accumulator.lastKey = key; } } @@ -189,7 +212,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: ToolCallAccumulator = { byIndex: {}, unindexed: [] }; + const toolCallAccumulator: ToolCallAccumulator = { calls: new Map(), keyByIndex: new Map() }; let buffer = ''; let spanEnded = false; @@ -201,7 +224,7 @@ export function instrumentWorkersAiStream( spanEnded = true; if (recordOutputs) { - const accumulatedToolCalls = [...Object.values(toolCallAccumulator.byIndex), ...toolCallAccumulator.unindexed]; + const accumulatedToolCalls = [...toolCallAccumulator.calls.values()]; 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 56e4e0a415de..60dac43e3e76 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 @@ -217,6 +217,54 @@ describe('instrumentWorkersAiStream', () => { expect(toolCalls.map((call: { id: string }) => call.id)).toEqual(['call_1', 'call_2']); }); + it('opens a new call when a different id arrives on an index that is already taken', 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":[{"index":0,"id":"call_2","type":"function","function":{"name":"getTime","arguments":"{"}}]}}]}\n\n', + 'data: {"choices":[{"delta":{"tool_calls":[{"index":0,"function":{"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']); + expect(toolCalls[1].function.arguments).toBe('{}'); + }); + + it('continues an indexed call with an argument fragment that has no index or id', async () => { + const { span, attributes } = createMockSpan(); + const chunks = [ + 'data: {"choices":[{"delta":{"tool_calls":[{"index":0,"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).toHaveLength(1); + expect(toolCalls[0].function).toEqual({ name: 'getRepoInfo', arguments: '{"owner":"cloudflare"}' }); + }); + + it('returns tool calls in the order they were opened when unindexed and indexed calls are mixed', async () => { + const { span, attributes } = createMockSpan(); + const chunks = [ + 'data: {"choices":[{"delta":{"tool_calls":[{"id":"call_1","type":"function","function":{"name":"getWeather","arguments":"{}"}}]}}]}\n\n', + 'data: {"choices":[{"delta":{"tool_calls":[{"index":0,"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 () => { From b2fade54d23ba2cf2793138abcda6f1fe572af8d Mon Sep 17 00:00:00 2001 From: Alex Castillo Date: Wed, 11 Nov 2026 09:48:22 -0400 Subject: [PATCH 3/3] fix(server-utils): share streamed tool call identity handling --- .../src/ai/core/streaming-tool-calls.ts | 51 +++++++++++++++++++ .../server-utils/src/ai/openai/streaming.ts | 46 +++-------------- .../src/ai/workers-ai/streaming.ts | 44 +--------------- .../openai-streaming-tool-calls.test.ts | 9 ++++ .../lib/tracing/workers-ai-streaming.test.ts | 44 ++++++++++++++-- 5 files changed, 109 insertions(+), 85 deletions(-) create mode 100644 packages/server-utils/src/ai/core/streaming-tool-calls.ts diff --git a/packages/server-utils/src/ai/core/streaming-tool-calls.ts b/packages/server-utils/src/ai/core/streaming-tool-calls.ts new file mode 100644 index 000000000000..e649caa57bf0 --- /dev/null +++ b/packages/server-utils/src/ai/core/streaming-tool-calls.ts @@ -0,0 +1,51 @@ +interface ToolCallKeyState { + calls: Map; + keyByIndex: Map; + lastKey?: string; +} + +/** + * Find the key of the call a tool call delta belongs to. A numeric `index` groups deltas, except that a different + * non-empty `id` arriving on an index that is already taken opens a new call (a known `id` continues its own call). + * Some providers omit the index and send each parallel call whole: 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. + */ +export function getToolCallKey( + toolCall: { index?: number | null; id?: string }, + name: string | undefined, + accumulator: ToolCallKeyState, +): string { + const { calls } = accumulator; + const findById = (id: string): string | undefined => { + for (const [key, call] of calls) { + if (call.id === id) return key; + } + return undefined; + }; + + if (typeof toolCall.index === 'number') { + const key = accumulator.keyByIndex.get(toolCall.index) ?? `index-${toolCall.index}`; + const known = toolCall.id ? findById(toolCall.id) : undefined; + if (known !== undefined) { + accumulator.keyByIndex.set(toolCall.index, known); + return known; + } + const existingId = calls.get(key)?.id; + if (existingId && toolCall.id && existingId !== toolCall.id) { + const next = findById(toolCall.id) ?? `index-${toolCall.index}-${calls.size}`; + accumulator.keyByIndex.set(toolCall.index, next); + return next; + } + return key; + } + if (toolCall.id) { + const known = findById(toolCall.id); + if (known !== undefined) { + return known; + } + } + if (toolCall.id || name || accumulator.lastKey === undefined) { + return `no-index-${calls.size}`; + } + return accumulator.lastKey; +} diff --git a/packages/server-utils/src/ai/openai/streaming.ts b/packages/server-utils/src/ai/openai/streaming.ts index 15686abfa072..98bb214c4146 100644 --- a/packages/server-utils/src/ai/openai/streaming.ts +++ b/packages/server-utils/src/ai/openai/streaming.ts @@ -1,5 +1,6 @@ import { captureException, SPAN_STATUS_ERROR } from '@sentry/core'; import type { Span } from '@sentry/core'; +import { getToolCallKey } from '../core/streaming-tool-calls'; import { endStreamSpan } from '../core/utils'; import { RESPONSE_EVENT_TYPES } from './constants'; import type { @@ -61,7 +62,11 @@ interface StreamingState { function processChatCompletionToolCalls(toolCalls: ChatCompletionToolCall[], state: StreamingState): void { for (const toolCall of toolCalls) { if (!toolCall.function) continue; - const index = getToolCallKey(toolCall, state); + const index = getToolCallKey(toolCall, toolCall.function.name, { + calls: state.chatCompletionToolCalls, + keyByIndex: (state.chatCompletionToolCallKeyByIndex ??= new Map()), + lastKey: state.lastChatCompletionToolCallKey, + }); // Initialize tool call if this is the first chunk for this index const existingToolCall = state.chatCompletionToolCalls.get(index); @@ -83,45 +88,6 @@ function processChatCompletionToolCalls(toolCalls: ChatCompletionToolCall[], sta } } -/** - * OpenAI-compatible providers such as Gemini's OpenAI endpoint send tool calls without an `index` - * (or with `null`), or reuse the same `index` for every call. A call without a usable `index` is - * matched by a known `id`; a new `id` or a function name opens a new call, and anything else - * continues the latest call. A different non-empty `id` on an `index` that is already taken also - * opens a new call. - */ -function getToolCallKey(toolCall: ChatCompletionToolCall, state: StreamingState): string { - const calls = state.chatCompletionToolCalls; - const findById = (id: string): string | undefined => { - for (const [key, call] of calls) { - if (call.id === id) return key; - } - return undefined; - }; - - if (typeof toolCall.index === 'number') { - const keyByIndex = (state.chatCompletionToolCallKeyByIndex ??= new Map()); - const key = keyByIndex.get(toolCall.index) ?? `index-${toolCall.index}`; - const existingId = calls.get(key)?.id; - if (existingId && toolCall.id && existingId !== toolCall.id) { - const next = findById(toolCall.id) ?? `index-${toolCall.index}-${calls.size}`; - keyByIndex.set(toolCall.index, next); - return next; - } - return key; - } - if (toolCall.id) { - const known = findById(toolCall.id); - if (known !== undefined) { - return known; - } - } - if (toolCall.id || toolCall.function?.name || state.lastChatCompletionToolCallKey === undefined) { - return `no-index-${calls.size}`; - } - return state.lastChatCompletionToolCallKey; -} - /** * Processes a single OpenAI ChatCompletionChunk event, updating the streaming state. * diff --git a/packages/server-utils/src/ai/workers-ai/streaming.ts b/packages/server-utils/src/ai/workers-ai/streaming.ts index 0f1636f3a1bd..7c75ef2da907 100644 --- a/packages/server-utils/src/ai/workers-ai/streaming.ts +++ b/packages/server-utils/src/ai/workers-ai/streaming.ts @@ -1,5 +1,6 @@ import { SPAN_STATUS_ERROR } from '@sentry/core'; import type { Span } from '@sentry/core'; +import { getToolCallKey } from '../core/streaming-tool-calls'; import { endStreamSpan, type StreamResponseState } from '../core/utils'; import type { WorkersAiUsage } from './types'; import { setOutputMessagesAttribute } from './utils'; @@ -39,47 +40,6 @@ interface ToolCallAccumulator { keyByIndex: Map; } -/** - * Find the key of the call a tool call delta belongs to. A numeric `index` groups deltas, except that a different - * non-empty `id` arriving on an index that is already taken opens a new call (a known `id` continues its own call). - * Some providers omit the index and send each parallel call whole: 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 getToolCallKey( - toolCall: WorkersAiStreamingToolCall, - name: string | undefined, - accumulator: ToolCallAccumulator, -): string { - const { calls } = accumulator; - const findById = (id: string): string | undefined => { - for (const [key, call] of calls) { - if (call.id === id) return key; - } - return undefined; - }; - - if (typeof toolCall.index === 'number') { - const key = accumulator.keyByIndex.get(toolCall.index) ?? `index-${toolCall.index}`; - const existingId = calls.get(key)?.id; - if (existingId && toolCall.id && existingId !== toolCall.id) { - const next = findById(toolCall.id) ?? `index-${toolCall.index}-${calls.size}`; - accumulator.keyByIndex.set(toolCall.index, next); - return next; - } - return key; - } - if (toolCall.id) { - const known = findById(toolCall.id); - if (known !== undefined) { - return known; - } - } - if (toolCall.id || name || accumulator.lastKey === undefined) { - return `no-index-${calls.size}`; - } - return accumulator.lastKey; -} - /** * Accumulate a fragmented OpenAI-compatible tool call (delivered across multiple * `choices[].delta.tool_calls` chunks) into the accumulator. @@ -100,7 +60,7 @@ function accumulateStreamingToolCalls(toolCalls: WorkersAiStreamingToolCall[], a if (!existing) { accumulator.calls.set(key, { - index: toolCall.index ?? accumulator.calls.size, + index: toolCall.index, id: toolCall.id, type: toolCall.type, function: { diff --git a/packages/server-utils/test/ai/lib/tracing/openai-streaming-tool-calls.test.ts b/packages/server-utils/test/ai/lib/tracing/openai-streaming-tool-calls.test.ts index cd39e060c9bc..36aca7b8a9c7 100644 --- a/packages/server-utils/test/ai/lib/tracing/openai-streaming-tool-calls.test.ts +++ b/packages/server-utils/test/ai/lib/tracing/openai-streaming-tool-calls.test.ts @@ -44,6 +44,15 @@ async function recordedToolCalls(chunks: ChatCompletionChunk[]): Promise { + it.each([undefined, null])('continues a known id when its index changes from %s to numeric', async index => { + const toolCalls = await recordedToolCalls([ + chunkWith([{ index, id: 'call_a', function: { name: 'lookup', arguments: '{"q":' } }]), + chunkWith([{ index: 0, id: 'call_a', function: { name: '', arguments: '"a"' } }]), + chunkWith([{ index: 0, function: { name: '', arguments: '}' } }]), + ]); + expect(toolCalls).toEqual([{ index, id: 'call_a', function: { name: 'lookup', arguments: '{"q":"a"}' } }]); + }); + it('keeps whole tool calls that carry an id and no index', async () => { const toolCalls = await recordedToolCalls([ chunkWith([ 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 60dac43e3e76..34ce6de6a5af 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 @@ -175,8 +175,8 @@ describe('instrumentWorkersAiStream', () => { 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"}' } }, + { id: 'call_1', type: 'function', function: { name: 'getWeather', arguments: '{"city":"Paris"}' } }, + { id: 'call_2', type: 'function', function: { name: 'getTime', arguments: '{"tz":"CET"}' } }, ]); }); @@ -194,7 +194,6 @@ describe('instrumentWorkersAiStream', () => { 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"}' }, @@ -202,6 +201,45 @@ describe('instrumentWorkersAiStream', () => { ]); }); + it.each([undefined, null])('continues a known id when its index changes from %s to numeric', async index => { + const { span, attributes } = createMockSpan(); + const deltas = [ + { index, id: 'call_1', function: { name: 'getWeather', arguments: '{"city":' } }, + { index: 0, id: 'call_1', function: { arguments: '"Paris"}' } }, + { index: 0, function: { arguments: ' ' } }, + ]; + const chunks = deltas.map( + delta => `data: ${JSON.stringify({ choices: [{ delta: { tool_calls: [delta] } }] })}\n\n`, + ); + await new Response(instrumentWorkersAiStream(streamFromChunks(chunks), span, true)).text(); + expect(JSON.parse(attributes[GEN_AI_RESPONSE_TOOL_CALLS] as string)).toEqual([ + { index, id: 'call_1', function: { name: 'getWeather', arguments: '{"city":"Paris"} ' } }, + ]); + expect(JSON.parse(attributes[GEN_AI_OUTPUT_MESSAGES] as string)).toEqual([ + { + role: 'assistant', + parts: [{ type: 'tool_call', id: 'call_1', name: 'getWeather', arguments: '{"city":"Paris"} ' }], + }, + ]); + }); + + it('continues an unindexed call by id after another call opens', async () => { + const { span, attributes } = createMockSpan(); + const deltas = [ + { index: null, id: 'call_1', function: { name: 'first', arguments: '{' } }, + { index: 0, id: 'call_2', function: { name: 'second', arguments: '{}' } }, + { id: 'call_1', function: { arguments: '}' } }, + ]; + const chunks = deltas.map( + delta => `data: ${JSON.stringify({ choices: [{ delta: { tool_calls: [delta] } }] })}\n\n`, + ); + await new Response(instrumentWorkersAiStream(streamFromChunks(chunks), span, true)).text(); + expect(JSON.parse(attributes[GEN_AI_RESPONSE_TOOL_CALLS] as string)).toEqual([ + { index: null, id: 'call_1', function: { name: 'first', arguments: '{}' } }, + { index: 0, id: 'call_2', function: { name: 'second', arguments: '{}' } }, + ]); + }); + 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 = [