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

Author: RulaKhaled

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

@@ -62,10 +62,17 @@ function onVercelAiSpanStart(span: Span): void {
 
 function vercelAiEventProcessor(event: Event): Event {
   if (event.type === 'transaction' && event.spans) {
+    // First pass: process all spans normally
     for (const span of event.spans) {
       // this mutates spans in-place
       processEndedVercelAiSpan(span);
     }
+
+    // Second pass: accumulate tokens for gen_ai.invoke_agent spans
+    // TODO: Determine how to handle token aggregation for tool call spans.
+    for (const span of event.spans) {
+      accumulateTokensFromChildSpans(span, event.spans);
+    }
   }
   return event;
 }
@@ -241,6 +248,48 @@ export function addVercelAiProcessors(client: Client): void {
   client.addEventProcessor(Object.assign(vercelAiEventProcessor, { id: 'VercelAiEventProcessor' }));
 }
 
+/**
+ * For the gen_ai.invoke_agent span, correctly iterate over child spans and aggregate:
+ * - Input tokens from client LLM child spans that include this attribute.
+ * - Output tokens from client LLM child spans that include this attribute.
+ * - Total tokens from client LLM child spans that include this attribute.
+ *
+ * Only immediate children of the invoke_agent span need to be considered,
+ * since aggregation will automatically occur for each parent span.
+ *
+ */
+function accumulateTokensFromChildSpans(spanJSON: SpanJSON, allSpans: SpanJSON[]): void {
+  if (spanJSON.op !== 'gen_ai.invoke_agent') {
+    return;
+  }
+  const childSpans = allSpans.filter(childSpan => childSpan.parent_span_id === spanJSON.span_id);
+
+  let totalInputTokens = 0;
+  let totalOutputTokens = 0;
+
+  for (const childSpan of childSpans) {
+    const inputTokens = childSpan.data[GEN_AI_USAGE_INPUT_TOKENS_ATTRIBUTE];
+    const outputTokens = childSpan.data[GEN_AI_USAGE_OUTPUT_TOKENS_ATTRIBUTE];
+
+    if (typeof inputTokens === 'number') {
+      totalInputTokens += inputTokens;
+    }
+    if (typeof outputTokens === 'number') {
+      totalOutputTokens += outputTokens;
+    }
+  }
+
+  if (totalInputTokens > 0) {
+    spanJSON.data[GEN_AI_USAGE_INPUT_TOKENS_ATTRIBUTE] = totalInputTokens;
+  }
+  if (totalOutputTokens > 0) {
+    spanJSON.data[GEN_AI_USAGE_OUTPUT_TOKENS_ATTRIBUTE] = totalOutputTokens;
+  }
+  if (totalInputTokens > 0 || totalOutputTokens > 0) {
+    spanJSON.data['gen_ai.usage.total_tokens'] = totalInputTokens + totalOutputTokens;
+  }
+}
+
 function addProviderMetadataToAttributes(attributes: SpanAttributes): void {
   const providerMetadata = attributes[AI_RESPONSE_PROVIDER_METADATA_ATTRIBUTE] as string | undefined;
   if (providerMetadata) {