feat(core): Accumulate tokens for gen_ai.invoke_agent spans from child LLM calls

dev-packages/node-integration-tests/suites/tracing/vercelai/test.ts

@@ -432,6 +432,9 @@ describe('Vercel AI integration', () => {
               'vercel.ai.settings.maxSteps': 1,
               'vercel.ai.streaming': false,
               'gen_ai.response.model': 'mock-model-id',
+              'gen_ai.usage.input_tokens': 15,
+              'gen_ai.usage.output_tokens': 25,
+              'gen_ai.usage.total_tokens': 40,
               'operation.name': 'ai.generateText',
               'sentry.op': 'gen_ai.invoke_agent',
               'sentry.origin': 'auto.vercelai.otel',
@@ -550,6 +553,9 @@ describe('Vercel AI integration', () => {
               'vercel.ai.settings.maxSteps': 1,
               'vercel.ai.streaming': false,
               'gen_ai.response.model': 'mock-model-id',
+              'gen_ai.usage.input_tokens': 15,
+              'gen_ai.usage.output_tokens': 25,
+              'gen_ai.usage.total_tokens': 40,
               'operation.name': 'ai.generateText',
               'sentry.op': 'gen_ai.invoke_agent',
               'sentry.origin': 'auto.vercelai.otel',

packages/core/src/utils/vercel-ai.ts

@@ -60,13 +60,34 @@ function onVercelAiSpanStart(span: Span): void {
   processGenerateSpan(span, name, attributes);
 }
 
+interface TokenSummary {
+  inputTokens: number;
+  outputTokens: number;
+}
+
 function vercelAiEventProcessor(event: Event): Event {
   if (event.type === 'transaction' && event.spans) {
+    // Map to accumulate token data by parent span ID
+    const tokenAccumulator: Map<string, TokenSummary> = new Map();
+
+    // First pass: process all spans and accumulate token data
     for (const span of event.spans) {
-      // this mutates spans in-place
       processEndedVercelAiSpan(span);
+
+      // Accumulate token data for parent spans
+      accumulateTokensForParent(span, tokenAccumulator);
+    }
+
+    // Second pass: apply accumulated token data to parent spans
+    for (const span of event.spans) {
+      if (span.op !== 'gen_ai.invoke_agent') {
+        continue;
+      }
+
+      applyAccumulatedTokens(span, tokenAccumulator);
     }
   }
+
   return event;
 }
 /**
@@ -241,6 +262,56 @@ export function addVercelAiProcessors(client: Client): void {
   client.addEventProcessor(Object.assign(vercelAiEventProcessor, { id: 'VercelAiEventProcessor' }));
 }
 
+/**
+ * Accumulates token data from a span to its parent in the token accumulator map.
+ * This function extracts token usage from the current span and adds it to the
+ * accumulated totals for its parent span.
+ */
+function accumulateTokensForParent(span: SpanJSON, tokenAccumulator: Map<string, TokenSummary>): void {
+  const parentSpanId = span.parent_span_id;
+  if (!parentSpanId) {
+    return;
+  }
+
+  const inputTokens = span.data[GEN_AI_USAGE_INPUT_TOKENS_ATTRIBUTE];
+  const outputTokens = span.data[GEN_AI_USAGE_OUTPUT_TOKENS_ATTRIBUTE];
+
+  if (typeof inputTokens === 'number' || typeof outputTokens === 'number') {
+    const existing = tokenAccumulator.get(parentSpanId) || { inputTokens: 0, outputTokens: 0 };
+
+    if (typeof inputTokens === 'number') {
+      existing.inputTokens += inputTokens;
+    }
+    if (typeof outputTokens === 'number') {
+      existing.outputTokens += outputTokens;
+    }
+
+    tokenAccumulator.set(parentSpanId, existing);
+  }
+}
+
+/**
+ * Applies accumulated token data to the `gen_ai.invoke_agent` span.
+ * Only immediate children of the `gen_ai.invoke_agent` span are considered,
+ * since aggregation will automatically occur for each parent span.
+ */
+function applyAccumulatedTokens(span: SpanJSON, tokenAccumulator: Map<string, TokenSummary>): void {
+  const accumulated = tokenAccumulator.get(span.span_id);
+  if (!accumulated) {
+    return;
+  }
+
+  if (accumulated.inputTokens > 0) {
+    span.data[GEN_AI_USAGE_INPUT_TOKENS_ATTRIBUTE] = accumulated.inputTokens;
+  }
+  if (accumulated.outputTokens > 0) {
+    span.data[GEN_AI_USAGE_OUTPUT_TOKENS_ATTRIBUTE] = accumulated.outputTokens;
+  }
+  if (accumulated.inputTokens > 0 || accumulated.outputTokens > 0) {
+    span.data['gen_ai.usage.total_tokens'] = accumulated.inputTokens + accumulated.outputTokens;
+  }
+}
+
 function addProviderMetadataToAttributes(attributes: SpanAttributes): void {
   const providerMetadata = attributes[AI_RESPONSE_PROVIDER_METADATA_ATTRIBUTE] as string | undefined;
   if (providerMetadata) {