Merge pull request #567 from koic/key_listen_streams_by_a_server_minted_identifier

Serve `subscriptions/listen` streams that carry the same request id

Author: koic

lib/mcp/server/transports/streamable_http_transport.rb

@@ -169,8 +169,10 @@ def initialize(
           @allowed_origins = Array(allowed_origins).map(&:downcase).freeze
           @pending_responses = {}
 
-          # Maps a `subscriptions/listen` request id to
-          # `{ stream: stream_object, filter: honored_subscription_filter, active: boolean, write_mutex: Mutex }` (SEP-2575).
+          # Maps a key the transport mints for each `subscriptions/listen` stream to
+          # `{ request_id: listen_request_id, stream: stream_object, filter: honored_subscription_filter, active: boolean,
+          # write_mutex: Mutex }` (SEP-2575). The request id is the client's, unique only among that client's own
+          # in-flight requests, so it stamps `subscriptionId` but cannot serve as the key: two clients may pick the same one.
           # In-process only; a multi-worker deployment needs an external event bus to fan notifications out across processes,
           # which is a follow-up.
           @listen_subscriptions = {}
@@ -925,19 +927,25 @@ def first
         # the legacy GET stream (`create_sse_body`).
         #
         # Registration and activation are split on purpose: the entry is inserted inactive
-        # (reserving the id and the cap slot atomically), the acknowledgement is written outside the lock,
+        # (reserving the cap slot atomically), the acknowledgement is written outside the lock,
         # and only then does the entry become eligible for delivery. A concurrent notification between
         # the insert and the acknowledgement write skips the inactive entry,
         # enforcing the SEP-2575 rule that no notification precedes the acknowledgement.
+        #
+        # The entry is keyed by an identifier minted here, not by the request id: that id is unique only among
+        # the requesting client's own in-flight requests, and two clients that pick the same one must each get
+        # their stream, stamped with the id they sent.
         def listen_sse_body(request_id, honored)
           ListenStreamBody.new do |stream|
+            subscription_key = SecureRandom.uuid
             rejected = false
             @mutex.synchronize do
-              if @listen_subscriptions.key?(request_id) ||
-                  (@max_listen_subscriptions && @listen_subscriptions.size >= @max_listen_subscriptions)
+              if @max_listen_subscriptions && @listen_subscriptions.size >= @max_listen_subscriptions
                 rejected = true
               else
-                @listen_subscriptions[request_id] = { stream: stream, filter: honored, active: false, write_mutex: Mutex.new }
+                @listen_subscriptions[subscription_key] = {
+                  request_id: request_id, stream: stream, filter: honored, active: false, write_mutex: Mutex.new,
+                }
               end
             end
 
@@ -955,10 +963,10 @@ def listen_sse_body(request_id, honored)
 
               begin
                 send_to_stream(stream, acknowledgement)
-                activate_listen_subscription(request_id)
-                start_listen_keepalive_thread(request_id)
+                activate_listen_subscription(subscription_key)
+                start_listen_keepalive_thread(subscription_key, request_id)
               rescue *STREAM_WRITE_ERRORS
-                remove_listen_subscription(request_id)
+                remove_listen_subscription(subscription_key)
                 close_stream_safely(stream)
               end
             end
@@ -967,9 +975,9 @@ def listen_sse_body(request_id, honored)
 
         # Marks a listen subscription eligible for delivery once its acknowledgement write has completed.
         # The entry may already be gone when the transport closed concurrently.
-        def activate_listen_subscription(request_id)
+        def activate_listen_subscription(subscription_key)
           @mutex.synchronize do
-            subscription = @listen_subscriptions[request_id]
+            subscription = @listen_subscriptions[subscription_key]
             subscription[:active] = true if subscription
           end
         end
@@ -978,37 +986,39 @@ def activate_listen_subscription(request_id)
         # connection is detected and its slot freed, rather than held until the next fan-out write.
         # Mirrors the legacy GET stream's `start_keepalive_thread`; a comment frame (not a data frame)
         # cannot corrupt an interleaved notification's JSON.
-        def start_listen_keepalive_thread(request_id)
+        def start_listen_keepalive_thread(subscription_key, request_id)
           return unless @listen_keepalive_interval
 
           Thread.new do
-            while listen_subscription_active?(request_id)
+            while listen_subscription_active?(subscription_key)
               sleep(@listen_keepalive_interval)
-              send_listen_keepalive_ping(request_id)
+              send_listen_keepalive_ping(subscription_key)
             end
           rescue *STREAM_WRITE_ERRORS
             # The peer went away; the ensure frees the slot. A dropped listen stream is the normal
             # way this loop ends, so it is not reported.
           rescue StandardError => e
+            # The request id is taken from the caller rather than the registry: a delivery failure may have
+            # removed the entry already, and the report should still name the stream.
             MCP.configuration.exception_reporter.call(e, { subscription_id: request_id })
           ensure
             stream = @mutex.synchronize do
-              subscription = @listen_subscriptions.delete(request_id)
+              subscription = @listen_subscriptions.delete(subscription_key)
               subscription && subscription[:stream]
             end
             close_stream_safely(stream) if stream
           end
         end
 
-        def listen_subscription_active?(request_id)
-          @mutex.synchronize { @listen_subscriptions.key?(request_id) }
+        def listen_subscription_active?(subscription_key)
+          @mutex.synchronize { @listen_subscriptions.key?(subscription_key) }
         end
 
         # Resolves the stream under the lock, then writes outside it so a stalled reader cannot block
         # every other subscription on `@mutex`. A write error propagates to end the keepalive loop.
-        def send_listen_keepalive_ping(request_id)
+        def send_listen_keepalive_ping(subscription_key)
           stream = @mutex.synchronize do
-            subscription = @listen_subscriptions[request_id]
+            subscription = @listen_subscriptions[subscription_key]
             subscription && subscription[:stream]
           end
           return unless stream
@@ -1054,7 +1064,7 @@ def deliver_to_listen_subscriptions(method, params)
           # The matching snapshot is taken under `@mutex`, but stream writes happen outside it:
           # a slow or stalled subscriber must not block the transport, matching the legacy delivery paths.
           matched = @mutex.synchronize do
-            @listen_subscriptions.filter_map do |request_id, subscription|
+            @listen_subscriptions.filter_map do |subscription_key, subscription|
               # An inactive entry has not finished writing its acknowledgement yet;
               # delivering to it would put a notification ahead of the acknowledgement.
               next unless subscription[:active]
@@ -1067,12 +1077,12 @@ def deliver_to_listen_subscriptions(method, params)
                 uris.is_a?(Array) && uris.include?(uri)
               end
 
-              [request_id, subscription] if hit
+              [subscription_key, subscription] if hit
             end
           end
 
-          matched.each do |request_id, subscription|
-            meta = { RequestEnvelope::SUBSCRIPTION_ID_META_KEY.to_sym => request_id }
+          matched.each do |subscription_key, subscription|
+            meta = { RequestEnvelope::SUBSCRIPTION_ID_META_KEY.to_sym => subscription[:request_id] }
             notification_params = (params || {}).merge(_meta: meta)
             notification = { jsonrpc: "2.0", method: method, params: notification_params }
 
@@ -1089,16 +1099,16 @@ def deliver_to_listen_subscriptions(method, params)
             rescue *STREAM_WRITE_ERRORS => e
               MCP.configuration.exception_reporter.call(
                 e,
-                { subscription_id: request_id, error: "Failed to send notification" },
+                { subscription_id: subscription[:request_id], error: "Failed to send notification" },
               )
-              remove_listen_subscription(request_id)
+              remove_listen_subscription(subscription_key)
               close_stream_safely(subscription[:stream])
             end
           end
         end
 
-        def remove_listen_subscription(request_id)
-          @mutex.synchronize { @listen_subscriptions.delete(request_id) }
+        def remove_listen_subscription(subscription_key)
+          @mutex.synchronize { @listen_subscriptions.delete(subscription_key) }
         end
 
         # Graceful teardown (SEP-2575): each open listen stream receives its `SubscriptionsListenResult` response
@@ -1110,7 +1120,7 @@ def teardown_listen_subscriptions
             subscriptions
           end
 
-          removed.each do |request_id, subscription|
+          removed.each_value do |subscription|
             # Marking the entry closed and writing the result under the stream's write mutex orders
             # this against in-flight deliveries: each one either lands before the result or observes
             # `closed` and skips, keeping the graceful result the stream's final message.
@@ -1120,13 +1130,13 @@ def teardown_listen_subscriptions
               begin
                 send_to_stream(subscription[:stream], {
                   jsonrpc: "2.0",
-                  id: request_id,
+                  id: subscription[:request_id],
                   result: {
                     # `SubscriptionsListenResult` is served at the transport layer and never
                     # passes through the dispatch path, so the REQUIRED 2026-07-28 `resultType` is
                     # stamped at its construction site.
                     resultType: ResultType::COMPLETE,
-                    _meta: { RequestEnvelope::SUBSCRIPTION_ID_META_KEY.to_sym => request_id },
+                    _meta: { RequestEnvelope::SUBSCRIPTION_ID_META_KEY.to_sym => subscription[:request_id] },
                   },
                 })
               rescue *STREAM_WRITE_ERRORS

test/mcp/server/transports/streamable_http_transport_test.rb

@@ -6113,7 +6113,7 @@ def string
           # the registry insert and the acknowledgement write, which happens outside the lock.
           io = StringIO.new
           @transport.instance_variable_get(:@listen_subscriptions)["listen-1"] = {
-            stream: io, filter: { toolsListChanged: true }, active: false, write_mutex: Mutex.new
+            request_id: "listen-1", stream: io, filter: { toolsListChanged: true }, active: false, write_mutex: Mutex.new
           }
 
           @server.notify_tools_list_changed
@@ -6128,13 +6128,14 @@ def string
 
         test "a delivery racing the graceful teardown cannot write after the final result" do
           io = open_listen_stream(id: "listen-1", notifications: { toolsListChanged: true })
-          entry = @transport.instance_variable_get(:@listen_subscriptions)["listen-1"]
+          registry = @transport.instance_variable_get(:@listen_subscriptions)
+          subscription_key, entry = registry.first
 
           @transport.close
 
           # Simulate an in-flight delivery that snapshotted the entry before teardown cleared
           # the registry: the closed flag set under the write mutex makes it a no-op.
-          @transport.instance_variable_get(:@listen_subscriptions)["listen-1"] = entry
+          registry[subscription_key] = entry
           @server.notify_tools_list_changed
 
           events = sse_events(io)
@@ -6257,16 +6258,54 @@ def string
           transport.close
         end
 
-        test "subscriptions/listen rejects a duplicate subscription id by closing the new stream" do
-          open_listen_stream(id: "listen-1", notifications: { toolsListChanged: true })
+        test "subscriptions/listen serves two streams that carry the same request id" do
+          # A request id is unique only among one client's own in-flight requests, and clients that number
+          # their requests with a counter reach the same small integers, so two clients' listen requests may
+          # carry the same id. Each gets its own stream, stamped with the id it sent.
+          first = open_listen_stream(id: 1, notifications: { toolsListChanged: true })
+          second = open_listen_stream(id: 1, notifications: { toolsListChanged: true })
 
-          duplicate = StringIO.new
-          response = @transport.handle_request(modern_rack_request(
-            modern_listen_body(id: "listen-1", params: { notifications: { toolsListChanged: true } }),
-          ))
-          response[2].call(duplicate)
+          @server.notify_tools_list_changed
+
+          [first, second].each do |io|
+            refute_predicate io, :closed?
+            events = sse_events(io)
+            assert_equal ["notifications/subscriptions/acknowledged", "notifications/tools/list_changed"], events.map { |event| event["method"] }
+            assert_equal [1, 1], events.map { |event| event.dig("params", "_meta", "io.modelcontextprotocol/subscriptionId") }
+          end
+        end
+
+        test "transport close sends each of two streams sharing a request id its own result" do
+          first = open_listen_stream(id: 1, notifications: { toolsListChanged: true })
+          second = open_listen_stream(id: 1, notifications: { toolsListChanged: true })
+
+          @transport.close
 
-          assert_predicate duplicate, :closed?
+          [first, second].each do |io|
+            result = sse_events(io).last
+            assert_equal 1, result["id"]
+            assert_equal "complete", result.dig("result", "resultType")
+            assert_predicate io, :closed?
+          end
+        end
+
+        test "a failed write on one of two streams sharing a request id frees only that stream" do
+          first = open_listen_stream(id: 1, notifications: { toolsListChanged: true })
+          second = open_listen_stream(id: 1, notifications: { toolsListChanged: true })
+          first.define_singleton_method(:write) { |_data| raise Errno::EPIPE }
+
+          @server.notify_tools_list_changed
+
+          assert_predicate first, :closed?
+          refute_predicate second, :closed?
+          registry = @transport.instance_variable_get(:@listen_subscriptions)
+          assert_equal 1, registry.size
+          assert_same second, registry.values.first[:stream]
+
+          # The entry that survived is the second stream's, so a further notification still reaches it.
+          @server.notify_tools_list_changed
+
+          assert_equal 2, sse_events(second).count { |event| event["method"] == "notifications/tools/list_changed" }
         end
 
         test "listen keepalive writes a comment frame outside the mutex" do
@@ -6281,7 +6320,9 @@ def string
             ping = data
           end
           stream.define_singleton_method(:flush) {}
-          @transport.instance_variable_get(:@listen_subscriptions)["listen-1"] = { stream: stream, filter: {}, write_mutex: Mutex.new }
+          @transport.instance_variable_get(:@listen_subscriptions)["listen-1"] = {
+            request_id: "listen-1", stream: stream, filter: {}, write_mutex: Mutex.new,
+          }
 
           @transport.send(:send_listen_keepalive_ping, "listen-1")