fix(instrumentation-bedrock): safely compute titan total tokens when output tokens are missing

[Yogesh0314]

Description

Fixes #1031.

For Amazon Titan embedding responses (e.g. amazon.titan-embed-text-v2:0), inputTextTokenCount is present in the response but output token counts are absent. Computing total tokens as inputTextTokenCount + titanOutputTokens resulted in number + undefined = NaN, which was serialized as "NaN" in exported span attributes.

Changes

• In packages/instrumentation-bedrock/src/instrumentation.ts:
  • Default missing output tokens to 0 when calculating GEN_AI_USAGE_TOTAL_TOKENS for Amazon Titan responses (titanInputTokens + (titanOutputTokens ?? 0)).
  • Only emit ATTR_GEN_AI_USAGE_OUTPUT_TOKENS if output tokens are defined.
  • Apply the same safeguard to the streaming response handler when totalOutputTextTokenCount is absent.
• Added comprehensive unit tests in packages/instrumentation-bedrock/tests/titan-tokens.test.ts.

Testing

• Ran all unit tests in @traceloop/instrumentation-bedrock (all passing).
• Verified code formatting with Prettier and linting with ESLint.

Summary by CodeRabbit

• Bug Fixes
  • Amazon Titan token usage now omits output-token counts when they aren’t available, while total-token counts continue to include input tokens.

[coderabbitai]

Navigate logical layers of code changes, visualize relationships, and explore their blast radius.

🧰 Additional context used

📚 Code guidelines (1)

CLAUDE.md — auto-discovered

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

• Configuration used: defaults
• Review profile: CHILL
• Plan: Advanced
• Run ID: ffec3218-fc6f-420f-9dd2-0e9a398ad886

📥 Commits

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

📒 Files selected for processing (2)

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

Included review availability: This review used your included allowance. Your plan provides up to 2 included reviews per hour; 1 remain after this review.

---

📝 Walkthrough

Walkthrough

Titan streaming and non-streaming responses now omit output-token usage attributes when output counts are unavailable. Total-token usage still includes the input count, treating a missing output count as zero. Tests cover both missing and present output counts.

Changes

Titan token usage

Layer / File(s)	Summary
Conditional token attributes and coverage packages/instrumentation-bedrock/src/instrumentation.ts, packages/instrumentation-bedrock/tests/titan-tokens.test.ts	Streaming and non-streaming response handling omits output-token attributes when their counts are missing. Total tokens use the input count when no output count is present. Tests cover missing and present counts in both response paths.

Priority: ➖ Normal

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

Change: Bug fix · Severity of issue fixed: Medium

Merge Risk: ⚪ Minimal · up to 6ea10

This change stops Titan embedding spans from reporting a NaN total token count when no output count is present. The output-token attribute is now left out in that case. No merge-blocking risk was found.

Security Architecture Review

Security architecture risk: ⚪ Minimal · up to 6ea10

The change corrects token-count reporting without adding access, privileges, sensitive content collection, or new dependencies. Responses containing both counts retain their existing behavior. No material security risk was identified in the changed flow.

Retained concerns
No architecture-level concerns identified.

Security review details

Security Blast Radius

• inferred — The changed flow remains confined to token-usage telemetry on existing Bedrock spans. The comparison introduces no new caller, credential use, privileged operation, or content-export path. This bounds the assessed change; it does not establish complete security coverage of the package or downstream consumers.

Trust Boundaries and Controls

• observed — Provider responses continue to enter the existing JSON-parsing and span-mapping path. Output-message collection remains conditional on _shouldSendPrompts, which uses the active context override or traceContent configuration. These controls and their defaults predate the PR and are unchanged.

🚥 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 summarizes the main change: safely computing Amazon Titan total tokens when output-token counts are missing.
Linked Issues check	✅ Passed	Issue #1031 requires safe Titan token totals for non-streaming and streaming responses, and omission of the output-token attribute when its count is absent. The reviewed change defaults the missing ou…
Out of Scope Changes check	✅ Passed	The changes are limited to Titan token accounting in the Bedrock instrumentation and focused tests for that behavior. The changes support issue #1031 and show no unrelated scope.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check. Docstring coverage is scoped to functions touched by this diff. Analyzed 0 functions across 2…

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create a new PR

• Autopilot · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

Warning

Some tools did not complete. Review the errors below.

🔧 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.

---

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

[CLAassistant]

All committers have signed the CLA.