Skip to content

fix: stop subtracting cache tokens from Bedrock Converse input_tokens (cherry-pick of upstream #832) - #29

Merged
samuelboland merged 1 commit into
mainfrom
dux-llmpatches-w1-03-converse-input-tokens
Sep 24, 2026
Merged

samuelboland merged 1 commit into
mainfrom
dux-llmpatches-w1-03-converse-input-tokens

Conversation

@staging-supernova-dx-appf-io

Copy link
Copy Markdown

What this does

Fixes double-subtraction of cache tokens from input_tokens in the Bedrock Converse protocol. RubyLLM::Protocols::Converse::Chat.input_tokens was computing max(inputTokens - cacheReadInputTokens - cacheWriteInputTokens, 0), but AWS Bedrock already reports inputTokens as excluding cached tokens — inputTokens, cacheReadInputTokens, and cacheWriteInputTokens are separate, additive buckets, not inputTokens inclusive of cache. The subtraction removed tokens never counted in the first place and floored the result to 0 whenever the cached prefix exceeded fresh input, which is the normal case in multi-turn conversations. Production data showed every assistant message with cache tokens recording input_tokens = 0 because of this bug.

This ships upstream's fix from crmne#832 (commit dc79ce62949d5648a838d4bdffda7499e55186d7), cherry-picked onto this fork's base. input_tokens(usage) now simply returns usage['inputTokens'] as-is, with a comment explaining why (linking AWS's prompt-caching docs). The streaming path (Protocols::Converse::Streaming#extract_input_tokens) delegates directly to Chat.input_tokens, so it picks up the fix automatically without any code change. Inspection of chat.rb and streaming.rb confirmed the only other uses of cacheReadInputTokens/cacheWriteInputTokens are pass-throughs populating cached_tokens/cache_creation_tokens on messages — no subtraction — so no other Converse code path needed the same fix.

The spec in spec/ruby_llm/protocols/converse/chat_spec.rb was updated: the old example asserting the subtraction behavior was replaced with one asserting AWS's inputTokens is exposed as-is (inputTokens=50, cacheReadInputTokens=40, cacheWriteInputTokens=10 → input_tokens == 50), and a new example covers the previously-broken case (inputTokens=3, cacheReadInputTokens=7714, cacheWriteInputTokens=327 → input_tokens == 3, no longer floored to 0). All other fork-only examples in that file were preserved unchanged.

Out of scope: the Mantle Responses protocol's token accounting is deliberately left untouched, since OpenAI's input_tokens there is inclusive of cached tokens and the subtraction is correct in that context. The corresponding Gemfile.lock bump on the supernova side follows separately and is not part of this PR.

Type of change

  • Bug fix
  • New feature
  • Breaking change
  • Documentation
  • Performance improvement

Scope check

  • I read the Contributing Guide
  • This aligns with RubyLLM's focus on LLM communication
  • This isn't application-specific logic that belongs in user code
  • This benefits most users, not just my specific use case

Required for new features

N/A — this is a bug fix, not a new feature.

Quality check

  • I ran overcommit --install and all hooks pass
  • I tested my changes thoroughly
    • For provider changes: Re-recorded VCR cassettes with bundle exec rake vcr:record[provider_name]
    • All tests pass: bundle exec rspec
  • I updated documentation if needed
  • I didn't modify auto-generated files manually (models.json, aliases.json)

bundle exec rspec spec/ruby_llm/protocols/converse/ spec/ruby_llm/providers/bedrock_spec.rb and bundle exec rubocop on the changed files were run and reported passing/clean. No VCR cassette re-recording was needed since no live API interaction changed — this is a purely arithmetic fix. The specs added for this change are precise, covering both the general non-subtraction behavior and the previously-broken high-cache-token case.

AI-generated code

  • I used AI tools to help write this code
  • I have reviewed and understand all generated code (required if above is checked)

API changes

  • Breaking change
  • New public methods/classes
  • Changed method signatures
  • No API changes

The method name and signature at the Chat.input_tokens boundary are unchanged; only the internal return-value calculation changed. Behaviorally, input_tokens on Bedrock Converse responses will now report higher, correct values instead of frequently reporting 0.

…mne#832)

Fixes crmne#828.

### Problem

`Converse::Chat#input_tokens` subtracts `cacheReadInputTokens` and
`cacheWriteInputTokens` from `inputTokens`. But AWS Bedrock's
`inputTokens` is **already** the non-cached count. The cache buckets are
reported separately, not folded in. Per the [AWS prompt-caching
docs](https://docs.aws.amazon.com/bedrock/latest/userguide/prompt-caching.html):

> When prompt caching is enabled, the `inputTokens` field represents
**only the non-cached input tokens**... `total input tokens =
inputTokens + cacheReadInputTokens + cacheWriteInputTokens`

So the subtraction removes tokens that were never there, understating
input on every cache hit and flooring it to `0` whenever the cached
prefix exceeds the fresh input (the common multi-turn case). Any cost or
usage tracking on `message.input_tokens` under-reports.

### Why pass-through is correct

The `Anthropic` provider already handles the same Claude models
correctly: `anthropic/chat.rb` passes `input_tokens` through raw and
reads the cache buckets separately. Bedrock serves those same models
with the same semantics, so this just brings Converse into line. The
same fix landed in LiteLLM
([#15292](BerriAI/litellm#15292)) and was
reported in LangSmith
([#1858](langchain-ai/langsmith-sdk#1858)).

### Change

- `input_tokens` returns `usage['inputTokens']` unchanged.
`cached_tokens` and `cache_creation_tokens` are still parsed separately,
so nothing is lost. Streaming delegates to this method, so both paths
are covered.
- Updated the converse chat spec. The previous test asserted the
subtraction on a payload that cannot occur under AWS semantics
(`inputTokens` already excludes cache), so its fixture is corrected to a
realistic one, and a regression is added for the floor-to-zero case.

Verified: `bundle exec rspec spec/ruby_llm/protocols/converse/` plus
`spec/ruby_llm/providers/bedrock_spec.rb` green (39 examples), rubocop
clean.

Co-authored-by: Carmine Paolino <carmine@paolino.me>

(cherry picked from commit dc79ce6)

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Co-authored-by: Sam Boland <sam.boland@appfolio.com>
@staging-supernova-dx-appf-io

Copy link
Copy Markdown
Author

🤖 Supernova Code Review — View trace

🔄 Re-run this review

— via automation "Code Review Opus 5.5 Medium" (run), installed by Sam Boland (sam.boland@appfolio.com)

@staging-supernova-dx-appf-io staging-supernova-dx-appf-io Bot left a comment

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

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

Reviewed the PR and found nothing I'd ask you to change. No inline comments.

What I checked:

  • lib/ruby_llm/protocols/converse/chat.rb: input_tokens now returns usage['inputTokens'] as-is. This matches how AWS reports it: inputTokens, cacheReadInputTokens and cacheWriteInputTokens are separate counts that add up. The old early return for a missing inputTokens is still covered, because the value just comes back as nil.
  • lib/ruby_llm/protocols/converse/streaming.rb (extract_input_tokens, lines 286–289) calls Chat.input_tokens, so the streaming path gets the fix too. cached_tokens and cache_creation_tokens were already copied straight from the response.
  • spec/ruby_llm/protocols/converse/streaming_spec.rb has no assertions that depended on the old subtraction.
  • The updated and new examples in chat_spec.rb test both cases: normal input with cache tokens, and the case where the cached prefix is larger than the new input and the result used to drop to 0.

🔄 Re-run this review

— via automation "Code Review Opus 5.5 Medium" (run), installed by Sam Boland (sam.boland@appfolio.com)

@samuelboland
samuelboland merged commit 8e0c88d into main Sep 24, 2026
24 of 25 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants