Skip to content

Commit 0cef024

Browse files
committed
Mark a removed subscriptions/listen stream closed
## Motivation and Context Removing a listen stream after a failed write or a finished keepalive did not mark its entry closed, as transport close does, so a delivery or keepalive that had taken the entry before the removal still wrote to the stream being closed. That write failed against the closing stream and was reported as a failed delivery. Removal now marks the entry closed, and those writes skip it under the stream's write mutex. The flag is set without taking the write mutex, which must never be taken inside the registry lock; the close that follows a removal takes the write mutex outside the registry lock, so a write already past its check lands before the stream closes, and every later one sees the flag. ## How Has This Been Tested? A new test in `test/mcp/server/transports/streamable_http_transport_test.rb` delivers to an entry taken before its removal and checks that nothing is written. Against the previous library the notification is written to the removed stream. The keepalive failure test also checks that the entry a failed ping removes is marked closed. Another blocks a delivery inside its write, removes the entry from a second thread, and checks that the stream is closed only after the write has completed. ## Breaking Changes None.
1 parent 402cb56 commit 0cef024

2 files changed

Lines changed: 100 additions & 11 deletions

File tree

‎lib/mcp/server/transports/streamable_http_transport.rb‎

Lines changed: 22 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1093,11 +1093,11 @@ def listen_sse_body(request_id, honored, acknowledgement)
10931093
if acknowledged
10941094
start_listen_keepalive_thread(subscription_key, request_id)
10951095
else
1096-
close_stream_safely(stream)
1096+
close_removed_listen_stream(subscription)
10971097
end
10981098
rescue *STREAM_WRITE_ERRORS
10991099
remove_listen_subscription(subscription_key)
1100-
close_stream_safely(stream)
1100+
close_removed_listen_stream(subscription)
11011101
end
11021102
end
11031103
end
@@ -1138,11 +1138,9 @@ def start_listen_keepalive_thread(subscription_key, request_id)
11381138
# removed the entry already, and the report should still name the stream.
11391139
MCP.configuration.exception_reporter.call(e, { subscription_id: request_id })
11401140
ensure
1141-
stream = @mutex.synchronize do
1142-
subscription = @listen_subscriptions.delete(subscription_key)
1143-
subscription && subscription[:stream]
1144-
end
1145-
close_stream_safely(stream) if stream
1141+
# Through the helper, so the entry is marked closed here as it is for every other removal.
1142+
subscription = remove_listen_subscription(subscription_key)
1143+
close_removed_listen_stream(subscription) if subscription
11461144
end
11471145
end
11481146

@@ -1254,20 +1252,35 @@ def deliver_to_listen_subscriptions(method, params)
12541252
{ subscription_id: subscription[:request_id], error: "Failed to send notification" },
12551253
)
12561254
remove_listen_subscription(subscription_key)
1257-
close_stream_safely(subscription[:stream])
1255+
close_removed_listen_stream(subscription)
12581256
end
12591257
end
12601258
end
12611259

12621260
def remove_listen_subscription(subscription_key)
12631261
@mutex.synchronize do
12641262
subscription = @listen_subscriptions.delete(subscription_key)
1265-
subscription[:keepalive_wakeup].signal if subscription
1263+
if subscription
1264+
# Marked closed as teardown does, so a delivery or keepalive that took the entry before the removal skips it under
1265+
# the write mutex instead of writing to a stream being closed. The flag is set without taking the write mutex,
1266+
# which must never be taken inside `@mutex`: a write that checks the flag from here on skips, and one already
1267+
# past its check lands before the stream closes, since `close_removed_listen_stream` waits for the write mutex.
1268+
subscription[:closed] = true
1269+
subscription[:keepalive_wakeup].signal
1270+
end
12661271

12671272
subscription
12681273
end
12691274
end
12701275

1276+
# Closes a subscription's stream while excluding concurrent writes. The write mutex, taken here outside
1277+
# `@mutex`, makes the close wait for a write already past its `closed` check, so that write completes or fails
1278+
# before the stream is closed, and a write that takes the mutex afterwards sees `closed` and skips. Teardown
1279+
# gets the same ordering by setting `closed` under the write mutex itself.
1280+
def close_removed_listen_stream(subscription)
1281+
subscription[:write_mutex].synchronize { close_stream_safely(subscription[:stream]) }
1282+
end
1283+
12711284
# Graceful teardown (SEP-2575): each open listen stream receives its `SubscriptionsListenResult` response
12721285
# before the stream closes.
12731286
def teardown_listen_subscriptions

‎test/mcp/server/transports/streamable_http_transport_test.rb‎

Lines changed: 78 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -6876,6 +6876,76 @@ def wait_until_blocked_or_done(thread)
68766876
assert_equal 2, sse_events(second).count { |event| event["method"] == "notifications/tools/list_changed" }
68776877
end
68786878

6879+
test "a delivery that took an entry before its removal does not write to the removed stream" do
6880+
io = open_listen_stream(id: "listen-1", notifications: { toolsListChanged: true })
6881+
registry = @transport.instance_variable_get(:@listen_subscriptions)
6882+
subscription_key, entry = registry.first
6883+
written = io.string.dup
6884+
6885+
# Holding the stream's write mutex lets the delivery take its snapshot under the registry lock and then
6886+
# block ahead of its write, which is where a removal can land between the two.
6887+
entry[:write_mutex].lock
6888+
delivery = Thread.new { @server.notify_tools_list_changed }
6889+
6890+
begin
6891+
wait_until_blocked_or_done(delivery)
6892+
@transport.send(:remove_listen_subscription, subscription_key)
6893+
ensure
6894+
entry[:write_mutex].unlock
6895+
end
6896+
6897+
refute_nil delivery.join(5)
6898+
6899+
assert entry[:closed]
6900+
assert_equal written, io.string
6901+
assert_empty registry
6902+
ensure
6903+
delivery&.join(5)
6904+
end
6905+
6906+
test "closing a removed stream waits for a delivery already past its closed check" do
6907+
io = open_listen_stream(id: "listen-1", notifications: { toolsListChanged: true })
6908+
registry = @transport.instance_variable_get(:@listen_subscriptions)
6909+
subscription_key, entry = registry.first
6910+
gate = Queue.new
6911+
entered_write = false
6912+
open_at_write = nil
6913+
# The write blocks once the delivery is inside it: past the `closed` check and holding the write mutex,
6914+
# which is the window a removal and the close that follows it used to race into.
6915+
io.define_singleton_method(:write) do |data|
6916+
entered_write = true
6917+
gate.pop
6918+
open_at_write = !closed?
6919+
super(data)
6920+
end
6921+
delivery = Thread.new { @server.notify_tools_list_changed }
6922+
deadline = Process.clock_gettime(Process::CLOCK_MONOTONIC) + 5
6923+
until entered_write
6924+
flunk("the delivery did not reach the write") if Process.clock_gettime(Process::CLOCK_MONOTONIC) > deadline
6925+
sleep(0.005)
6926+
end
6927+
6928+
closer = Thread.new do
6929+
@transport.send(:close_removed_listen_stream, @transport.send(:remove_listen_subscription, subscription_key))
6930+
end
6931+
wait_until_blocked_or_done(closer)
6932+
6933+
assert entry[:closed]
6934+
assert_empty registry
6935+
refute_predicate io, :closed?, "the stream was closed under a write still in progress"
6936+
6937+
gate << true
6938+
refute_nil delivery.join(5)
6939+
refute_nil closer.join(5)
6940+
assert open_at_write, "the write ran against a closed stream"
6941+
assert_predicate io, :closed?
6942+
assert_includes io.string, "notifications/tools/list_changed"
6943+
ensure
6944+
gate&.close
6945+
delivery&.join(5)
6946+
closer&.join(5)
6947+
end
6948+
68796949
test "listen keepalive writes a comment frame outside the mutex" do
68806950
# Register a subscription backed by a mutex-probing stream, then invoke the ping directly.
68816951
mutex = @transport.instance_variable_get(:@mutex)
@@ -6902,15 +6972,21 @@ def wait_until_blocked_or_done(thread)
69026972
test "listen keepalive frees the slot and closes the stream when the peer is gone" do
69036973
transport = StreamableHTTPTransport.new(@server, listen_keepalive_interval: 0.01)
69046974
io = open_listen_stream(id: "listen-1", notifications: { toolsListChanged: true }, transport: transport)
6975+
entry = transport.instance_variable_get(:@listen_subscriptions).values.first
69056976
io.define_singleton_method(:write) { |_data| raise Errno::ECONNRESET }
69066977

6978+
# The keepalive thread frees the slot first and closes the stream second, so wait for both.
69076979
deadline = Process.clock_gettime(Process::CLOCK_MONOTONIC) + 2
6908-
until transport.instance_variable_get(:@listen_subscriptions).empty?
6909-
flunk("keepalive did not free the slot for a dead peer") if Process.clock_gettime(Process::CLOCK_MONOTONIC) > deadline
6980+
until transport.instance_variable_get(:@listen_subscriptions).empty? && io.closed?
6981+
flunk("keepalive did not free the slot and close the stream for a dead peer") if Process.clock_gettime(Process::CLOCK_MONOTONIC) > deadline
69106982
sleep(0.01)
69116983
end
69126984

69136985
assert_predicate io, :closed?
6986+
6987+
# The removal a failed ping performs goes through the same helper as every other, so a delivery that took the entry before
6988+
# the ping failed skips it instead of writing to the closed stream.
6989+
assert entry[:closed], "the entry a failed ping removed was not marked closed"
69146990
ensure
69156991
transport.close
69166992
end

0 commit comments

Comments
 (0)