From b68d34be16621f7817ddd3400fc332c8bb9242f6 Mon Sep 17 00:00:00 2001 From: Justin Bowen Date: Fri, 14 Aug 2026 16:35:19 -0700 Subject: [PATCH 1/2] fix: record token usage for streaming generations MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A streamed generation reported zero tokens, zero cost, and no context pressure estimate. Three things had to be true at once for the usage to survive, and none of them were: 1. `api_prompt_execute` consumes the stream and returns nil, so there is no response body for `resolve_prompt` to read usage from — the non-streaming path's only source. 2. Chat Completions emits usage on a final chunk, and only when the request sets `stream_options: {include_usage: true}`. We never asked, so the chunk never came. `api_stream_usage_parameters` is a provider hook rather than a hardcoded flag: providers that report unconditionally, or not at all, are unaffected by the default empty hash. 3. That chunk arrives *after* `content.done`, which is where the OpenAI chunk handler called `process_prompt_finished` — building the response from a usage_stack the usage had not reached yet. Completion is now deferred to `stream_finished!`, which the base provider calls once the stream drains, and its result is returned from `resolve_prompt` so the generation (and any tool-call recursion inside it) still runs exactly once. `record_stream_usage` also converts the payload before handing it to `Usage.from_provider_usage`, which early-returns on anything that is not a Hash: the stainless gems hand back model objects like OpenAI::Models::CompletionUsage, and an unconverted object was silently dropped. All-zero and nil payloads are ignored so the `usage: null` on ordinary content chunks does not add empty entries to a stack that is summed with `reduce(:+)`. The usage chunk carries an empty `choices` array, so the handler reads usage before dereferencing `choices.first`. Verified against OpenRouter (anthropic/claude-sonnet-4.5): streaming went from `usage=nil` to input=14 output=4 total=18, matching the same prompt issued non-streaming. A real 30-page document-enrichment batch now reports 475/72 rather than 0/0. Co-Authored-By: Claude Opus 5 --- lib/active_agent/providers/_base_provider.rb | 80 +++++++++++++++- .../providers/open_ai/chat_provider.rb | 23 ++++- .../base_provider_stream_usage_test.rb | 96 +++++++++++++++++++ 3 files changed, 196 insertions(+), 3 deletions(-) create mode 100644 test/providers/base_provider_stream_usage_test.rb diff --git a/lib/active_agent/providers/_base_provider.rb b/lib/active_agent/providers/_base_provider.rb index 165fa8f7..5f4f6d27 100644 --- a/lib/active_agent/providers/_base_provider.rb +++ b/lib/active_agent/providers/_base_provider.rb @@ -54,6 +54,8 @@ 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 :max_tool_turns, :tool_turns # Tool-loop safety @@ -118,6 +120,8 @@ def initialize(kwargs = {}) self.context = kwargs self.message_stack = [] self.usage_stack = [] + self.stream_completion_pending = false + self.stream_completion_result = nil end # Generates prompt preview without executing the API call. @@ -193,6 +197,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 +241,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 +254,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 +286,27 @@ def api_prompt_execute(parameters) api_prompt_executer.create(**parameters) else api_prompt_executer.stream(**parameters.except(:stream)).each(¶meters[: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 +350,39 @@ 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(:+). + # + # @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? + + usage_stack.push(usage) + end + # Broadcasts stream open event. # # Fires once per request cycle, even during multi-turn tool calling. diff --git a/lib/active_agent/providers/open_ai/chat_provider.rb b/lib/active_agent/providers/open_ai/chat_provider.rb index b1485810..433c3152 100644 --- a/lib/active_agent/providers/open_ai/chat_provider.rb +++ b/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] 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" diff --git a/test/providers/base_provider_stream_usage_test.rb b/test/providers/base_provider_stream_usage_test.rb new file mode 100644 index 00000000..63e75aaa --- /dev/null +++ b/test/providers/base_provider_stream_usage_test.rb @@ -0,0 +1,96 @@ +# frozen_string_literal: true + +require "test_helper" +require_relative "../../lib/active_agent/providers/mock_provider" + +# Streamed token usage: a streaming request has no response body to read +# usage from — the provider emits it on a final chunk, after the content is +# done. Without recording that chunk and deferring the end-of-generation work +# until the stream drains, every streamed generation reports zero tokens (and +# so zero cost, and no context-pressure estimate downstream). +class BaseProviderStreamUsageTest < ActiveSupport::TestCase + # Stands in for a provider whose stream ends with a usage-bearing chunk. + class StreamingMockProvider < ActiveAgent::Providers::MockProvider + def self.name = "ActiveAgent::Providers::MockProvider" + + class_attribute :usage_payload, default: nil + + def api_stream_usage_parameters + { stream_options: { include_usage: true } } + end + + # Mimics the real shape: content arrives first and marks the generation + # complete, then a final chunk carries the usage. + def api_prompt_execute(parameters) + return super unless parameters[:stream] + + message_stack.push({ role: "assistant", content: "ok" }) + self.stream_completion_pending = true + record_stream_usage(self.class.usage_payload) + stream_finished! + nil + end + end + + def build(usage:) + StreamingMockProvider.usage_payload = usage + StreamingMockProvider.new(messages: [ { role: "user", content: "go" } ], stream: true) + end + + test "records usage delivered on a final streaming chunk" do + response = build(usage: { "prompt_tokens" => 14, "completion_tokens" => 4, "total_tokens" => 18 }).prompt + + assert_equal 14, response.usage.input_tokens + assert_equal 4, response.usage.output_tokens + end + + test "converts a provider model object rather than dropping it" do + # The stainless gems hand back model objects, not hashes; from_provider_usage + # only reads hashes, so an unconverted object is silently ignored. + object_usage = Struct.new(:to_h).new( + { "prompt_tokens" => 30, "completion_tokens" => 9, "total_tokens" => 39 } + ) + + response = build(usage: object_usage).prompt + + assert_equal 30, response.usage.input_tokens + assert_equal 9, response.usage.output_tokens + end + + test "ignores the empty usage sent on ordinary content chunks" do + provider = build(usage: nil) + provider.prompt + + assert_empty provider.usage_stack + + provider = build(usage: { "prompt_tokens" => 0, "completion_tokens" => 0, "total_tokens" => 0 }) + provider.prompt + + assert_empty provider.usage_stack, "an all-zero payload should not enter the stack" + end + + test "asks Chat Completions for usage when streaming" do + provider = ActiveAgent::Providers::OpenAI::ChatProvider.allocate + + assert_equal({ stream_options: { include_usage: true } }, + provider.send(:api_stream_usage_parameters)) + end + + test "the default provider asks for no extra streaming parameters" do + assert_empty build(usage: nil).send(:api_stream_usage_parameters) + end + + test "a deferred completion runs process_prompt_finished exactly once" do + provider = build(usage: { "prompt_tokens" => 5, "completion_tokens" => 1, "total_tokens" => 6 }) + + calls = 0 + provider.define_singleton_method(:process_prompt_finished) do |*args| + calls += 1 + super(*args) + end + + provider.prompt + + assert_equal 1, calls, "resolve_prompt must not re-finish a stream that already completed" + end +end From f99d431f0dc17e99356cf11dece6fa8b39c8d9c3 Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 15 Aug 2026 00:17:12 +0000 Subject: [PATCH 2/2] Fix the streaming-usage test failures, and the overcount behind one MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit CI was red on 14 tests. Three causes: Streaming lifecycle (OpenAI Chat, Gemini, Ollama). Completion is now deferred from content.done to stream_finished!, so the close event fires when the stream drains rather than on content.done. The tests drive chunks by hand and never drained, so no :close was ever emitted. They now call stream_finished! after the done event, and the OpenAI one asserts that content.done alone does not close — that is the new contract, not an incidental detail. Request building (OpenAI Chat, OpenRouter, Ollama). Streamed requests now carry stream_options: {include_usage: true}; the expected bodies and the recorded cassette requests did not. Both updated. The suite's own new test asserted that a default provider asks for no extra streaming parameters, against a mock that overrode exactly that method to ask for them. The override served no purpose — the mock's fake execute ignores the parameters — so it is gone. The Chat opt-in is asserted against a real serialized request instead (see below), which also drops that file's undeclared dependency on the openai gem. Fixing the request-building tests surfaced a defect worth its own note: usage payloads were pushed onto a stack summed with reduce(:+), but a streamed usage is a running total for the turn, not a delta. Chat Completions sends exactly one, so this was invisible — while Gemini's OpenAI-compatible endpoint, which now opts in by inheritance, repeats a cumulative usage on every chunk. A four-chunk reply reported 50 input tokens instead of 14. The turn now keeps a single entry that later payloads replace, reset per resolve_prompt so tool-calling turns still accumulate. Adds test/providers/open_ai/chat/streaming_usage_test.rb, which drives a real SSE body through the gem's stream helper — the one path CI never exercised. Before the PR, the usage chunk's empty choices array crashed the chunk handler outright (NoMethodError on nil.index), so it covers the nil guard, the recorded usage, the deferred close firing once, the opt-in on the wire, and a Gemini-shaped stream that repeats its usage. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01G8CBDTo2dd9Y4mYn2TDu5P --- lib/active_agent/providers/_base_provider.rb | 20 ++- .../test_agent_functions_with_streaming.yml | 4 +- .../test_agent_streaming.yml | 2 +- .../test_agent_functions_with_streaming.yml | 4 +- .../test_agent_streaming.yml | 2 +- .../test_agent_functions_with_streaming.yml | 4 +- .../test_agent_streaming.yml | 2 +- .../ollama/chat/native_format_test.rb | 5 +- .../open_ai/chat/native_format_test.rb | 5 +- .../open_router/native_format_test.rb | 5 +- .../base_provider_stream_usage_test.rb | 32 +++-- .../gemini/streaming_lifecycle_test.rb | 4 + .../ollama/streaming_lifecycle_test.rb | 4 + .../open_ai/chat/streaming_lifecycle_test.rb | 13 +- .../open_ai/chat/streaming_usage_test.rb | 117 ++++++++++++++++++ 15 files changed, 200 insertions(+), 23 deletions(-) create mode 100644 test/providers/open_ai/chat/streaming_usage_test.rb diff --git a/lib/active_agent/providers/_base_provider.rb b/lib/active_agent/providers/_base_provider.rb index 5f4f6d27..f88bfd68 100644 --- a/lib/active_agent/providers/_base_provider.rb +++ b/lib/active_agent/providers/_base_provider.rb @@ -58,6 +58,7 @@ class ProvidersError < StandardError; end :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 @@ -122,6 +123,7 @@ def initialize(kwargs = {}) 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. @@ -183,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) } @@ -363,6 +368,14 @@ def process_stream_chunk(api_response_chunk) # 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) @@ -380,7 +393,12 @@ def record_stream_usage(raw_usage) usage.input_tokens.to_i.zero? && usage.output_tokens.to_i.zero? - usage_stack.push(usage) + 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. diff --git a/test/fixtures/vcr_cassettes/integration/ollama/chat/native_format_test/test_agent_functions_with_streaming.yml b/test/fixtures/vcr_cassettes/integration/ollama/chat/native_format_test/test_agent_functions_with_streaming.yml index 666786d5..b21a3695 100644 --- a/test/fixtures/vcr_cassettes/integration/ollama/chat/native_format_test/test_agent_functions_with_streaming.yml +++ b/test/fixtures/vcr_cassettes/integration/ollama/chat/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":"qwen3","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 @@ -1765,7 +1765,7 @@ http_interactions: MA\"}","name":"get_current_weather"},"type":"function"}]},{"role":"tool","content":"{\"location\":\"Boston, MA\",\"unit\":\"fahrenheit\",\"temperature\":\"22\"}","tool_call_id":"call_e26xiwn7"}],"model":"qwen3","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 diff --git a/test/fixtures/vcr_cassettes/integration/ollama/chat/native_format_test/test_agent_streaming.yml b/test/fixtures/vcr_cassettes/integration/ollama/chat/native_format_test/test_agent_streaming.yml index 8d0c0040..c0703bbe 100644 --- a/test/fixtures/vcr_cassettes/integration/ollama/chat/native_format_test/test_agent_streaming.yml +++ b/test/fixtures/vcr_cassettes/integration/ollama/chat/native_format_test/test_agent_streaming.yml @@ -5,7 +5,7 @@ http_interactions: uri: http://127.0.0.1:11434/v1/chat/completions body: encoding: UTF-8 - string: '{"messages":[{"role":"developer","content":"You are a helpful assistant."},{"role":"user","content":"Hello!"}],"model":"deepseek-r1:latest","stream":true}' + string: '{"messages":[{"role":"developer","content":"You are a helpful assistant."},{"role":"user","content":"Hello!"}],"model":"deepseek-r1:latest","stream_options":{"include_usage":true},"stream":true}' headers: Accept-Encoding: - gzip;q=1.0,deflate;q=0.6,identity;q=0.3 diff --git a/test/fixtures/vcr_cassettes/integration/open_ai/chat/native_format_test/test_agent_functions_with_streaming.yml b/test/fixtures/vcr_cassettes/integration/open_ai/chat/native_format_test/test_agent_functions_with_streaming.yml index 798f087e..48d3f92a 100644 --- a/test/fixtures/vcr_cassettes/integration/open_ai/chat/native_format_test/test_agent_functions_with_streaming.yml +++ b/test/fixtures/vcr_cassettes/integration/open_ai/chat/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":"gpt-4.1","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 @@ -135,7 +135,7 @@ http_interactions: MA\"}","name":"get_current_weather"},"type":"function"}]},{"role":"tool","content":"{\"location\":\"Boston, MA\",\"unit\":\"fahrenheit\",\"temperature\":\"22\"}","tool_call_id":"call_jbujoPVdWfu8tZ7UjHPotN9G"}],"model":"gpt-4.1","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 diff --git a/test/fixtures/vcr_cassettes/integration/open_ai/chat/native_format_test/test_agent_streaming.yml b/test/fixtures/vcr_cassettes/integration/open_ai/chat/native_format_test/test_agent_streaming.yml index 085e6dda..44930da5 100644 --- a/test/fixtures/vcr_cassettes/integration/open_ai/chat/native_format_test/test_agent_streaming.yml +++ b/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 diff --git a/test/fixtures/vcr_cassettes/integration/open_router/native_format_test/test_agent_functions_with_streaming.yml b/test/fixtures/vcr_cassettes/integration/open_router/native_format_test/test_agent_functions_with_streaming.yml index 23acccb6..70351342 100644 --- a/test/fixtures/vcr_cassettes/integration/open_router/native_format_test/test_agent_functions_with_streaming.yml +++ b/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 diff --git a/test/fixtures/vcr_cassettes/integration/open_router/native_format_test/test_agent_streaming.yml b/test/fixtures/vcr_cassettes/integration/open_router/native_format_test/test_agent_streaming.yml index dc366aeb..ef9899ca 100644 --- a/test/fixtures/vcr_cassettes/integration/open_router/native_format_test/test_agent_streaming.yml +++ b/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 diff --git a/test/integration/ollama/chat/native_format_test.rb b/test/integration/ollama/chat/native_format_test.rb index 7eeac885..4fb8446d 100644 --- a/test/integration/ollama/chat/native_format_test.rb +++ b/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", diff --git a/test/integration/open_ai/chat/native_format_test.rb b/test/integration/open_ai/chat/native_format_test.rb index d548dbc9..fdcd640a 100644 --- a/test/integration/open_ai/chat/native_format_test.rb +++ b/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", diff --git a/test/integration/open_router/native_format_test.rb b/test/integration/open_router/native_format_test.rb index 3f0d97a7..3cbf510b 100644 --- a/test/integration/open_router/native_format_test.rb +++ b/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", diff --git a/test/providers/base_provider_stream_usage_test.rb b/test/providers/base_provider_stream_usage_test.rb index 63e75aaa..4c606e7a 100644 --- a/test/providers/base_provider_stream_usage_test.rb +++ b/test/providers/base_provider_stream_usage_test.rb @@ -15,10 +15,6 @@ def self.name = "ActiveAgent::Providers::MockProvider" class_attribute :usage_payload, default: nil - def api_stream_usage_parameters - { stream_options: { include_usage: true } } - end - # Mimics the real shape: content arrives first and marks the generation # complete, then a final chunk carries the usage. def api_prompt_execute(parameters) @@ -69,13 +65,33 @@ def build(usage:) assert_empty provider.usage_stack, "an all-zero payload should not enter the stack" end - test "asks Chat Completions for usage when streaming" do - provider = ActiveAgent::Providers::OpenAI::ChatProvider.allocate + test "a repeated cumulative usage replaces the turn's entry instead of stacking" do + # Gemini's OpenAI-compatible endpoint repeats a running total on every + # chunk rather than sending one usage at the end of the stream. + provider = build(usage: nil) + [ 6, 12, 18 ].each do |total| + provider.send(:record_stream_usage, + { "prompt_tokens" => total - 2, "completion_tokens" => 2, "total_tokens" => total }) + end + + assert_equal 1, provider.usage_stack.size + assert_equal 16, provider.usage_stack.first.input_tokens, "the last running total wins" + end + + test "each turn of a generation records its own usage" do + # Tool calling re-enters resolve_prompt per turn, and each turn streams a + # usage of its own — those accumulate rather than replacing one another. + provider = build(usage: { "prompt_tokens" => 5, "completion_tokens" => 1, "total_tokens" => 6 }) + + provider.prompt + response = provider.prompt - assert_equal({ stream_options: { include_usage: true } }, - provider.send(:api_stream_usage_parameters)) + assert_equal 2, provider.usage_stack.size + assert_equal 10, response.usage.input_tokens end + # Chat Completions' own opt-in is asserted against a real serialized request + # in Providers::OpenAI::Chat::StreamingUsageTest. test "the default provider asks for no extra streaming parameters" do assert_empty build(usage: nil).send(:api_stream_usage_parameters) end diff --git a/test/providers/gemini/streaming_lifecycle_test.rb b/test/providers/gemini/streaming_lifecycle_test.rb index 374a47ca..c7ebdf3c 100644 --- a/test/providers/gemini/streaming_lifecycle_test.rb +++ b/test/providers/gemini/streaming_lifecycle_test.rb @@ -132,8 +132,10 @@ def deep_to_h ) # Stub process_prompt_finished to just call broadcast_stream_close + # The generation finishes when the stream drains, not on content.done @provider.stub(:process_prompt_finished, ->(*_) { @provider.send(:broadcast_stream_close) }) do @provider.send(:process_stream_chunk, done_event) + @provider.send(:stream_finished!) end event_types = @stream_events.map { |e| e[:type] } @@ -183,8 +185,10 @@ def deep_to_h ) # Stub process_prompt_finished to just call broadcast_stream_close + # The generation finishes when the stream drains, not on content.done @provider.stub(:process_prompt_finished, ->(*_) { @provider.send(:broadcast_stream_close) }) do @provider.send(:process_stream_chunk, done_event) + @provider.send(:stream_finished!) end refute @provider.send(:streaming), "streaming should be false after close" diff --git a/test/providers/ollama/streaming_lifecycle_test.rb b/test/providers/ollama/streaming_lifecycle_test.rb index 4db2300b..3ffc6b22 100644 --- a/test/providers/ollama/streaming_lifecycle_test.rb +++ b/test/providers/ollama/streaming_lifecycle_test.rb @@ -93,8 +93,10 @@ def deep_to_h ) # Stub process_prompt_finished to just call broadcast_stream_close + # The generation finishes when the stream drains, not on content.done @provider.stub(:process_prompt_finished, ->(*_) { @provider.send(:broadcast_stream_close) }) do @provider.send(:process_stream_chunk, done_event) + @provider.send(:stream_finished!) end event_types = @stream_events.map { |e| e[:type] } @@ -140,8 +142,10 @@ def deep_to_h ) # Stub process_prompt_finished to just call broadcast_stream_close + # The generation finishes when the stream drains, not on content.done @provider.stub(:process_prompt_finished, ->(*_) { @provider.send(:broadcast_stream_close) }) do @provider.send(:process_stream_chunk, done_event) + @provider.send(:stream_finished!) end refute @provider.send(:streaming), "streaming should be false after close" diff --git a/test/providers/open_ai/chat/streaming_lifecycle_test.rb b/test/providers/open_ai/chat/streaming_lifecycle_test.rb index 3b656fbe..95119034 100644 --- a/test/providers/open_ai/chat/streaming_lifecycle_test.rb +++ b/test/providers/open_ai/chat/streaming_lifecycle_test.rb @@ -82,7 +82,7 @@ def deep_to_h assert_equal 1, open_events.size, "Expected only one :open event even after multiple chunks" end - test "content.done event triggers :close via process_prompt_finished" do + test "a drained stream, not content.done, triggers :close via process_prompt_finished" do # First send a chunk to trigger :open chunk = MockChunk.new( choices: [ MockChoice.new(index: 0, delta: MockDelta.new(content: "Hi", role: "assistant")) ] @@ -90,7 +90,7 @@ def deep_to_h chunk_event = MockChunkEvent.new(type: :chunk, chunk: chunk) @provider.send(:process_stream_chunk, chunk_event) - # Then send content.done event which triggers process_prompt_finished + # Then send content.done event, which only marks the generation complete done_event = MockContentDoneEvent.new( type: :"content.done", content: "Hi there!", @@ -101,6 +101,11 @@ def deep_to_h # This avoids the nil request issue while testing the streaming lifecycle @provider.stub(:process_prompt_finished, ->(*_) { @provider.send(:broadcast_stream_close) }) do @provider.send(:process_stream_chunk, done_event) + + assert_empty @stream_events.select { |e| e[:type] == :close }, + "content.done must not finish the generation: the usage chunk arrives after it" + + @provider.send(:stream_finished!) end close_events = @stream_events.select { |e| e[:type] == :close } @@ -133,8 +138,10 @@ def deep_to_h ) # Stub process_prompt_finished to just call broadcast_stream_close + # The generation finishes when the stream drains, not on content.done @provider.stub(:process_prompt_finished, ->(*_) { @provider.send(:broadcast_stream_close) }) do @provider.send(:process_stream_chunk, done_event) + @provider.send(:stream_finished!) end event_types = @stream_events.map { |e| e[:type] } @@ -183,8 +190,10 @@ def deep_to_h ) # Stub process_prompt_finished to just call broadcast_stream_close + # The generation finishes when the stream drains, not on content.done @provider.stub(:process_prompt_finished, ->(*_) { @provider.send(:broadcast_stream_close) }) do @provider.send(:process_stream_chunk, done_event) + @provider.send(:stream_finished!) end refute @provider.send(:streaming), "streaming should be false after close" diff --git a/test/providers/open_ai/chat/streaming_usage_test.rb b/test/providers/open_ai/chat/streaming_usage_test.rb new file mode 100644 index 00000000..14410825 --- /dev/null +++ b/test/providers/open_ai/chat/streaming_usage_test.rb @@ -0,0 +1,117 @@ +# frozen_string_literal: true + +require "test_helper" +require_relative "../../../../lib/active_agent/providers/open_ai/chat_provider" + +module Providers + module OpenAI + module Chat + # End-to-end proof that a streamed generation reports its tokens. + # + # Chat Completions sends usage on a final chunk that carries an empty + # choices array, after the content is done, and only when the request + # asked for it. Recording it has to survive the whole chain — the gem's + # stream helper, the chunk handler, and the completion deferred until + # the stream drains — so this drives real response bodies rather than + # hand-built chunk events. + class StreamingUsageTest < ActiveSupport::TestCase + include WebMock::API + + ENDPOINT = "https://api.openai.com/v1/chat/completions" + + USAGE = { prompt_tokens: 14, completion_tokens: 9, total_tokens: 23 } + + # The order a real Chat Completions stream arrives in: content, then + # the finish_reason chunk, then usage on a chunk with no choices. + CHUNKS = [ + { choices: [ { index: 0, delta: { role: "assistant", content: "" }, finish_reason: nil } ] }, + { choices: [ { index: 0, delta: { content: "Hi" }, finish_reason: nil } ] }, + { choices: [ { index: 0, delta: { content: " there!" }, finish_reason: nil } ] }, + { choices: [ { index: 0, delta: {}, finish_reason: "stop" } ] }, + { choices: [], usage: USAGE } + ] + + # Gemini's OpenAI-compatible endpoint ignores the spec and repeats a + # running usage total on every chunk instead of sending one at the end. + CUMULATIVE_CHUNKS = CHUNKS[0..3].each_with_index.map { |chunk, index| + chunk.merge(usage: USAGE.transform_values { |count| ((count * (index + 1)) / 4.0).ceil }) + } + [ CHUNKS.last ] + + def sse_body(chunks) + chunks.map { |chunk| + envelope = { + id: "chatcmpl-stream-usage", + object: "chat.completion.chunk", + created: 1_761_502_994, + model: "gpt-4o-mini" + } + + "data: #{envelope.merge(chunk).to_json}\n\n" + }.join + "data: [DONE]\n\n" + end + + def stub_stream(chunks) + stub_request(:post, ENDPOINT).to_return( + status: 200, + headers: { "Content-Type" => "text/event-stream" }, + body: sse_body(chunks) + ) + end + + setup do + @stream_events = [] + + stub_stream(CHUNKS) + end + + def stream_prompt + ActiveAgent::Providers::OpenAI::ChatProvider.new( + service: "OpenAI", + api_key: "test-api-key", + model: "gpt-4o-mini", + messages: [ { role: "user", content: "Hello" } ], + stream: true, + stream_broadcaster: ->(_message, _delta, event_type) { @stream_events << event_type } + ).prompt + end + + test "a streamed generation reports the usage from its final chunk" do + response = stream_prompt + + assert_equal 14, response.usage.input_tokens + assert_equal 9, response.usage.output_tokens + assert_equal 23, response.usage.total_tokens + end + + test "usage repeated on every chunk counts once, not once per chunk" do + stub_stream(CUMULATIVE_CHUNKS) + + response = stream_prompt + + assert_equal 14, response.usage.input_tokens + assert_equal 9, response.usage.output_tokens + assert_equal 23, response.usage.total_tokens + end + + test "the streamed content still arrives alongside the usage" do + assert_equal "Hi there!", stream_prompt.message.content + end + + test "the deferred completion still closes the stream exactly once" do + stream_prompt + + assert_equal 1, @stream_events.count(:close) + assert_equal :close, @stream_events.last, "the close must come after the last update" + end + + test "the request opts in to usage reporting" do + stream_prompt + + assert_requested :post, ENDPOINT, times: 1 do |request| + JSON.parse(request.body).dig("stream_options", "include_usage") == true + end + end + end + end + end +end