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

## Motivation and Context

A `subscriptions/listen` stream is registered before its acknowledgement is written,
and the acknowledgement was written outside any lock. Closing the transport snapshots the registry,
then writes each stream's `SubscriptionsListenResult` under the stream's write mutex and closes it.
When the close landed between the registration and the acknowledgement write, the client received
the result and the stream closed; the acknowledgement write then failed against the closed stream.
SEP-2575 requires the acknowledgement to be the first message on the stream, so such a stream ended
in a shape the specification does not allow.

The acknowledgement is now written under the same write mutex, and it is skipped when the transport has
already marked the entry closed. Closing the transport writes the result only to streams that were acknowledged;
a stream that never was is closed without one, an abrupt close the specification allows and the client answers
by sending `subscriptions/listen` again. Since both sides serialize on the one mutex, whenever both messages
are written the acknowledgement precedes the result, and no result reaches a stream that received no acknowledgement.

The flag that marks a stream acknowledged is set under the registry lock as well as the write mutex,
since the registry lock is the one a notification's delivery snapshot reads it under; nothing holds
the registry lock while waiting for a write mutex, so the two cannot deadlock.

## How Has This Been Tested?

New tests in `test/mcp/server/transports/streamable_http_transport_test.rb` close the transport while
an acknowledgement write is in progress and check that the stream carries the acknowledgement and then
the result, and close it with a registered but unacknowledged stream and check that no result is written.
Against the previous library the first stream carries only the result.

## Breaking Changes

None. A listen stream that the transport closes before its acknowledgement was written now ends without
a result instead of receiving one first.

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 })