fix(server-utils): Keep Workers AI streamed tool calls that have no index

packages/server-utils/src/ai/core/streaming-tool-calls.ts

@@ -0,0 +1,51 @@
+interface ToolCallKeyState {
+  calls: Map<string, { id?: string }>;
+  keyByIndex: Map<number, string>;
+  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;
+}

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.
  *

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<string, WorkersAiStreamingToolCall>;
+  /** 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<number, string>;
+}
+
 /**
  * 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<number, WorkersAiStreamingToolCall>,
-): 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<number, WorkersAiStreamingToolCall>,
+  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<number, WorkersAiStreamingToolCall> = {};
+  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);
       }

packages/server-utils/test/ai/lib/tracing/openai-streaming-tool-calls.test.ts

@@ -44,6 +44,15 @@ async function recordedToolCalls(chunks: ChatCompletionChunk[]): Promise<Array<R
 }
 
 describe('instrumentStream (openai chat completion tool calls)', () => {
+  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([

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 () => {