Skip to content

Perform SSE stream writes outside the session mutex #449

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
koic merged 1 commit into from
Jul 11, 2026
+200 −77
Merged

Perform SSE stream writes outside the session mutex #449

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
183 changes: 106 additions & 77 deletions lib/mcp/server/transports/streamable_http_transport.rb
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -196,80 +196,107 @@ def send_notification(method, params = nil, session_id: nil, related_request_id:
}
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 = []

result = @mutex.synchronize do
if session_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.
next false if @enable_json_response && related_request_id

# Send to specific session
if (session = @sessions[session_id])
stream = active_stream(session, related_request_id: related_request_id)
end
next false unless stream
# Resolve the target stream under the lock, then write outside it: a stalled SSE reader
# must not block every other session that needs `@mutex`.
stream = @mutex.synchronize do
next unless (session = @sessions[session_id])

if session_expired?(session)
cleanup_and_collect_stream(session_id, streams_to_close)
next
end

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)
cleanup_and_collect_stream(session_id, streams_to_close)
next false
expired_session_ids << session_id
next
end

begin
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" },
)
if related_request_id && session[:post_request_streams]&.key?(related_request_id)
session[:post_request_streams].delete(related_request_id)
streams_to_close << stream
else
cleanup_and_collect_stream(session_id, streams_to_close)
end
false
end
else
# Broadcast to all connected SSE sessions
sent_count = 0
failed_sessions = []
[session_id, stream]
end

@sessions.each do |sid, session|
next unless (stream = session[:get_sse_stream])
expired_session_ids.each do |session_id|
cleanup_and_collect_stream(session_id, streams_to_close)
end

if session_expired?(session)
failed_sessions << sid
next
end
collected
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
close_streams(streams_to_close)

# Clean up failed sessions
failed_sessions.each { |sid| cleanup_and_collect_stream(sid, streams_to_close) }
targets.count do |session_id, stream|
write_notification(stream, notification, session_id, nil)
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 = []

sent_count
@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
cleanup_and_collect_stream(session_id, streams_to_close)
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)
end

result
end

# Sends a server-to-client JSON-RPC request (e.g., `sampling/createMessage`) and
Expand Down Expand Up @@ -299,27 +326,25 @@ def send_request(method, params = nil, session_id: nil, related_request_id: nil,
request = { jsonrpc: "2.0", id: request_id, method: method }
request[:params] = params if params

sent = false

@mutex.synchronize do
# 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`.
stream = @mutex.synchronize do
unless (session = @sessions[session_id])
raise "Session not found: #{session_id}."
end

@pending_responses[request_id] = { queue: queue, session_id: session_id }

if (stream = active_stream(session, related_request_id: related_request_id))
begin
send_to_stream(stream, request)
sent = true
rescue *STREAM_WRITE_ERRORS
if related_request_id && session[:post_request_streams]&.key?(related_request_id)
session[:post_request_streams].delete(related_request_id)
close_stream_safely(stream)
else
cleanup_session_unsafe(session_id)
end
end
active_stream(session, related_request_id: related_request_id)
end

sent = false
if stream
begin
send_to_stream(stream, request)
sent = true
rescue *STREAM_WRITE_ERRORS
drop_broken_stream(session_id, stream, related_request_id)
end
end

Expand Down Expand Up @@ -1164,11 +1189,15 @@ def session_active_with_stream?(session_id)
end

def send_keepalive_ping(session_id)
@mutex.synchronize do
if @sessions[session_id] && @sessions[session_id][:get_sse_stream]
send_ping_to_stream(@sessions[session_id][:get_sse_stream])
end
# Resolve the stream under the lock, then write outside it so a stalled reader
# cannot block every other session on `@mutex`.
stream = @mutex.synchronize do
session = @sessions[session_id]
session && session[:get_sse_stream]
end
return unless stream

send_ping_to_stream(stream)
rescue *STREAM_WRITE_ERRORS => e
MCP.configuration.exception_reporter.call(
e,
Expand Down
94 changes: 94 additions & 0 deletions test/mcp/server/transports/streamable_http_transport_test.rb
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1335,6 +1335,62 @@ def string
assert @transport.instance_variable_get(:@sessions).key?(session_id2)
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
# Create and initialize a session.
init_request = create_rack_request(
Expand Down Expand Up @@ -5035,6 +5091,44 @@ def string

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)
default_accept = case method
when "POST"
Expand Down