Implement StreamableHTTP transport

Co-authored-by: Topher Bullock <topher.bullock@shopify.com>

lib/mcp.rb

@@ -13,6 +13,7 @@ require_relative "mcp/resource/contents"
 require_relative "mcp/resource/embedded"
 require_relative "mcp/resource_template"
 require_relative "mcp/server"
+require_relative "mcp/server/transports/streamable_http_transport"
 require_relative "mcp/server/transports/stdio_transport"
 require_relative "mcp/string_utils"
 require_relative "mcp/tool"

lib/mcp/server.rb

@@ -24,7 +24,7 @@ module MCP
     include Instrumentation
 
     attr_writer :capabilities
-    attr_accessor :name, :version, :tools, :prompts, :resources, :server_context, :configuration
+    attr_accessor :name, :version, :tools, :prompts, :resources, :server_context, :configuration, :transport
 
     def initialize(
       name: "model_context_protocol",
@@ -35,7 +35,8 @@ module MCP
       resource_templates: [],
       server_context: nil,
       configuration: nil,
-      capabilities: nil
+      capabilities: nil,
+      transport: nil
     )
       @name = name
       @version = version
@@ -68,6 +69,7 @@ module MCP
         Methods::COMPLETION_COMPLETE => ->(_) {},
         Methods::LOGGING_SET_LEVEL => ->(_) {},
       }
+      @transport = transport
     end
 
     def capabilities
@@ -96,6 +98,14 @@ module MCP
       @prompts[prompt.name_value] = prompt
     end
 
+    def handle_rack_request(request)
+      @transport.handle_request(request)
+    end
+
+    def close
+      @transport.close
+    end
+
     def resources_list_handler(&block)
       @capabilities.support_resources
       @handlers[Methods::RESOURCES_LIST] = block

lib/mcp/server/transports/stdio_transport.rb

@@ -8,6 +8,11 @@ module MCP
     module Transports
       class StdioTransport < Transport
         STATUS_INTERRUPTED = Signal.list["INT"] + 128
+        class << self
+          def register
+            super("stdio", self)
+          end
+        end
 
         def initialize(server)
           @server = server
@@ -37,6 +42,12 @@ module MCP
           $stdout.puts(json_message)
           $stdout.flush
         end
+
+        def send_response(message)
+          json_message = message.is_a?(String) ? message : JSON.generate(message)
+          $stdout.puts(json_message)
+          $stdout.flush
+        end
       end
     end
   end

lib/mcp/server/transports/streamable_http_transport.rb

@@ -0,0 +1,295 @@
+# frozen_string_literal: true
+
+require_relative "../../transport"
+require "json"
+require "securerandom"
+
+module MCP
+  class Server
+    module Transports
+      class StreamableHTTPTransport < Transport
+        class << self
+          def register
+            super("http", self)
+          end
+        end
+
+        def initialize(server)
+          super
+          # { session_id => { stream: stream_object }
+          @sessions = {}
+          @mutex = Mutex.new
+        end
+
+        def handle_request(request)
+          case request.env["REQUEST_METHOD"]
+          when "POST"
+            handle_post(request)
+          when "GET"
+            handle_get(request)
+          when "DELETE"
+            handle_delete(request)
+          else
+            [405, { "Content-Type" => "application/json" }, [{ error: "Method not allowed" }.to_json]]
+          end
+        end
+
+        def close
+          @mutex.synchronize do
+            @sessions.each_key { |session_id| cleanup_session_unsafe(session_id) }
+          end
+        end
+
+        def send_notification(notification, session_id: nil)
+          @mutex.synchronize do
+            if session_id
+              # Send to specific session
+              session = @sessions[session_id]
+              return false unless session && session[:stream]
+
+              begin
+                send_to_stream(session[:stream], notification)
+                true
+              rescue IOError, Errno::EPIPE => e
+                MCP.configuration.exception_reporter.call(
+                  e,
+                  { session_id: session_id, error: "Failed to send notification" },
+                )
+                cleanup_session_unsafe(session_id)
+                false
+              end
+            else
+              # Broadcast to all connected SSE sessions
+              sent_count = 0
+              failed_sessions = []
+
+              @sessions.each do |sid, session|
+                next unless session[:stream]
+
+                begin
+                  send_to_stream(session[:stream], notification)
+                  sent_count += 1
+                rescue IOError, Errno::EPIPE => 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_session_unsafe(sid) }
+
+              sent_count
+            end
+          end
+        end
+
+        private
+
+        def send_to_stream(stream, data)
+          message = data.is_a?(String) ? data : data.to_json
+          stream.write("data: #{message}\n\n")
+          stream.flush if stream.respond_to?(:flush)
+        end
+
+        def send_ping_to_stream(stream)
+          stream.write(": ping #{Time.now.iso8601}\n\n")
+          stream.flush if stream.respond_to?(:flush)
+        end
+
+        def handle_post(request)
+          body_string = request.body.read
+          session_id = extract_session_id(request)
+
+          body = parse_request_body(body_string)
+          return body unless body.is_a?(Hash) # Error response
+
+          if body["method"] == "initialize"
+            handle_initialization(body_string, body)
+          else
+            handle_regular_request(body_string, session_id)
+          end
+        rescue StandardError => e
+          ModelContextProtocol.configuration.exception_reporter.call(e, { request: body_string })
+          [500, { "Content-Type" => "application/json" }, [{ error: "Internal server error" }.to_json]]
+        end
+
+        def handle_get(request)
+          session_id = extract_session_id(request)
+
+          return missing_session_id_response unless session_id
+          return session_not_found_response unless session_exists?(session_id)
+
+          setup_sse_stream(session_id)
+        end
+
+        def handle_delete(request)
+          session_id = request.env["HTTP_MCP_SESSION_ID"]
+
+          return [
+            400,
+            { "Content-Type" => "application/json" },
+            [{ error: "Missing session ID" }.to_json],
+          ] unless session_id
+
+          cleanup_session(session_id)
+          [200, { "Content-Type" => "application/json" }, [{ success: true }.to_json]]
+        end
+
+        def cleanup_session(session_id)
+          @mutex.synchronize do
+            cleanup_session_unsafe(session_id)
+          end
+        end
+
+        def cleanup_session_unsafe(session_id)
+          session = @sessions[session_id]
+          return unless session
+
+          begin
+            session[:stream]&.close
+          rescue
+            nil
+          end
+          @sessions.delete(session_id)
+        end
+
+        def extract_session_id(request)
+          request.env["HTTP_MCP_SESSION_ID"]
+        end
+
+        def parse_request_body(body_string)
+          JSON.parse(body_string)
+        rescue JSON::ParserError
+          [400, { "Content-Type" => "application/json" }, [{ error: "Invalid JSON" }.to_json]]
+        end
+
+        def handle_initialization(body_string, body)
+          session_id = SecureRandom.uuid
+
+          @mutex.synchronize do
+            @sessions[session_id] = {
+              stream: nil,
+            }
+          end
+
+          response = @server.handle_json(body_string)
+
+          headers = {
+            "Content-Type" => "application/json",
+            "Mcp-Session-Id" => session_id,
+          }
+
+          [200, headers, [response]]
+        end
+
+        def handle_regular_request(body_string, session_id)
+          # If session ID is provided, but not in the sessions hash, return an error
+          if session_id && !@sessions.key?(session_id)
+            return [400, { "Content-Type" => "application/json" }, [{ error: "Invalid session ID" }.to_json]]
+          end
+
+          response = @server.handle_json(body_string)
+          stream = get_session_stream(session_id) if session_id
+
+          if stream
+            send_response_to_stream(stream, response, session_id)
+          else
+            [200, { "Content-Type" => "application/json" }, [response]]
+          end
+        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)
+          [200, { "Content-Type" => "application/json" }, [{ accepted: true }.to_json]]
+        rescue IOError, Errno::EPIPE => e
+          ModelContextProtocol.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
+
+        def missing_session_id_response
+          [400, { "Content-Type" => "application/json" }, [{ error: "Missing session ID" }.to_json]]
+        end
+
+        def session_not_found_response
+          [404, { "Content-Type" => "application/json" }, [{ error: "Session not found" }.to_json]]
+        end
+
+        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]
+        end
+
+        def create_sse_body(session_id)
+          proc do |stream|
+            store_stream_for_session(session_id, stream)
+            start_keepalive_thread(session_id)
+          end
+        end
+
+        def store_stream_for_session(session_id, stream)
+          @mutex.synchronize do
+            if @sessions[session_id]
+              @sessions[session_id][:stream] = stream
+            else
+              stream.close
+            end
+          end
+        end
+
+        def start_keepalive_thread(session_id)
+          Thread.new do
+            while session_active_with_stream?(session_id)
+              sleep(30)
+              send_keepalive_ping(session_id)
+            end
+          rescue StandardError => e
+            MCP.configuration.exception_reporter.call(e, { session_id: session_id })
+          ensure
+            cleanup_session(session_id)
+          end
+        end
+
+        def session_active_with_stream?(session_id)
+          @mutex.synchronize { @sessions.key?(session_id) && @sessions[session_id][:stream] }
+        end
+
+        def send_keepalive_ping(session_id)
+          @mutex.synchronize do
+            if @sessions[session_id] && @sessions[session_id][:stream]
+              send_ping_to_stream(@sessions[session_id][:stream])
+            end
+          end
+        rescue IOError, Errno::EPIPE => e
+          MCP.configuration.exception_reporter.call(
+            e,
+            { session_id: session_id, error: "Stream closed" },
+          )
+          raise # Re-raise to exit the keepalive loop
+        end
+      end
+    end
+  end
+end

lib/mcp/transport.rb

@@ -2,32 +2,57 @@
 
 module MCP
   class Transport
+    class << self
+      def register(name, transport_class)
+        @transports ||= {}
+        @transports[name] = transport_class
+      end
+
+      def get(name)
+        raise "Transport #{name} not found" unless @transports
+
+        @transports.fetch(name) { raise "Transport #{name} not found" }
+      end
+    end
+
+    # Initialize the transport with the server instance
     def initialize(server)
       @server = server
     end
 
+    # Send a response to the client
     def send_response(response)
       raise NotImplementedError, "Subclasses must implement send_response"
     end
 
+    # Open the transport connection
     def open
       raise NotImplementedError, "Subclasses must implement open"
     end
 
+    # Close the transport connection
     def close
       raise NotImplementedError, "Subclasses must implement close"
     end
 
-    private
+    # Handle a JSON request
+    # Returns a response that should be sent back to the client
+    def handle_json_request(request)
+      response = @server.handle_json(request)
+      send_response(response) if response
+    end
 
+    # Handle an incoming request
+    # Returns a response that should be sent back to the client
     def handle_request(request)
       response = @server.handle(request)
       send_response(response) if response
     end
 
-    def handle_json_request(request)
-      response = @server.handle_json(request)
-      send_response(response) if response
+    # Send a notification to the client
+    # Returns true if the notification was sent successfully
+    def send_notification(method, params = nil)
+      raise NotImplementedError, "Subclasses must implement send_notification"
     end
   end
 end