DUX-7485: Capture streaming usage tokens and context management for Bedrock InvokeModel

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
 

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
           )

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)
@@ -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&.to_i,
+            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
 

lib/ruby_llm/stream_accumulator.rb

@@ -8,19 +8,31 @@ 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
+      @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 +47,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
@@ -50,13 +64,17 @@ 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,
-          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
@@ -65,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|
@@ -153,11 +177,19 @@ 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
       @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)

lib/ruby_llm/tokens.rb

@@ -3,35 +3,49 @@
 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
       {
         input_tokens: input,
         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
 

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