Skip to content

Commit ffa8dd8

Browse files
authored
feat(core): Accumulate tokens for gen_ai.invoke_agent spans from child LLM calls (#17281)
## Problem Currently, `gen_ai.invoke_agent` spans (representing operations like `generateText()`) contain inaccurate token usage information. Users can only see token data on individual `gen_ai.generate_text` child spans, but the tokens are not accumulated across nested spans, making it difficult to track total token consumption for complete AI operations. ## Solution Implement token accumulation for `gen_ai.invoke_agent` spans by iterating over client LLM child spans and aggregating their token usage.
1 parent 94d0165 commit ffa8dd8

File tree

2 files changed

+78
-1
lines changed
  • dev-packages/node-integration-tests/suites/tracing/vercelai
  • packages/core/src/utils

2 files changed

+78
-1
lines changed

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

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -432,6 +432,9 @@ describe('Vercel AI integration', () => {
432432
'vercel.ai.settings.maxSteps': 1,
433433
'vercel.ai.streaming': false,
434434
'gen_ai.response.model': 'mock-model-id',
435+
'gen_ai.usage.input_tokens': 15,
436+
'gen_ai.usage.output_tokens': 25,
437+
'gen_ai.usage.total_tokens': 40,
435438
'operation.name': 'ai.generateText',
436439
'sentry.op': 'gen_ai.invoke_agent',
437440
'sentry.origin': 'auto.vercelai.otel',
@@ -550,6 +553,9 @@ describe('Vercel AI integration', () => {
550553
'vercel.ai.settings.maxSteps': 1,
551554
'vercel.ai.streaming': false,
552555
'gen_ai.response.model': 'mock-model-id',
556+
'gen_ai.usage.input_tokens': 15,
557+
'gen_ai.usage.output_tokens': 25,
558+
'gen_ai.usage.total_tokens': 40,
553559
'operation.name': 'ai.generateText',
554560
'sentry.op': 'gen_ai.invoke_agent',
555561
'sentry.origin': 'auto.vercelai.otel',

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

Lines changed: 72 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -60,13 +60,34 @@ function onVercelAiSpanStart(span: Span): void {
6060
processGenerateSpan(span, name, attributes);
6161
}
6262

63+
interface TokenSummary {
64+
inputTokens: number;
65+
outputTokens: number;
66+
}
67+
6368
function vercelAiEventProcessor(event: Event): Event {
6469
if (event.type === 'transaction' && event.spans) {
70+
// Map to accumulate token data by parent span ID
71+
const tokenAccumulator: Map<string, TokenSummary> = new Map();
72+
73+
// First pass: process all spans and accumulate token data
6574
for (const span of event.spans) {
66-
// this mutates spans in-place
6775
processEndedVercelAiSpan(span);
76+
77+
// Accumulate token data for parent spans
78+
accumulateTokensForParent(span, tokenAccumulator);
79+
}
80+
81+
// Second pass: apply accumulated token data to parent spans
82+
for (const span of event.spans) {
83+
if (span.op !== 'gen_ai.invoke_agent') {
84+
continue;
85+
}
86+
87+
applyAccumulatedTokens(span, tokenAccumulator);
6888
}
6989
}
90+
7091
return event;
7192
}
7293
/**
@@ -241,6 +262,56 @@ export function addVercelAiProcessors(client: Client): void {
241262
client.addEventProcessor(Object.assign(vercelAiEventProcessor, { id: 'VercelAiEventProcessor' }));
242263
}
243264

265+
/**
266+
* Accumulates token data from a span to its parent in the token accumulator map.
267+
* This function extracts token usage from the current span and adds it to the
268+
* accumulated totals for its parent span.
269+
*/
270+
function accumulateTokensForParent(span: SpanJSON, tokenAccumulator: Map<string, TokenSummary>): void {
271+
const parentSpanId = span.parent_span_id;
272+
if (!parentSpanId) {
273+
return;
274+
}
275+
276+
const inputTokens = span.data[GEN_AI_USAGE_INPUT_TOKENS_ATTRIBUTE];
277+
const outputTokens = span.data[GEN_AI_USAGE_OUTPUT_TOKENS_ATTRIBUTE];
278+
279+
if (typeof inputTokens === 'number' || typeof outputTokens === 'number') {
280+
const existing = tokenAccumulator.get(parentSpanId) || { inputTokens: 0, outputTokens: 0 };
281+
282+
if (typeof inputTokens === 'number') {
283+
existing.inputTokens += inputTokens;
284+
}
285+
if (typeof outputTokens === 'number') {
286+
existing.outputTokens += outputTokens;
287+
}
288+
289+
tokenAccumulator.set(parentSpanId, existing);
290+
}
291+
}
292+
293+
/**
294+
* Applies accumulated token data to the `gen_ai.invoke_agent` span.
295+
* Only immediate children of the `gen_ai.invoke_agent` span are considered,
296+
* since aggregation will automatically occur for each parent span.
297+
*/
298+
function applyAccumulatedTokens(span: SpanJSON, tokenAccumulator: Map<string, TokenSummary>): void {
299+
const accumulated = tokenAccumulator.get(span.span_id);
300+
if (!accumulated) {
301+
return;
302+
}
303+
304+
if (accumulated.inputTokens > 0) {
305+
span.data[GEN_AI_USAGE_INPUT_TOKENS_ATTRIBUTE] = accumulated.inputTokens;
306+
}
307+
if (accumulated.outputTokens > 0) {
308+
span.data[GEN_AI_USAGE_OUTPUT_TOKENS_ATTRIBUTE] = accumulated.outputTokens;
309+
}
310+
if (accumulated.inputTokens > 0 || accumulated.outputTokens > 0) {
311+
span.data['gen_ai.usage.total_tokens'] = accumulated.inputTokens + accumulated.outputTokens;
312+
}
313+
}
314+
244315
function addProviderMetadataToAttributes(attributes: SpanAttributes): void {
245316
const providerMetadata = attributes[AI_RESPONSE_PROVIDER_METADATA_ATTRIBUTE] as string | undefined;
246317
if (providerMetadata) {

0 commit comments

Comments
 (0)