Fix Bedrock streaming dropping all chunks when Faraday env is nil

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
 

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

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

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