Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
70 changes: 40 additions & 30 deletions lib/mcp/server/transports/streamable_http_transport.rb
Original file line number Diff line number Diff line change
Expand Up @@ -169,8 +169,10 @@ def initialize(
@allowed_origins = Array(allowed_origins).map(&:downcase).freeze
@pending_responses = {}

# Maps a `subscriptions/listen` request id to
# `{ stream: stream_object, filter: honored_subscription_filter, active: boolean, write_mutex: Mutex }` (SEP-2575).
# 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.
# 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 = {}
Expand Down Expand Up @@ -925,19 +927,25 @@ def first
# the legacy GET stream (`create_sse_body`).
#
# Registration and activation are split on purpose: the entry is inserted inactive
# (reserving the id and the cap slot atomically), the acknowledgement is written outside the lock,
# (reserving the cap slot atomically), the acknowledgement is written outside the 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 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
@mutex.synchronize do
if @listen_subscriptions.key?(request_id) ||
(@max_listen_subscriptions && @listen_subscriptions.size >= @max_listen_subscriptions)
if @max_listen_subscriptions && @listen_subscriptions.size >= @max_listen_subscriptions
rejected = true
else
@listen_subscriptions[request_id] = { stream: stream, filter: honored, active: false, write_mutex: Mutex.new }
@listen_subscriptions[subscription_key] = {
request_id: request_id, stream: stream, filter: honored, active: false, write_mutex: Mutex.new,
}
end
end

Expand All @@ -955,10 +963,10 @@ def listen_sse_body(request_id, honored)

begin
send_to_stream(stream, acknowledgement)
activate_listen_subscription(request_id)
start_listen_keepalive_thread(request_id)
activate_listen_subscription(subscription_key)
start_listen_keepalive_thread(subscription_key, request_id)
rescue *STREAM_WRITE_ERRORS
remove_listen_subscription(request_id)
remove_listen_subscription(subscription_key)
close_stream_safely(stream)
end
end
Expand All @@ -967,9 +975,9 @@ def listen_sse_body(request_id, honored)

# 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(request_id)
def activate_listen_subscription(subscription_key)
@mutex.synchronize do
subscription = @listen_subscriptions[request_id]
subscription = @listen_subscriptions[subscription_key]
subscription[:active] = true if subscription
end
end
Expand All @@ -978,37 +986,39 @@ def activate_listen_subscription(request_id)
# 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.
def start_listen_keepalive_thread(request_id)
def start_listen_keepalive_thread(subscription_key, request_id)
return unless @listen_keepalive_interval

Thread.new do
while listen_subscription_active?(request_id)
while listen_subscription_active?(subscription_key)
sleep(@listen_keepalive_interval)
send_listen_keepalive_ping(request_id)
send_listen_keepalive_ping(subscription_key)
end
rescue *STREAM_WRITE_ERRORS
# The peer went away; the ensure frees the slot. A dropped listen stream is the normal
# way this loop ends, so it is not reported.
rescue StandardError => e
# The request id is taken from the caller rather than the registry: a delivery failure may have
# removed the entry already, and the report should still name the stream.
MCP.configuration.exception_reporter.call(e, { subscription_id: request_id })
ensure
stream = @mutex.synchronize do
subscription = @listen_subscriptions.delete(request_id)
subscription = @listen_subscriptions.delete(subscription_key)
subscription && subscription[:stream]
end
close_stream_safely(stream) if stream
end
end

def listen_subscription_active?(request_id)
@mutex.synchronize { @listen_subscriptions.key?(request_id) }
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.
def send_listen_keepalive_ping(request_id)
def send_listen_keepalive_ping(subscription_key)
stream = @mutex.synchronize do
subscription = @listen_subscriptions[request_id]
subscription = @listen_subscriptions[subscription_key]
subscription && subscription[:stream]
end
return unless stream
Expand Down Expand Up @@ -1054,7 +1064,7 @@ def deliver_to_listen_subscriptions(method, params)
# The matching snapshot is taken under `@mutex`, but stream writes happen outside it:
# a slow or stalled subscriber must not block the transport, matching the legacy delivery paths.
matched = @mutex.synchronize do
@listen_subscriptions.filter_map do |request_id, subscription|
@listen_subscriptions.filter_map do |subscription_key, subscription|
# An inactive entry has not finished writing its acknowledgement yet;
# delivering to it would put a notification ahead of the acknowledgement.
next unless subscription[:active]
Expand All @@ -1067,12 +1077,12 @@ def deliver_to_listen_subscriptions(method, params)
uris.is_a?(Array) && uris.include?(uri)
end

[request_id, subscription] if hit
[subscription_key, subscription] if hit
end
end

matched.each do |request_id, subscription|
meta = { RequestEnvelope::SUBSCRIPTION_ID_META_KEY.to_sym => request_id }
matched.each do |subscription_key, subscription|
meta = { RequestEnvelope::SUBSCRIPTION_ID_META_KEY.to_sym => subscription[:request_id] }
notification_params = (params || {}).merge(_meta: meta)
notification = { jsonrpc: "2.0", method: method, params: notification_params }

Expand All @@ -1089,16 +1099,16 @@ def deliver_to_listen_subscriptions(method, params)
rescue *STREAM_WRITE_ERRORS => e
MCP.configuration.exception_reporter.call(
e,
{ subscription_id: request_id, error: "Failed to send notification" },
{ subscription_id: subscription[:request_id], error: "Failed to send notification" },
)
remove_listen_subscription(request_id)
remove_listen_subscription(subscription_key)
close_stream_safely(subscription[:stream])
end
end
end

def remove_listen_subscription(request_id)
@mutex.synchronize { @listen_subscriptions.delete(request_id) }
def remove_listen_subscription(subscription_key)
@mutex.synchronize { @listen_subscriptions.delete(subscription_key) }
end

# Graceful teardown (SEP-2575): each open listen stream receives its `SubscriptionsListenResult` response
Expand All @@ -1110,7 +1120,7 @@ def teardown_listen_subscriptions
subscriptions
end

removed.each do |request_id, subscription|
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.
Expand All @@ -1120,13 +1130,13 @@ def teardown_listen_subscriptions
begin
send_to_stream(subscription[:stream], {
jsonrpc: "2.0",
id: request_id,
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 => request_id },
_meta: { RequestEnvelope::SUBSCRIPTION_ID_META_KEY.to_sym => subscription[:request_id] },
},
})
rescue *STREAM_WRITE_ERRORS
Expand Down
65 changes: 53 additions & 12 deletions test/mcp/server/transports/streamable_http_transport_test.rb
Original file line number Diff line number Diff line change
Expand Up @@ -6113,7 +6113,7 @@ 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"] = {
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
}

@server.notify_tools_list_changed
Expand All @@ -6128,13 +6128,14 @@ def string

test "a delivery racing the graceful teardown cannot write after the final result" do
io = open_listen_stream(id: "listen-1", notifications: { toolsListChanged: true })
entry = @transport.instance_variable_get(:@listen_subscriptions)["listen-1"]
registry = @transport.instance_variable_get(:@listen_subscriptions)
subscription_key, entry = registry.first

@transport.close

# Simulate an in-flight delivery that snapshotted the entry before teardown cleared
# the registry: the closed flag set under the write mutex makes it a no-op.
@transport.instance_variable_get(:@listen_subscriptions)["listen-1"] = entry
registry[subscription_key] = entry
@server.notify_tools_list_changed

events = sse_events(io)
Expand Down Expand Up @@ -6257,16 +6258,54 @@ def string
transport.close
end

test "subscriptions/listen rejects a duplicate subscription id by closing the new stream" do
open_listen_stream(id: "listen-1", notifications: { toolsListChanged: true })
test "subscriptions/listen serves two streams that carry the same request id" do
# A request id is unique only among one client's own in-flight requests, and clients that number
# their requests with a counter reach the same small integers, so two clients' listen requests may
# carry the same id. Each gets its own stream, stamped with the id it sent.
first = open_listen_stream(id: 1, notifications: { toolsListChanged: true })
second = open_listen_stream(id: 1, notifications: { toolsListChanged: true })

duplicate = StringIO.new
response = @transport.handle_request(modern_rack_request(
modern_listen_body(id: "listen-1", params: { notifications: { toolsListChanged: true } }),
))
response[2].call(duplicate)
@server.notify_tools_list_changed

[first, second].each do |io|
refute_predicate io, :closed?
events = sse_events(io)
assert_equal ["notifications/subscriptions/acknowledged", "notifications/tools/list_changed"], events.map { |event| event["method"] }
assert_equal [1, 1], events.map { |event| event.dig("params", "_meta", "io.modelcontextprotocol/subscriptionId") }
end
end

test "transport close sends each of two streams sharing a request id its own result" do
first = open_listen_stream(id: 1, notifications: { toolsListChanged: true })
second = open_listen_stream(id: 1, notifications: { toolsListChanged: true })

@transport.close

assert_predicate duplicate, :closed?
[first, second].each do |io|
result = sse_events(io).last
assert_equal 1, result["id"]
assert_equal "complete", result.dig("result", "resultType")
assert_predicate io, :closed?
end
end

test "a failed write on one of two streams sharing a request id frees only that stream" do
first = open_listen_stream(id: 1, notifications: { toolsListChanged: true })
second = open_listen_stream(id: 1, notifications: { toolsListChanged: true })
first.define_singleton_method(:write) { |_data| raise Errno::EPIPE }

@server.notify_tools_list_changed

assert_predicate first, :closed?
refute_predicate second, :closed?
registry = @transport.instance_variable_get(:@listen_subscriptions)
assert_equal 1, registry.size
assert_same second, registry.values.first[:stream]

# The entry that survived is the second stream's, so a further notification still reaches it.
@server.notify_tools_list_changed

assert_equal 2, sse_events(second).count { |event| event["method"] == "notifications/tools/list_changed" }
end

test "listen keepalive writes a comment frame outside the mutex" do
Expand All @@ -6281,7 +6320,9 @@ def string
ping = data
end
stream.define_singleton_method(:flush) {}
@transport.instance_variable_get(:@listen_subscriptions)["listen-1"] = { stream: stream, filter: {}, write_mutex: Mutex.new }
@transport.instance_variable_get(:@listen_subscriptions)["listen-1"] = {
request_id: "listen-1", stream: stream, filter: {}, write_mutex: Mutex.new,
}

@transport.send(:send_listen_keepalive_ping, "listen-1")

Expand Down
Loading