Merge pull request #535 from koic/serialize_listen_stream_writes_on_teardown

Serialize `subscriptions/listen` writes so the graceful result is the final message
This commit is contained in:
Koichi ITO 2026-08-27 01:12:33 +09:00 • committed by GitHub
commit 6b6b641655
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 55 additions and 22 deletions

View file

@ -169,7 +169,8 @@ module MCP
@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 }` (SEP-2575).
# Maps a `subscriptions/listen` request id to
# `{ stream: stream_object, filter: honored_subscription_filter, active: boolean, write_mutex: Mutex }` (SEP-2575).
# 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 = {}
@ -936,7 +937,7 @@ module MCP
(@max_listen_subscriptions && @listen_subscriptions.size >= @max_listen_subscriptions)
rejected = true
else
@listen_subscriptions[request_id] = { stream: stream, filter: honored, active: false }
@listen_subscriptions[request_id] = { stream: stream, filter: honored, active: false, write_mutex: Mutex.new }
end
end
@ -1066,24 +1067,32 @@ module MCP
uris.is_a?(Array) && uris.include?(uri)
end
[request_id, subscription[:stream]] if hit
[request_id, subscription] if hit
end
end
matched.each do |request_id, stream|
matched.each do |request_id, subscription|
meta = { RequestEnvelope::SUBSCRIPTION_ID_META_KEY.to_sym => request_id }
notification_params = (params || {}).merge(_meta: meta)
notification = { jsonrpc: "2.0", method: method, params: notification_params }
begin
send_to_stream(stream, notification)
# The per-stream write mutex orders this write against a concurrent graceful teardown:
# once teardown has marked the entry closed and written its `SubscriptionsListenResult`,
# a delivery that snapshotted the entry before the registry was cleared skips it instead
# of writing after the final message.
subscription[:write_mutex].synchronize do
next if subscription[:closed]
send_to_stream(subscription[:stream], notification)
end
rescue *STREAM_WRITE_ERRORS => e
MCP.configuration.exception_reporter.call(
e,
{ subscription_id: request_id, error: "Failed to send notification" },
)
remove_listen_subscription(request_id)
close_stream_safely(stream)
close_stream_safely(subscription[:stream])
end
end
end
@ -1102,20 +1111,27 @@ module MCP
end
removed.each do |request_id, subscription|
begin
send_to_stream(subscription[:stream], {
jsonrpc: "2.0",
id: 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 },
},
})
rescue *STREAM_WRITE_ERRORS
nil
# 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.
subscription[:write_mutex].synchronize do
subscription[:closed] = true
begin
send_to_stream(subscription[:stream], {
jsonrpc: "2.0",
id: 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 },
},
})
rescue *STREAM_WRITE_ERRORS
nil
end
end
close_stream_safely(subscription[:stream])
end

View file

@ -6113,7 +6113,7 @@ module MCP
# 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
stream: io, filter: { toolsListChanged: true }, active: false, write_mutex: Mutex.new
}
@server.notify_tools_list_changed
@ -6126,6 +6126,23 @@ module MCP
assert_equal ["notifications/tools/list_changed"], sse_events(io).map { |event| event["method"] }
end
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"]
@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
@server.notify_tools_list_changed
events = sse_events(io)
assert_equal "complete", events.last.dig("result", "resultType")
refute(events.any? { |event| event["method"] == "notifications/tools/list_changed" })
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 })
@ -6264,7 +6281,7 @@ module MCP
ping = data
end
stream.define_singleton_method(:flush) {}
@transport.instance_variable_get(:@listen_subscriptions)["listen-1"] = { stream: stream, filter: {} }
@transport.instance_variable_get(:@listen_subscriptions)["listen-1"] = { stream: stream, filter: {}, write_mutex: Mutex.new }
@transport.send(:send_listen_keepalive_ping, "listen-1")