Merge pull request #521 from koic/reject_duplicate_in_flight_request_ids

Reject duplicate in-flight JSON-RPC request IDs

lib/mcp/server.rb

@@ -619,8 +619,24 @@ module MCP
 
       # `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 @@ module MCP
           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

lib/mcp/server/transports/streamable_http_transport.rb

@@ -450,8 +450,13 @@ module MCP
 
           @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 @@ module MCP
             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 @@ module MCP
               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 @@ module MCP
               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 @@ module MCP
           )
         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)
 

lib/mcp/server_session.rb

@@ -58,19 +58,41 @@ module MCP
       @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)

test/mcp/server/transports/streamable_http_transport_test.rb

@@ -3431,6 +3431,110 @@ module MCP
           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 @@ module MCP
           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"

test/mcp/server_cancellation_test.rb

@@ -5,6 +5,7 @@ require "test_helper"
 module MCP
   class ServerCancellationTest < ActiveSupport::TestCase
     include InstrumentationTestHelper
+    include InitializeParamsTestHelper
 
     class MockTransport < Transport
       attr_reader :requests, :notifications, :cancelled_request_ids
@@ -333,6 +334,74 @@ module MCP
       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: {} })