fix: Capture provider managed tool calls and capture right provider when using Vercel AI gateway

.changeset/pi-tool-current-span.md

@@ -2,4 +2,4 @@
 "braintrust": patch
 ---
 
-Keep Pi Coding Agent tool spans current during execution so nested spans attach to the tool span.
+fix: Keep Pi Coding Agent tool spans current during execution so nested spans attach to the tool span.

.changeset/social-pumas-smoke.md

@@ -0,0 +1,5 @@
+---
+"braintrust": patch
+---
+
+fix: Capture provider managed tool calls and capture right provider when using Vercel AI gateway

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

@@ -1081,6 +1081,60 @@ export function defineAISDKInstrumentationAssertions(options: {
       );
     }
 
+    if (options.sdkMajorVersion >= 7) {
+      for (const operationName of ["generateText", "streamText"]) {
+        test(
+          `captures provider-executed tools in ${operationName}`,
+          testConfig,
+          () => {
+            const operation = findLatestSpan(
+              events,
+              `ai-sdk-provider-tool-${operationName}-operation`,
+            );
+            const parent = latestEvent(
+              findChildSpans(events, operationName, operation?.span.id),
+            );
+            expect(parent).toBeDefined();
+            const tools = findChildSpans(events, "web_search", parent?.span.id);
+            expect(tools).toHaveLength(1);
+            const tool = tools[0];
+            expect(tool.span.type).toBe("tool");
+            expect(tool.input).toBeDefined();
+            expect(tool.output).toBeDefined();
+            expect(tool.row.error).toBeUndefined();
+            expect(tool.metadata).toMatchObject({
+              providerExecuted: true,
+              toolCallId: expect.any(String),
+            });
+            expect(collectToolCallNames(parent?.output)).toContain(
+              "web_search",
+            );
+            expect(collectToolResultNames(parent?.output)).toContain(
+              "web_search",
+            );
+            expect(parent?.output).toMatchObject({
+              steps: expect.arrayContaining([
+                expect.objectContaining({
+                  content: expect.arrayContaining([
+                    expect.objectContaining({
+                      type: "tool-call",
+                      toolCallId: tool.metadata?.toolCallId,
+                      input: tool.input,
+                    }),
+                    expect.objectContaining({
+                      type: "tool-result",
+                      toolCallId: tool.metadata?.toolCallId,
+                      output: tool.output,
+                    }),
+                  ]),
+                }),
+              ]),
+            });
+          },
+        );
+      }
+    }
+
     test("captures trace for generateText() with tools", testConfig, () => {
       const root = findLatestSpan(events, ROOT_NAME);
       const trace = findToolTrace(events);

e2e/scenarios/ai-sdk-instrumentation/scenario.impl.mjs

@@ -564,6 +564,33 @@ async function runAISDKInstrumentationScenario(
         await runDirectModelCalls();
       }
 
+      if (sdkMajorVersion >= 7) {
+        for (const operation of ["generateText", "streamText"]) {
+          await runOperation(
+            `ai-sdk-provider-tool-${operation}-operation`,
+            `provider-tool-${operation}`,
+            async () => {
+              const result = await instrumentedAI[operation]({
+                model: openai.responses("gpt-4.1-mini"),
+                prompt:
+                  "Search the web for the official Braintrust website. Reply with its URL only.",
+                tools: {
+                  web_search: openai.tools.webSearch({
+                    searchContextSize: "low",
+                  }),
+                },
+                toolChoice: { type: "tool", toolName: "web_search" },
+                maxOutputTokens: 128,
+              });
+              if (operation === "streamText") {
+                for await (const _chunk of result.fullStream) {
+                }
+              }
+            },
+          );
+        }
+      }
+
       if (options.supportsEvaluate) {
         const evaluate =
           options.evaluate ?? instrumentedAI.experimental_evaluate;

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

@@ -32,6 +32,7 @@ import {
   processAISDKOutput as processAISDKOutputActual,
   processAISDKGenerateImageOutput,
   extractTokenMetrics,
+  serializeModelWithProvider,
 } from "./ai-sdk-plugin";
 import iso from "../../isomorph";
 import { serializeAISDKToolsForLogging } from "../../wrappers/ai-sdk/tool-serialization";
@@ -666,6 +667,34 @@ describe("AI SDK utility functions", () => {
       });
     });
 
+    it.each([
+      ["openai/gpt-4", "openai", "gpt-4"],
+      ["anthropic/claude-3", "anthropic", "claude-3"],
+      ["openai/text-embedding-3-small", "openai", "text-embedding-3-small"],
+    ])(
+      "should resolve the gateway provider for %s",
+      (modelId, provider, model) => {
+        expect(
+          serializeModelWithProvider({ modelId, provider: "gateway" }),
+        ).toEqual({
+          model,
+          provider,
+        });
+      },
+    );
+
+    it.each(["gpt-4", "/gpt-4", "openai/", ""])(
+      "should retain gateway without a usable provider prefix in %s",
+      (modelId) => {
+        expect(
+          serializeModelWithProvider({ modelId, provider: "gateway" }),
+        ).toEqual({
+          model: modelId,
+          provider: "gateway",
+        });
+      },
+    );
+
     it("should prefer explicit provider over parsed provider", () => {
       const result = serializeModelWithProvider({
         modelId: "anthropic/claude-3",
@@ -678,6 +707,7 @@ describe("AI SDK utility functions", () => {
     });
 
     it("should handle null/undefined model", () => {
+      // @ts-expect-error Exercise malformed input outside the declared SDK types.
       const result1 = serializeModelWithProvider(null);
       expect(result1).toEqual({
         model: undefined,
@@ -1883,25 +1913,6 @@ describe("AI SDK utility functions", () => {
 
 // Helper functions exported for testing
 // These would normally be private but we're testing them through the module
-function serializeModelWithProvider(model: any): {
-  model: string;
-  provider?: string;
-} {
-  const modelId = typeof model === "string" ? model : model?.modelId;
-  const explicitProvider =
-    typeof model === "object" ? model?.provider : undefined;
-
-  if (!modelId) {
-    return { model: modelId, provider: explicitProvider };
-  }
-
-  const parsed = parseGatewayModelString(modelId);
-  return {
-    model: parsed.model,
-    provider: explicitProvider || parsed.provider,
-  };
-}
-
 function parseGatewayModelString(modelString: string): {
   model: string;
   provider?: string;

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

@@ -4354,7 +4354,10 @@ export function serializeModelWithProvider(model: AISDKModel | undefined): {
   const parsed = parseGatewayModelString(modelId);
   return {
     model: parsed.model,
-    provider: explicitProvider || parsed.provider,
+    provider:
+      explicitProvider === "gateway"
+        ? parsed.provider || explicitProvider
+        : explicitProvider || parsed.provider,
   };
 }