Skip to content

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

Merged
RulaKhaled merged 3 commits into
developfrom
rolaabuhasna/js-662-token-attributes-on-gen_aiinvoke_agent-span
Aug 4, 2025
Merged

feat(core): Accumulate tokens for gen_ai.invoke_agent spans from child LLM calls#17281
RulaKhaled merged 3 commits into
developfrom
rolaabuhasna/js-662-token-attributes-on-gen_aiinvoke_agent-span

Conversation

@RulaKhaled

Copy link
Copy Markdown
Collaborator

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.

@linear

linearBot commented Aug 1, 2025

Copy link
Copy Markdown

cursor[bot]

This comment was marked as outdated.

@github-actions

github-actionsBot commented Aug 1, 2025

Copy link
Copy Markdown
Contributor

size-limit report 📦

PathSize% ChangeChange
@sentry/browser23.76 kB--
@sentry/browser - with treeshaking flags22.35 kB--
@sentry/browser (incl. Tracing)39.43 kB--
@sentry/browser (incl. Tracing, Replay)77.52 kB--
@sentry/browser (incl. Tracing, Replay) - with treeshaking flags67.39 kB--
@sentry/browser (incl. Tracing, Replay with Canvas)82.22 kB--
@sentry/browser (incl. Tracing, Replay, Feedback)94.35 kB--
@sentry/browser (incl. Feedback)40.45 kB--
@sentry/browser (incl. sendFeedback)28.45 kB--
@sentry/browser (incl. FeedbackAsync)33.34 kB--
@sentry/react25.5 kB--
@sentry/react (incl. Tracing)41.4 kB--
@sentry/vue28.2 kB--
@sentry/vue (incl. Tracing)41.23 kB--
@sentry/svelte23.79 kB--
CDN Bundle25.28 kB--
CDN Bundle (incl. Tracing)39.3 kB--
CDN Bundle (incl. Tracing, Replay)75.38 kB--
CDN Bundle (incl. Tracing, Replay, Feedback)80.82 kB--
CDN Bundle - uncompressed73.86 kB--
CDN Bundle (incl. Tracing) - uncompressed116.35 kB--
CDN Bundle (incl. Tracing, Replay) - uncompressed230.56 kB--
CDN Bundle (incl. Tracing, Replay, Feedback) - uncompressed243.37 kB--
@sentry/nextjs (client)43.46 kB--
@sentry/sveltekit (client)39.87 kB--
@sentry/node-core47.53 kB-0.01%-1 B 🔽
@sentry/node146.36 kB+0.12%+168 B 🔺
@sentry/node - without tracing91.63 kB--
@sentry/aws-serverless103.08 kB+0.01%+1 B 🔺

View base workflow run

@mydeamydea changed the title feat(core): Accumulate tokens for gen_ai.invoke_agent spans from child LLM callsfeat(core): Accumulate tokens for gen_ai.invoke_agent spans from child LLM callsAug 4, 2025
Comment threadpackages/core/src/utils/vercel-ai.ts Outdated
// 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);

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

hmm, can we combine this into a single pass somehow? 🤔 that would be more efficient I suppose.

@RulaKhaledRulaKhaledAug 4, 2025

Copy link
Copy Markdown
CollaboratorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

we can but it's safer to leave them separate because we are processing child spans to accumulate the tokens, we need to make sure they are all transformed completely to the attribute we are expecting, otherwise we could end up with a parent span first that have child spans that still use vercel naming so we end up not processing it.

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

hmm overall this is quite expensive, looking at it :D because we loop over all spans a lot:

  1. Once in the first pass
  2. Once in the second pass
    a. for each second pass, we iterate over all spans again when filtering
    b. then finally once more over the filtered spans (which is a subset already, but still)

IMHO it's probably worth it to find a way to streamline this to get to max. 2 passes. I would propose something like this:

  1. We create an object to hold the summarized tokens. e.g. something like this: const tmpAmounts: Map<string, TokenSummary> = new Map();
  2. We pass this into the first pass function as second argument, so that can mutate it
  3. The function that transforms attributes, at the very end, updates the passed in map by incrementing the token summary for the map item of the parentSpanId, e.g. something like this (simplified...):
constparentSpanId=spanToJSON(span).parent_span_id;if(parentSpanId){constparentSummary=tmpAmounts.get(parentSpanId)||{total: 0};parentSummary.total+=amount;tmpAmounts.set(parentSpanId,parentSummary);}
  1. Finally, after the first pass, we iterate over tmpAmounts and update the parent spans with the stored summaries.

I think this or something like this should reduce overhead considerably 🤔

Copy link
Copy Markdown
CollaboratorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

very nice, i'll try it out.


// Second pass: apply accumulated token data to parent spans
for (const span of event.spans) {
if (span.op !== 'gen_ai.invoke_agent') {

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
if(span.op!=='gen_ai.invoke_agent'){
if(span.op!=='gen_ai.invoke_agent'||!tokenAccumulator.has(span.span_id)){

should also be a reasonable shortcut?

@mydeamydea left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

sweet!

Sign up for freeto join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants

@RulaKhaled@mydea