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 37aefd9a6046..7c75ef2da907 100644 --- a/packages/server-utils/src/ai/workers-ai/streaming.ts +++ b/packages/server-utils/src/ai/workers-ai/streaming.ts @@ -1,11 +1,12 @@ 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'; interface WorkersAiStreamingToolCall { - index?: number; + index?: number | null; id?: string; type?: string; function?: { name?: string; arguments?: string }; @@ -28,14 +29,22 @@ interface WorkersAiStreamChunk { usage?: WorkersAiUsage & { prompt_tokens?: number; completion_tokens?: number; total_tokens?: number }; } +/** + * Tool calls accumulated from a stream, keyed in the order they were opened (so output order is arrival order). + */ +interface ToolCallAccumulator { + 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; +} + /** * 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,19 +55,19 @@ function accumulateStreamingToolCalls( continue; } - const index = toolCall.index ?? 0; - const existing = accumulator[index]; + const key = getToolCallKey(toolCall, name, accumulator); + const existing = accumulator.calls.get(key); if (!existing) { - accumulator[index] = { - index, + accumulator.calls.set(key, { + index: toolCall.index, id: toolCall.id, type: toolCall.type, function: { name, arguments: args ?? '', }, - }; + }); } else if (existing.function) { if (name && !existing.function.name) { existing.function.name = name; @@ -67,6 +76,7 @@ function accumulateStreamingToolCalls( existing.function.arguments = `${existing.function.arguments ?? ''}${args}`; } } + accumulator.lastKey = key; } } @@ -81,7 +91,7 @@ function processLine( line: string, state: StreamResponseState, recordOutputs: boolean, - toolCallAccumulator: Record, + toolCallAccumulator: ToolCallAccumulator, ): void { const trimmed = line.trim(); if (!trimmed.startsWith('data:')) { @@ -162,7 +172,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 = { calls: new Map(), keyByIndex: new Map() }; let buffer = ''; let spanEnded = false; @@ -174,7 +184,7 @@ export function instrumentWorkersAiStream( spanEnded = true; if (recordOutputs) { - const accumulatedToolCalls = Object.values(toolCallAccumulator); + const accumulatedToolCalls = [...toolCallAccumulator.calls.values()]; if (accumulatedToolCalls.length > 0) { state.toolCalls.push(...accumulatedToolCalls); } 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 433a0774cd20..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 @@ -162,6 +162,147 @@ 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([ + { id: 'call_1', type: 'function', function: { name: 'getWeather', arguments: '{"city":"Paris"}' } }, + { 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([ + { + id: 'call_1', + type: 'function', + function: { name: 'getRepoInfo', arguments: '{"owner":"cloudflare"}' }, + }, + ]); + }); + + 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 = [ + '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']); + }); + + 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 () => {