fix: record token usage for streaming generations

lib/active_agent/providers/_base_provider.rb

@@ -54,8 +54,11 @@ class ProvidersError < StandardError; end
       attr_internal :options, :context, :trace_id,   # Setup
                     :request, :message_stack,        # Runtime
                     :stream_broadcaster, :streaming, # Callback (Streams)
+                    :stream_completion_pending,      # Callback (Streams)
+                    :stream_completion_result,       # Callback (Streams)
                     :tools_function,                 # Callback (Tools)
                     :usage_stack,                    # Usage Tracking
+                    :stream_usage_index,             # Usage Tracking (Streams)
                     :max_tool_turns, :tool_turns     # Tool-loop safety
 
       # Upper bound on tool-calling round-trips within one generation. A
@@ -118,6 +121,9 @@ def initialize(kwargs = {})
         self.context            = kwargs
         self.message_stack      = []
         self.usage_stack        = []
+        self.stream_completion_pending = false
+        self.stream_completion_result  = nil
+        self.stream_usage_index        = nil
       end
 
       # Generates prompt preview without executing the API call.
@@ -179,6 +185,9 @@ def instrument(name, payload = {}, &block)
       #
       # @return [ActiveAgent::Providers::Common::PromptResponse]
       def resolve_prompt
+        # Each turn streams its own usage; see record_stream_usage.
+        self.stream_usage_index = nil
+
         api_parameters = api_request_build(prepare_prompt_request, prompt_request_type)
         api_response = instrument("prompt.provider.active_agent") do |payload|
           raw_response = with_exception_handling { api_prompt_execute(api_parameters) }
@@ -193,6 +202,12 @@ def resolve_prompt
           raw_response
         end
 
+        # A stream that deferred its completion has already run
+        # process_prompt_finished (including any tool-call recursion) from
+        # stream_finished!, and holds the response it produced. Calling it a
+        # second time here would re-run that work.
+        return stream_completion_result if stream_completion_result
+
         process_prompt_finished(api_response)
       end
 
@@ -231,7 +246,11 @@ def prepare_prompt_request
       # @return [Hash] API request parameters
       def api_request_build(request, request_type)
         parameters          = request_type.serialize(request)
-        parameters[:stream] = process_stream if request.try(:stream)
+
+        if request.try(:stream)
+          parameters[:stream] = process_stream
+          parameters.deep_merge!(api_stream_usage_parameters)
+        end
 
         if options.extra_headers.present?
           parameters[:request_options] = { extra_headers: options.extra_headers }.deep_merge(parameters[:request_options] || {})
@@ -240,6 +259,20 @@ def api_request_build(request, request_type)
         parameters
       end
 
+      # Extra request parameters needed to make the provider report token usage
+      # while streaming.
+      #
+      # A streaming request returns its usage on a final chunk rather than in a
+      # response body, and several providers only send that chunk when the
+      # request asks for it. Providers that need such a flag override this;
+      # the default asks for nothing, so a provider that reports usage
+      # unconditionally — or not at all — is unaffected.
+      #
+      # @return [Hash]
+      def api_stream_usage_parameters
+        {}
+      end
+
       # @return [Proc] for each response chunk
       def process_stream
         proc do |api_response_chunk|
@@ -258,10 +291,27 @@ def api_prompt_execute(parameters)
           api_prompt_executer.create(**parameters)
         else
           api_prompt_executer.stream(**parameters.except(:stream)).each(&parameters[:stream])
+          stream_finished!
           nil
         end
       end
 
+      # Runs the deferred end-of-generation work once a stream has drained.
+      #
+      # A chunk handler that completed the message marks the generation
+      # pending rather than finishing inline, because providers emit their
+      # usage chunk after the content is done — finishing inline would build
+      # the response before usage is recorded. Handlers that finish inline
+      # leave the flag unset and this is a no-op.
+      #
+      # @return [Object, nil] result of process_prompt_finished, if deferred
+      def stream_finished!
+        return unless stream_completion_pending
+
+        self.stream_completion_pending = false
+        self.stream_completion_result  = process_prompt_finished
+      end
+
       # Returns provider-specific API executer for prompt requests.
       #
       # Since all currently implemented providers use stainless gems, subclasses
@@ -305,6 +355,52 @@ def process_stream_chunk(api_response_chunk)
         fail NotImplementedError, "Subclass expected to implement"
       end
 
+      # Records token usage carried on a streaming chunk.
+      #
+      # The non-streaming path pushes onto usage_stack in resolve_prompt, from
+      # the response body. A streaming request has no response body — the
+      # provider streams chunks and api_prompt_execute returns nil — so usage
+      # would otherwise be lost, and every streamed generation reports zero
+      # tokens and zero cost. Chunk handlers call this when they see usage so
+      # the two paths converge on the same usage_stack.
+      #
+      # Ignores blank and all-zero payloads: providers send `usage: null` on
+      # ordinary content chunks, and pushing those would add empty entries to
+      # a stack that is summed with reduce(:+).
+      #
+      # A streamed usage payload is a running total for the turn, not a delta,
+      # so the turn keeps a single entry that later payloads replace. Chat
+      # Completions sends exactly one, on a final chunk, but Gemini's
+      # OpenAI-compatible endpoint repeats a cumulative usage on every chunk —
+      # summing those would report a turn's tokens many times over. Tool
+      # calling still accumulates across turns: each turn enters resolve_prompt
+      # and starts a fresh entry.
+      #
+      # @param raw_usage [Hash, Object, nil] provider-shaped usage payload
+      # @return [void]
+      def record_stream_usage(raw_usage)
+        return if raw_usage.blank?
+
+        # from_provider_usage only reads hashes, and the stainless gems hand
+        # back model objects (e.g. OpenAI::Models::CompletionUsage), so an
+        # unconverted object is silently dropped.
+        raw_usage = raw_usage.deep_to_h if raw_usage.respond_to?(:deep_to_h)
+        raw_usage = raw_usage.to_h if !raw_usage.is_a?(Hash) && raw_usage.respond_to?(:to_h)
+
+        usage = Common::Usage.from_provider_usage(raw_usage)
+        return if usage.blank?
+        return if usage.total_tokens.to_i.zero? &&
+                  usage.input_tokens.to_i.zero? &&
+                  usage.output_tokens.to_i.zero?
+
+        if stream_usage_index
+          usage_stack[stream_usage_index] = usage
+        else
+          self.stream_usage_index = usage_stack.length
+          usage_stack.push(usage)
+        end
+      end
+
       # Broadcasts stream open event.
       #
       # Fires once per request cycle, even during multi-turn tool calling.

lib/active_agent/providers/open_ai/chat_provider.rb

@@ -26,6 +26,15 @@ def self.prompt_request_type
 
         protected
 
+        # Chat Completions reports a streamed request's token usage on a final
+        # chunk, and only when the request opts in.
+        #
+        # @return [Hash]
+        # @see Base#api_stream_usage_parameters
+        def api_stream_usage_parameters
+          { stream_options: { include_usage: true } }
+        end
+
         # @return [OpenAI::Client::Completions] the API client for chat completions
         # @see Base#api_prompt_executer
         def api_prompt_executer
@@ -100,7 +109,12 @@ def process_stream_chunk(api_response_event)
           # Called Multiple Times: [Chunk<T>, T]<Content, ToolsCall>
           case api_response_event.type
           when :chunk
+            # The usage chunk carries usage and an empty choices array, so it
+            # must be read before choices.first is dereferenced below.
+            record_stream_usage(api_response_event.chunk.try(:usage))
+
             api_message = api_response_event.chunk.choices.first
+            return if api_message.nil?
 
             # If we have a delta, we need to update a message in the stack
             message = find_or_create_message(api_message.index)
@@ -118,8 +132,13 @@ def process_stream_chunk(api_response_event)
             # Returns the full content when complete
             # => {type: :"content.done", content: "Hi there! How can I help you today?", parsed: nil}
 
-            # Once we are finished, close out and run tooling callbacks (Recursive)
-            process_prompt_finished
+            # Close out and run tooling callbacks (Recursive) — but not here.
+            # The usage chunk is emitted *after* content.done, and finishing
+            # at this point builds the response from a usage_stack that has
+            # not received it yet, so a streamed generation reports zero
+            # tokens. Deferred to stream_finished!, which the base provider
+            # calls once the stream has drained.
+            self.stream_completion_pending = true
           when :"tool_calls.function.arguments.delta"
             # => {type: :"tool_calls.function.arguments.delta", name: "get_current_weather", index: 0, arguments: "", parsed: nil, arguments_delta: ""}
           when :"tool_calls.function.arguments.done"

test/fixtures/vcr_cassettes/integration/open_ai/chat/native_format_test/test_agent_streaming.yml

@@ -5,7 +5,7 @@ http_interactions:
     uri: https://api.openai.com/v1/chat/completions
     body:
       encoding: UTF-8
-      string: '{"messages":[{"role":"developer","content":"You are a helpful assistant."},{"role":"user","content":"Hello!"}],"model":"gpt-5","stream":true}'
+      string: '{"messages":[{"role":"developer","content":"You are a helpful assistant."},{"role":"user","content":"Hello!"}],"model":"gpt-5","stream_options":{"include_usage":true},"stream":true}'
     headers:
       Accept-Encoding:
       - gzip;q=1.0,deflate;q=0.6,identity;q=0.3

test/fixtures/vcr_cassettes/integration/open_router/native_format_test/test_agent_functions_with_streaming.yml

@@ -8,7 +8,7 @@ http_interactions:
       string: '{"messages":[{"role":"user","content":"What is the weather like in
         Boston today?"}],"model":"google/gemini-2.0-flash-001","tool_choice":"auto","tools":[{"type":"function","function":{"name":"get_current_weather","description":"Get
         the current weather in a given location","parameters":{"type":"object","properties":{"location":{"type":"string","description":"The
-        city and state, e.g. San Francisco, CA"},"unit":{"type":"string","enum":["celsius","fahrenheit"]}},"required":["location"]}}}],"stream":true}'
+        city and state, e.g. San Francisco, CA"},"unit":{"type":"string","enum":["celsius","fahrenheit"]}},"required":["location"]}}}],"stream_options":{"include_usage":true},"stream":true}'
     headers:
       Accept-Encoding:
       - gzip;q=1.0,deflate;q=0.6,identity;q=0.3
@@ -96,7 +96,7 @@ http_interactions:
         MA\"}","name":"get_current_weather"},"type":"function"}]},{"role":"tool","content":"{\"location\":\"Boston,
         MA\",\"unit\":\"fahrenheit\",\"temperature\":\"22\"}","tool_call_id":"tool_0_get_current_weather_PWzyW2sI3uKsGniIxpGn"}],"model":"google/gemini-2.0-flash-001","tool_choice":"auto","tools":[{"type":"function","function":{"name":"get_current_weather","description":"Get
         the current weather in a given location","parameters":{"type":"object","properties":{"location":{"type":"string","description":"The
-        city and state, e.g. San Francisco, CA"},"unit":{"type":"string","enum":["celsius","fahrenheit"]}},"required":["location"]}}}],"stream":true}'
+        city and state, e.g. San Francisco, CA"},"unit":{"type":"string","enum":["celsius","fahrenheit"]}},"required":["location"]}}}],"stream_options":{"include_usage":true},"stream":true}'
     headers:
       Accept-Encoding:
       - gzip;q=1.0,deflate;q=0.6,identity;q=0.3

test/fixtures/vcr_cassettes/integration/open_router/native_format_test/test_agent_streaming.yml

@@ -5,7 +5,7 @@ http_interactions:
     uri: https://openrouter.ai/api/v1/chat/completions
     body:
       encoding: UTF-8
-      string: '{"messages":[{"role":"developer","content":"You are a helpful assistant."},{"role":"user","content":"Hello!"}],"model":"openrouter/auto","stream":true}'
+      string: '{"messages":[{"role":"developer","content":"You are a helpful assistant."},{"role":"user","content":"Hello!"}],"model":"openrouter/auto","stream_options":{"include_usage":true},"stream":true}'
     headers:
       Accept-Encoding:
       - gzip;q=1.0,deflate;q=0.6,identity;q=0.3

test/integration/ollama/chat/native_format_test.rb

@@ -79,6 +79,9 @@ def image_input
                 content: "Hello!"
               }
             ],
+            # Chat Completions only reports a streamed request's token usage
+            # when the request opts in.
+            stream_options: { include_usage: true },
             stream: true
           }
           def streaming
@@ -215,7 +218,7 @@ def web_search
           ###############################################################
           # Extended Example
           ###############################################################
-          FUNCTIONS_WITH_STREAMING = FUNCTIONS.merge(stream: true)
+          FUNCTIONS_WITH_STREAMING = FUNCTIONS.merge(stream_options: { include_usage: true }, stream: true)
           def functions_with_streaming
             prompt(
               model: "qwen3",

test/integration/open_ai/chat/native_format_test.rb

@@ -100,6 +100,9 @@ def image_input
                 content: "Hello!"
               }
             ],
+            # Chat Completions only reports a streamed request's token usage
+            # when the request opts in.
+            stream_options: { include_usage: true },
             stream: true
           }
           def streaming
@@ -238,7 +241,7 @@ def web_search
           ###############################################################
           # Extended Example
           ###############################################################
-          FUNCTIONS_WITH_STREAMING = FUNCTIONS.merge(stream: true)
+          FUNCTIONS_WITH_STREAMING = FUNCTIONS.merge(stream_options: { include_usage: true }, stream: true)
           def functions_with_streaming
             prompt(
               model: "gpt-4.1",

test/integration/open_router/native_format_test.rb

@@ -96,6 +96,9 @@ def image_input
               content: "Hello!"
             }
           ],
+          # Chat Completions only reports a streamed request's token usage
+          # when the request opts in.
+          stream_options: { include_usage: true },
           stream: true
         }
         def streaming
@@ -310,7 +313,7 @@ def structured_output
           )
         end
 
-        FUNCTIONS_WITH_STREAMING = FUNCTIONS.merge(stream: true)
+        FUNCTIONS_WITH_STREAMING = FUNCTIONS.merge(stream_options: { include_usage: true }, stream: true)
         def functions_with_streaming
           prompt(
             model: "google/gemini-2.0-flash-001",