feat: Track tool ids for ai SDK v6

.changeset/ai-sdk-v6-tool-call-ids.md

@@ -0,0 +1,5 @@
+---
+"braintrust": patch
+---
+
+feat: Track tool ids for ai SDK v6

e2e/scenarios/ai-sdk-instrumentation/assertions.ts

@@ -292,16 +292,11 @@ function findRerankTrace(events: CapturedLogEvent[]) {
 function findToolTrace(events: CapturedLogEvent[]) {
   const operation = findLatestSpan(events, "ai-sdk-tool-operation");
   const parent = findParentSpan(events, "generateText", operation?.span.id);
-  const toolSpans = findAllSpans(events, "get_weather").filter(
-    (event) => event.span.rootId === operation?.span.rootId,
-  );
-  const modelChildren = events
-    .filter((event) => event.span.rootId === operation?.span.rootId)
-    .filter((event) => {
-      const name = event.span.name ?? "";
-      return name === "doGenerate" || name === "doStream";
-    })
-    .filter((event) => event.span.parentIds[0] !== parent?.span.id);
+  const toolSpans = findChildSpans(events, "get_weather", parent?.span.id);
+  const modelChildren = [
+    ...findChildSpans(events, "doGenerate", parent?.span.id),
+    ...findChildSpans(events, "doStream", parent?.span.id),
+  ];
 
   return {
     modelChildren,
@@ -1010,6 +1005,31 @@ export function defineAISDKInstrumentationAssertions(options: {
         expect(trace.toolSpans.length).toBeGreaterThanOrEqual(1);
         expect(trace.toolSpans[0]?.input).toBeDefined();
         expect(trace.toolSpans[0]?.output).toBeDefined();
+        if (options.sdkMajorVersion >= 6) {
+          for (const toolSpan of trace.toolSpans) {
+            const toolCallId = toolSpan.metadata?.toolCallId;
+            expect(toolCallId).toEqual(expect.any(String));
+            const messages = expect.arrayContaining([
+              expect.objectContaining({
+                role: "tool",
+                content: expect.arrayContaining([
+                  expect.objectContaining({
+                    type: "tool-result",
+                    toolName: toolSpan.span.name,
+                    toolCallId,
+                  }),
+                ]),
+              }),
+            ]);
+            expect(trace.parent?.output).toMatchObject({
+              steps: expect.arrayContaining([
+                expect.objectContaining({
+                  response: expect.objectContaining({ messages }),
+                }),
+              ]),
+            });
+          }
+        }
         expect(collectToolCallNames(trace.parent?.output)).toContain(
           "get_weather",
         );
@@ -1339,7 +1359,12 @@ export function defineAISDKInstrumentationAssertions(options: {
       await matchSpanTreeSnapshot(events, spanSnapshotPath, {
         normalize: {
           additionalProviderIdKeys: ["callId"],
-          omittedKeys: ["id", "performance", "prompt_cache_key", "toolCallId"],
+          omittedKeys: [
+            "id",
+            "performance",
+            "prompt_cache_key",
+            ...(options.sdkMajorVersion < 6 ? ["toolCallId"] : []),
+          ],
         },
       });
     });

e2e/scenarios/cloudflare-ai-chat-instrumentation/__snapshots__/cloudflare-ai-chat-v0-auto.span-tree.txt

@@ -743,6 +743,9 @@ span_tree:
 │                 "condition": "sunny",
 │                 "marker": "WEATHER_TOOL_EXECUTED"
 │               }
+│               metadata: {
+│                 "toolCallId": "<toolCallId:1>"
+│               }
 └── cloudflare-ai-chat-error-root [task]
     metadata: {
       "mode": "auto",

e2e/scenarios/cloudflare-ai-chat-instrumentation/__snapshots__/cloudflare-ai-chat-v0-latest-auto.span-tree.txt

@@ -743,6 +743,9 @@ span_tree:
 │                 "condition": "sunny",
 │                 "marker": "WEATHER_TOOL_EXECUTED"
 │               }
+│               metadata: {
+│                 "toolCallId": "<toolCallId:1>"
+│               }
 └── cloudflare-ai-chat-error-root [task]
     metadata: {
       "mode": "auto",

e2e/scenarios/cloudflare-ai-chat-instrumentation/__snapshots__/cloudflare-ai-chat-v0-latest-manual.span-tree.txt

@@ -743,6 +743,9 @@ span_tree:
 │                 "condition": "sunny",
 │                 "marker": "WEATHER_TOOL_EXECUTED"
 │               }
+│               metadata: {
+│                 "toolCallId": "<toolCallId:1>"
+│               }
 └── cloudflare-ai-chat-error-root [task]
     metadata: {
       "mode": "manual",

e2e/scenarios/cloudflare-ai-chat-instrumentation/__snapshots__/cloudflare-ai-chat-v0-manual.span-tree.txt

@@ -743,6 +743,9 @@ span_tree:
 │                 "condition": "sunny",
 │                 "marker": "WEATHER_TOOL_EXECUTED"
 │               }
+│               metadata: {
+│                 "toolCallId": "<toolCallId:1>"
+│               }
 └── cloudflare-ai-chat-error-root [task]
     metadata: {
       "mode": "manual",

js/src/instrumentation/plugins/ai-sdk-plugin.streaming.test.ts

@@ -50,6 +50,61 @@ describe("AI SDK streaming instrumentation", () => {
     _exportsForTestingOnly.clearTestBackgroundLogger();
   });
 
+  describe.each(["promise", "generator"])("%s tool execution", (kind) => {
+    test.each([
+      { options: { toolCallId: "tool-1" }, expectedId: "tool-1" },
+      { options: undefined, expectedId: undefined },
+      { options: null, expectedId: undefined },
+      { options: {}, expectedId: undefined },
+      { options: { toolCallId: 123 }, expectedId: undefined },
+    ])("logs toolCallId from $options", async ({ options, expectedId }) => {
+      const input = { location: "Paris" };
+      const output = { temperature: 22 };
+      const params = {
+        tools: {
+          get_weather: {
+            execute:
+              kind === "generator"
+                ? async function* (..._args: unknown[]) {
+                    yield { temperature: 20 };
+                    yield output;
+                  }
+                : async (..._args: unknown[]) => output,
+          },
+        },
+      };
+
+      await aiSDKChannels.generateText.tracePromise(
+        async () => {
+          const result = params.tools.get_weather.execute(input, options);
+          if (Symbol.asyncIterator in result) {
+            const values = [];
+            for await (const value of result) {
+              values.push(value);
+            }
+            expect(values).toEqual([{ temperature: 20 }, output]);
+          } else {
+            expect(await result).toEqual(output);
+          }
+          return { text: "done" };
+        },
+        { arguments: [params] } as any,
+      );
+
+      const spans = (await backgroundLogger.drain()) as any[];
+      const toolSpans = spans.filter(
+        (span) => span.span_attributes?.type === "tool",
+      );
+      expect(toolSpans).toHaveLength(1);
+      expect(toolSpans[0]).toMatchObject({ input, output });
+      if (expectedId === undefined) {
+        expect(toolSpans[0].metadata ?? {}).not.toHaveProperty("toolCallId");
+      } else {
+        expect(toolSpans[0].metadata).toMatchObject({ toolCallId: expectedId });
+      }
+    });
+  });
+
   test("generateText child span logs missing usage diagnostic when output has no usage", async () => {
     expect(await backgroundLogger.drain()).toHaveLength(0);
 

js/src/instrumentation/plugins/ai-sdk-plugin.ts

@@ -2625,6 +2625,15 @@ function prepareAISDKChildTracing(
         return result;
       }
 
+      const executionOptions = args[1];
+      const toolCallId = isObject(executionOptions)
+        ? executionOptions.toolCallId
+        : undefined;
+      const toolInput = {
+        input: serializeToolExecutionInput(args),
+        ...(typeof toolCallId === "string" ? { metadata: { toolCallId } } : {}),
+      };
+
       if (isAsyncGenerator(result)) {
         return (async function* () {
           const span = activeEntry.parentSpan.startSpan(
@@ -2638,7 +2647,7 @@ function prepareAISDKChildTracing(
               INSTRUMENTATION_NAMES.AI_SDK,
             ),
           );
-          span.log({ input: serializeToolExecutionInput(args) });
+          span.log(toolInput);
 
           try {
             let lastValue: unknown;
@@ -2658,7 +2667,7 @@ function prepareAISDKChildTracing(
 
       return activeEntry.parentSpan.traced(
         async (span) => {
-          span.log({ input: serializeToolExecutionInput(args) });
+          span.log(toolInput);
           const awaitedResult = await result;
           span.log({ output: awaitedResult });
           return awaitedResult;