diff --git a/lib/mcp/server/transports/streamable_http_transport.rb b/lib/mcp/server/transports/streamable_http_transport.rb index a4a1ab2d..57e2f6c0 100644 --- a/lib/mcp/server/transports/streamable_http_transport.rb +++ b/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 diff --git a/test/mcp/server/transports/streamable_http_transport_test.rb b/test/mcp/server/transports/streamable_http_transport_test.rb index 118f2ae3..185bcb04 100644 --- a/test/mcp/server/transports/streamable_http_transport_test.rb +++ b/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.