Perform SSE stream writes outside the session mutex

## Motivation and Context

`StreamableHTTPTransport` guards its `@sessions` and `@pending_responses` maps with one `@mutex`.
The send paths wrote to an SSE stream while holding that lock: the targeted and
broadcast branches of `send_notification`, the server-to-client `send_request`,
and `send_keepalive_ping` all called `send_to_stream`/`send_ping_to_stream`
(a `stream.write` plus `stream.flush`) inside `@mutex.synchronize`.

A stream write is blocking I/O with no timeout: if a client reads its SSE stream slowly,
the server-side `write` waits until the socket buffer drains.
Performing that write under `@mutex` serializes every other session on the same lock (head-of-line blocking):
new `initialize` inserts, session validation, `handle_response`, the reaper, and delivery to other sessions
all wait behind one slow reader. The broadcast path holds the lock across every session,
so one slow stream also delays delivery to the sessions after it.

The cleanup paths already release the lock before touching a stream: they collect `streams_to_close`
and call `close` outside `synchronize`. This change applies the same discipline to the writes.

- Each send path resolves its target stream under `@mutex` (and, for `send_request`, registers
  the pending response there), then releases the lock before writing.
- A write error still drops the broken stream. That teardown is centralized in a new `drop_broken_stream`,
  which mutates `@sessions` under `@mutex` and closes the affected streams outside it.
  A request-scoped stream is dropped on its own; a GET SSE failure removes the whole session,
  matching the previous behavior.
- `send_notification` splits into `deliver_targeted_notification` and `deliver_broadcast_notification`;
  the return values (true/false for a session, the delivered count for a broadcast) are unchanged.

## How Has This Been Tested?

New tests in `test/mcp/server/transports/streamable_http_transport_test.rb` install a probe stream that records,
via `mutex.try_lock`, whether `@mutex` is held during each write, and assert it is not held for the targeted notification,
broadcast, `send_request`, and keepalive paths. Each new test fails against the previous code and passes now.
The existing send and cleanup tests still pass.

## Breaking Changes

None. Delivery semantics and return values are unchanged; the lock is simply no longer held while writing to a stream.
This commit is contained in:
Koichi ITO 2026-07-11 17:14:34 +09:00
commit 5b935eb031
2 changed files with 200 additions and 77 deletions

View file

@ -196,80 +196,107 @@ module MCP
} }
notification[:params] = params if params notification[:params] = params if params
if session_id
deliver_targeted_notification(notification, session_id, related_request_id)
else
deliver_broadcast_notification(notification)
end
end
def deliver_targeted_notification(notification, session_id, related_request_id)
# JSON response mode returns a single JSON object as the POST response,
# so request-scoped notifications (e.g. progress, log) cannot be delivered
# alongside it. Session-scoped standalone notifications
# (e.g. `resources/updated`, `elicitation/complete`) still flow via GET SSE.
return false if @enable_json_response && related_request_id
streams_to_close = [] streams_to_close = []
result = @mutex.synchronize do # Resolve the target stream under the lock, then write outside it: a stalled SSE reader
if session_id # must not block every other session that needs `@mutex`.
# JSON response mode returns a single JSON object as the POST response, stream = @mutex.synchronize do
# so request-scoped notifications (e.g. progress, log) cannot be delivered next unless (session = @sessions[session_id])
# alongside it. Session-scoped standalone notifications
# (e.g. `resources/updated`, `elicitation/complete`) still flow via GET SSE.
next false if @enable_json_response && related_request_id
# Send to specific session if session_expired?(session)
if (session = @sessions[session_id]) cleanup_and_collect_stream(session_id, streams_to_close)
stream = active_stream(session, related_request_id: related_request_id) next
end end
next false unless stream
active_stream(session, related_request_id: related_request_id)
end
close_streams(streams_to_close)
return false unless stream
write_notification(stream, notification, session_id, related_request_id)
end
def deliver_broadcast_notification(notification)
streams_to_close = []
# Snapshot the connected streams under the lock, then write to each outside it.
targets = @mutex.synchronize do
expired_session_ids = []
collected = @sessions.filter_map do |session_id, session|
next unless (stream = session[:get_sse_stream])
if session_expired?(session) if session_expired?(session)
cleanup_and_collect_stream(session_id, streams_to_close) expired_session_ids << session_id
next false next
end end
begin [session_id, stream]
send_to_stream(stream, notification) end
true
rescue *STREAM_WRITE_ERRORS => e expired_session_ids.each do |session_id|
MCP.configuration.exception_reporter.call( cleanup_and_collect_stream(session_id, streams_to_close)
e, end
{ session_id: session_id, error: "Failed to send notification" },
) collected
if related_request_id && session[:post_request_streams]&.key?(related_request_id) end
session[:post_request_streams].delete(related_request_id)
streams_to_close << stream close_streams(streams_to_close)
else
cleanup_and_collect_stream(session_id, streams_to_close) targets.count do |session_id, stream|
end write_notification(stream, notification, session_id, nil)
false end
end end
# Writes a notification to an SSE stream without holding `@mutex`. On a write error,
# drops the broken stream and returns false; on success returns true.
def write_notification(stream, notification, session_id, related_request_id)
send_to_stream(stream, notification)
true
rescue *STREAM_WRITE_ERRORS => e
MCP.configuration.exception_reporter.call(e, { session_id: session_id, error: "Failed to send notification" })
drop_broken_stream(session_id, stream, related_request_id)
false
end
# Removes a stream that failed to accept a write. A request-scoped stream is dropped on its own;
# a session-scoped (GET SSE) failure tears down the whole session. The `@sessions` mutation runs
# under `@mutex`, and the affected streams are closed outside it.
def drop_broken_stream(session_id, stream, related_request_id)
streams_to_close = []
@mutex.synchronize do
session = @sessions[session_id]
if related_request_id && session&.dig(:post_request_streams, related_request_id)
session[:post_request_streams].delete(related_request_id)
streams_to_close << stream
else else
# Broadcast to all connected SSE sessions cleanup_and_collect_stream(session_id, streams_to_close)
sent_count = 0
failed_sessions = []
@sessions.each do |sid, session|
next unless (stream = session[:get_sse_stream])
if session_expired?(session)
failed_sessions << sid
next
end
begin
send_to_stream(stream, notification)
sent_count += 1
rescue *STREAM_WRITE_ERRORS => e
MCP.configuration.exception_reporter.call(
e,
{ session_id: sid, error: "Failed to send notification" },
)
failed_sessions << sid
end
end
# Clean up failed sessions
failed_sessions.each { |sid| cleanup_and_collect_stream(sid, streams_to_close) }
sent_count
end end
end end
streams_to_close.each do |stream| close_streams(streams_to_close)
end
def close_streams(streams)
streams.each do |stream|
close_stream_safely(stream) close_stream_safely(stream)
end end
result
end end
# Sends a server-to-client JSON-RPC request (e.g., `sampling/createMessage`) and # Sends a server-to-client JSON-RPC request (e.g., `sampling/createMessage`) and
@ -299,27 +326,25 @@ module MCP
request = { jsonrpc: "2.0", id: request_id, method: method } request = { jsonrpc: "2.0", id: request_id, method: method }
request[:params] = params if params request[:params] = params if params
sent = false # Register the pending response and resolve the stream under the lock, but perform
# the write outside it so a stalled reader cannot block every other session on `@mutex`.
@mutex.synchronize do stream = @mutex.synchronize do
unless (session = @sessions[session_id]) unless (session = @sessions[session_id])
raise "Session not found: #{session_id}." raise "Session not found: #{session_id}."
end end
@pending_responses[request_id] = { queue: queue, session_id: session_id } @pending_responses[request_id] = { queue: queue, session_id: session_id }
if (stream = active_stream(session, related_request_id: related_request_id)) active_stream(session, related_request_id: related_request_id)
begin end
send_to_stream(stream, request)
sent = true sent = false
rescue *STREAM_WRITE_ERRORS if stream
if related_request_id && session[:post_request_streams]&.key?(related_request_id) begin
session[:post_request_streams].delete(related_request_id) send_to_stream(stream, request)
close_stream_safely(stream) sent = true
else rescue *STREAM_WRITE_ERRORS
cleanup_session_unsafe(session_id) drop_broken_stream(session_id, stream, related_request_id)
end
end
end end
end end
@ -1164,11 +1189,15 @@ module MCP
end end
def send_keepalive_ping(session_id) def send_keepalive_ping(session_id)
@mutex.synchronize do # Resolve the stream under the lock, then write outside it so a stalled reader
if @sessions[session_id] && @sessions[session_id][:get_sse_stream] # cannot block every other session on `@mutex`.
send_ping_to_stream(@sessions[session_id][:get_sse_stream]) stream = @mutex.synchronize do
end session = @sessions[session_id]
session && session[:get_sse_stream]
end end
return unless stream
send_ping_to_stream(stream)
rescue *STREAM_WRITE_ERRORS => e rescue *STREAM_WRITE_ERRORS => e
MCP.configuration.exception_reporter.call( MCP.configuration.exception_reporter.call(
e, e,

View file

@ -1335,6 +1335,62 @@ module MCP
assert @transport.instance_variable_get(:@sessions).key?(session_id2) assert @transport.instance_variable_get(:@sessions).key?(session_id2)
end end
test "send_notification to a session writes outside the mutex" do
session_id = initialize_test_session
writes = install_mutex_probe_stream(session_id)
result = @transport.send_notification("test", { message: "hi" }, session_id: session_id)
assert result
assert_equal 1, writes[:total]
assert_equal 0, writes[:under_mutex], "notification write must not hold @mutex"
end
test "send_notification broadcast writes outside the mutex" do
session_id1 = initialize_test_session(id: "1")
session_id2 = initialize_test_session(id: "2")
writes1 = install_mutex_probe_stream(session_id1)
writes2 = install_mutex_probe_stream(session_id2)
sent_count = @transport.send_notification("broadcast", { message: "hi" }, **{})
assert_equal 2, sent_count
assert_equal 0, writes1[:under_mutex], "broadcast write must not hold @mutex"
assert_equal 0, writes2[:under_mutex], "broadcast write must not hold @mutex"
end
test "send_keepalive_ping writes outside the mutex" do
session_id = initialize_test_session
writes = install_mutex_probe_stream(session_id)
@transport.send(:send_keepalive_ping, session_id)
assert_equal 1, writes[:total]
assert_equal 0, writes[:under_mutex], "keepalive write must not hold @mutex"
end
test "send_request writes outside the mutex" do
session_id = initialize_test_session
writes = install_mutex_probe_stream(session_id)
result_queue = Queue.new
thread = Thread.new do
result_queue.push(@transport.send_request("ping", nil, session_id: session_id))
end
# Wait until the request has been written to the probe stream.
Thread.pass until writes[:total].positive?
request_id = @transport.instance_variable_get(:@pending_responses).keys.first
@transport.send(:handle_response, { jsonrpc: "2.0", id: request_id, result: {} }, session_id: session_id)
result_queue.pop
thread.join
assert_equal 1, writes[:total]
assert_equal 0, writes[:under_mutex], "server-to-client request write must not hold @mutex"
end
test "send_keepalive_ping handles Errno::ECONNRESET gracefully" do test "send_keepalive_ping handles Errno::ECONNRESET gracefully" do
# Create and initialize a session. # Create and initialize a session.
init_request = create_rack_request( init_request = create_rack_request(
@ -5035,6 +5091,44 @@ module MCP
private private
def initialize_test_session(id: "init")
init_request = create_rack_request(
"POST",
"/",
{ "CONTENT_TYPE" => "application/json" },
{ jsonrpc: "2.0", method: "initialize", id: id }.to_json,
)
@transport.handle_request(init_request)[1]["Mcp-Session-Id"]
end
# Installs a stream that records, on every write, whether `@mutex` was held at that moment.
# `mutex.try_lock` fails while the transport still holds the lock, so a nonzero `writes[:under_mutex]`
# means a blocking write is being performed under the lock.
def install_mutex_probe_stream(session_id, related_request_id: nil)
mutex = @transport.instance_variable_get(:@mutex)
writes = { total: 0, under_mutex: 0 }
stream = Object.new
stream.define_singleton_method(:write) do |_data|
writes[:total] += 1
if mutex.try_lock
mutex.unlock
else
writes[:under_mutex] += 1
end
end
stream.define_singleton_method(:flush) {}
stream.define_singleton_method(:close) {}
session = @transport.instance_variable_get(:@sessions)[session_id]
if related_request_id
(session[:post_request_streams] ||= {})[related_request_id] = stream
else
session[:get_sse_stream] = stream
end
writes
end
def create_rack_request(method, path, headers, body = nil) def create_rack_request(method, path, headers, body = nil)
default_accept = case method default_accept = case method
when "POST" when "POST"