fix: Record tools in metadata for AI SDK and Google gen ai (#1570)

Record tools for AI SDK in metadata for ai sdk and google gen ai

Ref:
https://braintrustdata.slack.com/archives/C083GCUTVDZ/p1773612247369699

Author: lforst

e2e/scenarios/wrap-ai-sdk-generation-traces/__snapshots__/scenario.test.ts.snap

@@ -323,6 +323,13 @@ test.each(wrapAISDKScenarios)(
 
       expect(toolParent?.input).toBeDefined();
       expect(toolParent?.output).toBeDefined();
+      expect(
+        (toolParent?.row.metadata as { tools?: unknown } | undefined)?.tools,
+      ).toMatchObject({
+        get_weather: {
+          description: "Get the weather for a location",
+        },
+      });
       if (supportsToolExecution) {
         expect(toolModelSpans.length).toBeGreaterThanOrEqual(2);
         expect(toolSpans.length).toBeGreaterThanOrEqual(1);
@@ -416,6 +423,7 @@ test.each(wrapAISDKScenarios)(
                   "provider",
                   "model",
                   "operation",
+                  "tools",
                   "braintrust",
                   "scenario",
                 ]),

e2e/scenarios/wrap-ai-sdk-generation-traces/scenario.test.ts

@@ -808,28 +808,6 @@ exports[`wrap-google-genai-content-traces captures generate, attachment, stream,
             "mode": "ANY",
           },
         },
-        "tools": [
-          {
-            "functionDeclarations": [
-              {
-                "description": "Get the current weather in a given location",
-                "name": "get_weather",
-                "parametersJsonSchema": {
-                  "properties": {
-                    "location": {
-                      "description": "The city and state or city and country",
-                      "type": "string",
-                    },
-                  },
-                  "required": [
-                    "location",
-                  ],
-                  "type": "object",
-                },
-              },
-            ],
-          },
-        ],
       },
       "contents": {
         "text": "Use the get_weather function for Paris, France. Do not answer from memory.",
@@ -849,6 +827,28 @@ exports[`wrap-google-genai-content-traces captures generate, attachment, stream,
           "mode": "ANY",
         },
       },
+      "tools": [
+        {
+          "functionDeclarations": [
+            {
+              "description": "Get the current weather in a given location",
+              "name": "get_weather",
+              "parametersJsonSchema": {
+                "properties": {
+                  "location": {
+                    "description": "The city and state or city and country",
+                    "type": "string",
+                  },
+                },
+                "required": [
+                  "location",
+                ],
+                "type": "object",
+              },
+            },
+          ],
+        },
+      ],
     },
     "metrics": {
       "start": 0,

e2e/scenarios/wrap-google-genai-content-traces/__snapshots__/scenario.test.ts.snap

@@ -147,23 +147,30 @@ test("wrap-google-genai-content-traces captures generate, attachment, stream, ea
 
     expect(JSON.stringify(attachmentSpan?.input)).toContain("file.png");
 
-    const toolInput = toolSpan?.input as
+    const toolMetadata = toolSpan?.row.metadata as
       | {
-          config?: {
-            tools?: Array<{
-              functionDeclarations?: Array<{ name?: string }>;
-            }>;
-          };
+          tools?: Array<{
+            functionDeclarations?: Array<{ name?: string }>;
+          }>;
         }
       | undefined;
     expect(
-      toolInput?.config?.tools?.some((tool) =>
+      toolMetadata?.tools?.some((tool) =>
         tool.functionDeclarations?.some(
           (declaration) => declaration.name === "get_weather",
         ),
       ),
     ).toBe(true);
 
+    const toolInput = toolSpan?.input as
+      | {
+          config?: {
+            tools?: Array<unknown>;
+          };
+        }
+      | undefined;
+    expect(toolInput?.config?.tools).toBeUndefined();
+
     const toolOutput = toolSpan?.output as
       | {
           candidates?: Array<{

e2e/scenarios/wrap-google-genai-content-traces/scenario.test.ts

@@ -10,6 +10,7 @@ vi.mock("../../isomorph", () => ({
 import { AISDKPlugin } from "./ai-sdk-plugin";
 import { Attachment } from "../../logger";
 import iso from "../../isomorph";
+import { serializeAISDKToolsForLogging } from "../../wrappers/ai-sdk/tool-serialization";
 
 const mockNewTracingChannel = iso.newTracingChannel as ReturnType<typeof vi.fn>;
 
@@ -840,6 +841,36 @@ describe("AI SDK utility functions", () => {
       expect(metadata.model).toBeUndefined();
       expect(metadata.braintrust.integration_name).toBe("ai-sdk");
     });
+
+    it("should put tools in metadata", () => {
+      const params = {
+        model: "gpt-4",
+        tools: {
+          echo: {
+            description: "Echo the message",
+            execute: async () => "ok",
+            parameters: {
+              type: "object",
+              properties: {
+                message: {
+                  type: "string",
+                },
+              },
+            },
+          },
+        },
+      };
+      const metadata = extractMetadataFromParams(params);
+      expect(metadata.tools).toMatchObject({
+        echo: {
+          description: "Echo the message",
+          parameters: {
+            type: "object",
+          },
+        },
+      });
+      expect(metadata.tools).not.toHaveProperty("echo.execute");
+    });
   });
 
   describe("processAISDKOutput", () => {
@@ -1189,6 +1220,10 @@ function extractMetadataFromParams(params: any): Record<string, any> {
   if (provider) {
     metadata.provider = provider;
   }
+  const tools = serializeAISDKToolsForLogging(params.tools);
+  if (tools) {
+    metadata.tools = tools;
+  }
 
   return metadata;
 }

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

@@ -3,6 +3,7 @@ import { traceStreamingChannel, unsubscribeAll } from "../core/channel-tracing";
 import { SpanTypeAttribute } from "../../../util/index";
 import { getCurrentUnixTimestamp } from "../../util";
 import { processInputAttachments } from "../../wrappers/attachment-utils";
+import { serializeAISDKToolsForLogging } from "../../wrappers/ai-sdk/tool-serialization";
 import { aiSDKChannels } from "./ai-sdk-channels";
 import type {
   AISDKCallParams,
@@ -234,8 +235,15 @@ export class AISDKPlugin extends BasePlugin {
 function processAISDKInput(params: AISDKCallParams): unknown {
   if (!params) return params;
 
-  // Use the attachment processing from the manual wrapper
-  return processInputAttachments(params);
+  // Use the attachment processing from the manual wrapper, but keep tool
+  // definitions in metadata to match the OpenAI wrapper convention.
+  const input = processInputAttachments(params);
+  if (!input || typeof input !== "object" || Array.isArray(input)) {
+    return input;
+  }
+
+  const { tools: _tools, ...rest } = input as Record<string, unknown>;
+  return rest;
 }
 
 /**
@@ -260,6 +268,10 @@ function extractMetadataFromParams(
   if (provider) {
     metadata.provider = provider;
   }
+  const tools = serializeAISDKToolsForLogging(params.tools);
+  if (tools) {
+    metadata.tools = tools;
+  }
 
   return metadata;
 }

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

@@ -1396,6 +1396,30 @@ describe("ai sdk client unit tests", TEST_SUITE_OPTIONS, () => {
     // Should have at least 2 spans: the main generateText span and the tool execution span
     expect(spans.length).toBeGreaterThanOrEqual(2);
 
+    const generateTextSpan = spans.find(
+      (span: any) => span.span_attributes?.name === "generateText",
+    );
+    expect(generateTextSpan).toBeDefined();
+    expect((generateTextSpan as any).input.tools).toMatchObject({
+      calculate: {
+        description: "Perform a mathematical calculation",
+        inputSchema: {
+          type: "object",
+        },
+      },
+    });
+    expect((generateTextSpan as any).metadata.tools).toMatchObject({
+      calculate: {
+        description: "Perform a mathematical calculation",
+      },
+    });
+    expect((generateTextSpan as any).input.tools).not.toHaveProperty(
+      "calculate.execute",
+    );
+    expect((generateTextSpan as any).metadata.tools).not.toHaveProperty(
+      "calculate.execute",
+    );
+
     // Find the tool execution span
     // eslint-disable-next-line @typescript-eslint/consistent-type-assertions, @typescript-eslint/no-explicit-any
     const toolSpan = spans.find(