model-context-protocol-rb 0.3.4 → 0.4.0

data/lib/model_context_protocol/server/streamable_http_transport.rb

@@ -20,6 +20,11 @@ module ModelContextProtocol
 
       transport_options = @configuration.transport_options
       @redis = transport_options[:redis_client]
+      @require_sessions = transport_options.fetch(:require_sessions, false)
+      @default_protocol_version = transport_options.fetch(:default_protocol_version, "2025-03-26")
+      @session_protocol_versions = {}  # Track protocol versions per session
+      @validate_origin = transport_options.fetch(:validate_origin, true)
+      @allowed_origins = transport_options.fetch(:allowed_origins, ["http://localhost", "https://localhost", "http://127.0.0.1", "https://127.0.0.1"])
 
       @session_store = ModelContextProtocol::Server::SessionStore.new(
         @redis,
@@ -29,6 +34,7 @@ module ModelContextProtocol
       @server_instance = "#{Socket.gethostname}-#{Process.pid}-#{SecureRandom.hex(4)}"
       @local_streams = {}
       @notification_queue = []
+      @sse_event_counter = 0
 
       setup_redis_subscriber
     end
@@ -36,20 +42,19 @@ module ModelContextProtocol
     def handle
       @configuration.logger.connect_transport(self)
 
-      request = @configuration.transport_options[:request]
-      response = @configuration.transport_options[:response]
+      env = @configuration.transport_options[:env]
 
-      unless request && response
-        raise ArgumentError, "StreamableHTTP transport requires request and response objects in transport_options"
+      unless env
+        raise ArgumentError, "StreamableHTTP transport requires Rack env hash in transport_options"
       end
 
-      case request.method
+      case env["REQUEST_METHOD"]
       when "POST"
-        handle_post_request(request)
+        handle_post_request(env)
       when "GET"
-        handle_sse_request(request, response)
+        handle_sse_request(env)
       when "DELETE"
-        handle_delete_request(request)
+        handle_delete_request(env)
       else
         error_response = ErrorResponse[id: nil, error: {code: -32601, message: "Method not allowed"}]
         {json: error_response.serialized, status: 405}
@@ -72,16 +77,90 @@ module ModelContextProtocol
 
     private
 
-    def handle_post_request(request)
-      body_string = request.body.read
+    def validate_headers(env)
+      if @validate_origin
+        origin = env["HTTP_ORIGIN"]
+        if origin && !@allowed_origins.any? { |allowed| origin.start_with?(allowed) }
+          error_response = ErrorResponse[id: nil, error: {code: -32600, message: "Origin not allowed"}]
+          return {json: error_response.serialized, status: 403}
+        end
+      end
+
+      accept_header = env["HTTP_ACCEPT"]
+      if accept_header
+        unless accept_header.include?("application/json") || accept_header.include?("text/event-stream")
+          error_response = ErrorResponse[id: nil, error: {code: -32600, message: "Invalid Accept header. Must include application/json or text/event-stream"}]
+          return {json: error_response.serialized, status: 400}
+        end
+      end
+
+      protocol_version = env["HTTP_MCP_PROTOCOL_VERSION"]
+      if protocol_version
+        # Check if this matches a known negotiated version
+        valid_versions = @session_protocol_versions.values.compact.uniq
+        unless valid_versions.empty? || valid_versions.include?(protocol_version)
+          error_response = ErrorResponse[id: nil, error: {code: -32600, message: "Invalid MCP protocol version: #{protocol_version}. Expected one of: #{valid_versions.join(", ")}"}]
+          return {json: error_response.serialized, status: 400}
+        end
+      end
+
+      nil
+    end
+
+    def determine_message_type(body)
+      if body.key?("method") && body.key?("id")
+        :request
+      elsif body.key?("method") && !body.key?("id")
+        :notification
+      elsif body.key?("id") && body.key?("result") || body.key?("error")
+        :response
+      else
+        :unknown
+      end
+    end
+
+    def create_initialization_sse_stream_proc(response_data)
+      proc do |stream|
+        event_id = next_event_id
+        send_sse_event(stream, response_data, event_id)
+      end
+    end
+
+    def create_request_sse_stream_proc(response_data)
+      proc do |stream|
+        event_id = next_event_id
+        send_sse_event(stream, response_data, event_id)
+      end
+    end
+
+    def next_event_id
+      @sse_event_counter += 1
+      "#{@server_instance}-#{@sse_event_counter}"
+    end
+
+    def send_sse_event(stream, data, event_id = nil)
+      if event_id
+        stream.write("id: #{event_id}\n")
+      end
+      message = data.is_a?(String) ? data : data.to_json
+      stream.write("data: #{message}\n\n")
+      stream.flush if stream.respond_to?(:flush)
+    end
+
+    def handle_post_request(env)
+      validation_error = validate_headers(env)
+      return validation_error if validation_error
+
+      body_string = env["rack.input"].read
       body = JSON.parse(body_string)
-      session_id = request.headers["Mcp-Session-Id"]
+      session_id = env["HTTP_MCP_SESSION_ID"]
+      accept_header = env["HTTP_ACCEPT"] || ""
 
       case body["method"]
       when "initialize"
-        handle_initialization(body)
+        handle_initialization(body, accept_header)
       else
-        handle_regular_request(body, session_id)
+        handle_regular_request(body, session_id, accept_header)
       end
     rescue JSON::ParserError
       error_response = ErrorResponse[id: "", error: {code: -32700, message: "Parse error"}]
@@ -96,51 +175,122 @@ module ModelContextProtocol
       {json: error_response.serialized, status: 500}
     end
 
-    def handle_initialization(body)
-      session_id = SecureRandom.uuid
-
-      @session_store.create_session(session_id, {
-        server_instance: @server_instance,
-        context: @configuration.context || {},
-        created_at: Time.now.to_f
-      })
-
+    def handle_initialization(body, accept_header)
       result = @router.route(body)
       response = Response[id: body["id"], result: result.serialized]
+      response_headers = {}
+
+      negotiated_protocol_version = result.serialized[:protocolVersion] || result.serialized["protocolVersion"]
+
+      if @require_sessions
+        session_id = SecureRandom.uuid
+        @session_store.create_session(session_id, {
+          server_instance: @server_instance,
+          context: @configuration.context || {},
+          created_at: Time.now.to_f,
+          negotiated_protocol_version: negotiated_protocol_version
+        })
+        response_headers["Mcp-Session-Id"] = session_id
+        @session_protocol_versions[session_id] = negotiated_protocol_version
+      else
+        @session_protocol_versions[:default] = negotiated_protocol_version
+      end
 
-      {
-        json: response.serialized,
-        status: 200,
-        headers: {"Mcp-Session-Id" => session_id}
-      }
+      if accept_header.include?("text/event-stream") && !accept_header.include?("application/json")
+        response_headers.merge!({
+          "Content-Type" => "text/event-stream",
+          "Cache-Control" => "no-cache",
+          "Connection" => "keep-alive"
+        })
+
+        {
+          stream: true,
+          headers: response_headers,
+          stream_proc: create_initialization_sse_stream_proc(response.serialized)
+        }
+      else
+        response_headers["Content-Type"] = "application/json"
+        {
+          json: response.serialized,
+          status: 200,
+          headers: response_headers
+        }
+      end
     end
 
-    def handle_regular_request(body, session_id)
-      unless session_id && @session_store.session_exists?(session_id)
-        error_response = ErrorResponse[id: body["id"], error: {code: -32600, message: "Invalid or missing session ID"}]
-        return {json: error_response.serialized, status: 400}
+    def handle_regular_request(body, session_id, accept_header)
+      if @require_sessions
+        unless session_id && @session_store.session_exists?(session_id)
+          if session_id && !@session_store.session_exists?(session_id)
+            error_response = ErrorResponse[id: body["id"], error: {code: -32600, message: "Session terminated"}]
+            return {json: error_response.serialized, status: 404}
+          else
+            error_response = ErrorResponse[id: body["id"], error: {code: -32600, message: "Invalid or missing session ID"}]
+            return {json: error_response.serialized, status: 400}
+          end
+        end
       end
 
-      result = @router.route(body)
-      response = Response[id: body["id"], result: result.serialized]
+      message_type = determine_message_type(body)
 
-      if @session_store.session_has_active_stream?(session_id)
-        deliver_to_session_stream(session_id, response.serialized)
-        {json: {accepted: true}, status: 200}
-      else
-        {json: response.serialized, status: 200}
+      case message_type
+      when :notification, :response
+        if session_id && @session_store.session_has_active_stream?(session_id)
+          deliver_to_session_stream(session_id, body)
+        end
+        {json: {}, status: 202}
+
+      when :request
+        result = @router.route(body)
+        response = Response[id: body["id"], result: result.serialized]
+
+        if session_id && @session_store.session_has_active_stream?(session_id)
+          deliver_to_session_stream(session_id, response.serialized)
+          return {json: {accepted: true}, status: 200}
+        end
+
+        if accept_header.include?("text/event-stream") && !accept_header.include?("application/json")
+          {
+            stream: true,
+            headers: {
+              "Content-Type" => "text/event-stream",
+              "Cache-Control" => "no-cache",
+              "Connection" => "keep-alive"
+            },
+            stream_proc: create_request_sse_stream_proc(response.serialized)
+          }
+        else
+          {
+            json: response.serialized,
+            status: 200,
+            headers: {"Content-Type" => "application/json"}
+          }
+        end
       end
     end
 
-    def handle_sse_request(request, response)
-      session_id = request.headers["Mcp-Session-Id"]
-
-      unless session_id && @session_store.session_exists?(session_id)
-        error_response = ErrorResponse[id: nil, error: {code: -32600, message: "Invalid or missing session ID"}]
+    def handle_sse_request(env)
+      accept_header = env["HTTP_ACCEPT"] || ""
+      unless accept_header.include?("text/event-stream")
+        error_response = ErrorResponse[id: nil, error: {code: -32600, message: "Accept header must include text/event-stream"}]
         return {json: error_response.serialized, status: 400}
       end
 
-      @session_store.mark_stream_active(session_id, @server_instance)
+      session_id = env["HTTP_MCP_SESSION_ID"]
+      last_event_id = env["HTTP_LAST_EVENT_ID"]
+
+      if @require_sessions
+        unless session_id && @session_store.session_exists?(session_id)
+          if session_id && !@session_store.session_exists?(session_id)
+            error_response = ErrorResponse[id: nil, error: {code: -32600, message: "Session terminated"}]
+            return {json: error_response.serialized, status: 404}
+          else
+            error_response = ErrorResponse[id: nil, error: {code: -32600, message: "Invalid or missing session ID"}]
+            return {json: error_response.serialized, status: 400}
+          end
+        end
+        @session_store.mark_stream_active(session_id, @server_instance)
+      end
 
       {
         stream: true,
@@ -149,12 +299,12 @@ module ModelContextProtocol
           "Cache-Control" => "no-cache",
           "Connection" => "keep-alive"
         },
-        stream_proc: create_sse_stream_proc(session_id)
+        stream_proc: create_sse_stream_proc(session_id, last_event_id)
       }
     end
 
-    def handle_delete_request(request)
-      session_id = request.headers["Mcp-Session-Id"]
+    def handle_delete_request(env)
+      session_id = env["HTTP_MCP_SESSION_ID"]
 
       if session_id
         cleanup_session(session_id)
@@ -163,11 +313,15 @@ module ModelContextProtocol
       {json: {success: true}, status: 200}
     end
 
-    def create_sse_stream_proc(session_id)
+    def create_sse_stream_proc(session_id, last_event_id = nil)
       proc do |stream|
-        register_local_stream(session_id, stream)
+        register_local_stream(session_id, stream) if session_id
 
-        flush_notifications_to_stream(stream)
+        if last_event_id
+          replay_messages_after_event_id(stream, session_id, last_event_id)
+        else
+          flush_notifications_to_stream(stream)
+        end
 
         start_keepalive_thread(session_id, stream)
 
@@ -176,7 +330,7 @@ module ModelContextProtocol
           sleep 0.1
         end
       ensure
-        cleanup_local_stream(session_id)
+        cleanup_local_stream(session_id) if session_id
       end
     end
 
@@ -226,9 +380,12 @@ module ModelContextProtocol
     end
 
     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)
+      event_id = next_event_id
+      send_sse_event(stream, data, event_id)
+    end
+
+    def replay_messages_after_event_id(stream, session_id, last_event_id)
+      flush_notifications_to_stream(stream)
     end
 
     def deliver_to_session_stream(session_id, data)

data/lib/model_context_protocol/server/tool.rb

@@ -2,6 +2,11 @@ require "json-schema"
 
 module ModelContextProtocol
   class Server::Tool
+    # Raised when output schema validation fails.
+    class OutputSchemaValidationError < StandardError; end
+
+    include ModelContextProtocol::Server::ContentHelpers
+
     attr_reader :arguments, :context, :logger
 
     def initialize(arguments, logger, context = {})
@@ -15,58 +20,58 @@ module ModelContextProtocol
       raise NotImplementedError, "Subclasses must implement the call method"
     end
 
-    TextResponse = Data.define(:text) do
-      def serialized
-        {content: [{type: "text", text:}], isError: false}
-      end
-    end
-    private_constant :TextResponse
-
-    ImageResponse = Data.define(:data, :mime_type) do
-      def initialize(data:, mime_type: "image/png")
-        super
-      end
-
+    Response = Data.define(:content) do
       def serialized
-        {content: [{type: "image", data:, mimeType: mime_type}], isError: false}
+        serialized_contents = content.map(&:serialized)
+        {content: serialized_contents, isError: false}
       end
     end
-    private_constant :ImageResponse
-
-    ResourceResponse = Data.define(:uri, :text, :mime_type) do
-      def initialize(uri:, text:, mime_type: "text/plain")
-        super
-      end
+    private_constant :Response
 
+    StructuredContentResponse = Data.define(:structured_content, :tool) do
       def serialized
-        {content: [{type: "resource", resource: {uri:, mimeType: mime_type, text:}}], isError: false}
+        json_text = JSON.generate(structured_content)
+        text_content = ModelContextProtocol::Server::Content::Text[
+          meta: nil,
+          annotations: nil,
+          text: json_text
+        ]
+
+        validation_errors = JSON::Validator.fully_validate(
+          tool.class.definition[:outputSchema], structured_content
+        )
+
+        if validation_errors.empty?
+          {
+            structuredContent: structured_content,
+            content: [text_content.serialized],
+            isError: false
+          }
+        else
+          raise OutputSchemaValidationError, validation_errors.join(", ")
+        end
       end
     end
-    private_constant :ResourceResponse
+    private_constant :StructuredContentResponse
 
-    ToolErrorResponse = Data.define(:text) do
+    ErrorResponse = Data.define(:error) do
       def serialized
-        {content: [{type: "text", text:}], isError: true}
+        {content: [{type: "text", text: error}], isError: true}
       end
     end
-    private_constant :ToolErrorResponse
-
-    private def respond_with(type, **options)
-      case [type, options]
-      in [:text, {text:}]
-        TextResponse[text:]
-      in [:image, {data:, mime_type:}]
-        ImageResponse[data:, mime_type:]
-      in [:image, {data:}]
-        ImageResponse[data:]
-      in [:resource, {mime_type:, text:, uri:}]
-        ResourceResponse[mime_type:, text:, uri:]
-      in [:resource, {text:, uri:}]
-        ResourceResponse[text:, uri:]
-      in [:error, {text:}]
-        ToolErrorResponse[text:]
+    private_constant :ErrorResponse
+
+    private def respond_with(**kwargs)
+      case [kwargs]
+      in [{content:}]
+        content_array = content.is_a?(Array) ? content : [content]
+        Response[content: content_array]
+      in [{structured_content:}]
+        StructuredContentResponse[structured_content:, tool: self]
+      in [{error:}]
+        ErrorResponse[error:]
       else
-        raise ModelContextProtocol::Server::ResponseArgumentsError, "Invalid arguments: #{type}, #{options}"
+        raise ModelContextProtocol::Server::ResponseArgumentsError, "Invalid arguments: #{kwargs.inspect}"
       end
     end
 
@@ -75,39 +80,46 @@ module ModelContextProtocol
     end
 
     class << self
-      attr_reader :name, :description, :input_schema
+      attr_reader :name, :description, :title, :input_schema, :output_schema
 
-      def with_metadata(&block)
-        metadata_dsl = MetadataDSL.new
-        metadata_dsl.instance_eval(&block)
+      def define(&block)
+        definition_dsl = DefinitionDSL.new
+        definition_dsl.instance_eval(&block)
 
-        @name = metadata_dsl.name
-        @description = metadata_dsl.description
-        @input_schema = metadata_dsl.input_schema
+        @name = definition_dsl.name
+        @description = definition_dsl.description
+        @title = definition_dsl.title
+        @input_schema = definition_dsl.input_schema
+        @output_schema = definition_dsl.output_schema
       end
 
       def inherited(subclass)
         subclass.instance_variable_set(:@name, @name)
         subclass.instance_variable_set(:@description, @description)
+        subclass.instance_variable_set(:@title, @title)
         subclass.instance_variable_set(:@input_schema, @input_schema)
+        subclass.instance_variable_set(:@output_schema, @output_schema)
       end
 
       def call(arguments, logger, context = {})
         new(arguments, logger, context).call
       rescue JSON::Schema::ValidationError => validation_error
         raise ModelContextProtocol::Server::ParameterValidationError, validation_error.message
-      rescue ModelContextProtocol::Server::ResponseArgumentsError => response_arguments_error
-        raise response_arguments_error
+      rescue OutputSchemaValidationError, ModelContextProtocol::Server::ResponseArgumentsError => tool_error
+        raise tool_error, tool_error.message
       rescue => error
-        ToolErrorResponse[text: error.message]
+        ErrorResponse[error: error.message]
       end
 
-      def metadata
-        {name: @name, description: @description, inputSchema: @input_schema}
+      def definition
+        result = {name: @name, description: @description, inputSchema: @input_schema}
+        result[:title] = @title if @title
+        result[:outputSchema] = @output_schema if @output_schema
+        result
       end
     end
 
-    class MetadataDSL
+    class DefinitionDSL
       def name(value = nil)
         @name = value if value
         @name
@@ -118,10 +130,20 @@ module ModelContextProtocol
         @description
       end
 
+      def title(value = nil)
+        @title = value if value
+        @title
+      end
+
       def input_schema(&block)
         @input_schema = instance_eval(&block) if block_given?
         @input_schema
       end
+
+      def output_schema(&block)
+        @output_schema = instance_eval(&block) if block_given?
+        @output_schema
+      end
     end
   end
 end