diff --git a/lib/mcp/server/transports/streamable_http_transport.rb b/lib/mcp/server/transports/streamable_http_transport.rb index 2cedd643..fd45b332 100644 --- a/lib/mcp/server/transports/streamable_http_transport.rb +++ b/lib/mcp/server/transports/streamable_http_transport.rb @@ -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 @@ -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 @@ -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, diff --git a/test/mcp/server/transports/streamable_http_transport_test.rb b/test/mcp/server/transports/streamable_http_transport_test.rb index 8bf02ede..0bf51f1d 100644 --- a/test/mcp/server/transports/streamable_http_transport_test.rb +++ b/test/mcp/server/transports/streamable_http_transport_test.rb @@ -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( @@ -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"