diff --git a/lib/mcp/server/transports/streamable_http_transport.rb b/lib/mcp/server/transports/streamable_http_transport.rb index 57e2f6c0..186f7287 100644 --- a/lib/mcp/server/transports/streamable_http_transport.rb +++ b/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", diff --git a/test/mcp/server/transports/streamable_http_transport_test.rb b/test/mcp/server/transports/streamable_http_transport_test.rb index 185bcb04..ce8e0ae7 100644 --- a/test/mcp/server/transports/streamable_http_transport_test.rb +++ b/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 })