diff --git a/lib/mcp/client/http.rb b/lib/mcp/client/http.rb index 5c4ab6eb..4e728748 100644 --- a/lib/mcp/client/http.rb +++ b/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,16 @@ 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. This is the behavior every + # method had before the abort shortcut. It is the only case that still needs + # the finished response object. + 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 +420,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.5 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 +449,7 @@ 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 + capture_session_info(method, response ? response.headers : stream.response_headers, body) capture_mcp_param_declarations(method, params, body) body @@ -963,11 +978,11 @@ 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 # 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 diff --git a/test/mcp/client/http_test.rb b/test/mcp/client/http_test.rb index 64028da8..3ed866d7 100644 --- a/test/mcp/client/http_test.rb +++ b/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(