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

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<number, WorkersAiStreamingToolCall>;
+  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<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 +68,24 @@ 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: {
           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;
@@ -81,7 +108,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 +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<number, WorkersAiStreamingToolCall> = {};
+  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);
       }

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