fix(server-utils): Read AI SDK cache token counts in Vercel AI channel subscriber - #24350
fix(server-utils): Read AI SDK cache token counts in Vercel AI channel subscriber#24350sergical wants to merge 4 commits into
Conversation
…l subscriber The `ai:telemetry` channel subscriber only derived `gen_ai.usage.cache_read.input_tokens` and `gen_ai.usage.cache_creation.input_tokens` from provider-keyed `providerMetadata`. Through the Vercel AI Gateway that object is keyed `gateway`, so the cache counts the AI SDK already normalizes into its usage object were dropped. The v10 OTel-based integration mapped them. Read them from the usage object as well: `cachedInputTokens` (v5), `inputTokenDetails.cacheReadTokens` / `cacheWriteTokens` (v6), and the `inputTokens.cacheRead` / `cacheWrite` objects of v7 model-call usage. Provider-derived values are applied afterwards and still win. Co-authored-by: Claude <claude@anthropic.com>
![sentry[bot]](https://avatars.githubusercontent.com/in/12637?s=60&v=4){kind=small}
size-limit report 📦
|
Co-authored-by: Claude <claude@anthropic.com>
logaretm
left a comment
logaretm
left a comment
There was a problem hiding this comment.
clanker flagged some stuff, plus got a question on undefined possibly now wiping attributes but doubt it could happen.
Could that be tested?
| GEN_AI_USAGE_OUTPUT_TOKENS, | ||
| GEN_AI_USAGE_TOTAL_TOKENS, | ||
| GEN_AI_USAGE_CACHE_READ_INPUT_TOKENS, | ||
| GEN_AI_USAGE_CACHE_CREATION_INPUT_TOKENS, |
There was a problem hiding this comment.
Adding the two cache keys here drops them from providerAttributes on every root operation, but I think v6/v7 have replacements for them so it's not an issue there.
On v4/v5 cache_creation would get dropped from the root span, is that desirable?
There was a problem hiding this comment.
I think so. On a root span the other last-step keys (output_tokens, total_tokens) are already dropped for the same reason, and a last-step cache write count on a span whose input/read counts are aggregated across steps reads as the whole call's writes. On v4/v5 the SDK reports no cache writes at all, so the root span ends up with reads only, which matches what the SDK itself exposes. Model-call spans keep the providerMetadata value in every version.
| if (totalTokens !== undefined) { | ||
| span.setAttribute(GEN_AI_USAGE_TOTAL_TOKENS, totalTokens); | ||
| } | ||
| span.setAttributes(cacheTokenAttributes(usage)); |
There was a problem hiding this comment.
Q: Now that undefined could be passed in here, is it possible that it wipes previously set attributes? I don't think that case could happen but maybe something to check as well.
There was a problem hiding this comment.
Good catch, it did: setAttributes with an undefined value removed an attribute that was already on the span, so a v5 usage object (no cache write count) wiped a cache_creation set earlier. Fixed in fa57569 by only setting the counts that are reported, with a test for the existing-attribute case.
| * Cache token counts as the AI SDK normalizes them: v5 `cachedInputTokens`, v6 `inputTokenDetails`, | ||
| * v7 `inputTokens.{cacheRead,cacheWrite}`. | ||
| */ | ||
| function cacheTokenAttributes(usage: Record<string, unknown>): Record<string, number> { |
There was a problem hiding this comment.
Since it could be number | undefined this might fail now in the tests.
There was a problem hiding this comment.
Resolved by the same change: the helper now returns { cacheRead?, cacheWrite? } and the callers skip undefined, so nothing typed number receives it.
Co-authored-by: Claude <claude@anthropic.com>
2 participants
The v11 channel subscriber only maps cache token counts out of
providerMetadata, keyed by provider name. Calls through the Vercel AI Gateway carry agatewaykey instead, sogen_ai.usage.cache_read.input_tokensandgen_ai.usage.cache_creation.input_tokensare never set for them. The OTel processor read the SDK's normalized counts (ai.usage.cachedInputTokens) and was removed in #23384 together with its test, so this went out untested.The subscriber now reads the SDK's normalized usage in all three shapes:
usage.cachedInputTokensusage.inputTokenDetails.{cacheReadTokens,cacheWriteTokens}usage.inputTokens.{cacheRead,cacheWrite}A provider-reported count from
providerMetadatastill wins on model-call spans. On root operations the cache keys are now last-step-only, so the aggregated SDK counts are kept.Created with Claude Code