Merge pull request #595 from koic/bound_the_listen_teardown_on_transport_close

Bound the listen teardown on transport close

Author: koic

docs/_server/subscriptions.md

@@ -88,5 +88,7 @@ transport = MCP::Server::Transports::StreamableHTTPTransport.new(
 
 A stream stays open until the client closes the connection; on graceful shutdown via `transport.close`,
 each open stream receives its `subscriptions/listen` response before closing.
+`close` waits at most 5 seconds in all for those responses to be written: a stream whose client stopped reading
+is closed without one when that time is up, and the client re-sends its `subscriptions/listen` request.
 
 See [Notifications](/server/notifications/) for the notification types themselves and their session scoping.

lib/mcp/server/transports/streamable_http_transport.rb

@@ -361,6 +361,9 @@ def initialize(
         SESSION_REAP_INTERVAL = 60
         # Seconds `close` waits for a reap in progress to finish closing its streams before it kills the reaper.
         REAPER_JOIN_TIMEOUT = 5
+        # How long `close` waits in all for the closing results of the open `subscriptions/listen` streams
+        # before closing the streams still open without one.
+        LISTEN_TEARDOWN_TIMEOUT = 5
 
         # Loopback hosts always accepted by DNS rebinding protection. A locally bound MCP server (the canonical pattern) is
         # protected out of the box; non-loopback deployments widen the list via `allowed_hosts:`.
@@ -1453,8 +1456,14 @@ def close_removed_listen_stream(subscription)
         end
 
         # Graceful teardown (SEP-2575): each open listen stream receives its `SubscriptionsListenResult` response
-        # before the stream closes.
-        def teardown_listen_subscriptions
+        # before the stream closes. The results are written on a helper thread that is waited for at most `timeout`
+        # seconds in all, because a peer that stopped reading stalls the write of its result, or a delivery already
+        # holding the stream's write mutex, for as long as its socket buffers stay full. Once the time is up, every stream
+        # still open is closed from here: closing a Ruby IO ends a write stalled on it with an `IOError` the helper rescues,
+        # while a stream object of another kind may keep the helper blocked, which no longer holds `close`. Either way
+        # the stream goes without its result, like one whose acknowledgement was never written. The helper is left
+        # to finish on its own rather than killed, so a write mutex is never dropped in the middle of a write.
+        def teardown_listen_subscriptions(timeout: LISTEN_TEARDOWN_TIMEOUT)
           removed = @mutex.synchronize do
             subscriptions = @listen_subscriptions.dup
             @listen_subscriptions.clear
@@ -1465,35 +1474,72 @@ def teardown_listen_subscriptions
 
             subscriptions
           end
+          return if removed.empty?
 
-          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 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
+          # Whichever side takes an entry out of `pending` closes its stream: the helper after the write, or `close` once
+          # the time is up. Claiming before closing keeps a stream from being closed by both in either order.
+          pending = removed.dup
+          pending_mutex = Mutex.new
+          writer = Thread.new do
+            removed.each do |subscription_key, subscription|
+              write_listen_closing_result(subscription)
+              close_stream_safely(subscription[:stream]) if pending_mutex.synchronize { pending.delete(subscription_key) }
+            end
+          end
 
-              # 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
+            writer.join(timeout)
+          rescue StandardError
+            # A helper that died with an exception has stopped writing; the streams it did not finish are closed below.
+          end
 
-              begin
-                send_to_stream(subscription[:stream], {
-                  jsonrpc: "2.0",
-                  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 => subscription[:request_id] },
-                  },
-                })
-              rescue *STREAM_WRITE_ERRORS
-                nil
-              end
+          left = pending_mutex.synchronize do
+            subscriptions = pending.values
+            pending.clear
+            subscriptions
+          end
+          left.each { |subscription| close_stream_safely(subscription[:stream]) }
+        end
+
+        # Marking the entry closed and writing the result under the stream's write mutex orders
+        # 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.
+        # The stream is left open for the caller to close.
+        def write_listen_closing_result(subscription)
+          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",
+                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 => subscription[:request_id] },
+                },
+              })
+            rescue *STREAM_WRITE_ERRORS
+              nil
             end
-            close_stream_safely(subscription[:stream])
+          end
+        rescue StandardError => e
+          # Reported rather than raised: on the helper thread an error would end the teardown of every stream after this one,
+          # and `close` is past waiting for it by the time a stalled write fails. A reporter that raises is dropped for
+          # the same reason.
+          begin
+            MCP.configuration.exception_reporter.call(
+              e,
+              { subscription_id: subscription[:request_id], error: "Failed to close listen stream" },
+            )
+          rescue StandardError
+            nil
           end
         end
 

test/mcp/server/transports/streamable_http_transport_test.rb

@@ -6927,6 +6927,129 @@ def wait_until_blocked_or_done(thread)
           assert_equal events, sse_events(io)
         end
 
+        test "listen teardown returns once its timeout passed behind a held write mutex" do
+          # A delivery stalled in its write holds the stream's write mutex for as long as the peer leaves
+          # the socket buffers full; the test holds the mutex in that delivery's place.
+          io = open_listen_stream(id: "listen-1", notifications: { toolsListChanged: true })
+          closes = 0
+          io.define_singleton_method(:close) do
+            closes += 1
+            super()
+          end
+          registry = @transport.instance_variable_get(:@listen_subscriptions)
+          entry = registry.values.first
+          before = Thread.list
+
+          entry[:write_mutex].lock
+          begin
+            teardown = Thread.new { @transport.send(:teardown_listen_subscriptions, timeout: 0.1) }
+
+            refute_nil teardown.join(5), "the teardown did not return while the write mutex stayed held"
+            assert_predicate io, :closed?
+            assert_empty registry
+          ensure
+            entry[:write_mutex].unlock
+          end
+
+          # Released, the helper finds the stream closed and ends without a result written or a second close.
+          wait_until { (Thread.list - before).empty? }
+          assert_nil sse_events(io).find { |event| event.dig("result", "resultType") == "complete" }
+          assert_equal 1, closes
+        end
+
+        test "listen teardown ends a stalled closing result write by closing its stream" do
+          gate = Queue.new
+          reached = Queue.new
+          stalled = Class.new(StringIO) do
+            attr_accessor :stall
+            attr_reader :closes
+
+            # Blocks like a write to a socket whose peer stopped reading, until the stream is closed from another thread.
+            define_method(:write) do |data|
+              return super(data) unless stall
+
+              reached.push(true)
+              gate.pop
+              raise IOError, "closed stream"
+            end
+
+            define_method(:close) do
+              @closes = (@closes || 0) + 1
+              gate.close
+              super()
+            end
+          end.new
+          response = @transport.handle_request(modern_rack_request(
+            modern_listen_body(id: "listen-1", params: { notifications: { toolsListChanged: true } }),
+          ))
+          response[2].call(stalled)
+          stalled.stall = true
+          before = Thread.list
+
+          begin
+            teardown = Thread.new { @transport.send(:teardown_listen_subscriptions, timeout: 0.5) }
+
+            # The helper is parked in the stalled write, with the stream still open, before the timeout runs out.
+            wait_until { !reached.empty? }
+            refute_predicate stalled, :closed?
+
+            refute_nil teardown.join(5), "the teardown did not return while the result write stalled"
+            assert_predicate stalled, :closed?
+          ensure
+            # A failed step must not leave the helper parked in the stalled write for the rest of the suite.
+            gate.close
+          end
+
+          # The close ended the write, so the helper finished on its own, without closing the stream again.
+          wait_until { (Thread.list - before).empty? }
+          assert_nil sse_events(stalled).find { |event| event.dig("result", "resultType") == "complete" }
+          assert_equal 1, stalled.closes
+        end
+
+        test "transport close reports a closing result write failing unexpectedly and still closes the other streams" do
+          failing = open_listen_stream(id: "listen-1", notifications: { toolsListChanged: true })
+          healthy = open_listen_stream(id: "listen-2", notifications: { toolsListChanged: true })
+          failing.define_singleton_method(:write) { |_data| raise "boom" }
+          reported = []
+          original_reporter = MCP.configuration.exception_reporter
+          MCP.configuration.exception_reporter = ->(error, context) { reported << [error, context] }
+
+          begin
+            @transport.close
+          ensure
+            MCP.configuration.exception_reporter = original_reporter
+          end
+
+          assert_predicate failing, :closed?
+          assert_predicate healthy, :closed?
+          assert_equal "complete", sse_events(healthy).last.dig("result", "resultType")
+          assert_equal 1, reported.size
+          assert_equal "boom", reported.first[0].message
+          assert_equal({ subscription_id: "listen-1", error: "Failed to close listen stream" }, reported.first[1])
+        end
+
+        test "transport close tears down even when the exception reporter raises on a failed closing result write" do
+          failing = open_listen_stream(id: "listen-1", notifications: { toolsListChanged: true })
+          healthy = open_listen_stream(id: "listen-2", notifications: { toolsListChanged: true })
+          failing.define_singleton_method(:write) { |_data| raise "boom" }
+          session_stream = StringIO.new
+          open_sse_session(@transport, io: session_stream)
+          original_reporter = MCP.configuration.exception_reporter
+          MCP.configuration.exception_reporter = ->(_error, _context) { raise "reporter failed" }
+
+          begin
+            @transport.close
+          ensure
+            MCP.configuration.exception_reporter = original_reporter
+          end
+
+          assert_predicate failing, :closed?
+          assert_predicate healthy, :closed?
+          assert_equal "complete", sse_events(healthy).last.dig("result", "resultType")
+          assert_predicate session_stream, :closed?
+          assert_empty @transport.instance_variable_get(:@sessions)
+        end
+
         test "subscriptions/listen honoring reads the capability flags, not capability presence" do
           # A server declaring `tools` without `listChanged: true` promises no
           # list-changed delivery, so the acknowledgement omits the type.