diff --git a/lib/ruby_llm/protocols/converse/streaming.rb b/lib/ruby_llm/protocols/converse/streaming.rb index c54c5e990..afc32201d 100644 --- a/lib/ruby_llm/protocols/converse/streaming.rb +++ b/lib/ruby_llm/protocols/converse/streaming.rb @@ -22,23 +22,22 @@ def stream_response(payload, additional_headers = {}, &block) decoder = event_stream_decoder body = JSON.generate(payload) + faraday_v1 = Faraday::VERSION.start_with?('1') + on_data = RubyLLM::Streaming::FaradayHandlers.build( + faraday_v1: faraday_v1, + on_chunk: ->(chunk, _env) { parse_stream_chunk(decoder, chunk, accumulator, &block) }, + on_failed_response: ->(chunk, env) { handle_failed_stream(chunk, env) } + ) + response = @connection.post(stream_url, payload) do |req| req.headers.merge!(@provider.sign_headers('POST', stream_url, body)) req.headers.merge!(additional_headers) unless additional_headers.empty? req.headers['Accept'] = 'application/vnd.amazon.eventstream' - if Faraday::VERSION.start_with?('1') - req.options[:on_data] = proc do |chunk, _size| - parse_stream_chunk(decoder, chunk, accumulator, &block) - end + if faraday_v1 + req.options[:on_data] = on_data else - req.options.on_data = proc do |chunk, _bytes, env| - if env&.status == 200 - parse_stream_chunk(decoder, chunk, accumulator, &block) - else - handle_failed_stream(chunk, env) - end - end + req.options.on_data = on_data end end diff --git a/lib/ruby_llm/streaming.rb b/lib/ruby_llm/streaming.rb index 9c5451160..016f253a4 100644 --- a/lib/ruby_llm/streaming.rb +++ b/lib/ruby_llm/streaming.rb @@ -171,10 +171,13 @@ def v1_on_data(on_chunk) def v2_on_data(on_chunk, on_failed_response) proc do |chunk, _bytes, env| - if env&.status == 200 - on_chunk.call(chunk, env) - else + # Some adapters do not expose response status during on_data. + status = env&.status + + if status && status != 200 on_failed_response.call(chunk, env) + else + on_chunk.call(chunk, env) end end end diff --git a/spec/ruby_llm/protocols/converse/streaming_spec.rb b/spec/ruby_llm/protocols/converse/streaming_spec.rb index 0df43480b..29cd169bb 100644 --- a/spec/ruby_llm/protocols/converse/streaming_spec.rb +++ b/spec/ruby_llm/protocols/converse/streaming_spec.rb @@ -7,6 +7,7 @@ Object.new.tap do |object| object.extend(described_class) object.instance_variable_set(:@model, instance_double(RubyLLM::Model, id: 'bedrock-test-model')) + object.define_singleton_method(:escape_model_id) { |id| id.to_s.gsub('/', '%2F') } end end @@ -98,4 +99,63 @@ expect(message.thinking.text).to eq('thinking text') expect(message.thinking.signature).to eq('thinking-signature') end + + describe 'on_data routing in stream_response' do + # Captures the on_data proc that stream_response installs on the Faraday + # request, by faking @connection/@provider/req just enough to run the block. + let(:captured_on_data) do + captured = nil + request_options = Object.new + request_options.define_singleton_method(:on_data=) { |proc| captured = proc } + req = Object.new + req.define_singleton_method(:headers) { {} } + req.define_singleton_method(:options) { request_options } + + connection = instance_double(RubyLLM::Connection) + allow(connection).to receive(:post) do |_url, _payload, &block| + block.call(req) + nil + end + + provider = double('provider', sign_headers: {}) # rubocop:disable RSpec/VerifiedDoubles + + streaming.instance_variable_set(:@connection, connection) + streaming.instance_variable_set(:@provider, provider) + allow(streaming).to receive_messages(event_stream_decoder: :decoder, parse_stream_chunk: nil, + handle_failed_stream: nil) + allow(RubyLLM::StreamAccumulator).to receive(:new).and_return( + instance_double(RubyLLM::StreamAccumulator, to_message: RubyLLM::Message.new(role: :assistant, content: '')) + ) + + streaming.send(:stream_response, {}) + captured + end + + before { stub_const('Faraday::VERSION', '2.0.0') } + + # Faraday 2 with the net_http adapter passes a nil env during streaming + # (status unknown). Chunks must still be parsed, not dropped as failures. + it 'parses the chunk when env is nil (status not yet known)' do + captured_on_data.call('event-frame', 11, nil) + + expect(streaming).to have_received(:parse_stream_chunk).with(:decoder, 'event-frame', anything) + expect(streaming).not_to have_received(:handle_failed_stream) + end + + it 'parses the chunk when env reports a 200 status' do + env = Struct.new(:status).new(200) + captured_on_data.call('event-frame', 11, env) + + expect(streaming).to have_received(:parse_stream_chunk).with(:decoder, 'event-frame', anything) + expect(streaming).not_to have_received(:handle_failed_stream) + end + + it 'routes the chunk to failure handling when env reports a non-200 status' do + env = Struct.new(:status).new(403) + captured_on_data.call('error-frame', 11, env) + + expect(streaming).to have_received(:handle_failed_stream).with('error-frame', env) + expect(streaming).not_to have_received(:parse_stream_chunk) + end + end end diff --git a/spec/ruby_llm/streaming_spec.rb b/spec/ruby_llm/streaming_spec.rb index 2ab9c2cad..ba4eb1d96 100644 --- a/spec/ruby_llm/streaming_spec.rb +++ b/spec/ruby_llm/streaming_spec.rb @@ -63,4 +63,64 @@ expect(response.status).to eq(429) expect(response.body).to eq(parsed_error) end + + # Faraday 2 with the net_http adapter invokes on_data with a nil env (the + # status is not yet known mid-stream). The handler must process such chunks + # normally rather than treating them as a failed response and discarding them. + it 'processes chunks when env is nil (status not yet known)' do + yielded_chunks = [] + handler = test_obj.send(:handle_stream) { |chunk| yielded_chunks << chunk } + + handler.call("data: {\"x\":\"ok\"}\n\n", 0, nil) + + expect(yielded_chunks).to eq(['chunk:ok']) + end + + describe RubyLLM::Streaming::FaradayHandlers do + describe '.v2_on_data' do + it 'routes the chunk to on_chunk when env is nil (status unknown)' do + on_chunk_calls = [] + on_failed_calls = [] + handler = described_class.v2_on_data( + ->(chunk, faraday_env) { on_chunk_calls << [chunk, faraday_env] }, + ->(chunk, faraday_env) { on_failed_calls << [chunk, faraday_env] } + ) + + handler.call('frame', 5, nil) + + expect(on_chunk_calls).to eq([['frame', nil]]) + expect(on_failed_calls).to be_empty + end + + it 'routes the chunk to on_chunk when env reports a 200 status' do + on_chunk_calls = [] + on_failed_calls = [] + ok_env = Struct.new(:status).new(200) + handler = described_class.v2_on_data( + ->(chunk, faraday_env) { on_chunk_calls << [chunk, faraday_env] }, + ->(chunk, faraday_env) { on_failed_calls << [chunk, faraday_env] } + ) + + handler.call('frame', 5, ok_env) + + expect(on_chunk_calls).to eq([['frame', ok_env]]) + expect(on_failed_calls).to be_empty + end + + it 'routes the chunk to on_failed_response when env reports a non-200 status' do + on_chunk_calls = [] + on_failed_calls = [] + err_env = Struct.new(:status).new(403) + handler = described_class.v2_on_data( + ->(chunk, faraday_env) { on_chunk_calls << [chunk, faraday_env] }, + ->(chunk, faraday_env) { on_failed_calls << [chunk, faraday_env] } + ) + + handler.call('error-frame', 11, err_env) + + expect(on_failed_calls).to eq([['error-frame', err_env]]) + expect(on_chunk_calls).to be_empty + end + end + end end