fix(instrumentation-bedrock): avoid NaN total_tokens for Titan embedding responses

[BrontoStephen]

Summary

Fixes #1031.

Titan embedding responses (e.g. amazon.titan-embed-text-v2:0) carry inputTextTokenCount but no output-token field — there is no generated text. Both Amazon response paths computed total_tokens as input + undefined, which evaluates to NaN and gets serialized as the string "NaN" on the exported span. In our production telemetry, 100% of spans carrying total_tokens from this path had the value "NaN".

Changes

• Non-streaming Amazon path: titanOutputTokens (response.results?.[0]?.tokenCount) doesn't exist for embeddings. Now omits gen_ai.usage.output_tokens when absent and computes total_tokens = input + (output ?? 0).
• Streaming final-chunk path: same guard for totalOutputTextTokenCount.
• Embeddings are billed on input tokens only, so total_tokens = input is the semantically correct value for them.

Tests

Added tests/titan-embedding-tokens.test.ts (same direct-handler style as cache-token-fold-in.test.ts) covering:

• embedding response → numeric total_tokens equal to input, output_tokens omitted, never NaN
• Titan text response → unchanged input + output behavior (regression)
• streaming final chunk with and without output counts

nx test @traceloop/instrumentation-bedrock: 82 passing (78 existing + 4 new). Lint clean.

🤖 Generated with Claude Code

Summary by CodeRabbit

• Bug Fixes

  • Fixed token usage reporting for Amazon Bedrock Titan streaming and embedding responses when output token counts are unavailable.
  • Prevented total token counts from displaying invalid NaN values.
  • Ensured total tokens accurately reflect available input and output token data.

• Tests

  • Added coverage for streaming, embedding, and text response token accounting.

[CLAassistant]

All committers have signed the CLA.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: f725e8a4-afdd-4eb7-8a61-23ee56081838

📥 Commits

Reviewing files that changed from the base of the PR and between 1faf69d and 1dc41ce.

📒 Files selected for processing (2)

• packages/instrumentation-bedrock/src/instrumentation.ts
• packages/instrumentation-bedrock/tests/titan-embedding-tokens.test.ts

---

📝 Walkthrough

Walkthrough

Changes

Titan token accounting

Layer / File(s)	Summary
Guard Titan token attributes packages/instrumentation-bedrock/src/instrumentation.ts, packages/instrumentation-bedrock/tests/titan-embedding-tokens.test.ts	Streaming and non-streaming Titan responses omit unavailable output-token attributes and calculate totals with missing output treated as zero. Tests cover embedding, text, and streaming responses.

Estimated code review effort: 2 (Simple) | ~10 minutes

Suggested reviewers: dvirski

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly matches the main fix: preventing NaN total_tokens for Titan embedding responses in Bedrock instrumentation.
Linked Issues check	✅ Passed	The code and tests address #1031 by handling missing output tokens in non-streaming and streaming Titan embedding paths while preserving text behavior.
Out of Scope Changes check	✅ Passed	The changes stay focused on the Titan token-count fix and its tests, with no unrelated code paths or features introduced.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

Warning

There were issues while running some tools. Please review the errors and either fix the tool's configuration or disable the tool if it's a critical failure.

🔧 ESLint

If the error stems from missing dependencies, add them to the package.json file. For unrecoverable errors (e.g., due to private dependencies), disable the tool in the CodeRabbit configuration.

ESLint install timed out. The project may have too many dependencies for the sandbox.

---

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

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}