mcp 0.10.0 → 0.11.0

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

@@ -1,13 +1,18 @@
 # frozen_string_literal: true
 
 require "json"
-require "securerandom"
 require_relative "../../transport"
 
 module MCP
   class Server
     module Transports
       class StreamableHTTPTransport < Transport
+        SSE_HEADERS = {
+          "Content-Type" => "text/event-stream",
+          "Cache-Control" => "no-cache",
+          "Connection" => "keep-alive",
+        }.freeze
+
         def initialize(server, stateless: false, session_idle_timeout: nil)
           super(server)
           # Maps `session_id` to `{ stream: stream_object, server_session: ServerSession, last_active_at: float_from_monotonic_clock }`.
@@ -16,6 +21,7 @@ module MCP
 
           @stateless = stateless
           @session_idle_timeout = session_idle_timeout
+          @pending_responses = {}
 
           if @session_idle_timeout
             if @stateless
@@ -50,12 +56,17 @@ module MCP
           @reaper_thread&.kill
           @reaper_thread = nil
 
-          @mutex.synchronize do
-            @sessions.each_key { |session_id| cleanup_session_unsafe(session_id) }
+          removed_sessions = @mutex.synchronize do
+            @sessions.each_key.filter_map { |session_id| cleanup_session_unsafe(session_id) }
+          end
+
+          removed_sessions.each do |session|
+            close_stream_safely(session[:stream])
+            close_post_request_streams(session)
           end
         end
 
-        def send_notification(method, params = nil, session_id: nil)
+        def send_notification(method, params = nil, session_id: nil, related_request_id: nil)
           # Stateless mode doesn't support notifications
           raise "Stateless mode does not support notifications" if @stateless
 
@@ -65,26 +76,35 @@ module MCP
           }
           notification[:params] = params if params
 
-          @mutex.synchronize do
+          streams_to_close = []
+
+          result = @mutex.synchronize do
             if session_id
               # Send to specific session
-              session = @sessions[session_id]
-              return false unless session && session[:stream]
+              if (session = @sessions[session_id])
+                stream = active_stream(session, related_request_id: related_request_id)
+              end
+              next false unless stream
 
               if session_expired?(session)
-                cleanup_session_unsafe(session_id)
-                return false
+                cleanup_and_collect_stream(session_id, streams_to_close)
+                next false
               end
 
               begin
-                send_to_stream(session[:stream], notification)
+                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" },
                 )
-                cleanup_session_unsafe(session_id)
+                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
@@ -93,7 +113,7 @@ module MCP
               failed_sessions = []
 
               @sessions.each do |sid, session|
-                next unless session[:stream]
+                next unless (stream = session[:stream])
 
                 if session_expired?(session)
                   failed_sessions << sid
@@ -101,7 +121,7 @@ module MCP
                 end
 
                 begin
-                  send_to_stream(session[:stream], notification)
+                  send_to_stream(stream, notification)
                   sent_count += 1
                 rescue *STREAM_WRITE_ERRORS => e
                   MCP.configuration.exception_reporter.call(
@@ -113,11 +133,93 @@ module MCP
               end
 
               # Clean up failed sessions
-              failed_sessions.each { |sid| cleanup_session_unsafe(sid) }
+              failed_sessions.each { |sid| cleanup_and_collect_stream(sid, streams_to_close) }
 
               sent_count
             end
           end
+
+          streams_to_close.each do |stream|
+            close_stream_safely(stream)
+          end
+
+          result
+        end
+
+        # Sends a server-to-client JSON-RPC request (e.g., `sampling/createMessage`) and
+        # blocks until the client responds.
+        #
+        # Uses a `Queue` for cross-thread synchronization. This method creates a `Queue`,
+        # sends the request via SSE stream, then blocks on `queue.pop`.
+        # When the client POSTs a response, `handle_response` matches it by `request_id`
+        # and pushes the result onto the queue, unblocking this thread.
+        def send_request(method, params = nil, session_id: nil, related_request_id: nil)
+          if @stateless
+            raise "Stateless mode does not support server-to-client requests."
+          end
+
+          unless session_id
+            raise "session_id is required for server-to-client requests."
+          end
+
+          request_id = generate_request_id
+          queue = Queue.new
+
+          request = { jsonrpc: "2.0", id: request_id, method: method }
+          request[:params] = params if params
+
+          sent = false
+
+          @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
+            end
+          end
+
+          # TODO: Replace with event store + replay when resumability is implemented.
+          # Resumability is a separate MCP specification feature (SSE event IDs, Last-Event-ID replay,
+          # event store management) independent of sampling.
+          # See: https://modelcontextprotocol.io/specification/latest/basic/transports#resumability-and-redelivery
+          #
+          # The TypeScript and Python SDKs buffer messages and replay on reconnect.
+          # Until then, raise to prevent queue.pop from blocking indefinitely.
+          unless sent
+            raise "No active stream for #{method} request."
+          end
+
+          response = queue.pop
+
+          if response.is_a?(Hash) && response.key?(:error)
+            raise StandardError, "Client returned an error for #{method} request (code: #{response[:error][:code]}): #{response[:error][:message]}"
+          end
+
+          if response == :session_closed
+            raise "SSE session closed while waiting for #{method} response."
+          end
+
+          response
+        ensure
+          if request_id
+            @mutex.synchronize do
+              @pending_responses.delete(request_id)
+            end
+          end
         end
 
         private
@@ -136,22 +238,17 @@ module MCP
         def reap_expired_sessions
           return unless @session_idle_timeout
 
-          expired_streams = @mutex.synchronize do
-            @sessions.each_with_object([]) do |(session_id, session), streams|
-              next unless session_expired?(session)
+          removed_sessions = @mutex.synchronize do
+            @sessions.each_key.filter_map do |session_id|
+              next unless session_expired?(@sessions[session_id])
 
-              streams << session[:stream] if session[:stream]
-              @sessions.delete(session_id)
+              cleanup_session_unsafe(session_id)
             end
           end
 
-          expired_streams.each do |stream|
-            # Closing outside the mutex is safe because expired sessions are already
-            # removed from `@sessions` above, so other threads will not find them
-            # and will not attempt to close the same stream.
-            stream.close
-          rescue
-            nil
+          removed_sessions.each do |session|
+            close_stream_safely(session[:stream])
+            close_post_request_streams(session)
           end
         end
 
@@ -176,15 +273,19 @@ module MCP
           body = parse_request_body(body_string)
           return body unless body.is_a?(Hash) # Error response
 
-          if body["method"] == "initialize"
+          if body[:method] == "initialize"
             handle_initialization(body_string, body)
           else
             return missing_session_id_response if !@stateless && !session_id
 
-            if notification?(body) || response?(body)
+            if notification?(body)
               handle_accepted
+            elsif response?(body)
+              return session_not_found_response if !@stateless && !session_exists?(session_id)
+
+              handle_response(body, session_id: session_id)
             else
-              handle_regular_request(body_string, session_id)
+              handle_regular_request(body_string, session_id, related_request_id: body[:id])
             end
           end
         rescue StandardError => e
@@ -228,21 +329,51 @@ module MCP
         end
 
         def cleanup_session(session_id)
-          @mutex.synchronize do
+          session = @mutex.synchronize do
             cleanup_session_unsafe(session_id)
           end
+
+          if session
+            close_stream_safely(session[:stream])
+            close_post_request_streams(session)
+          end
         end
 
+        # Removes a session from `@sessions` and returns it. Does not close the stream.
+        # Callers must close the stream outside the mutex to avoid holding the lock during
+        # potentially blocking I/O.
         def cleanup_session_unsafe(session_id)
-          session = @sessions[session_id]
-          return unless session
+          session = @sessions.delete(session_id)
 
-          begin
-            session[:stream]&.close
-          rescue
-            nil
+          # Unblock threads waiting on pending responses for this session.
+          @pending_responses.each_value do |pending_response|
+            if pending_response[:session_id] == session_id
+              pending_response[:queue].push(:session_closed)
+            end
+          end
+
+          session
+        end
+
+        def cleanup_and_collect_stream(session_id, streams_to_close)
+          return unless (removed = cleanup_session_unsafe(session_id))
+
+          streams_to_close << removed[:stream]
+          removed[:post_request_streams]&.each_value { |stream| streams_to_close << stream }
+        end
+
+        def close_stream_safely(stream)
+          stream&.close
+        rescue StandardError
+          # Ignore close-related errors from already closed/broken streams.
+        end
+
+        def close_post_request_streams(session)
+          return unless (post_request_streams = session[:post_request_streams])
+
+          post_request_streams.each_value do |stream|
+            close_stream_safely(stream)
           end
-          @sessions.delete(session_id)
         end
 
         def extract_session_id(request)
@@ -277,17 +408,35 @@ module MCP
         end
 
         def parse_request_body(body_string)
-          JSON.parse(body_string)
+          JSON.parse(body_string, symbolize_names: true)
         rescue JSON::ParserError, TypeError
           [400, { "Content-Type" => "application/json" }, [{ error: "Invalid JSON" }.to_json]]
         end
 
         def notification?(body)
-          !body["id"] && !!body["method"]
+          !body[:id] && !!body[:method]
         end
 
         def response?(body)
-          !!body["id"] && !body["method"]
+          !!body[:id] && !body[:method]
+        end
+
+        # Verifies that the response came from the expected session to prevent
+        # cross-session response injection if request IDs are ever leaked.
+        def handle_response(body, session_id:)
+          request_id = body[:id]
+          @mutex.synchronize do
+            if (pending_response = @pending_responses[request_id]) && pending_response[:session_id] == session_id
+              if body.key?(:error)
+                error = body[:error]
+                pending_response[:queue].push(error: { code: error[:code], message: error[:message] })
+              else
+                pending_response[:queue].push(body[:result])
+              end
+            end
+          end
+
+          handle_accepted
         end
 
         def handle_initialization(body_string, body)
@@ -326,9 +475,8 @@ module MCP
           [202, {}, []]
         end
 
-        def handle_regular_request(body_string, session_id)
+        def handle_regular_request(body_string, session_id, related_request_id: nil)
           server_session = nil
-          stream = nil
 
           unless @stateless
             if session_id
@@ -338,57 +486,106 @@ module MCP
               @mutex.synchronize do
                 session = @sessions[session_id]
                 server_session = session[:server_session] if session
-                stream = session[:stream] if session
               end
             end
           end
 
-          response = if server_session
-            server_session.handle_json(body_string)
+          if session_id && !@stateless
+            handle_request_with_sse_response(body_string, session_id, server_session, related_request_id: related_request_id)
           else
-            @server.handle_json(body_string)
+            response = dispatch_handle_json(body_string, server_session)
+            [200, { "Content-Type" => "application/json" }, [response]]
+          end
+        end
+
+        # Returns the POST response as an SSE stream so the server can send
+        # JSON-RPC requests and notifications during request processing.
+        # https://modelcontextprotocol.io/specification/2025-11-25/basic/transports#sending-messages-to-the-server
+        def handle_request_with_sse_response(body_string, session_id, server_session, related_request_id: nil)
+          body = proc do |stream|
+            @mutex.synchronize do
+              session = @sessions[session_id]
+              if session && related_request_id
+                session[:post_request_streams] ||= {}
+                session[:post_request_streams][related_request_id] = stream
+              end
+            end
+
+            begin
+              response = dispatch_handle_json(body_string, server_session)
+
+              send_to_stream(stream, response) if response
+            ensure
+              if related_request_id
+                @mutex.synchronize do
+                  session = @sessions[session_id]
+                  session[:post_request_streams]&.delete(related_request_id) if session
+                end
+              end
+
+              begin
+                stream.close
+              rescue StandardError
+                # Ignore close-related errors from already closed/broken streams.
+              end
+            end
           end
 
-          if stream
-            send_response_to_stream(stream, response, session_id)
+          [200, SSE_HEADERS, body]
+        end
+
+        # Returns the SSE stream available for server-to-client messages.
+        # When `related_request_id` is given, returns only the POST response
+        # stream for that request (no fallback to GET SSE). This prevents
+        # request-scoped messages from leaking to the wrong stream.
+        # When `related_request_id` is nil, returns the GET SSE stream.
+        def active_stream(session, related_request_id: nil)
+          if related_request_id
+            session.dig(:post_request_streams, related_request_id)
           else
-            [200, { "Content-Type" => "application/json" }, [response]]
+            session[:stream]
+          end
+        end
+
+        def dispatch_handle_json(body_string, server_session)
+          if server_session
+            server_session.handle_json(body_string)
+          else
+            @server.handle_json(body_string)
           end
         end
 
         def validate_and_touch_session(session_id)
-          @mutex.synchronize do
-            return session_not_found_response unless (session = @sessions[session_id])
-            return unless @session_idle_timeout
+          removed = nil
+
+          response = @mutex.synchronize do
+            next session_not_found_response unless (session = @sessions[session_id])
+            next unless @session_idle_timeout
 
             if session_expired?(session)
-              cleanup_session_unsafe(session_id)
-              return session_not_found_response
+              removed = cleanup_session_unsafe(session_id)
+              next session_not_found_response
             end
 
             session[:last_active_at] = Process.clock_gettime(Process::CLOCK_MONOTONIC)
+            nil
           end
 
-          nil
+          if removed
+            close_stream_safely(removed[:stream])
+
+            removed[:post_request_streams]&.each_value do |stream|
+              close_stream_safely(stream)
+            end
+          end
+
+          response
         end
 
         def get_session_stream(session_id)
           @mutex.synchronize { @sessions[session_id]&.fetch(:stream, nil) }
         end
 
-        def send_response_to_stream(stream, response, session_id)
-          message = JSON.parse(response)
-          send_to_stream(stream, message)
-          handle_accepted
-        rescue *STREAM_WRITE_ERRORS => e
-          MCP.configuration.exception_reporter.call(
-            e,
-            { session_id: session_id, error: "Stream closed during response" },
-          )
-          cleanup_session(session_id)
-          [200, { "Content-Type" => "application/json" }, [response]]
-        end
-
         def session_exists?(session_id)
           @mutex.synchronize { @sessions.key?(session_id) }
         end
@@ -416,13 +613,7 @@ module MCP
         def setup_sse_stream(session_id)
           body = create_sse_body(session_id)
 
-          headers = {
-            "Content-Type" => "text/event-stream",
-            "Cache-Control" => "no-cache",
-            "Connection" => "keep-alive",
-          }
-
-          [200, headers, body]
+          [200, SSE_HEADERS, body]
         end
 
         def create_sse_body(session_id)