fix(core): Instrument invoke_agent root span, and move vercel ai logic to a folder

Author: RulaKhaled

packages/core/src/utils/vercel-ai.ts renamed to packages/core/src/utils/vercel-ai/index.ts

@@ -1,8 +1,10 @@
-import type { Client } from '../client';
-import { SEMANTIC_ATTRIBUTE_SENTRY_OP, SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN } from '../semanticAttributes';
-import type { Event } from '../types-hoist/event';
-import type { Span, SpanAttributes, SpanAttributeValue, SpanJSON, SpanOrigin } from '../types-hoist/span';
-import { spanToJSON } from './spanUtils';
+import type { Client } from '../../client';
+import { SEMANTIC_ATTRIBUTE_SENTRY_OP, SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN } from '../../semanticAttributes';
+import type { Event } from '../../types-hoist/event';
+import type { Span, SpanAttributes, SpanAttributeValue, SpanJSON, SpanOrigin } from '../../types-hoist/span';
+import { spanToJSON } from '../spanUtils';
+import type { TokenSummary } from './types';
+import { accumulateTokensForParent, applyAccumulatedTokens } from './utils';
 import type { ProviderMetadata } from './vercel-ai-attributes';
 import {
   AI_MODEL_ID_ATTRIBUTE,
@@ -60,11 +62,6 @@ 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
@@ -86,6 +83,12 @@ function vercelAiEventProcessor(event: Event): Event {
 
       applyAccumulatedTokens(span, tokenAccumulator);
     }
+
+    // Also apply to root when it is the invoke_agent pipeline
+    const trace = event.contexts?.trace;
+    if (trace && trace.op === 'gen_ai.invoke_agent') {
+      applyAccumulatedTokens(trace, tokenAccumulator);
+    }
   }
 
   return event;
@@ -262,56 +265,6 @@ 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) {

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

@@ -0,0 +1,4 @@
+export interface TokenSummary {
+  inputTokens: number;
+  outputTokens: number;
+}

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

@@ -0,0 +1,57 @@
+import type { TraceContext } from '../../types-hoist/context';
+import type { SpanJSON } from '../../types-hoist/span';
+import { GEN_AI_USAGE_INPUT_TOKENS_ATTRIBUTE, GEN_AI_USAGE_OUTPUT_TOKENS_ATTRIBUTE } from '../gen-ai-attributes';
+import type { TokenSummary } from './types';
+
+/**
+ * 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.
+ */
+export 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.
+ */
+export function applyAccumulatedTokens(
+  spanOrTrace: SpanJSON | TraceContext,
+  tokenAccumulator: Map<string, TokenSummary>,
+): void {
+  const accumulated = tokenAccumulator.get(spanOrTrace.span_id);
+  if (!accumulated || !spanOrTrace.data) {
+    return;
+  }
+
+  if (accumulated.inputTokens > 0) {
+    spanOrTrace.data[GEN_AI_USAGE_INPUT_TOKENS_ATTRIBUTE] = accumulated.inputTokens;
+  }
+  if (accumulated.outputTokens > 0) {
+    spanOrTrace.data[GEN_AI_USAGE_OUTPUT_TOKENS_ATTRIBUTE] = accumulated.outputTokens;
+  }
+  if (accumulated.inputTokens > 0 || accumulated.outputTokens > 0) {
+    spanOrTrace.data['gen_ai.usage.total_tokens'] = accumulated.inputTokens + accumulated.outputTokens;
+  }
+}

packages/core/src/utils/vercel-ai-attributes.ts renamed to packages/core/src/utils/vercel-ai/vercel-ai-attributes.ts

@@ -89,7 +89,7 @@ export class SentryVercelAiInstrumentation extends InstrumentationBase {
    * Initializes the instrumentation by defining the modules to be patched.
    */
   public init(): InstrumentationModuleDefinition {
-    const module = new InstrumentationNodeModuleDefinition('ai', ['>=3.0.0 <5'], this._patch.bind(this));
+    const module = new InstrumentationNodeModuleDefinition('ai', ['>=3.0.0 <6'], this._patch.bind(this));
     return module;
   }