End a listen stream's keepalive thread when its slot is freed

## Motivation and Context

Each `subscriptions/listen` stream has a keepalive thread that sleeps for the configured interval
and then writes a comment frame, so a silently dropped peer is noticed and its slot freed.
The thread checked for its entry only between sleeps, so when another path removed the entry,
a failed notification write or the transport closing, the thread slept out the rest of its interval,
up to 15 seconds by default, before it noticed and exited.
How many such threads existed therefore depended on how quickly clients opened and dropped listen streams,
not on `max_listen_subscriptions`, which bounds only the entries. The TypeScript SDK stops its keepalive timer
the moment it tears a stream down.

The wait between pings is now a condition variable wait under the registry lock, and whatever removes
an entry signals it under the same lock. The presence check and the wait cannot be separated by a removal,
so a thread waiting out its interval wakes and exits as soon as its entry goes, whichever path removed it;
one already past the wait, in a ping, finishes that write first. Pings still happen outside the lock,
and a dropped peer is still detected by the ping's write failure as before.

The ping itself is now written under the stream's write mutex, as notification delivery and the closing result are,
and skipped once the transport has marked the entry closed. It used to be written outside that mutex,
so a comment frame could follow the closing result, or land between the bytes of a notification on a stream
that does not serialize its writes itself.

## How Has This Been Tested?

New tests in `test/mcp/server/transports/streamable_http_transport_test.rb` open listen streams with a 30 second interval,
free their slots through a failed delivery write and through `close`, and check that every keepalive thread ends within
a bound far below that interval. Against the previous library the threads are still alive at that bound.

Two more tests hold a stream's write mutex and check that the ping waits for it, and mark an entry closed
and check that the ping writes nothing.

## Breaking Changes

None.

Author: koic

lib/mcp/server/transports/streamable_http_transport.rb

@@ -171,8 +171,10 @@ def initialize(
 
           # 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.
+          # write_mutex: Mutex, keepalive_wakeup: ConditionVariable }` (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. Whoever removes an entry signals `keepalive_wakeup` under `@mutex`,
+          # so the stream's keepalive thread ends with its slot instead of sleeping out its interval.
           # 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 = {}
@@ -944,7 +946,7 @@ def listen_sse_body(request_id, honored)
                 rejected = true
               else
                 @listen_subscriptions[subscription_key] = {
-                  request_id: request_id, stream: stream, filter: honored, active: false, write_mutex: Mutex.new,
+                  request_id: request_id, stream: stream, filter: honored, active: false, write_mutex: Mutex.new, keepalive_wakeup: ConditionVariable.new
                 }
               end
             end
@@ -986,12 +988,27 @@ def activate_listen_subscription(subscription_key)
         # 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.
+        #
+        # The wait between pings is a condition variable wait under `@mutex`, not a plain sleep:
+        # the presence check and the wait happen under the same lock that removals signal from,
+        # so a removal cannot slip in between them and a thread waiting out its interval wakes
+        # at once when its entry goes, whichever path removed it. A thread already past the wait,
+        # in a ping, finishes that write first and then finds its entry gone.
         def start_listen_keepalive_thread(subscription_key, request_id)
           return unless @listen_keepalive_interval
 
           Thread.new do
-            while listen_subscription_active?(subscription_key)
-              sleep(@listen_keepalive_interval)
+            loop do
+              registered = @mutex.synchronize do
+                subscription = @listen_subscriptions[subscription_key]
+                next false unless subscription
+
+                subscription[:keepalive_wakeup].wait(@mutex, @listen_keepalive_interval)
+
+                @listen_subscriptions.key?(subscription_key)
+              end
+              break unless registered
+
               send_listen_keepalive_ping(subscription_key)
             end
           rescue *STREAM_WRITE_ERRORS
@@ -1010,20 +1027,20 @@ def start_listen_keepalive_thread(subscription_key, request_id)
           end
         end
 
-        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.
+        # Resolves the entry under the registry lock, then writes outside it so a stalled reader cannot
+        # block every other subscription on `@mutex`. The write itself holds the stream's write mutex,
+        # like notification delivery and the closing result: the comment frame then cannot land between
+        # the bytes of a notification or after the closing result, and once teardown has marked
+        # the entry closed the ping is skipped. A write error propagates to end the keepalive loop.
         def send_listen_keepalive_ping(subscription_key)
-          stream = @mutex.synchronize do
-            subscription = @listen_subscriptions[subscription_key]
-            subscription && subscription[:stream]
-          end
-          return unless stream
+          subscription = @mutex.synchronize { @listen_subscriptions[subscription_key] }
+          return unless subscription
 
-          send_ping_to_stream(stream)
+          subscription[:write_mutex].synchronize do
+            next if subscription[:closed]
+
+            send_ping_to_stream(subscription[:stream])
+          end
         end
 
         # Per SEP-2575, the server MUST NOT send notification types the client has not requested,
@@ -1108,7 +1125,12 @@ def deliver_to_listen_subscriptions(method, params)
         end
 
         def remove_listen_subscription(subscription_key)
-          @mutex.synchronize { @listen_subscriptions.delete(subscription_key) }
+          @mutex.synchronize do
+            subscription = @listen_subscriptions.delete(subscription_key)
+            subscription[:keepalive_wakeup].signal if subscription
+
+            subscription
+          end
         end
 
         # Graceful teardown (SEP-2575): each open listen stream receives its `SubscriptionsListenResult` response
@@ -1117,6 +1139,11 @@ def teardown_listen_subscriptions
           removed = @mutex.synchronize do
             subscriptions = @listen_subscriptions.dup
             @listen_subscriptions.clear
+
+            subscriptions.each_value do |subscription|
+              subscription[:keepalive_wakeup].signal
+            end
+
             subscriptions
           end
 

test/mcp/server/transports/streamable_http_transport_test.rb

@@ -6113,7 +6113,12 @@ 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"] = {
-            request_id: "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,
+            keepalive_wakeup: ConditionVariable.new,
           }
 
           @server.notify_tools_list_changed
@@ -6321,7 +6326,7 @@ def string
           end
           stream.define_singleton_method(:flush) {}
           @transport.instance_variable_get(:@listen_subscriptions)["listen-1"] = {
-            request_id: "listen-1", stream: stream, filter: {}, write_mutex: Mutex.new,
+            request_id: "listen-1", stream: stream, filter: {}, write_mutex: Mutex.new, keepalive_wakeup: ConditionVariable.new,
           }
 
           @transport.send(:send_listen_keepalive_ping, "listen-1")
@@ -6347,6 +6352,113 @@ def string
           transport.close
         end
 
+        test "listen keepalive ends with its slot when a delivery write fails" do
+          # A long interval: were the thread still sleeping it out after the slot is freed, the join below would time out.
+          transport = StreamableHTTPTransport.new(@server, listen_keepalive_interval: 30)
+          before = Thread.list
+          io = open_listen_stream(id: "listen-1", notifications: { toolsListChanged: true }, transport: transport)
+          keepalive_threads = Thread.list - before
+          assert_equal 1, keepalive_threads.size
+          wait_until_asleep(keepalive_threads)
+          io.define_singleton_method(:write) do |_data|
+            raise Errno::EPIPE
+          end
+
+          transport.send_notification("notifications/tools/list_changed", nil, **{})
+
+          assert_empty transport.instance_variable_get(:@listen_subscriptions)
+          assert(keepalive_threads.all? { |thread| thread.join(5) }, "the keepalive thread outlived its freed slot")
+          assert_predicate io, :closed?
+        ensure
+          # Whatever the assertions did, the streams and their threads must not outlive the test.
+          transport.close
+
+          keepalive_threads.each do |thread|
+            thread.join(5)
+          end
+        end
+
+        test "listen keepalive ends with its slot on transport close" do
+          transport = StreamableHTTPTransport.new(@server, listen_keepalive_interval: 30)
+          before = Thread.list
+          open_listen_stream(id: "listen-1", notifications: { toolsListChanged: true }, transport: transport)
+          open_listen_stream(id: "listen-2", notifications: { toolsListChanged: true }, transport: transport)
+          keepalive_threads = Thread.list - before
+          assert_equal 2, keepalive_threads.size
+          wait_until_asleep(keepalive_threads)
+
+          transport.close
+
+          assert(keepalive_threads.all? { |thread| thread.join(5) }, "a keepalive thread outlived the transport")
+        ensure
+          transport.close
+
+          keepalive_threads.each do |thread|
+            thread.join(5)
+          end
+        end
+
+        test "listen keepalive writes under the stream's write mutex" do
+          # Holding the mutex a delivery or the closing result would hold makes the ping wait its turn,
+          # so a comment frame cannot land between the bytes of another message.
+          transport = StreamableHTTPTransport.new(@server, listen_keepalive_interval: 30)
+          before = Thread.list
+          io = open_listen_stream(id: "listen-1", notifications: { toolsListChanged: true }, transport: transport)
+          keepalive_threads = Thread.list - before
+          subscription_key, subscription = transport.instance_variable_get(:@listen_subscriptions).first
+          written_before = io.string.dup
+
+          subscription[:write_mutex].lock
+          ping_thread = Thread.new { transport.send(:send_listen_keepalive_ping, subscription_key) }
+          wait_until_asleep([ping_thread])
+          assert_equal written_before, io.string, "the ping must wait for the stream's write mutex"
+          subscription[:write_mutex].unlock
+
+          assert(ping_thread.join(5), "the ping did not finish once the mutex was released")
+          assert_match(/\A: ping /, io.string.delete_prefix(written_before))
+        ensure
+          subscription[:write_mutex].unlock if subscription && subscription[:write_mutex].owned?
+
+          transport.close
+
+          keepalive_threads.each do |thread|
+            thread.join(5)
+          end
+        end
+
+        test "listen keepalive does not write once the transport marked the stream closed" do
+          # Teardown marks the entry closed and writes the result under the write mutex; a ping that resolved
+          # the entry just before must find the flag and skip, or a comment frame follows the final message.
+          transport = StreamableHTTPTransport.new(@server, listen_keepalive_interval: 30)
+          before = Thread.list
+          io = open_listen_stream(id: "listen-1", notifications: { toolsListChanged: true }, transport: transport)
+          keepalive_threads = Thread.list - before
+          subscription_key, subscription = transport.instance_variable_get(:@listen_subscriptions).first
+          written_before = io.string.dup
+          subscription[:closed] = true
+
+          transport.send(:send_listen_keepalive_ping, subscription_key)
+
+          assert_equal written_before, io.string
+        ensure
+          transport.close
+
+          keepalive_threads.each do |thread|
+            thread.join(5)
+          end
+        end
+
+        # A freshly started keepalive thread may not have reached its wait yet; the removal has to land while
+        # the thread is asleep for the test to say anything about waking it. The bound is generous for
+        # a starved CI runner while staying far below the 30 second interval these tests use.
+        def wait_until_asleep(threads)
+          deadline = Process.clock_gettime(Process::CLOCK_MONOTONIC) + 5
+          until threads.all? { |thread| thread.status == "sleep" }
+            flunk("keepalive threads did not reach their wait") if Process.clock_gettime(Process::CLOCK_MONOTONIC) > deadline
+            sleep(0.005)
+          end
+        end
+
         test "listen keepalive is not started when the interval is nil" do
           # The set of threads, not their count: `Thread.list` is process-wide, so a thread another test left running
           # that finishes between the two samples moves the count in the direction this assertion does not care about.