From 20abb58dc21a92f15d17a9c622bf20f9cb17a2e6 Mon Sep 17 00:00:00 2001 From: RulaKhaled Date: Fri, 1 Aug 2025 11:55:54 +0200 Subject: [PATCH 1/3] feat(core): Accumulate tokens for gen_ai.invoke_agent spans from child LLM calls --- .../suites/tracing/vercelai/test.ts | 6 +++ packages/core/src/utils/vercel-ai.ts | 49 +++++++++++++++++++ 2 files changed, 55 insertions(+) diff --git a/dev-packages/node-integration-tests/suites/tracing/vercelai/test.ts b/dev-packages/node-integration-tests/suites/tracing/vercelai/test.ts index 5353f53f42e3..94fd0dde8486 100644 --- a/dev-packages/node-integration-tests/suites/tracing/vercelai/test.ts +++ b/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', diff --git a/packages/core/src/utils/vercel-ai.ts b/packages/core/src/utils/vercel-ai.ts index 9c20d49ea157..c693ffd27414 100644 --- a/packages/core/src/utils/vercel-ai.ts +++ b/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) { From 0ce72f0661b4a64aca17d8cd516492c0666023a1 Mon Sep 17 00:00:00 2001 From: RulaKhaled Date: Fri, 1 Aug 2025 12:00:43 +0200 Subject: [PATCH 2/3] fix comment --- packages/core/src/utils/vercel-ai.ts | 11 +++++------ 1 file changed, 5 insertions(+), 6 deletions(-) diff --git a/packages/core/src/utils/vercel-ai.ts b/packages/core/src/utils/vercel-ai.ts index c693ffd27414..e3a944f81971 100644 --- a/packages/core/src/utils/vercel-ai.ts +++ b/packages/core/src/utils/vercel-ai.ts @@ -249,14 +249,13 @@ export function addVercelAiProcessors(client: Client): void { } /** - * 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. + * For the gen_ai.invoke_agent span, iterate over child spans and aggregate tokens: + * - Input tokens from client LLM child spans that include `gen_ai.usage.input_tokens` attribute. + * - Output tokens from client LLM child spans that include `gen_ai.usage.output_tokens` attribute. + * - Total tokens from client LLM child spans that include `gen_ai.usage.total_tokens` attribute. * - * Only immediate children of the invoke_agent span need to be considered, + * Only immediate children of the `gen_ai.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') { From 94d3f8d0757dadd62b9b54a994004c6a68b38631 Mon Sep 17 00:00:00 2001 From: RulaKhaled Date: Mon, 4 Aug 2025 10:58:54 +0200 Subject: [PATCH 3/3] make this 2n faster --- packages/core/src/utils/vercel-ai.ts | 79 ++++++++++++++++++---------- 1 file changed, 51 insertions(+), 28 deletions(-) diff --git a/packages/core/src/utils/vercel-ai.ts b/packages/core/src/utils/vercel-ai.ts index e3a944f81971..4ef437a1b922 100644 --- a/packages/core/src/utils/vercel-ai.ts +++ b/packages/core/src/utils/vercel-ai.ts @@ -60,20 +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) { - // First pass: process all spans normally + // Map to accumulate token data by parent span ID + const tokenAccumulator: Map = 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: accumulate tokens for gen_ai.invoke_agent spans - // TODO: Determine how to handle token aggregation for tool call spans. + // Second pass: apply accumulated token data to parent spans for (const span of event.spans) { - accumulateTokensFromChildSpans(span, event.spans); + if (span.op !== 'gen_ai.invoke_agent') { + continue; + } + + applyAccumulatedTokens(span, tokenAccumulator); } } + return event; } /** @@ -249,43 +263,52 @@ export function addVercelAiProcessors(client: Client): void { } /** - * For the gen_ai.invoke_agent span, iterate over child spans and aggregate tokens: - * - Input tokens from client LLM child spans that include `gen_ai.usage.input_tokens` attribute. - * - Output tokens from client LLM child spans that include `gen_ai.usage.output_tokens` attribute. - * - Total tokens from client LLM child spans that include `gen_ai.usage.total_tokens` attribute. - * - * Only immediate children of the `gen_ai.invoke_agent` span need to be considered, - * since aggregation will automatically occur for each parent span. + * 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 accumulateTokensFromChildSpans(spanJSON: SpanJSON, allSpans: SpanJSON[]): void { - if (spanJSON.op !== 'gen_ai.invoke_agent') { +function accumulateTokensForParent(span: SpanJSON, tokenAccumulator: Map): void { + const parentSpanId = span.parent_span_id; + if (!parentSpanId) { return; } - const childSpans = allSpans.filter(childSpan => childSpan.parent_span_id === spanJSON.span_id); - let totalInputTokens = 0; - let totalOutputTokens = 0; + const inputTokens = span.data[GEN_AI_USAGE_INPUT_TOKENS_ATTRIBUTE]; + const outputTokens = span.data[GEN_AI_USAGE_OUTPUT_TOKENS_ATTRIBUTE]; - 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' || typeof outputTokens === 'number') { + const existing = tokenAccumulator.get(parentSpanId) || { inputTokens: 0, outputTokens: 0 }; if (typeof inputTokens === 'number') { - totalInputTokens += inputTokens; + existing.inputTokens += inputTokens; } if (typeof outputTokens === 'number') { - totalOutputTokens += outputTokens; + 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): void { + const accumulated = tokenAccumulator.get(span.span_id); + if (!accumulated) { + return; } - if (totalInputTokens > 0) { - spanJSON.data[GEN_AI_USAGE_INPUT_TOKENS_ATTRIBUTE] = totalInputTokens; + if (accumulated.inputTokens > 0) { + span.data[GEN_AI_USAGE_INPUT_TOKENS_ATTRIBUTE] = accumulated.inputTokens; } - if (totalOutputTokens > 0) { - spanJSON.data[GEN_AI_USAGE_OUTPUT_TOKENS_ATTRIBUTE] = totalOutputTokens; + if (accumulated.outputTokens > 0) { + span.data[GEN_AI_USAGE_OUTPUT_TOKENS_ATTRIBUTE] = accumulated.outputTokens; } - if (totalInputTokens > 0 || totalOutputTokens > 0) { - spanJSON.data['gen_ai.usage.total_tokens'] = totalInputTokens + totalOutputTokens; + if (accumulated.inputTokens > 0 || accumulated.outputTokens > 0) { + span.data['gen_ai.usage.total_tokens'] = accumulated.inputTokens + accumulated.outputTokens; } }