diff --git a/lib/active_agent/providers/_base_provider.rb b/lib/active_agent/providers/_base_provider.rb index 165fa8f7..f88bfd68 100644 --- a/lib/active_agent/providers/_base_provider.rb +++ b/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(¶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 +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. 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/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 new file mode 100644 index 00000000..4c606e7a --- /dev/null +++ b/test/providers/base_provider_stream_usage_test.rb @@ -0,0 +1,112 @@ +# 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 + + # 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 "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 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 + + 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 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