fix: return from initialize as soon as its SSE response arrives

Author: atesgoral

lib/mcp/client/http.rb

@@ -120,7 +120,7 @@ class MessageTooLargeError < StandardError; end
       # the incrementally parsed JSON-RPC response, the last received SSE event id (for `Last-Event-ID`),
       # and the server's `retry:` reconnection delay. Non-SSE bodies accumulate in `buffer` for the JSON path.
       class SSEStream
-        attr_reader :buffer, :response, :last_event_id, :retry_ms
+        attr_reader :buffer, :response, :response_headers, :last_event_id, :retry_ms
         attr_accessor :abortable
 
         def initialize(abortable:, max_message_bytes: MAX_MESSAGE_BYTES, on_request: nil)
@@ -132,6 +132,7 @@ def initialize(abortable:, max_message_bytes: MAX_MESSAGE_BYTES, on_request: nil
           @response = nil
           @last_event_id = nil
           @retry_ms = nil
+          @response_headers = nil
         end
 
         # Whether the server sent a SEP-1699 priming event (any event carrying an id),
@@ -146,9 +147,10 @@ def primed?
         # abort the stream so an over-limit tail cannot mask a response that already arrived.
         def on_data
           proc do |chunk, _received_bytes, env|
+            @response_headers ||= env.response_headers if env
             if event_stream?(env)
               feed(chunk)
-              raise StreamAbort if @abortable && @response
+              raise StreamAbort if @abortable && @response && !error_status?(env)
 
               if parser_buffered_bytes > @max_message_bytes
                 raise MessageTooLargeError, "Server SSE event exceeds #{@max_message_bytes} bytes without completing"
@@ -163,6 +165,15 @@ def on_data
           end
         end
 
+        # An error status must reach Faraday's `raise_error` middleware (which drives the
+        # OAuth and step-up retries in `send_request`), so the stream is read to EOF instead
+        # of aborting when the server answered with one — the pre-abort behavior for every
+        # method, and the only case where the finished response object is still load-bearing.
+        def error_status?(env)
+          status = env.status if env.respond_to?(:status)
+          status.is_a?(Integer) && status >= 400
+        end
+
         # A fresh parser for a new HTTP connection (reconnection GET),
         # so a partial line from the previous stream cannot corrupt the next one.
         def reset_parser!
@@ -408,10 +419,13 @@ def send_request(request:)
 
         begin
           # The response is consumed incrementally so that an SSE stream the server holds open
-          # (or closes early per SEP-1699) can be handled; `initialize` streams are read to EOF
-          # so the response object (and its `Mcp-Session-Id` header) is always available for capture.
+          # (or closes early per SEP-1699) can be handled. Every request, `initialize` included,
+          # aborts once its response event arrives: `SSEStream` has already captured the response
+          # headers from the streaming `env` by then, so `Mcp-Session-Id` survives the abort.
+          # When `env` is unavailable (Faraday < 2.1 or an adapter without `on_data` support)
+          # nothing aborts and the stream is still read to EOF, as before.
           stream = SSEStream.new(
-            abortable: method.to_s != MCP::Methods::INITIALIZE,
+            abortable: true,
             max_message_bytes: @max_message_bytes,
             on_request: ->(message) { dispatch_server_request(message) },
           )
@@ -434,7 +448,8 @@ def send_request(request:)
           # anything is learned from the body, and only for a request: a notification awaits no response.
           reject_invalid_response!(body, method, params) if request_id
 
-          capture_session_info(method, response, body) if response
+          response_headers = response ? response.headers : stream.response_headers
+          capture_session_info(method, response_headers, body) if response_headers
           capture_mcp_param_declarations(method, params, body)
 
           body
@@ -963,11 +978,12 @@ def refresh_token_available?
         !(tokens["refresh_token"] || tokens[:refresh_token]).to_s.empty?
       end
 
-      def capture_session_info(method, response, body)
+      def capture_session_info(method, headers, body)
         return unless method.to_s == Methods::INITIALIZE
+        return unless headers
 
         # Faraday normalizes header names to lowercase.
-        session_id = response.headers[SESSION_ID_HEADER.downcase]
+        session_id = headers[SESSION_ID_HEADER.downcase]
         @session_id ||= session_id unless session_id.to_s.empty?
         @protocol_version ||= body.is_a?(Hash) ? body.dig("result", "protocolVersion") : nil
       end

test/mcp/client/http_test.rb

@@ -1914,6 +1914,133 @@ def test_captures_session_id_and_protocol_version_on_initialize
         assert_equal("2025-11-25", client.protocol_version)
       end
 
+      def test_initialize_returns_promptly_and_captures_session_id_when_sse_stream_stays_open
+        require "socket"
+
+        hold_seconds = 5
+        seen_session_ids = []
+        WebMock.disable!
+        server = TCPServer.new("127.0.0.1", 0)
+        connection_threads = []
+        server_thread = Thread.new do
+          2.times do |index|
+            socket = server.accept
+            connection_threads << Thread.new(socket, index) do |conn, conn_index|
+              request_headers = {}
+              conn.gets
+              while (line = conn.gets) && line != "\r\n"
+                name, value = line.split(": ", 2)
+                request_headers[name.downcase] = value&.strip
+              end
+              content_length = request_headers["content-length"].to_i
+              conn.read(content_length) if content_length.positive?
+              seen_session_ids << request_headers["mcp-session-id"]
+              if conn_index.zero?
+                event = 'data: {"jsonrpc":"2.0","id":"1","result":{"protocolVersion":"2025-11-25"}}' + "\n\n"
+                conn.write(
+                  "HTTP/1.1 200 OK\r\n" \
+                    "Content-Type: text/event-stream\r\n" \
+                    "Mcp-Session-Id: session-held-open\r\n" \
+                    "Connection: close\r\n\r\n",
+                )
+                conn.write(event)
+                conn.flush
+                sleep(hold_seconds)
+              else
+                body = { jsonrpc: "2.0", result: { tools: [] } }.to_json
+                conn.write(
+                  "HTTP/1.1 200 OK\r\n" \
+                    "Content-Type: application/json\r\n" \
+                    "Content-Length: #{body.bytesize}\r\n" \
+                    "Connection: close\r\n\r\n#{body}",
+                )
+              end
+              conn.close
+            end
+          end
+        rescue IOError, SystemCallError
+          nil
+        end
+        begin
+          held_client = HTTP.new(url: "http://127.0.0.1:#{server.addr[1]}/")
+          started = Process.clock_gettime(Process::CLOCK_MONOTONIC)
+          held_client.send_request(request: { jsonrpc: "2.0", id: "1", method: "initialize" })
+          elapsed = Process.clock_gettime(Process::CLOCK_MONOTONIC) - started
+          assert_operator(elapsed, :<, hold_seconds / 2.0, "initialize waited for the held-open SSE stream to close")
+          assert_equal("session-held-open", held_client.session_id)
+          assert_equal("2025-11-25", held_client.protocol_version)
+          held_client.send_request(request: { jsonrpc: "2.0", id: "2", method: "tools/list" })
+          server_thread.join(5)
+          assert_equal("session-held-open", seen_session_ids[1])
+        ensure
+          connection_threads.each(&:kill)
+          connection_threads.each { |thread| thread.join(1) }
+          server_thread.kill
+          server.close unless server.closed?
+        end
+      ensure
+        WebMock.enable!
+      end
+
+      def test_initialize_retries_through_oauth_when_sse_answer_is_401
+        require "socket"
+
+        WebMock.disable!
+        server = TCPServer.new("127.0.0.1", 0)
+        connection_threads = []
+        server_thread = Thread.new do
+          2.times do |index|
+            socket = server.accept
+            connection_threads << Thread.new(socket, index) do |conn, conn_index|
+              conn.gets
+              headers = {}
+              while (line = conn.gets) && line != "\r\n"
+                name, value = line.split(": ", 2)
+                headers[name.downcase] = value&.strip
+              end
+              length = headers["content-length"].to_i
+              conn.read(length) if length.positive?
+              if conn_index.zero?
+                event = 'data: {"jsonrpc":"2.0","id":"1","error":{"code":-32001,"message":"unauthorized"}}' + "\n\n"
+                conn.write(
+                  "HTTP/1.1 401 Unauthorized\r\n" \
+                    "Content-Type: text/event-stream\r\n" \
+                    "Connection: close\r\n\r\n#{event}",
+                )
+              else
+                event = 'data: {"jsonrpc":"2.0","id":"1","result":{"protocolVersion":"2025-11-25"}}' + "\n\n"
+                conn.write(
+                  "HTTP/1.1 200 OK\r\n" \
+                    "Content-Type: text/event-stream\r\n" \
+                    "Mcp-Session-Id: session-after-retry\r\n" \
+                    "Connection: close\r\n\r\n#{event}",
+                )
+              end
+              conn.close
+            end
+          end
+        rescue IOError, SystemCallError
+          nil
+        end
+        begin
+          oauth_client = HTTP.new(url: "http://127.0.0.1:#{server.addr[1]}/", oauth: Struct.new(:access_token).new)
+          oauth_calls = 0
+          oauth_client.define_singleton_method(:run_oauth_flow!) { |unauthorized_error:| oauth_calls += 1 }
+          oauth_client.send_request(request: { jsonrpc: "2.0", id: "1", method: "initialize" })
+          server_thread.join(5)
+          assert_equal(1, oauth_calls, "the 401 must reach the OAuth retry instead of being returned as a response")
+          assert_equal("session-after-retry", oauth_client.session_id)
+          assert_equal("2025-11-25", oauth_client.protocol_version)
+        ensure
+          connection_threads.each(&:kill)
+          connection_threads.each { |thread| thread.join(1) }
+          server_thread.kill
+          server.close unless server.closed?
+        end
+      ensure
+        WebMock.enable!
+      end
+
       def test_includes_session_and_protocol_version_headers_after_initialize
         stub_request(:post, url)
           .to_return(