From 6d40d4cb78b3370849eaea2b36acb044f1ed3a16 Mon Sep 17 00:00:00 2001 From: Koichi ITO Date: Tue, 13 Oct 2026 00:57:00 +0900 Subject: [PATCH] Close a failed listen stream only from the path that removed its entry ## Motivation and Context When a `subscriptions/listen` stream's write fails, in notification delivery or in the acknowledgement write, `StreamableHTTPTransport` removes the stream's registry entry and closes the stream. The close ran whether or not the removal was this path's own: when the keepalive loop, a token expiry, or the transport teardown had removed the entry a moment earlier, that path closes the stream too, and the failed write closed it a second time. Every other removal path closes the stream only when it was the one to remove the entry. The two write failure paths now do the same, so a stream is closed exactly once, by the path that removed it. An acknowledgement skipped because the teardown marked the entry closed first went the same way, closing a stream the teardown closes itself; that close is gone too. ## How Has This Been Tested? A new test in `test/mcp/server/transports/streamable_http_transport_test.rb` fails a delivery write on a stream whose entry another path removed just before, and counts the stream's closes: the delivery closes it once against the previous library, and leaves it to the remover now. A second test has the teardown mark an entry closed while its acknowledgement waits on the write mutex, and counts one close instead of the two the previous library made. ## Breaking Changes None. --- .../transports/streamable_http_transport.rb | 16 ++-- .../streamable_http_transport_test.rb | 80 +++++++++++++++++++ 2 files changed, 87 insertions(+), 9 deletions(-) diff --git a/lib/mcp/server/transports/streamable_http_transport.rb b/lib/mcp/server/transports/streamable_http_transport.rb index 826c69b0..9d5b29ce 100644 --- a/lib/mcp/server/transports/streamable_http_transport.rb +++ b/lib/mcp/server/transports/streamable_http_transport.rb @@ -1204,14 +1204,11 @@ def listen_sse_body(request_id, honored, acknowledgement, auth_info = nil) true end - if acknowledged - start_listen_keepalive_thread(subscription_key, request_id) - else - close_removed_listen_stream(subscription) - end + # An acknowledgement skipped because the entry was marked closed leaves the stream to the path that marked + # it, the teardown, which closes the stream itself once it holds the write mutex. + start_listen_keepalive_thread(subscription_key, request_id) if acknowledged rescue *STREAM_WRITE_ERRORS - remove_listen_subscription(subscription_key) - close_removed_listen_stream(subscription) + close_removed_listen_stream(subscription) if remove_listen_subscription(subscription_key) end end end @@ -1422,8 +1419,9 @@ def deliver_to_listen_subscriptions(method, params) e, { subscription_id: subscription[:request_id], error: "Failed to send notification" }, ) - remove_listen_subscription(subscription_key) - close_removed_listen_stream(subscription) + # Closed only when the removal is this path's own, as every other removal path does: the keepalive loop, + # a token expiry, or the teardown that removed the entry first closes the stream itself. + close_removed_listen_stream(subscription) if remove_listen_subscription(subscription_key) end end end diff --git a/test/mcp/server/transports/streamable_http_transport_test.rb b/test/mcp/server/transports/streamable_http_transport_test.rb index 8f0dafdb..72aa8c3e 100644 --- a/test/mcp/server/transports/streamable_http_transport_test.rb +++ b/test/mcp/server/transports/streamable_http_transport_test.rb @@ -6887,6 +6887,56 @@ def string assert_predicate io, :closed? end + test "transport close that marks an entry closed before its acknowledgement closes the stream once" do + # The acknowledgement waits on the stream's write mutex while the teardown marks the entry closed and closes the stream; + # the skipped acknowledgement then leaves the close to the teardown. The write mutex the body creates is swapped at + # registration for one the test can hold the acknowledgement on while letting the teardown through. + reached = Queue.new + release = Queue.new + first_call = true + write_mutex = Object.new + write_mutex.define_singleton_method(:synchronize) do |&block| + if first_call + first_call = false + reached.push(true) + release.pop + end + block.call + end + + registry = Class.new(Hash) do + define_method(:[]=) do |key, subscription| + subscription[:write_mutex] = write_mutex + super(key, subscription) + end + end.new + @transport.instance_variable_set(:@listen_subscriptions, registry) + + io = StringIO.new + closes = 0 + io.define_singleton_method(:close) do + closes += 1 + super() + 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) } + + begin + wait_until { !reached.empty? } + @transport.close + ensure + release.push(true) + end + + assert(body_thread.join(5), "the listen body did not finish") + assert_empty sse_events(io) + assert_predicate io, :closed? + assert_equal 1, closes + 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) @@ -7061,6 +7111,36 @@ def wait_until_blocked_or_done(thread) assert_equal 2, sse_events(second).count { |event| event["method"] == "notifications/tools/list_changed" } end + test "a delivery whose write fails leaves closing the stream to the path that removed its entry" do + io = open_listen_stream(id: "listen-1", notifications: { toolsListChanged: true }) + registry = @transport.instance_variable_get(:@listen_subscriptions) + subscription_key, entry = registry.first + transport = @transport + closes = 0 + + io.define_singleton_method(:close) do + closes += 1 + super() + end + + # The peer went away: the keepalive loop removed the entry a moment before this write failed on the same stream. + io.define_singleton_method(:write) do |_data| + transport.send(:remove_listen_subscription, subscription_key) + raise Errno::EPIPE + end + + @server.notify_tools_list_changed + + assert_equal 0, closes, "the delivery closed a stream whose removal belonged to another path" + assert_empty registry + + # The path that removed the entry closes the stream, once. + @transport.send(:close_removed_listen_stream, entry) + + assert_equal 1, closes + assert_predicate io, :closed? + end + test "a delivery that took an entry before its removal does not write to the removed stream" do io = open_listen_stream(id: "listen-1", notifications: { toolsListChanged: true }) registry = @transport.instance_variable_get(:@listen_subscriptions)