fix: include ruby_llm cache tokens in total_tokens

lib/active_agent/providers/ruby_llm_provider.rb

@@ -185,13 +185,17 @@ def process_stream_chunk(chunk)
 
         # Counts arrive in separate chunks (Anthropic sends input at the start
         # and output at the end). They are cumulative for this turn, so merge
-        # reported fields rather than adding each chunk's counts.
+        # reported fields rather than adding each chunk's counts. The total is
+        # stamped after the merge: a later chunk often carries only output, and
+        # its own total would drop the cache already counted.
         if (usage = ruby_llm_usage(chunk))
-          (@stream_response[:usage] ||= {}).merge!(usage) do |field, previous, current|
+          merged = (@stream_response[:usage] ||= {})
+          merged.merge!(usage) do |field, previous, current|
             # Chat Completions synthesizes zero when cache writes are omitted.
             # That cannot erase writes already counted in this turn.
             field == :cache_creation_tokens && current == 0 ? previous : current
           end
+          merged[:total_tokens] = ruby_llm_total_tokens(merged)
         end
         if chunk.respond_to?(:finish_reason) && chunk.finish_reason
           @stream_response[:stop_reason] = ruby_llm_stop_reason(chunk, tool_calls: message[:tool_calls])
@@ -517,7 +521,10 @@ def normalize_ruby_llm_response(response, model_id)
         end
 
         usage = ruby_llm_usage(response)
-        hash[:usage] = usage if usage
+        if usage
+          usage[:total_tokens] = ruby_llm_total_tokens(usage)
+          hash[:usage] = usage
+        end
 
         hash[:model] = model_id if model_id
 
@@ -592,7 +599,8 @@ def ruby_llm_stop_reason(response, tool_calls: response.tool_calls)
       # none: the message has no Tokens then in ruby_llm 1.x. RubyLLM counts
       # cached tokens apart from the input, so input_tokens leaves them out.
       # A count the provider left out is left out here too, and Usage reads it
-      # as zero.
+      # as zero. Total is not set here: a stream merges these hashes field by
+      # field and stamps the total once the merge is done.
       #
       # @param response [RubyLLM::Message]
       # @return [Hash, nil] Common::Usage attributes
@@ -619,6 +627,18 @@ def ruby_llm_usage(response)
           }.compact
         end
       end
+
+      # Every token the response processed. ruby_llm's input excludes cache
+      # reads and writes: OpenAI and Gemini subtract them from the prompt
+      # count, and Anthropic already reports them apart from input_tokens.
+      # Output already includes thinking when that is billed as output, so
+      # reasoning_tokens is not added again.
+      #
+      # @param usage [Hash] fields from {#ruby_llm_usage}, after a stream merge
+      # @return [Integer]
+      def ruby_llm_total_tokens(usage)
+        %i[input_tokens output_tokens cached_tokens cache_creation_tokens].sum { |key| usage[key] || 0 }
+      end
     end
   end
 end

test/providers/ruby_llm/response_test.rb

@@ -48,6 +48,8 @@ class RubyLLMResponseTest < ActiveSupport::TestCase
     assert_equal 30, usage.cached_tokens
     assert_equal 4, usage.cache_creation_tokens
     assert_equal 2, usage.reasoning_tokens
+    # input excludes cache, and thinking is already inside output when billed.
+    assert_equal 49, usage.total_tokens
   end
 
   # ruby_llm 1.16 has no token counts at all on such a message (its tokens is
@@ -83,6 +85,8 @@ class RubyLLMResponseTest < ActiveSupport::TestCase
       assert_equal 3, usage.public_send(field)
       assert_equal 0, usage.input_tokens
       assert_equal 0, usage.output_tokens
+      # A lone thinking count is a breakdown of output, not an extra bucket.
+      assert_equal(count == :thinking ? 0 : 3, usage.total_tokens)
     end
   end
 
@@ -142,6 +146,7 @@ class RubyLLMResponseTest < ActiveSupport::TestCase
     assert_equal 7, usage.input_tokens
     assert_equal 4, usage.output_tokens
     assert_equal 5, usage.cached_tokens
+    assert_equal 16, usage.total_tokens
   end
 
   test "reports the tokens Anthropic counted, cache reads and writes included" do
@@ -155,6 +160,7 @@ class RubyLLMResponseTest < ActiveSupport::TestCase
     assert_equal 4, usage.output_tokens
     assert_equal 5, usage.cached_tokens
     assert_equal 2, usage.cache_creation_tokens
+    assert_equal 23, usage.total_tokens
   end
 
   # --- Stop reason ---

test/providers/ruby_llm/streaming_test.rb

@@ -239,6 +239,7 @@ def message_delta(stop_reason)
       assert_not_nil response.usage
       assert_equal(stream.name == :chat_completions ? 7 : 12, response.usage.input_tokens)
       assert_equal 4, response.usage.output_tokens
+      assert_equal 16, response.usage.total_tokens
       assert_equal "end_turn", response.finish_reason
       assert_equal "Done.", response.messages.last.content
       assert_equal 1, response.messages.count { |message| message.role == "assistant" }
@@ -414,6 +415,7 @@ def message_delta(stop_reason)
       assert_equal 5, response.usage.output_tokens
       assert_equal 5, response.usage.cached_tokens
       assert_equal 1, response.usage.reasoning_tokens
+      assert_equal 27, response.usage.total_tokens
       assert_equal "end_turn", response.finish_reason
     end
   end
@@ -435,6 +437,7 @@ def message_delta(stop_reason)
     assert_equal 7, usage.input_tokens
     assert_equal 4, usage.output_tokens
     assert_equal 5, usage.cache_creation_tokens
+    assert_equal 16, usage.total_tokens
   end
 
   test "updates positive cumulative cache writes and starts a fresh count on each tool turn" do
@@ -453,6 +456,7 @@ def message_delta(stop_reason)
       assert_equal 17, response.usage.input_tokens
       assert_equal 5, response.usage.output_tokens
       assert_equal 9, response.usage.cache_creation_tokens
+      assert_equal 31, response.usage.total_tokens
       assert_equal [ 9, 0 ], response.usages.map(&:cache_creation_tokens)
     end
   end