fix(sdk): proper formatting for vercel AI SDK tool calls (#736)

Author: nirga

packages/traceloop-sdk/src/lib/tracing/ai-sdk-transformations.ts

@@ -1,15 +1,19 @@
-import { ReadableSpan } from "@opentelemetry/sdk-trace-node";
+import { ReadableSpan, Span } from "@opentelemetry/sdk-trace-node";
 import { SpanAttributes } from "@traceloop/ai-semantic-conventions";
 
+const AI_GENERATE_TEXT = "ai.generateText";
 const AI_GENERATE_TEXT_DO_GENERATE = "ai.generateText.doGenerate";
 const AI_GENERATE_OBJECT_DO_GENERATE = "ai.generateObject.doGenerate";
 const AI_STREAM_TEXT_DO_STREAM = "ai.streamText.doStream";
 const HANDLED_SPAN_NAMES: Record<string, string> = {
-  [AI_GENERATE_TEXT_DO_GENERATE]: "ai.generateText.generate",
-  [AI_GENERATE_OBJECT_DO_GENERATE]: "ai.generateObject.generate",
-  [AI_STREAM_TEXT_DO_STREAM]: "ai.streamText.stream",
+  [AI_GENERATE_TEXT]: "run.ai",
+  [AI_GENERATE_TEXT_DO_GENERATE]: "text.generate",
+  [AI_GENERATE_OBJECT_DO_GENERATE]: "object.generate",
+  [AI_STREAM_TEXT_DO_STREAM]: "text.stream",
 };
 
+const TOOL_SPAN_NAME = "ai.toolCall";
+
 const AI_RESPONSE_TEXT = "ai.response.text";
 const AI_RESPONSE_OBJECT = "ai.response.object";
 const AI_RESPONSE_TOOL_CALLS = "ai.response.toolCalls";
@@ -19,6 +23,7 @@ const AI_USAGE_PROMPT_TOKENS = "ai.usage.promptTokens";
 const AI_USAGE_COMPLETION_TOKENS = "ai.usage.completionTokens";
 const AI_MODEL_PROVIDER = "ai.model.provider";
 const AI_PROMPT_TOOLS = "ai.prompt.tools";
+const AI_TELEMETRY_METADATA_PREFIX = "ai.telemetry.metadata.";
 const TYPE_TEXT = "text";
 const TYPE_TOOL_CALL = "tool_call";
 const ROLE_ASSISTANT = "assistant";
@@ -47,14 +52,6 @@ const VENDOR_MAPPING: Record<string, string> = {
   openrouter: "OpenRouter",
 };
 
-export const transformAiSdkSpanName = (span: ReadableSpan): void => {
-  // Unfortunately, the span name is not writable as this is not the intended behavior
-  // but it is a workaround to set the correct span name
-  if (span.name in HANDLED_SPAN_NAMES) {
-    (span as any).name = HANDLED_SPAN_NAMES[span.name];
-  }
-};
-
 const transformResponseText = (attributes: Record<string, any>): void => {
   if (AI_RESPONSE_TEXT in attributes) {
     attributes[`${SpanAttributes.LLM_COMPLETIONS}.0.content`] =
@@ -367,9 +364,41 @@ const transformVendor = (attributes: Record<string, any>): void => {
   }
 };
 
-export const transformAiSdkAttributes = (
-  attributes: Record<string, any>,
-): void => {
+const transformTelemetryMetadata = (attributes: Record<string, any>): void => {
+  const metadataAttributes: Record<string, string> = {};
+  const keysToDelete: string[] = [];
+
+  // Find all ai.telemetry.metadata.* attributes
+  for (const [key, value] of Object.entries(attributes)) {
+    if (key.startsWith(AI_TELEMETRY_METADATA_PREFIX)) {
+      const metadataKey = key.substring(AI_TELEMETRY_METADATA_PREFIX.length);
+
+      // Always mark for deletion since it's a telemetry metadata attribute
+      keysToDelete.push(key);
+
+      if (metadataKey && value != null) {
+        // Convert value to string for association properties
+        const stringValue = typeof value === "string" ? value : String(value);
+        metadataAttributes[metadataKey] = stringValue;
+
+        // Also set as traceloop association property attribute
+        attributes[
+          `${SpanAttributes.TRACELOOP_ASSOCIATION_PROPERTIES}.${metadataKey}`
+        ] = stringValue;
+      }
+    }
+  }
+
+  // Remove original ai.telemetry.metadata.* attributes
+  keysToDelete.forEach((key) => {
+    delete attributes[key];
+  });
+
+  // Note: Context setting for child span inheritance should be done before span creation,
+  // not during transformation. Use `withTelemetryMetadataContext` function for context propagation.
+};
+
+export const transformLLMSpans = (attributes: Record<string, any>): void => {
   transformResponseText(attributes);
   transformResponseObject(attributes);
   transformResponseToolCalls(attributes);
@@ -379,16 +408,40 @@ export const transformAiSdkAttributes = (
   transformCompletionTokens(attributes);
   calculateTotalTokens(attributes);
   transformVendor(attributes);
+  transformTelemetryMetadata(attributes);
+};
+
+const transformToolCalls = (span: ReadableSpan): void => {
+  if (
+    span.attributes["ai.toolCall.args"] &&
+    span.attributes["ai.toolCall.result"]
+  ) {
+    span.attributes[SpanAttributes.TRACELOOP_ENTITY_INPUT] =
+      span.attributes["ai.toolCall.args"];
+    delete span.attributes["ai.toolCall.args"];
+    span.attributes[SpanAttributes.TRACELOOP_ENTITY_OUTPUT] =
+      span.attributes["ai.toolCall.result"];
+    delete span.attributes["ai.toolCall.result"];
+  }
 };
 
 const shouldHandleSpan = (span: ReadableSpan): boolean => {
-  return span.name in HANDLED_SPAN_NAMES;
+  return span.instrumentationScope?.name === "ai";
+};
+
+export const transformAiSdkSpanNames = (span: Span): void => {
+  if (span.name === TOOL_SPAN_NAME) {
+    span.updateName(`${span.attributes["ai.toolCall.name"] as string}.tool`);
+  }
+  if (span.name in HANDLED_SPAN_NAMES) {
+    span.updateName(HANDLED_SPAN_NAMES[span.name]);
+  }
 };
 
-export const transformAiSdkSpan = (span: ReadableSpan): void => {
+export const transformAiSdkSpanAttributes = (span: ReadableSpan): void => {
   if (!shouldHandleSpan(span)) {
     return;
   }
-  transformAiSdkSpanName(span);
-  transformAiSdkAttributes(span.attributes);
+  transformLLMSpans(span.attributes);
+  transformToolCalls(span);
 };

packages/traceloop-sdk/src/lib/tracing/span-processor.ts

@@ -2,9 +2,10 @@ import {
   SimpleSpanProcessor,
   BatchSpanProcessor,
   SpanProcessor,
+  Span,
   ReadableSpan,
 } from "@opentelemetry/sdk-trace-node";
-import { Span, context } from "@opentelemetry/api";
+import { context } from "@opentelemetry/api";
 import { OTLPTraceExporter } from "@opentelemetry/exporter-trace-otlp-proto";
 import { SpanExporter } from "@opentelemetry/sdk-trace-base";
 import {
@@ -13,7 +14,10 @@ import {
   WORKFLOW_NAME_KEY,
 } from "./tracing";
 import { SpanAttributes } from "@traceloop/ai-semantic-conventions";
-import { transformAiSdkSpan } from "./ai-sdk-transformations";
+import {
+  transformAiSdkSpanAttributes,
+  transformAiSdkSpanNames,
+} from "./ai-sdk-transformations";
 import { parseKeyPairsIntoRecord } from "./baggage-utils";
 
 export const ALL_INSTRUMENTATION_LIBRARIES = "all" as const;
@@ -155,6 +159,8 @@ const onSpanStart = (span: Span): void => {
       );
     }
   }
+
+  transformAiSdkSpanNames(span);
 };
 
 /**
@@ -220,7 +226,7 @@ const onSpanEnd = (
     }
 
     // Apply AI SDK transformations (if needed)
-    transformAiSdkSpan(span);
+    transformAiSdkSpanAttributes(span);
 
     // Ensure OTLP transformer compatibility
     const compatibleSpan = ensureSpanCompatibility(span);

packages/traceloop-sdk/test/ai-sdk-integration.test.ts

@@ -105,17 +105,15 @@ describe("Test AI SDK Integration with Recording", function () {
     const spans = memoryExporter.getFinishedSpans();
 
     const generateTextSpan = spans.find(
-      (span) =>
-        span.name === "ai.generateText.generate" ||
-        span.name === "ai.generateText.doGenerate",
+      (span) => span.name === "text.generate",
     );
 
     assert.ok(result);
     assert.ok(result.text);
     assert.ok(generateTextSpan);
 
-    // Verify span name
-    assert.strictEqual(generateTextSpan.name, "ai.generateText.generate");
+    // Verify span name (should be transformed from ai.generateText.doGenerate to text.generate)
+    assert.strictEqual(generateTextSpan.name, "text.generate");
 
     // Verify vendor
     assert.strictEqual(generateTextSpan.attributes["gen_ai.system"], "OpenAI");
@@ -174,17 +172,16 @@ describe("Test AI SDK Integration with Recording", function () {
     // Find the Google span specifically (should have workflow name test_google_workflow)
     const generateTextSpan = spans.find(
       (span) =>
-        (span.name === "ai.generateText.generate" ||
-          span.name === "ai.generateText.doGenerate") &&
+        span.name === "text.generate" &&
         span.attributes["traceloop.workflow.name"] === "test_google_workflow",
     );
 
     assert.ok(result);
     assert.ok(result.text);
     assert.ok(generateTextSpan, "Could not find Google generateText span");
 
-    // Verify span name
-    assert.strictEqual(generateTextSpan.name, "ai.generateText.generate");
+    // Verify span name (should be transformed from ai.generateText.doGenerate to text.generate)
+    assert.strictEqual(generateTextSpan.name, "text.generate");
 
     // Verify vendor
     assert.strictEqual(generateTextSpan.attributes["gen_ai.system"], "Google");
@@ -236,9 +233,7 @@ describe("Test AI SDK Integration with Recording", function () {
     assert.ok(result.text);
 
     const spans = memoryExporter.getFinishedSpans();
-    const aiSdkSpan = spans.find((span) =>
-      span.name.startsWith("ai.generateText"),
-    );
+    const aiSdkSpan = spans.find((span) => span.name === "text.generate");
 
     assert.ok(aiSdkSpan);