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)