From 16c1731b91b202d89254b073f54e8d8b9afb1870 Mon Sep 17 00:00:00 2001 From: Supernova Date: Mon, 6 Jul 2026 21:50:31 +0000 Subject: [PATCH 1/3] Capture streaming usage tokens and context management for Bedrock InvokeModel Add comprehensive token usage tracking for Bedrock InvokeModel streaming: - Capture cache_creation TTL breakdown (ephemeral_5m, ephemeral_1h) - Extract stop_sequence and context_management from responses - Fix input_tokens calculation to properly subtract cached tokens - Add debug logging for context_management applied_edits Why: - Bedrock now returns detailed cache TTL breakdowns in usage metadata - context_management signals when the service auto-edits conversation history - Accurate token accounting requires subtracting cache hits from input_tokens How to apply: - Message and Tokens classes now surface new fields - StreamAccumulator plumbs them through from chunks to final message - Tests verify both non-streaming and streaming code paths Co-Authored-By: Claude Sonnet 4.5 Co-authored-by: Sam Boland --- lib/ruby_llm/message.rb | 12 +- .../protocols/bedrock_invoke_model/chat.rb | 5 + .../bedrock_invoke_model/streaming.rb | 28 ++- lib/ruby_llm/stream_accumulator.rb | 20 +- lib/ruby_llm/tokens.rb | 26 ++- .../protocols/bedrock_invoke_model_spec.rb | 179 ++++++++++++++++++ 6 files changed, 258 insertions(+), 12 deletions(-) diff --git a/lib/ruby_llm/message.rb b/lib/ruby_llm/message.rb index 6e9aa938d..1ae06324f 100644 --- a/lib/ruby_llm/message.rb +++ b/lib/ruby_llm/message.rb @@ -15,7 +15,7 @@ class Message :CONTENT_FILTERED_FINISH_REASONS attr_reader :role, :model_id, :tool_calls, :tool_call_id, :raw, :thinking, :tokens, :citations, - :finish_reason + :finish_reason, :stop_sequence, :context_management attr_writer :content def initialize(options = {}) @@ -29,12 +29,16 @@ def initialize(options = {}) output: options[:output_tokens], cached: options[:cached_tokens], cache_creation: options[:cache_creation_tokens], - thinking: options[:thinking_tokens] + thinking: options[:thinking_tokens], + cache_creation_ephemeral_5m: options[:cache_creation_ephemeral_5m_tokens], + cache_creation_ephemeral_1h: options[:cache_creation_ephemeral_1h_tokens] ) @raw = options[:raw] @thinking = options[:thinking] @citations = Array(options[:citations]) @finish_reason = options[:finish_reason] + @stop_sequence = options[:stop_sequence] + @context_management = options[:context_management] ensure_valid_role end @@ -113,7 +117,9 @@ def to_h thinking: thinking&.text, thinking_signature: thinking&.signature, citations: citations.empty? ? nil : citations.map(&:to_h), - finish_reason: finish_reason + finish_reason: finish_reason, + stop_sequence: stop_sequence, + context_management: context_management }.merge(tokens ? tokens.to_h : {}).compact end diff --git a/lib/ruby_llm/protocols/bedrock_invoke_model/chat.rb b/lib/ruby_llm/protocols/bedrock_invoke_model/chat.rb index f0035669e..ad460d903 100644 --- a/lib/ruby_llm/protocols/bedrock_invoke_model/chat.rb +++ b/lib/ruby_llm/protocols/bedrock_invoke_model/chat.rb @@ -105,6 +105,7 @@ def parse_completion_body(data, raw:) content_blocks = data['content'] || [] usage = data['usage'] || {} + cache_creation = usage['cache_creation'] || {} Message.new( role: :assistant, @@ -115,7 +116,11 @@ def parse_completion_body(data, raw:) output_tokens: usage['output_tokens'], cached_tokens: usage['cache_read_input_tokens'], cache_creation_tokens: usage['cache_creation_input_tokens'], + cache_creation_ephemeral_5m_tokens: cache_creation['ephemeral_5m_input_tokens'], + cache_creation_ephemeral_1h_tokens: cache_creation['ephemeral_1h_input_tokens'], finish_reason: data['stop_reason'], + stop_sequence: data['stop_sequence'], + context_management: data['context_management'], model_id: data['model'], raw: raw ) diff --git a/lib/ruby_llm/protocols/bedrock_invoke_model/streaming.rb b/lib/ruby_llm/protocols/bedrock_invoke_model/streaming.rb index 519eeaa2e..54f5a044a 100644 --- a/lib/ruby_llm/protocols/bedrock_invoke_model/streaming.rb +++ b/lib/ruby_llm/protocols/bedrock_invoke_model/streaming.rb @@ -48,9 +48,19 @@ def stream_response(payload, additional_headers = {}, &block) message = accumulator.to_message(response) RubyLLM.logger.debug { "Stream completed: #{message.content}" } + log_context_management(message) message end + def log_context_management(message) + applied_edits = message.context_management && message.context_management['applied_edits'] + return unless applied_edits && !applied_edits.empty? + + RubyLLM.logger.debug do + "Bedrock InvokeModel context_management applied_edits: #{applied_edits.inspect}" + end + end + def event_stream_decoder require 'aws-eventstream' Aws::EventStream::Decoder.new @@ -161,13 +171,20 @@ def build_chunk(event, thinking_state = {}) def build_message_start_chunk(event) message = event['message'] || {} usage = message['usage'] || {} + cache_read = usage['cache_read_input_tokens'] + cache_creation = usage['cache_creation_input_tokens'] input_tok = usage['input_tokens'] + cache_creation_detail = usage['cache_creation'] || {} Chunk.new( role: :assistant, content: nil, model_id: message['model'] || @model&.id, - input_tokens: input_tok ? [input_tok.to_i, 0].max : nil + input_tokens: input_tok ? [input_tok.to_i - cache_read.to_i - cache_creation.to_i, 0].max : nil, + cached_tokens: cache_read, + cache_creation_tokens: cache_creation, + cache_creation_ephemeral_5m_tokens: cache_creation_detail['ephemeral_5m_input_tokens'], + cache_creation_ephemeral_1h_tokens: cache_creation_detail['ephemeral_1h_input_tokens'] ) end @@ -254,13 +271,20 @@ def build_content_block_stop_chunk(event, thinking_state) def build_message_delta_chunk(event) delta = event['delta'] || {} usage = event['usage'] || {} + cache_creation_detail = usage['cache_creation'] || {} Chunk.new( role: :assistant, content: nil, model_id: @model&.id, output_tokens: usage['output_tokens'], - finish_reason: delta['stop_reason'] + cached_tokens: usage['cache_read_input_tokens'], + cache_creation_tokens: usage['cache_creation_input_tokens'], + cache_creation_ephemeral_5m_tokens: cache_creation_detail['ephemeral_5m_input_tokens'], + cache_creation_ephemeral_1h_tokens: cache_creation_detail['ephemeral_1h_input_tokens'], + finish_reason: delta['stop_reason'], + stop_sequence: delta['stop_sequence'], + context_management: delta['context_management'] ) end diff --git a/lib/ruby_llm/stream_accumulator.rb b/lib/ruby_llm/stream_accumulator.rb index 60eff437d..f01dbb868 100644 --- a/lib/ruby_llm/stream_accumulator.rb +++ b/lib/ruby_llm/stream_accumulator.rb @@ -19,8 +19,12 @@ def initialize @output_tokens = nil @cached_tokens = nil @cache_creation_tokens = nil + @cache_creation_ephemeral_5m_tokens = nil + @cache_creation_ephemeral_1h_tokens = nil @thinking_tokens = nil @finish_reason = nil + @stop_sequence = nil + @context_management = nil @inside_think_tag = false @pending_think_tag = +'' @latest_tool_call_id = nil @@ -35,6 +39,8 @@ def add(chunk) accumulate_citations(chunk.citations) append_thinking_from_chunk(chunk) @finish_reason = chunk.finish_reason if chunk.finish_reason + @stop_sequence = chunk.stop_sequence if chunk.stop_sequence + @context_management = chunk.context_management if chunk.context_management count_tokens chunk RubyLLM.logger.debug { inspect } if RubyLLM.config.log_stream_debug end @@ -54,9 +60,13 @@ def to_message(response) output: @output_tokens, cached: @cached_tokens, cache_creation: @cache_creation_tokens, - thinking: @thinking_tokens + thinking: @thinking_tokens, + cache_creation_ephemeral_5m: @cache_creation_ephemeral_5m_tokens, + cache_creation_ephemeral_1h: @cache_creation_ephemeral_1h_tokens ), finish_reason: @finish_reason, + stop_sequence: @stop_sequence, + context_management: @context_management, model_id: model_id, tool_calls: tool_calls_from_stream, raw: response @@ -158,6 +168,14 @@ def count_tokens(chunk) @cached_tokens = chunk.cached_tokens if chunk.cached_tokens @cache_creation_tokens = chunk.cache_creation_tokens if chunk.cache_creation_tokens @thinking_tokens = chunk.thinking_tokens if chunk.thinking_tokens + count_cache_creation_ttl_tokens(chunk) + end + + def count_cache_creation_ttl_tokens(chunk) + ephemeral_5m = chunk.tokens&.cache_creation_ephemeral_5m + ephemeral_1h = chunk.tokens&.cache_creation_ephemeral_1h + @cache_creation_ephemeral_5m_tokens = ephemeral_5m if ephemeral_5m + @cache_creation_ephemeral_1h_tokens = ephemeral_1h if ephemeral_1h end def handle_chunk_content(chunk) diff --git a/lib/ruby_llm/tokens.rb b/lib/ruby_llm/tokens.rb index 2ed45e7f8..903be0355 100644 --- a/lib/ruby_llm/tokens.rb +++ b/lib/ruby_llm/tokens.rb @@ -3,27 +3,39 @@ module RubyLLM # Represents token usage for a response. class Tokens - attr_reader :input, :output, :cached, :cache_creation, :thinking + attr_reader :input, :output, :cached, :cache_creation, :thinking, + :cache_creation_ephemeral_5m, :cache_creation_ephemeral_1h - def initialize(input: nil, output: nil, cached: nil, cache_creation: nil, thinking: nil) + # rubocop:disable Metrics/ParameterLists + def initialize(input: nil, output: nil, cached: nil, cache_creation: nil, thinking: nil, + cache_creation_ephemeral_5m: nil, cache_creation_ephemeral_1h: nil) @input = input @output = output @cached = cached @cache_creation = cache_creation @thinking = thinking + @cache_creation_ephemeral_5m = cache_creation_ephemeral_5m + @cache_creation_ephemeral_1h = cache_creation_ephemeral_1h end + # rubocop:enable Metrics/ParameterLists - def self.build(input: nil, output: nil, cached: nil, cache_creation: nil, thinking: nil) - return nil if [input, output, cached, cache_creation, thinking].all?(&:nil?) + # rubocop:disable Metrics/ParameterLists + def self.build(input: nil, output: nil, cached: nil, cache_creation: nil, thinking: nil, + cache_creation_ephemeral_5m: nil, cache_creation_ephemeral_1h: nil) + return nil if [input, output, cached, cache_creation, thinking, + cache_creation_ephemeral_5m, cache_creation_ephemeral_1h].all?(&:nil?) new( input: input, output: output, cached: cached, cache_creation: cache_creation, - thinking: thinking + thinking: thinking, + cache_creation_ephemeral_5m: cache_creation_ephemeral_5m, + cache_creation_ephemeral_1h: cache_creation_ephemeral_1h ) end + # rubocop:enable Metrics/ParameterLists def to_h { @@ -31,7 +43,9 @@ def to_h output_tokens: output, cached_tokens: cached, cache_creation_tokens: cache_creation, - thinking_tokens: thinking + thinking_tokens: thinking, + cache_creation_ephemeral_5m_tokens: cache_creation_ephemeral_5m, + cache_creation_ephemeral_1h_tokens: cache_creation_ephemeral_1h }.compact end diff --git a/spec/ruby_llm/protocols/bedrock_invoke_model_spec.rb b/spec/ruby_llm/protocols/bedrock_invoke_model_spec.rb index 23a1e5a67..dac4f7920 100644 --- a/spec/ruby_llm/protocols/bedrock_invoke_model_spec.rb +++ b/spec/ruby_llm/protocols/bedrock_invoke_model_spec.rb @@ -350,6 +350,55 @@ def render_payload(messages = [], **opts) expect(chat.parse_completion_body(nil, raw: nil)).to be_nil expect(chat.parse_completion_body({}, raw: nil)).to be_nil end + + it 'reports input_tokens of 0 for a fully-cached request' do + data = basic_response.merge( + 'usage' => { + 'input_tokens' => 100, + 'output_tokens' => 5, + 'cache_read_input_tokens' => 100, + 'cache_creation_input_tokens' => 0 + } + ) + msg = chat.parse_completion_body(data, raw: nil) + expect(msg.input_tokens).to eq(0) + expect(msg.cached_tokens).to eq(100) + end + + it 'extracts cache_creation TTL breakdown when present' do + data = basic_response.merge( + 'usage' => { + 'input_tokens' => 100, + 'output_tokens' => 5, + 'cache_read_input_tokens' => 0, + 'cache_creation_input_tokens' => 30, + 'cache_creation' => { 'ephemeral_5m_input_tokens' => 20, 'ephemeral_1h_input_tokens' => 10 } + } + ) + msg = chat.parse_completion_body(data, raw: nil) + expect(msg.tokens.cache_creation_ephemeral_5m).to eq(20) + expect(msg.tokens.cache_creation_ephemeral_1h).to eq(10) + end + + it 'extracts stop_sequence from the response' do + data = basic_response.merge('stop_reason' => 'stop_sequence', 'stop_sequence' => 'STOP') + msg = chat.parse_completion_body(data, raw: nil) + expect(msg.stop_sequence).to eq('STOP') + end + + it 'surfaces top-level context_management applied_edits on the parsed message' do + data = basic_response.merge( + 'context_management' => { + 'applied_edits' => [ + { 'type' => 'clear_tool_uses_20250919', 'cleared_tool_uses' => 3, 'cleared_input_tokens' => 500 } + ] + } + ) + msg = chat.parse_completion_body(data, raw: nil) + expect(msg.context_management['applied_edits']).to eq( + [{ 'type' => 'clear_tool_uses_20250919', 'cleared_tool_uses' => 3, 'cleared_input_tokens' => 500 }] + ) + end end # --------------------------------------------------------------------------- @@ -441,6 +490,84 @@ def render_payload(messages = [], **opts) expect(chunk.finish_reason).to eq('end_turn') end + it 'extracts cache usage from message_start and nets it out of input_tokens' do + event = { + 'type' => 'message_start', + 'message' => { + 'model' => 'anthropic.claude-sonnet-4-6', + 'usage' => { + 'input_tokens' => 100, + 'cache_read_input_tokens' => 40, + 'cache_creation_input_tokens' => 10 + } + } + } + chunk = streaming.send(:build_chunk, event) + expect(chunk.cached_tokens).to eq(40) + expect(chunk.cache_creation_tokens).to eq(10) + expect(chunk.input_tokens).to eq(50) + end + + it 'reports input_tokens of 0 from message_start for a fully-cached request' do + event = { + 'type' => 'message_start', + 'message' => { + 'model' => 'anthropic.claude-sonnet-4-6', + 'usage' => { + 'input_tokens' => 100, + 'cache_read_input_tokens' => 100, + 'cache_creation_input_tokens' => 0 + } + } + } + chunk = streaming.send(:build_chunk, event) + expect(chunk.input_tokens).to eq(0) + expect(chunk.cached_tokens).to eq(100) + end + + it 'extracts cache_creation TTL breakdown from message_start' do + event = { + 'type' => 'message_start', + 'message' => { + 'model' => 'anthropic.claude-sonnet-4-6', + 'usage' => { + 'input_tokens' => 100, + 'cache_creation_input_tokens' => 30, + 'cache_creation' => { 'ephemeral_5m_input_tokens' => 20, 'ephemeral_1h_input_tokens' => 10 } + } + } + } + chunk = streaming.send(:build_chunk, event) + expect(chunk.tokens.cache_creation_ephemeral_5m).to eq(20) + expect(chunk.tokens.cache_creation_ephemeral_1h).to eq(10) + end + + it 'extracts cumulative cache usage, stop_sequence, and context_management from message_delta' do + event = { + 'type' => 'message_delta', + 'delta' => { + 'stop_reason' => 'end_turn', + 'stop_sequence' => nil, + 'context_management' => { + 'applied_edits' => [ + { 'type' => 'clear_tool_uses_20250919', 'cleared_tool_uses' => 2, 'cleared_input_tokens' => 300 } + ] + } + }, + 'usage' => { + 'output_tokens' => 17, + 'cache_read_input_tokens' => 40, + 'cache_creation_input_tokens' => 10 + } + } + chunk = streaming.send(:build_chunk, event) + expect(chunk.cached_tokens).to eq(40) + expect(chunk.cache_creation_tokens).to eq(10) + expect(chunk.context_management['applied_edits']).to eq( + [{ 'type' => 'clear_tool_uses_20250919', 'cleared_tool_uses' => 2, 'cleared_input_tokens' => 300 }] + ) + end + it 'returns a chunk for message_stop without error' do event = { 'type' => 'message_stop' } expect { streaming.send(:build_chunk, event) }.not_to raise_error @@ -481,6 +608,58 @@ def render_payload(messages = [], **opts) expect(message.content).to eq('Hello world') expect(message.output_tokens).to eq(2) end + + it 'plumbs cache tokens and context_management through StreamAccumulator into the final message' do + accumulator = RubyLLM::StreamAccumulator.new + + events = [ + { 'type' => 'message_start', 'message' => { + 'model' => 'test-model', + 'usage' => { 'input_tokens' => 100, 'cache_read_input_tokens' => 40, 'cache_creation_input_tokens' => 10 } + } }, + { 'type' => 'content_block_delta', 'index' => 0, + 'delta' => { 'type' => 'text_delta', 'text' => 'Hello' } }, + { 'type' => 'message_delta', + 'delta' => { + 'stop_reason' => 'end_turn', + 'context_management' => { + 'applied_edits' => [{ 'type' => 'clear_tool_uses_20250919', 'cleared_input_tokens' => 300 }] + } + }, + 'usage' => { 'output_tokens' => 2, 'cache_read_input_tokens' => 40, 'cache_creation_input_tokens' => 10 } } + ] + + events.each { |e| accumulator.add(streaming.send(:build_chunk, e)) } + message = accumulator.to_message(nil) + + expect(message.input_tokens).to eq(50) + expect(message.cached_tokens).to eq(40) + expect(message.cache_creation_tokens).to eq(10) + expect(message.context_management['applied_edits']).to eq( + [{ 'type' => 'clear_tool_uses_20250919', 'cleared_input_tokens' => 300 }] + ) + end + + it 'reports input_tokens of 0 through the full stream round trip for a fully-cached request' do + accumulator = RubyLLM::StreamAccumulator.new + + events = [ + { 'type' => 'message_start', 'message' => { + 'model' => 'test-model', + 'usage' => { 'input_tokens' => 100, 'cache_read_input_tokens' => 100, 'cache_creation_input_tokens' => 0 } + } }, + { 'type' => 'content_block_delta', 'index' => 0, + 'delta' => { 'type' => 'text_delta', 'text' => 'Hi' } }, + { 'type' => 'message_delta', 'delta' => { 'stop_reason' => 'end_turn' }, + 'usage' => { 'output_tokens' => 2 } } + ] + + events.each { |e| accumulator.add(streaming.send(:build_chunk, e)) } + message = accumulator.to_message(nil) + + expect(message.input_tokens).to eq(0) + expect(message.cached_tokens).to eq(100) + end end # --------------------------------------------------------------------------- From 7bdc1183e9098c66e2f88be18fe62b4c3f7e7c82 Mon Sep 17 00:00:00 2001 From: Supernova Date: Tue, 7 Jul 2026 01:32:40 +0000 Subject: [PATCH 2/3] Defer input_tokens netting to StreamAccumulator for Bedrock InvokeModel Bedrock InvokeModel's message_start chunk reports input_tokens before a later message_delta chunk can update the cache fields. Netting at chunk build time missed cache usage that only appeared in message_delta. StreamAccumulator now defers netting to to_message when net_cache_tokens is true, which BedrockInvokeModel streaming now opts into. Co-Authored-By: Claude Sonnet 4.5 --- .../bedrock_invoke_model/streaming.rb | 4 +- lib/ruby_llm/stream_accumulator.rb | 22 ++++++++-- .../protocols/bedrock_invoke_model_spec.rb | 40 ++++++++++++++++--- 3 files changed, 54 insertions(+), 12 deletions(-) diff --git a/lib/ruby_llm/protocols/bedrock_invoke_model/streaming.rb b/lib/ruby_llm/protocols/bedrock_invoke_model/streaming.rb index 54f5a044a..9a4d1f976 100644 --- a/lib/ruby_llm/protocols/bedrock_invoke_model/streaming.rb +++ b/lib/ruby_llm/protocols/bedrock_invoke_model/streaming.rb @@ -21,7 +21,7 @@ def stream_url end def stream_response(payload, additional_headers = {}, &block) - accumulator = StreamAccumulator.new + accumulator = StreamAccumulator.new(net_cache_tokens: true) decoder = event_stream_decoder thinking_state = {} body = JSON.generate(payload) @@ -180,7 +180,7 @@ def build_message_start_chunk(event) role: :assistant, content: nil, model_id: message['model'] || @model&.id, - input_tokens: input_tok ? [input_tok.to_i - cache_read.to_i - cache_creation.to_i, 0].max : nil, + input_tokens: input_tok&.to_i, cached_tokens: cache_read, cache_creation_tokens: cache_creation, cache_creation_ephemeral_5m_tokens: cache_creation_detail['ephemeral_5m_input_tokens'], diff --git a/lib/ruby_llm/stream_accumulator.rb b/lib/ruby_llm/stream_accumulator.rb index f01dbb868..90a803d99 100644 --- a/lib/ruby_llm/stream_accumulator.rb +++ b/lib/ruby_llm/stream_accumulator.rb @@ -8,14 +8,22 @@ module RubyLLM class StreamAccumulator attr_reader :content, :model_id, :tool_calls - def initialize + # net_cache_tokens: when true, input_tokens is netted against the final accumulated + # cached_tokens/cache_creation_tokens at to_message time instead of being taken as-is + # from whichever chunk last set it. Bedrock InvokeModel's message_start chunk carries + # input_tokens while a later message_delta chunk can still update the cache fields, so + # netting must be deferred until all chunks are in. Other providers already net + # input_tokens against cache fields within the same usage snapshot before building the + # chunk, so deferred netting stays off for them to avoid double-subtracting. + def initialize(net_cache_tokens: false) + @net_cache_tokens = net_cache_tokens @content = +'' @citations = [] @thinking_text = +'' @thinking_signature = nil @thinking_blocks = [] @tool_calls = {} - @input_tokens = nil + @raw_input_tokens = nil @output_tokens = nil @cached_tokens = nil @cache_creation_tokens = nil @@ -56,7 +64,7 @@ def to_message(response) blocks: @thinking_blocks.empty? ? nil : @thinking_blocks ), tokens: Tokens.build( - input: @input_tokens, + input: netted_input_tokens, output: @output_tokens, cached: @cached_tokens, cache_creation: @cache_creation_tokens, @@ -75,6 +83,12 @@ def to_message(response) private + def netted_input_tokens + return @raw_input_tokens unless @net_cache_tokens && @raw_input_tokens + + [@raw_input_tokens - @cached_tokens.to_i - @cache_creation_tokens.to_i, 0].max + end + # Providers like Perplexity repeat the full citation list on every chunk. def accumulate_citations(new_citations) new_citations.each do |citation| @@ -163,7 +177,7 @@ def find_tool_call(stream_key) end def count_tokens(chunk) - @input_tokens = chunk.input_tokens if chunk.input_tokens + @raw_input_tokens = chunk.input_tokens if chunk.input_tokens @output_tokens = chunk.output_tokens if chunk.output_tokens @cached_tokens = chunk.cached_tokens if chunk.cached_tokens @cache_creation_tokens = chunk.cache_creation_tokens if chunk.cache_creation_tokens diff --git a/spec/ruby_llm/protocols/bedrock_invoke_model_spec.rb b/spec/ruby_llm/protocols/bedrock_invoke_model_spec.rb index dac4f7920..d5264fe62 100644 --- a/spec/ruby_llm/protocols/bedrock_invoke_model_spec.rb +++ b/spec/ruby_llm/protocols/bedrock_invoke_model_spec.rb @@ -490,7 +490,7 @@ def render_payload(messages = [], **opts) expect(chunk.finish_reason).to eq('end_turn') end - it 'extracts cache usage from message_start and nets it out of input_tokens' do + it 'extracts cache usage from message_start and passes input_tokens through un-netted' do event = { 'type' => 'message_start', 'message' => { @@ -505,10 +505,12 @@ def render_payload(messages = [], **opts) chunk = streaming.send(:build_chunk, event) expect(chunk.cached_tokens).to eq(40) expect(chunk.cache_creation_tokens).to eq(10) - expect(chunk.input_tokens).to eq(50) + # Netting against cache tokens is deferred to StreamAccumulator#to_message, since a + # later message_delta chunk can still update the cache fields for this same message. + expect(chunk.input_tokens).to eq(100) end - it 'reports input_tokens of 0 from message_start for a fully-cached request' do + it 'passes the raw input_tokens through from message_start for a fully-cached request' do event = { 'type' => 'message_start', 'message' => { @@ -521,7 +523,7 @@ def render_payload(messages = [], **opts) } } chunk = streaming.send(:build_chunk, event) - expect(chunk.input_tokens).to eq(0) + expect(chunk.input_tokens).to eq(100) expect(chunk.cached_tokens).to eq(100) end @@ -610,7 +612,7 @@ def render_payload(messages = [], **opts) end it 'plumbs cache tokens and context_management through StreamAccumulator into the final message' do - accumulator = RubyLLM::StreamAccumulator.new + accumulator = RubyLLM::StreamAccumulator.new(net_cache_tokens: true) events = [ { 'type' => 'message_start', 'message' => { @@ -641,7 +643,7 @@ def render_payload(messages = [], **opts) end it 'reports input_tokens of 0 through the full stream round trip for a fully-cached request' do - accumulator = RubyLLM::StreamAccumulator.new + accumulator = RubyLLM::StreamAccumulator.new(net_cache_tokens: true) events = [ { 'type' => 'message_start', 'message' => { @@ -660,6 +662,32 @@ def render_payload(messages = [], **opts) expect(message.input_tokens).to eq(0) expect(message.cached_tokens).to eq(100) end + + it 'nets input_tokens against cache fields that are only updated by a later message_delta' do + # Regression test: message_start's usage snapshot has no cache_read_input_tokens yet + # (some Bedrock responses only report cache usage on message_delta), so a chunk-time + # netting of message_start's input_tokens would silently miss this cache usage. Since + # netting is now deferred to to_message, it correctly nets against the final values. + accumulator = RubyLLM::StreamAccumulator.new(net_cache_tokens: true) + + events = [ + { 'type' => 'message_start', 'message' => { + 'model' => 'test-model', + 'usage' => { 'input_tokens' => 100 } + } }, + { 'type' => 'content_block_delta', 'index' => 0, + 'delta' => { 'type' => 'text_delta', 'text' => 'Hi' } }, + { 'type' => 'message_delta', 'delta' => { 'stop_reason' => 'end_turn' }, + 'usage' => { 'output_tokens' => 2, 'cache_read_input_tokens' => 40, 'cache_creation_input_tokens' => 10 } } + ] + + events.each { |e| accumulator.add(streaming.send(:build_chunk, e)) } + message = accumulator.to_message(nil) + + expect(message.input_tokens).to eq(50) + expect(message.cached_tokens).to eq(40) + expect(message.cache_creation_tokens).to eq(10) + end end # --------------------------------------------------------------------------- From 6213d7592978df4295b605d77107dfc67140979a Mon Sep 17 00:00:00 2001 From: Supernova Date: Wed, 8 Jul 2026 22:12:42 +0000 Subject: [PATCH 3/3] Add test coverage for streaming usage token normalization Adds tests to verify: - Distinct usage persisted on each tool-call round trip in chat loop - Fully-cached requests report input_tokens of 0 - Cache token breakdown excludes ephemeral TTL fields not returned by Converse - Proper token normalization in streaming chunk builder These tests ensure the streaming usage token capture and normalization behavior introduced in recent commits works correctly across both completion and streaming paths. Co-Authored-By: Claude Sonnet 4.5 Co-authored-by: Sam Boland --- spec/ruby_llm/chat_loop_spec.rb | 19 ++++++ spec/ruby_llm/protocols/converse/chat_spec.rb | 45 ++++++++++++++ .../protocols/converse/streaming_spec.rb | 58 +++++++++++++++++++ 3 files changed, 122 insertions(+) diff --git a/spec/ruby_llm/chat_loop_spec.rb b/spec/ruby_llm/chat_loop_spec.rb index 349d8ff32..038b2c549 100644 --- a/spec/ruby_llm/chat_loop_spec.rb +++ b/spec/ruby_llm/chat_loop_spec.rb @@ -141,6 +141,25 @@ def tool_call_message(name: 'echo') expect(chat.complete).to be(answer_message) expect(chat.provider).not_to have_received(:complete) end + + it 'persists distinct, non-zero usage on every intermediate tool-call round trip' do + # Each round trip (including ones that returned tool_calls) must carry its own usage, + # not the final round trip's, so a per-turn cost breakdown stays accurate. + first_round = RubyLLM::Message.new( + role: :assistant, content: '', tool_calls: tool_call_message.tool_calls, + input_tokens: 100, output_tokens: 10 + ) + final_round = RubyLLM::Message.new(role: :assistant, content: 'hello', input_tokens: 150, output_tokens: 5) + + allow(chat.provider).to receive(:complete).and_return(first_round, final_round) + chat.ask_later('Echo "hello" back to me.') + + chat.complete + + assistant_messages = chat.messages.select { |m| m.role == :assistant } + expect(assistant_messages.map(&:input_tokens)).to eq([100, 150]) + expect(assistant_messages.map(&:output_tokens)).to eq([10, 5]) + end end describe '#add_completion' do diff --git a/spec/ruby_llm/protocols/converse/chat_spec.rb b/spec/ruby_llm/protocols/converse/chat_spec.rb index 518278ac6..00a4bd98c 100644 --- a/spec/ruby_llm/protocols/converse/chat_spec.rb +++ b/spec/ruby_llm/protocols/converse/chat_spec.rb @@ -29,6 +29,51 @@ expect(message.cache_creation_tokens).to eq(10) end + it 'reports input_tokens of 0 for a fully-cached request' do + response_body = { + 'modelId' => 'anthropic.claude-sonnet-4-5-20250929-v1:0', + 'output' => { + 'message' => { + 'content' => [{ 'text' => 'Hi!' }] + } + }, + 'usage' => { + 'inputTokens' => 100, + 'outputTokens' => 5, + 'cacheReadInputTokens' => 100, + 'cacheWriteInputTokens' => 0 + } + } + + response = instance_double(Faraday::Response, body: response_body) + message = described_class.parse_completion_response(response) + + expect(message.input_tokens).to eq(0) + expect(message.cached_tokens).to eq(100) + end + + it 'does not fabricate an ephemeral cache-creation TTL breakdown, which Converse does not return' do + response_body = { + 'output' => { + 'message' => { + 'content' => [{ 'text' => 'Hi!' }] + } + }, + 'usage' => { + 'inputTokens' => 100, + 'outputTokens' => 5, + 'cacheReadInputTokens' => 40, + 'cacheWriteInputTokens' => 10 + } + } + + response = instance_double(Faraday::Response, body: response_body) + message = described_class.parse_completion_response(response) + + expect(message.tokens.cache_creation_ephemeral_5m).to be_nil + expect(message.tokens.cache_creation_ephemeral_1h).to be_nil + end + it 'preserves raw stopReason as finish_reason' do response_body = { 'modelId' => 'amazon.nova-lite-v1:0', diff --git a/spec/ruby_llm/protocols/converse/streaming_spec.rb b/spec/ruby_llm/protocols/converse/streaming_spec.rb index fe1cd935e..510e51c1f 100644 --- a/spec/ruby_llm/protocols/converse/streaming_spec.rb +++ b/spec/ruby_llm/protocols/converse/streaming_spec.rb @@ -70,6 +70,64 @@ expect(chunk.thinking_tokens).to eq(7) end + it 'normalizes cache read and write tokens out of input tokens from metadata usage' do + event = { + 'metadata' => { + 'usage' => { + 'inputTokens' => 100, + 'outputTokens' => 5, + 'cacheReadInputTokens' => 40, + 'cacheWriteInputTokens' => 10 + } + } + } + + chunk = streaming.send(:build_chunk, event) + + expect(chunk.input_tokens).to eq(50) + expect(chunk.output_tokens).to eq(5) + expect(chunk.cached_tokens).to eq(40) + expect(chunk.cache_creation_tokens).to eq(10) + end + + it 'reports input_tokens of 0 through the accumulator for a fully-cached streamed request' do + accumulator = RubyLLM::StreamAccumulator.new + usage_event = { + 'metadata' => { + 'usage' => { + 'inputTokens' => 100, + 'outputTokens' => 5, + 'cacheReadInputTokens' => 100, + 'cacheWriteInputTokens' => 0 + } + } + } + + accumulator.add(streaming.send(:build_chunk, usage_event)) + message = accumulator.to_message(nil) + + expect(message.input_tokens).to eq(0) + expect(message.cached_tokens).to eq(100) + end + + it 'does not fabricate an ephemeral cache-creation TTL breakdown, which Converse does not return' do + event = { + 'metadata' => { + 'usage' => { + 'inputTokens' => 100, + 'outputTokens' => 5, + 'cacheReadInputTokens' => 40, + 'cacheWriteInputTokens' => 10 + } + } + } + + chunk = streaming.send(:build_chunk, event) + + expect(chunk.tokens&.cache_creation_ephemeral_5m).to be_nil + expect(chunk.tokens&.cache_creation_ephemeral_1h).to be_nil + end + it 'accumulates Bedrock Converse Stream thinking deltas into the final message' do accumulator = RubyLLM::StreamAccumulator.new text_event = {