diff --git a/lib/mcp/server.rb b/lib/mcp/server.rb index ee397cec..ea2a3c04 100644 --- a/lib/mcp/server.rb +++ b/lib/mcp/server.rb @@ -619,8 +619,24 @@ def handle_request(request, method, session: nil, related_request_id: nil) # `initialize` MUST NOT be cancelled (MCP spec 2025-11-25, cancellation item 2), # so do not track it in the in-flight registry. - cancellation = if related_request_id && method != Methods::INITIALIZE - session&.register_in_flight(related_request_id) + cancellation = nil + if related_request_id && method != Methods::INITIALIZE && session + cancellation = session.register_in_flight(related_request_id) + + # The spec puts the uniqueness obligation on the sender - "The request ID MUST NOT have been previously used by + # the requestor within the same session" - and says nothing about what a receiver does with a duplicate. + # Answering one is the only option that stays correct: the id routes request-scoped messages back to + # the request that caused them, and the transport's rule is that those messages "SHOULD relate to + # the originating client request", which a second live request under the same id makes impossible to honor + # for either of them. Refused the same way a duplicate `initialize` is, and for the same reason: + # so that a repeated id cannot silently displace state negotiated by the first one. + if cancellation.nil? + raise RequestHandlerError.new( + "Invalid Request: request id #{related_request_id.inspect} is already in flight", + request, + error_type: :invalid_request, + ) + end end ->(params) { @@ -727,7 +743,10 @@ def handle_request(request, method, session: nil, related_request_id: nil) reported_exception = wrapped raise wrapped ensure - session&.unregister_in_flight(related_request_id) if related_request_id + # `cancellation` is non-nil exactly when this request claimed the id above, so this also keeps `initialize` + # (which never registers) from evicting an in-flight registration under a reused id when the duplicate-`initialize` + # refusal raises out of the handler. + session&.unregister_in_flight(related_request_id, cancellation: cancellation) if related_request_id && cancellation end } end diff --git a/lib/mcp/server/transports/streamable_http_transport.rb b/lib/mcp/server/transports/streamable_http_transport.rb index e420dc0d..8d889773 100644 --- a/lib/mcp/server/transports/streamable_http_transport.rb +++ b/lib/mcp/server/transports/streamable_http_transport.rb @@ -450,8 +450,13 @@ def drop_broken_stream(session_id, stream, related_request_id) @mutex.synchronize do session = @sessions[session_id] - if related_request_id && session&.dig(:post_request_streams, related_request_id) - session[:post_request_streams].delete(related_request_id) + if related_request_id + # Unregister only our own stream: removing on the id alone would drop whichever stream currently holds it, + # which is not necessarily the one that failed. The failed stream is closed either way, and a request-scoped + # failure never reaches the session teardown below. + registered = session&.dig(:post_request_streams, related_request_id) + session[:post_request_streams].delete(related_request_id) if registered.equal?(stream) + streams_to_close << stream else cleanup_and_collect_stream(session_id, streams_to_close) @@ -1595,6 +1600,14 @@ def handle_regular_request(body_string, session_id, related_request_id: nil) end end + # `Server` refuses a duplicate id as well, but only once the request reaches it. The SSE branch below + # registers this request's stream under that id first, so without this check the colliding request + # would take over the routing entry for the moment it takes to be rejected, and its `ensure` would then + # clear the entry the original request still needs. + if related_request_id && server_session&.in_flight?(related_request_id) + return request_id_conflict_response + end + if session_id && !@stateless && !@enable_json_response handle_request_with_sse_response(body_string, session_id, server_session, related_request_id: related_request_id) else @@ -1618,7 +1631,11 @@ def handle_request_with_sse_response(body_string, session_id, server_session, re session = @sessions[session_id] if session && related_request_id session[:post_request_streams] ||= {} - session[:post_request_streams][related_request_id] = stream + + # Claim the id only while it is free. `handle_regular_request` already refused the colliding request, + # so reaching an occupied slot means a race got past that check; leaving the first stream in place keeps + # its messages going where they belong. + session[:post_request_streams][related_request_id] ||= stream end end @@ -1630,7 +1647,11 @@ def handle_request_with_sse_response(body_string, session_id, server_session, re if related_request_id @mutex.synchronize do session = @sessions[session_id] - session[:post_request_streams]&.delete(related_request_id) if session + # Only retire our own registration: a request that never claimed the id, or one whose claim has + # already been replaced, must not unregister the stream that owns it. + registered = session&.dig(:post_request_streams, related_request_id) + + session[:post_request_streams].delete(related_request_id) if registered.equal?(stream) end end @@ -1849,6 +1870,16 @@ def session_already_connected_response ) end + # The POST counterpart of the GET conflict above. A request id already in flight cannot be given + # a stream of its own, because the id is what routes request-scoped messages back. + def request_id_conflict_response + json_rpc_error_response( + status: 409, + code: JsonRpcHandler::ErrorCode::INVALID_REQUEST, + message: "Conflict: Request id is already in flight for this session", + ) + end + def setup_sse_stream(session_id) body = create_sse_body(session_id) diff --git a/lib/mcp/server_session.rb b/lib/mcp/server_session.rb index d12b18c2..cfee2ab8 100644 --- a/lib/mcp/server_session.rb +++ b/lib/mcp/server_session.rb @@ -58,19 +58,41 @@ def lock_era!(era) @era = era end - # Registers a `Cancellation` token for an in-flight request. + # Registers a `Cancellation` token for an in-flight request, or returns `nil` when `request_id` is already in flight. + # The request id is the only key that routes request-scoped notifications, server-to-client requests, + # and `notifications/cancelled` back to the request that caused them, so a second live request under the same id + # has no destination of its own. Rather than let the newcomer displace the registration, report the collision + # and leave the first request intact; the caller turns that into an Invalid Request. def register_in_flight(request_id) return if request_id.nil? cancellation = Cancellation.new(request_id: request_id) - @in_flight_mutex.synchronize { @in_flight[request_id] = cancellation } - cancellation + registered = @in_flight_mutex.synchronize do + next false if @in_flight.key?(request_id) + + @in_flight[request_id] = cancellation + true + end + + registered ? cancellation : nil end - def unregister_in_flight(request_id) + # Removes an in-flight registration. Passing the `Cancellation` that `register_in_flight` returned removes + # the entry only while it is still that one, so a request can never evict a registration it does not own. + def unregister_in_flight(request_id, cancellation: nil) return if request_id.nil? - @in_flight_mutex.synchronize { @in_flight.delete(request_id) } + @in_flight_mutex.synchronize do + next if cancellation && !@in_flight[request_id].equal?(cancellation) + + @in_flight.delete(request_id) + end + end + + # Whether `request_id` is currently in flight, so a transport can refuse a colliding request + # before registering any state of its own for it. + def in_flight?(request_id) + !lookup_in_flight(request_id).nil? end def lookup_in_flight(request_id) diff --git a/test/mcp/server/transports/streamable_http_transport_test.rb b/test/mcp/server/transports/streamable_http_transport_test.rb index 26b89060..2cc42de0 100644 --- a/test/mcp/server/transports/streamable_http_transport_test.rb +++ b/test/mcp/server/transports/streamable_http_transport_test.rb @@ -3431,6 +3431,110 @@ def string assert_equal "ok", result[:content][:text] end + test "refuses a POST reusing an in-flight request id and keeps the original stream" do + server = Server.new(name: "test", tools: [], prompts: [], resources: []) + transport = StreamableHTTPTransport.new(server) + server.transport = transport + + gate = Queue.new + server.define_tool(name: "victim_tool") do |server_context:| + server_context.report_progress(1, message: "first-frame") + gate.pop + server_context.report_progress(2, message: "second-frame") + Tool::Response.new([{ type: "text", text: "done" }]) + end + + session_id, server_session = start_session(transport) + + victim = transport.handle_request(colliding_tool_call(session_id, "req-1")) + victim_stream = TestStream.new + victim_thread = Thread.new { victim[2].call(victim_stream) } + sleep(0.01) until server_session.lookup_in_flight("req-1") + + attacker = transport.handle_request(colliding_tool_call(session_id, "req-1")) + + assert_equal 409, attacker[0] + assert_equal( + JsonRpcHandler::ErrorCode::INVALID_REQUEST, + JSON.parse(attacker[2][0]).dig("error", "code"), + ) + + gate << :go + victim_thread.join + + # The registration survived the refusal, so the frame emitted after it still reaches + # the request that asked for it. + assert_includes victim_stream.string, "first-frame" + assert_includes victim_stream.string, "second-frame" + end + + test "a finishing request leaves a stream another request registered under the same id" do + server = Server.new(name: "test", tools: [], prompts: [], resources: []) + transport = StreamableHTTPTransport.new(server) + server.transport = transport + + gate = Queue.new + server.define_tool(name: "victim_tool") do |server_context:| + gate.pop + Tool::Response.new([{ type: "text", text: "done" }]) + end + + session_id, server_session = start_session(transport) + + victim = transport.handle_request(colliding_tool_call(session_id, "req-1")) + victim_stream = TestStream.new + victim_thread = Thread.new { victim[2].call(victim_stream) } + sleep(0.01) until server_session.lookup_in_flight("req-1") + + # Stand in for a stream that won the registration in a race the pre-dispatch refusal normally prevents. + # Finishing the other request must not unregister it. + foreign_stream = TestStream.new + sessions = transport.instance_variable_get(:@sessions) + sessions[session_id][:post_request_streams]["req-1"] = foreign_stream + + gate << :go + victim_thread.join + + assert_same foreign_stream, sessions[session_id][:post_request_streams]["req-1"] + end + + test "a broken request-scoped stream drops only itself, even when it is not the registered one" do + server = Server.new(name: "test", tools: [], prompts: [], resources: []) + transport = StreamableHTTPTransport.new(server) + server.transport = transport + + session_id, = start_session(transport) + sessions = transport.instance_variable_get(:@sessions) + registered = TestStream.new + sessions[session_id][:post_request_streams] = { "req-1" => registered } + + transport.send(:drop_broken_stream, session_id, TestStream.new, "req-1") + + assert sessions.key?(session_id), "a request-scoped failure must not tear down the session" + assert_same registered, sessions[session_id][:post_request_streams]["req-1"] + end + + test "allows a request id to be reused once the earlier request has finished" do + server = Server.new(name: "test", tools: [], prompts: [], resources: []) + transport = StreamableHTTPTransport.new(server) + server.transport = transport + + server.define_tool(name: "victim_tool") do |server_context:| + Tool::Response.new([{ type: "text", text: "done" }]) + end + + session_id, = start_session(transport) + + 2.times do + response = transport.handle_request(colliding_tool_call(session_id, "req-1")) + + assert_equal 200, response[0] + stream = TestStream.new + response[2].call(stream) + assert_includes stream.string, "done" + end + end + test "JSON response mode returns accepted when cancellation suppresses response" do server = Server.new(name: "test", tools: [], prompts: [], resources: []) transport = StreamableHTTPTransport.new(server, enable_json_response: true) @@ -6199,6 +6303,36 @@ def install_mutex_probe_stream(session_id, related_request_id: nil) writes end + # Initializes `transport` and returns its session id together with the `ServerSession`, + # which the in-flight request id tests poll to know when a request has really started. + def start_session(transport) + request = create_rack_request( + "POST", + "/", + { "CONTENT_TYPE" => "application/json" }, + { jsonrpc: "2.0", method: "initialize", id: "init", params: initialize_params }.to_json, + ) + session_id = transport.handle_request(request)[1]["mcp-session-id"] + + [session_id, transport.instance_variable_get(:@sessions)[session_id][:server_session]] + end + + # A `tools/call` for `victim_tool` under a caller-chosen request id, with a progress token so + # the tool's `report_progress` has somewhere to go. + def colliding_tool_call(session_id, request_id) + create_rack_request( + "POST", + "/", + { "CONTENT_TYPE" => "application/json", "HTTP_MCP_SESSION_ID" => session_id }, + { + jsonrpc: "2.0", + id: request_id, + method: "tools/call", + params: { name: "victim_tool", arguments: {}, _meta: { progressToken: "tok" } }, + }.to_json, + ) + end + def create_rack_request(method, path, headers, body = nil) default_accept = case method when "POST" diff --git a/test/mcp/server_cancellation_test.rb b/test/mcp/server_cancellation_test.rb index a91bf1ae..9fc28cf9 100644 --- a/test/mcp/server_cancellation_test.rb +++ b/test/mcp/server_cancellation_test.rb @@ -5,6 +5,7 @@ module MCP class ServerCancellationTest < ActiveSupport::TestCase include InstrumentationTestHelper + include InitializeParamsTestHelper class MockTransport < Transport attr_reader :requests, :notifications, :cancelled_request_ids @@ -333,6 +334,74 @@ def handle_request(request); end assert_includes @mock_transport.cancelled_request_ids, "req-9" end + test "register_in_flight refuses an id that is already in flight" do + first = @session.register_in_flight("req-1") + + assert first + assert_nil @session.register_in_flight("req-1"), "a second live request cannot share the id" + assert_same first, @session.lookup_in_flight("req-1"), "the first registration must survive" + end + + test "unregister_in_flight leaves a registration it does not own" do + owner = @session.register_in_flight("req-1") + @session.unregister_in_flight("req-1", cancellation: Cancellation.new(request_id: "req-1")) + + assert_same owner, @session.lookup_in_flight("req-1") + + @session.unregister_in_flight("req-1", cancellation: owner) + + assert_nil @session.lookup_in_flight("req-1") + end + + test "in_flight? reports whether an id is registered" do + refute @session.in_flight?("req-1") + + @session.register_in_flight("req-1") + + assert @session.in_flight?("req-1") + end + + test "a request reusing an in-flight id is answered with Invalid Request" do + @server.define_tool(name: "slow") do |server_context:| + sleep(0.2) + Tool::Response.new([{ type: "text", text: "ok" }]) + end + + request = { + jsonrpc: "2.0", + id: "req-1", + method: Methods::TOOLS_CALL, + params: { name: "slow", arguments: {} }, + } + + in_flight = Thread.new { @session.handle(request) } + sleep(0.01) until @session.lookup_in_flight("req-1") + + duplicate = @session.handle(request) + + assert_equal( + JsonRpcHandler::ErrorCode::INVALID_REQUEST, + duplicate.dig(:error, :code) || duplicate.dig("error", "code"), + ) + in_flight.join + end + + test "a refused duplicate initialize leaves an in-flight registration under its reused id" do + @session.handle(jsonrpc: "2.0", id: "init", method: Methods::INITIALIZE, params: initialize_params) + + owner = @session.register_in_flight("req-1") + + # `initialize` bypasses the duplicate-id refusal (it is never in flight itself), so its + # rejection path is the one place a reused id reaches a handler while the id is still live. + duplicate = @session.handle(jsonrpc: "2.0", id: "req-1", method: Methods::INITIALIZE, params: initialize_params) + + assert_equal( + JsonRpcHandler::ErrorCode::INVALID_REQUEST, + duplicate.dig(:error, :code) || duplicate.dig("error", "code"), + ) + assert_same owner, @session.lookup_in_flight("req-1"), "the registration must survive the refused initialize" + end + test "parent cancellation propagates to nested server-to-client requests" do @session.instance_variable_set(:@client_capabilities, { elicitation: {} })