Skip to content

fix(instrumentations): add cache token span attributes for OpenAI, Bedrock, Together AI, and Vertex AI#1026

Merged
dvirski merged 2 commits into
mainfrom
dr/fix(instrumentations)-emit-cache-tokens
Jun 18, 2026
Merged

fix(instrumentations): add cache token span attributes for OpenAI, Bedrock, Together AI, and Vertex AI#1026
dvirski merged 2 commits into
mainfrom
dr/fix(instrumentations)-emit-cache-tokens

Conversation

@dvirski

@dvirski dvirski commented Jun 15, 2026

Copy link
Copy Markdown
Contributor

Summary by CodeRabbit

  • Bug Fixes

    • Improved cache token metrics tracking and validation across Anthropic, Google Generative AI, and Bedrock integrations with stricter null/undefined checks.
    • Enhanced totalInputTokens calculation to accurately include cache-related token values.
  • Tests

    • Added test recording for Bedrock cache token handling scenarios.

@coderabbitai

coderabbitai Bot commented Jun 15, 2026

Copy link
Copy Markdown
Contributor

Review Change Stack

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info
⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: 9a515825-4a0c-41ba-8e36-143b96791a74

📥 Commits

Reviewing files that changed from the base of the PR and between 45a1278 and ce12fa3.

📒 Files selected for processing (4)
  • packages/instrumentation-anthropic/src/instrumentation.ts
  • packages/instrumentation-bedrock/recordings/Test-Anthropic-with-AWS-Bedrock-Instrumentation_3427560822/should-set-cache-tokens-in-span-for-Anthropic-messages-API-with-cached-tokens_3921807056/recording.har
  • packages/instrumentation-bedrock/src/instrumentation.ts
  • packages/instrumentation-google-generativeai/src/instrumentation.ts
✅ Files skipped from review due to trivial changes (2)
  • packages/instrumentation-bedrock/recordings/Test-Anthropic-with-AWS-Bedrock-Instrumentation_3427560822/should-set-cache-tokens-in-span-for-Anthropic-messages-API-with-cached-tokens_3921807056/recording.har
  • packages/instrumentation-google-generativeai/src/instrumentation.ts
🚧 Files skipped from review as they are similar to previous changes (2)
  • packages/instrumentation-anthropic/src/instrumentation.ts
  • packages/instrumentation-bedrock/src/instrumentation.ts

📝 Walkthrough

Walkthrough

Cache token span attribute guards are tightened from loose != null to strict !== null && !== undefined checks in the Anthropic and Google GenerativeAI instrumentations. The Bedrock non-streaming path is updated to fold cache_read_input_tokens and cache_creation_input_tokens into totalInputTokens using nullish coalescing, and a new HAR recording is added for the Bedrock cached-tokens test.

Changes

Cache Token Guard and Calculation Fixes

Layer / File(s) Summary
Cache token guard tightening across packages
packages/instrumentation-anthropic/src/instrumentation.ts, packages/instrumentation-google-generativeai/src/instrumentation.ts
_endSpan in both packages replaces loose != null checks with explicit !== null && !== undefined guards before emitting cache creation and cache read input token span attributes.
Bedrock: totalInputTokens fold and HAR recording
packages/instrumentation-bedrock/src/instrumentation.ts, packages/instrumentation-bedrock/recordings/.../recording.har
Non-streaming Anthropic usage mapping drops || 0 pre-defaults and uses ?? 0 to include cache token fields in totalInputTokens, while still emitting cache semconv attributes only when the fields are present. A new HAR recording captures the Bedrock Anthropic cached-tokens response for tests.

Estimated code review effort

🎯 1 (Trivial) | ⏱️ ~5 minutes

Possibly related PRs

  • traceloop/openllmetry-js#1028: Directly modifies the same packages/instrumentation-bedrock/src/instrumentation.ts Anthropic non-streaming cache token accounting logic that this PR also changes.

Suggested reviewers

  • max-deygin-traceloop

Poem

🐇 Null checks once loose, now tight as a drum,
Cache tokens are folded with ?? 0 sum.
The HAR records faithfully what Bedrock said,
No phantom zeros, just values instead.
Hop hop, consistency spreads through the code! 🌿

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name Status Explanation Resolution
Title check ⚠️ Warning The PR title states it adds cache token span attributes to multiple instrumentations (OpenAI, Bedrock, Together AI, Vertex AI), but the provided changeset only shows modifications to Bedrock, Anthropic, and Google GenerativeAI instrumentations, with no evidence of OpenAI, Together AI, or Vertex AI changes. Update the PR title to accurately reflect which instrumentations were modified (e.g., 'fix(instrumentations): add/tighten cache token span attributes for Bedrock, Anthropic, and Google GenerativeAI')
✅ Passed checks (4 passed)
Check name Status Explanation
Description Check ✅ Passed Check skipped - CodeRabbit’s high-level summary is enabled.
Docstring Coverage ✅ Passed No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check ✅ Passed Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check ✅ Passed Check skipped because no linked issues were found for this pull request.

✏️ Tip: You can configure your own custom pre-merge checks in the settings.

✨ Finishing Touches
📝 Generate docstrings
  • Create stacked PR
  • Commit on current branch
🧪 Generate unit tests (beta)
  • Create PR with unit tests
  • Commit unit tests in branch dr/fix(instrumentations)-emit-cache-tokens

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

Comment @coderabbitai help to get the list of available commands and usage tips.

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

Actionable comments posted: 1

🧹 Nitpick comments (1)
packages/instrumentation-together/src/instrumentation.ts (1)

525-533: ⚡ Quick win

Consider extending the Completion type instead of using a weak type cast.

The current implementation uses as unknown as Record<string, unknown> to access cached_tokens, which bypasses TypeScript's type safety. Following the pattern established in the OpenAI instrumentation (lines 117-119 in packages/instrumentation-openai/src/instrumentation.ts), consider defining a local type extension at the top of the file:

interface CompletionUsageWithCache extends Completion.CompletionUsage {
  cached_tokens?: number;
}

Then access it with better type safety:

-        const cachedTokens = (
-          result.usage as unknown as Record<string, unknown>
-        ).cached_tokens;
-        if (cachedTokens) {
+        const cachedTokens = (result.usage as CompletionUsageWithCache)
+          ?.cached_tokens;
+        if (typeof cachedTokens === 'number') {
           span.setAttribute(
             ATTR_GEN_AI_USAGE_CACHE_READ_INPUT_TOKENS,
-            cachedTokens as number,
+            cachedTokens,
           );
         }

This approach provides better type safety and is more maintainable.

🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@packages/instrumentation-together/src/instrumentation.ts` around lines 525 -
533, The code uses a weak type cast (as unknown as Record<string, unknown>) to
access cached_tokens from the result.usage object, which bypasses TypeScript's
type safety. Define a local interface extension at the top of the
instrumentation.ts file that extends Completion.CompletionUsage with an optional
cached_tokens property (following the pattern in the OpenAI instrumentation),
then replace the weak type cast with this properly typed interface when
accessing cached_tokens before the span.setAttribute call for
ATTR_GEN_AI_USAGE_CACHE_READ_INPUT_TOKENS. This improves type safety and
maintainability without changing the functional behavior.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@packages/instrumentation-bedrock/src/instrumentation.ts`:
- Around line 706-717: Cache token attributes are currently being guarded with
truthy checks, which silently drops valid zero values from provider responses.
In packages/instrumentation-bedrock/src/instrumentation.ts at lines 706-717,
replace the truthy checks for usage["cache_read_input_tokens"] and
usage["cache_creation_input_tokens"] with null/undefined checks using != null or
!== undefined operators to preserve zero values in
ATTR_GEN_AI_USAGE_CACHE_READ_INPUT_TOKENS and
ATTR_GEN_AI_USAGE_CACHE_CREATION_INPUT_TOKENS attributes. In
packages/instrumentation-vertexai/src/vertexai-instrumentation.ts at lines
256-260, similarly replace the truthy check for cachedContentTokenCount with a
!== undefined or != null check when setting
ATTR_GEN_AI_USAGE_CACHE_READ_INPUT_TOKENS to ensure zero values are faithfully
emitted.

---

Nitpick comments:
In `@packages/instrumentation-together/src/instrumentation.ts`:
- Around line 525-533: The code uses a weak type cast (as unknown as
Record<string, unknown>) to access cached_tokens from the result.usage object,
which bypasses TypeScript's type safety. Define a local interface extension at
the top of the instrumentation.ts file that extends Completion.CompletionUsage
with an optional cached_tokens property (following the pattern in the OpenAI
instrumentation), then replace the weak type cast with this properly typed
interface when accessing cached_tokens before the span.setAttribute call for
ATTR_GEN_AI_USAGE_CACHE_READ_INPUT_TOKENS. This improves type safety and
maintainability without changing the functional behavior.
🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

  • Push a commit to this branch (recommended)
  • Create a new PR with the fixes

ℹ️ Review info
⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: 10f66fbf-e369-4454-b281-f604f43ee0e1

📥 Commits

Reviewing files that changed from the base of the PR and between 28c4a7a and c495dc1.

⛔ Files ignored due to path filters (1)
  • pnpm-lock.yaml is excluded by !**/pnpm-lock.yaml
📒 Files selected for processing (10)
  • packages/instrumentation-bedrock/src/instrumentation.ts
  • packages/instrumentation-bedrock/tests/anthropic.test.ts
  • packages/instrumentation-openai/src/instrumentation.ts
  • packages/instrumentation-openai/test/instrumentation.test.ts
  • packages/instrumentation-together/package.json
  • packages/instrumentation-together/src/instrumentation.ts
  • packages/instrumentation-together/test/instrumentation.test.ts
  • packages/instrumentation-vertexai/package.json
  • packages/instrumentation-vertexai/src/vertexai-instrumentation.ts
  • packages/instrumentation-vertexai/tests/gemini.test.ts

Comment thread packages/instrumentation-bedrock/src/instrumentation.ts Outdated
@dvirski dvirski changed the title feat(instrumentations): add cache token span attributes for OpenAI, Bedrock, Together AI, and Vertex AI fix(instrumentations): add cache token span attributes for OpenAI, Bedrock, Together AI, and Vertex AI Jun 17, 2026
@dvirski dvirski force-pushed the dr/fix(instrumentations)-emit-cache-tokens branch from 45a1278 to ce12fa3 Compare June 18, 2026 14:43
@dvirski dvirski merged commit 1faf69d into main Jun 18, 2026
8 of 9 checks passed
@dvirski dvirski deleted the dr/fix(instrumentations)-emit-cache-tokens branch June 18, 2026 14:49
Sign up for free to 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