fix(openai,openrouter,huggingface): keep streamed tool calls apart when the provider omits index

.changeset/quiet-owls-count.md

@@ -0,0 +1,5 @@
+---
+"braintrust": patch
+---
+
+fix(openai,openrouter,huggingface): keep streamed tool calls apart when the provider omits `index`

js/src/instrumentation/plugins/huggingface-plugin.test.ts

@@ -89,6 +89,98 @@ describe("aggregateChatCompletionChunks", () => {
     });
   });
 
+  it("keeps tool calls apart when the provider omits index", () => {
+    const chunk = (toolCall: Record<string, unknown>) => ({
+      choices: [{ delta: { tool_calls: [toolCall] } }],
+    });
+    const result = aggregateChatCompletionChunks([
+      chunk({
+        id: "call_1",
+        type: "function",
+        function: { name: "get_weather", arguments: '{"city":' },
+      }),
+      chunk({ function: { arguments: '"Paris"}' } }),
+      chunk({
+        id: "call_2",
+        type: "function",
+        function: { name: "get_time", arguments: '{"tz":"CET"}' },
+      }),
+    ]);
+
+    expect(result).toMatchObject({
+      choices: [
+        {
+          message: {
+            tool_calls: [
+              {
+                id: "call_1",
+                type: "function",
+                function: {
+                  name: "get_weather",
+                  arguments: '{"city":"Paris"}',
+                },
+              },
+              {
+                id: "call_2",
+                type: "function",
+                function: { name: "get_time", arguments: '{"tz":"CET"}' },
+              },
+            ],
+          },
+        },
+      ],
+    });
+  });
+
+  it.each([undefined, null])(
+    "merges interleaved tool calls by ID when index is %s",
+    (index) => {
+      const chunks = [
+        {
+          id: "call_1",
+          type: "function",
+          function: { name: "get_weather", arguments: '{"city":' },
+        },
+        {
+          id: "call_2",
+          type: "function",
+          function: { name: "get_time", arguments: '{"tz":' },
+        },
+        { id: "call_1", function: { arguments: '"Paris"}' } },
+        { id: "call_2", function: { arguments: '"CET"}' } },
+      ].map((toolCall) => ({
+        choices: [
+          { index: 0, delta: { tool_calls: [{ ...toolCall, index }] } },
+        ],
+      }));
+      const result = aggregateChatCompletionChunks(chunks);
+
+      expect(result).toMatchObject({
+        choices: [
+          {
+            message: {
+              tool_calls: [
+                {
+                  id: "call_1",
+                  type: "function",
+                  function: {
+                    name: "get_weather",
+                    arguments: '{"city":"Paris"}',
+                  },
+                },
+                {
+                  id: "call_2",
+                  type: "function",
+                  function: { name: "get_time", arguments: '{"tz":"CET"}' },
+                },
+              ],
+            },
+          },
+        ],
+      });
+    },
+  );
+
   it("preserves full tool calls emitted in a single streamed chunk", () => {
     expect(
       aggregateChatCompletionChunks([

js/src/instrumentation/plugins/huggingface-plugin.ts

@@ -560,7 +560,7 @@ function mergeChatToolCallDeltas(
     const toolIndex =
       typeof toolDelta.index === "number" && toolDelta.index >= 0
         ? toolDelta.index
-        : 0;
+        : resolveMissingToolCallIndex(toolCallsByIndex, toolDelta);
     const existing = toolCallsByIndex.get(toolIndex);
 
     if (!existing) {
@@ -572,6 +572,31 @@ function mergeChatToolCallDeltas(
   }
 }
 
+// Some providers omit `index`: match existing ids before starting a new call.
+// A delta without an id continues the latest one.
+function resolveMissingToolCallIndex(
+  toolCallsByIndex: Map<number, Record<string, unknown>>,
+  toolDelta: Record<string, unknown>,
+): number {
+  if (typeof toolDelta.id === "string") {
+    for (const [index, toolCall] of toolCallsByIndex) {
+      if (toolCall.id === toolDelta.id) {
+        return index;
+      }
+    }
+  }
+  if (toolCallsByIndex.size === 0) {
+    return 0;
+  }
+  const lastIndex = Math.max(...toolCallsByIndex.keys());
+  const lastId = toolCallsByIndex.get(lastIndex)?.id;
+  return typeof toolDelta.id === "string" &&
+    typeof lastId === "string" &&
+    lastId !== toolDelta.id
+    ? lastIndex + 1
+    : lastIndex;
+}
+
 function createChatToolCall(
   toolDelta: Record<string, unknown>,
 ): Record<string, unknown> {

js/src/instrumentation/plugins/openai-plugin.test.ts

@@ -1028,6 +1028,114 @@ describe("aggregateChatCompletionChunks", () => {
         '{"a":1}',
       );
     });
+
+    it("should keep parallel tool calls apart when the provider omits index", () => {
+      const toolCallChunk = (toolCall: Record<string, unknown>) => ({
+        choices: [{ index: 0, delta: { tool_calls: [toolCall] } }],
+      });
+      const chunks = [
+        toolCallChunk({
+          id: "call_1",
+          type: "function",
+          function: { name: "get_weather", arguments: '{"city":' },
+        }),
+        toolCallChunk({ function: { arguments: '"Paris"}' } }),
+        toolCallChunk({
+          id: "call_2",
+          type: "function",
+          function: { name: "get_time", arguments: '{"tz":"CET"}' },
+        }),
+      ];
+
+      const result = aggregateChatCompletionChunks(chunks as any);
+      const toolCalls = (result.output as any[])[0].message.tool_calls;
+
+      expect(toolCalls).toEqual([
+        {
+          id: "call_1",
+          type: "function",
+          function: { name: "get_weather", arguments: '{"city":"Paris"}' },
+        },
+        {
+          id: "call_2",
+          type: "function",
+          function: { name: "get_time", arguments: '{"tz":"CET"}' },
+        },
+      ]);
+    });
+
+    it.each([undefined, null])(
+      "merges interleaved tool calls by ID when index is %s",
+      (index) => {
+        const chunks = [
+          {
+            id: "call_1",
+            type: "function",
+            function: { name: "get_weather", arguments: '{"city":' },
+          },
+          {
+            id: "call_2",
+            type: "function",
+            function: { name: "get_time", arguments: '{"tz":' },
+          },
+          { id: "call_1", function: { arguments: '"Paris"}' } },
+          { id: "call_2", function: { arguments: '"CET"}' } },
+        ].map((toolCall) => ({
+          choices: [
+            { index: 0, delta: { tool_calls: [{ ...toolCall, index }] } },
+          ],
+        }));
+        // @ts-expect-error Compatible providers may omit or null the required index.
+        const result = aggregateChatCompletionChunks(chunks);
+
+        expect(result.output[0].message.tool_calls).toMatchObject([
+          {
+            id: "call_1",
+            type: "function",
+            function: { name: "get_weather", arguments: '{"city":"Paris"}' },
+          },
+          {
+            id: "call_2",
+            type: "function",
+            function: { name: "get_time", arguments: '{"tz":"CET"}' },
+          },
+        ]);
+      },
+    );
+
+    it("should keep every tool call when one chunk carries several without index", () => {
+      const chunks = [
+        {
+          choices: [
+            {
+              index: 0,
+              delta: {
+                tool_calls: [
+                  {
+                    id: "call_1",
+                    type: "function",
+                    function: { name: "a", arguments: "{}" },
+                  },
+                  {
+                    id: "call_2",
+                    type: "function",
+                    function: { name: "b", arguments: "{}" },
+                  },
+                ],
+              },
+            },
+          ],
+        },
+      ];
+
+      const result = aggregateChatCompletionChunks(chunks as any);
+      const toolCalls = (result.output as any[])[0].message.tool_calls;
+
+      expect(toolCalls.map((toolCall: any) => toolCall.id)).toEqual([
+        "call_1",
+        "call_2",
+      ]);
+    });
   });
 
   describe("finish reason", () => {

js/src/instrumentation/plugins/openai-plugin.ts

@@ -640,6 +640,33 @@ function toChatChoice(choice: AggregatedChatChoice): OpenAIChatChoice {
   };
 }
 
+// Some OpenAI-compatible providers omit `index` on streamed tool call deltas.
+// Match existing ids before starting a new call; a delta without an id
+// continues the latest one.
+function resolveToolCallIndex(
+  toolCallsByIndex: AggregatedChatChoice["toolCallsByIndex"],
+  toolDelta: { index?: number | null; id?: string },
+): number {
+  if (typeof toolDelta.index === "number") {
+    return toolDelta.index;
+  }
+  if (toolDelta.id) {
+    for (const [index, toolCall] of toolCallsByIndex) {
+      if (toolCall.id === toolDelta.id) {
+        return index;
+      }
+    }
+  }
+  if (toolCallsByIndex.size === 0) {
+    return 0;
+  }
+  const lastIndex = Math.max(...toolCallsByIndex.keys());
+  const lastId = toolCallsByIndex.get(lastIndex)?.id;
+  return toolDelta.id && lastId && lastId !== toolDelta.id
+    ? lastIndex + 1
+    : lastIndex;
+}
+
 /**
  * Aggregate chat completion chunks into a single response.
  * Combines role (first), content, reasoning_content and audio transcripts (concatenated),
@@ -740,15 +767,18 @@ export function aggregateChatCompletionChunks(
 
       if (delta.tool_calls) {
         for (const toolDelta of delta.tool_calls) {
-          let aggregatedToolCall = aggregatedChoice.toolCallsByIndex.get(
-            toolDelta.index,
+          const toolIndex = resolveToolCallIndex(
+            aggregatedChoice.toolCallsByIndex,
+            toolDelta,
           );
+          let aggregatedToolCall =
+            aggregatedChoice.toolCallsByIndex.get(toolIndex);
           if (!aggregatedToolCall) {
             aggregatedToolCall = {
               function: { arguments: "" },
             };
             aggregatedChoice.toolCallsByIndex.set(
-              toolDelta.index,
+              toolIndex,
               aggregatedToolCall,
             );
           }

js/src/instrumentation/plugins/openrouter-plugin.test.ts

@@ -196,6 +196,75 @@ describe("OpenRouter Plugin", () => {
       });
     });
 
+    it("should keep tool calls apart when the provider omits index", () => {
+      const chunk = (toolCall: Record<string, unknown>) => ({
+        choices: [{ delta: { tool_calls: [toolCall] } }],
+      });
+      const result = aggregateOpenRouterChatChunks([
+        chunk({
+          id: "call_1",
+          type: "function",
+          function: { name: "lookup", arguments: '{"city":' },
+        }),
+        chunk({ function: { arguments: '"Vienna"}' } }),
+        chunk({
+          id: "call_2",
+          type: "function",
+          function: { name: "weather", arguments: '{"tz":"CET"}' },
+        }),
+      ]);
+
+      expect(
+        (result.output[0].message?.tool_calls as any[]).map((toolCall) => [
+          toolCall.id,
+          toolCall.function.name,
+          toolCall.function.arguments,
+        ]),
+      ).toEqual([
+        ["call_1", "lookup", '{"city":"Vienna"}'],
+        ["call_2", "weather", '{"tz":"CET"}'],
+      ]);
+    });
+
+    it.each([undefined, null])(
+      "merges interleaved tool calls by ID when index is %s",
+      (index) => {
+        const chunks = [
+          {
+            id: "call_1",
+            type: "function",
+            function: { name: "get_weather", arguments: '{"city":' },
+          },
+          {
+            id: "call_2",
+            type: "function",
+            function: { name: "get_time", arguments: '{"tz":' },
+          },
+          { id: "call_1", function: { arguments: '"Paris"}' } },
+          { id: "call_2", function: { arguments: '"CET"}' } },
+        ].map((toolCall) => ({
+          choices: [
+            { index: 0, delta: { tool_calls: [{ ...toolCall, index }] } },
+          ],
+        }));
+        // @ts-expect-error Providers may emit null indexes despite the SDK type.
+        const result = aggregateOpenRouterChatChunks(chunks);
+
+        expect(result.output[0].message?.tool_calls).toMatchObject([
+          {
+            id: "call_1",
+            type: "function",
+            function: { name: "get_weather", arguments: '{"city":"Paris"}' },
+          },
+          {
+            id: "call_2",
+            type: "function",
+            function: { name: "get_time", arguments: '{"tz":"CET"}' },
+          },
+        ]);
+      },
+    );
+
     it("should return an empty assistant message for empty chunks", () => {
       expect(aggregateOpenRouterChatChunks([])).toEqual({
         output: [