Merge pull request #574 from koic/keep_the_listen_acknowledgement_first_on_transport_close

Keep the acknowledgement first when the transport closes during a listen registration

Author: koic

lib/mcp/server/transports/streamable_http_transport.rb

@@ -929,29 +929,33 @@ def first
         # the legacy GET stream (`create_sse_body`).
         #
         # Registration and activation are split on purpose: the entry is inserted inactive
-        # (reserving the cap slot atomically), the acknowledgement is written outside the lock,
+        # (reserving the cap slot atomically), the acknowledgement is written outside the registry 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 acknowledgement write holds the stream's write mutex, which `teardown_listen_subscriptions` also takes
+        # before it marks an entry closed. The two therefore cannot interleave: the acknowledgement either lands
+        # before the result, or, if the transport closed first, is not written at all and the stream just closes,
+        # so no stream ever carries a result ahead of its 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
+            subscription = nil
             @mutex.synchronize do
-              if @max_listen_subscriptions && @listen_subscriptions.size >= @max_listen_subscriptions
-                rejected = true
-              else
-                @listen_subscriptions[subscription_key] = {
+              unless @max_listen_subscriptions && @listen_subscriptions.size >= @max_listen_subscriptions
+                subscription = {
                   request_id: request_id, stream: stream, filter: honored, active: false, write_mutex: Mutex.new, keepalive_wakeup: ConditionVariable.new
                 }
+                @listen_subscriptions[subscription_key] = subscription
               end
             end
 
-            if rejected
+            if subscription.nil?
               close_stream_safely(stream)
             else
               acknowledgement = {
@@ -964,9 +968,27 @@ def listen_sse_body(request_id, honored)
               }
 
               begin
-                send_to_stream(stream, acknowledgement)
-                activate_listen_subscription(subscription_key)
-                start_listen_keepalive_thread(subscription_key, request_id)
+                acknowledged = subscription[:write_mutex].synchronize do
+                  next false if subscription[:closed]
+
+                  send_to_stream(stream, acknowledgement)
+
+                  # Set on the entry itself, not through the registry: a concurrent close may already have cleared the registry
+                  # while its result write waits on this mutex, and that write must still find the stream acknowledged.
+                  # Set under the registry lock as well, since that is the lock the delivery snapshot reads the flag under.
+                  # This is the one place a write mutex is held while the registry lock is taken; it stays deadlock-free only
+                  # as long as no path takes a write mutex inside `@mutex.synchronize`, so resolve entries under `@mutex`,
+                  # release it, then write.
+                  @mutex.synchronize { subscription[:active] = true }
+
+                  true
+                end
+
+                if acknowledged
+                  start_listen_keepalive_thread(subscription_key, request_id)
+                else
+                  close_stream_safely(stream)
+                end
               rescue *STREAM_WRITE_ERRORS
                 remove_listen_subscription(subscription_key)
                 close_stream_safely(stream)
@@ -975,15 +997,6 @@ def listen_sse_body(request_id, honored)
           end
         end
 
-        # 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(subscription_key)
-          @mutex.synchronize do
-            subscription = @listen_subscriptions[subscription_key]
-            subscription[:active] = true if subscription
-          end
-        end
-
         # Periodically writes an SSE keepalive comment frame to a listen stream so a silently dropped
         # 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)
@@ -1149,11 +1162,15 @@ def teardown_listen_subscriptions
 
           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.
+            # this against in-flight deliveries and against the acknowledgement write: each one either lands
+            # before the result or observes `closed` and skips, keeping the graceful result the stream's final message.
             subscription[:write_mutex].synchronize do
               subscription[:closed] = true
 
+              # A stream whose acknowledgement was never written gets no result either: SEP-2575 makes
+              # the acknowledgement the first message, so the stream closes abruptly and the client re-sends.
+              next unless subscription[:active]
+
               begin
                 send_to_stream(subscription[:stream], {
                   jsonrpc: "2.0",

test/mcp/server/transports/streamable_http_transport_test.rb

@@ -6125,7 +6125,7 @@ def string
 
           assert_empty sse_events(io)
 
-          @transport.send(:activate_listen_subscription, "listen-1")
+          @transport.instance_variable_get(:@listen_subscriptions)["listen-1"][:active] = true
           @server.notify_tools_list_changed
 
           assert_equal ["notifications/tools/list_changed"], sse_events(io).map { |event| event["method"] }
@@ -6149,6 +6149,79 @@ def string
           refute(events.any? { |event| event["method"] == "notifications/tools/list_changed" })
         end
 
+        test "transport close during the acknowledgement write still sends the acknowledgement first" do
+          # The stream blocks its first write, the acknowledgement, until released, so the close arrives while
+          # that write is in progress and has to queue behind it on the stream's write mutex.
+          reached = Queue.new
+          release = Queue.new
+          io = StringIO.new
+          first_write = true
+
+          io.define_singleton_method(:write) do |data|
+            if first_write
+              first_write = false
+              reached.push(true)
+              release.pop
+            end
+            super(data)
+          end
+
+          response = @transport.handle_request(modern_rack_request(
+            modern_listen_body(id: "listen-1", params: { notifications: { toolsListChanged: true } }),
+          ))
+          body_thread = Thread.new { response[2].call(io) }
+          close_thread = nil
+
+          begin
+            wait_until { !reached.empty? }
+            close_thread = Thread.new { @transport.close }
+            wait_until_blocked_or_done(close_thread)
+          ensure
+            # Released whatever the waits above did, so a failed wait cannot leave the body holding
+            # the stream's write mutex with the close queued behind it.
+            release.push(true)
+          end
+
+          assert(body_thread.join(5), "the listen body did not finish")
+          assert(close_thread.join(5), "the transport close did not finish")
+
+          events = sse_events(io)
+
+          assert_equal 2, events.size
+          assert_equal "notifications/subscriptions/acknowledged", events[0]["method"]
+          assert_equal "complete", events[1].dig("result", "resultType")
+          assert_predicate io, :closed?
+        end
+
+        test "transport close before the acknowledgement closes the stream without a result" do
+          # An entry in the registered-but-not-yet-acknowledged state: with no acknowledgement written,
+          # a result would be the stream's first message, which SEP-2575 forbids.
+          io = StringIO.new
+          @transport.instance_variable_get(:@listen_subscriptions)["listen-1"] = {
+            request_id: "listen-1",
+            stream: io,
+            filter: { toolsListChanged: true },
+            active: false,
+            write_mutex: Mutex.new,
+            keepalive_wakeup: ConditionVariable.new,
+          }
+
+          @transport.close
+
+          assert_empty sse_events(io)
+          assert_predicate io, :closed?
+        end
+
+        # Waits until `thread` is either blocked (asleep on a lock or queue) or finished, so the step after it
+        # runs against a thread that has made its move.
+        def wait_until_blocked_or_done(thread)
+          deadline = Process.clock_gettime(Process::CLOCK_MONOTONIC) + 5
+          until thread.status == "sleep" || thread.status == false
+            flunk("thread did not block or finish") if Process.clock_gettime(Process::CLOCK_MONOTONIC) > deadline
+            sleep(0.005)
+          end
+        end
+
         test "subscriptions/listen streams for different subscriptions receive their own subscriptionId" do
           first = open_listen_stream(id: "listen-1", notifications: { toolsListChanged: true })
           second = open_listen_stream(id: "listen-2", notifications: { toolsListChanged: true })