model-context-protocol-rb 0.3.2 → 0.3.4

data/lib/model_context_protocol/server/prompt.rb

@@ -1,10 +1,12 @@
 module ModelContextProtocol
   class Server::Prompt
-    attr_reader :params
+    attr_reader :arguments, :context, :logger
 
-    def initialize(params)
-      validate!(params)
-      @params = params
+    def initialize(arguments, logger, context = {})
+      validate!(arguments)
+      @arguments = arguments
+      @context = context
+      @logger = logger
     end
 
     def call
@@ -22,18 +24,18 @@ module ModelContextProtocol
       Response[messages:, description: self.class.description]
     end
 
-    private def validate!(params = {})
-      arguments = self.class.arguments || []
-      required_args = arguments.select { |arg| arg[:required] }.map { |arg| arg[:name] }
-      valid_arg_names = arguments.map { |arg| arg[:name] }
+    private def validate!(arguments = {})
+      defined_arguments = self.class.defined_arguments || []
+      required_args = defined_arguments.select { |arg| arg[:required] }.map { |arg| arg[:name].to_sym }
+      valid_arg_names = defined_arguments.map { |arg| arg[:name].to_sym }
 
-      missing_args = required_args - params.keys
+      missing_args = required_args - arguments.keys
       unless missing_args.empty?
         missing_args_list = missing_args.join(", ")
         raise ArgumentError, "Missing required arguments: #{missing_args_list}"
       end
 
-      extra_args = params.keys - valid_arg_names
+      extra_args = arguments.keys - valid_arg_names
       unless extra_args.empty?
         extra_args_list = extra_args.join(", ")
         raise ArgumentError, "Unexpected arguments: #{extra_args_list}"
@@ -41,10 +43,10 @@ module ModelContextProtocol
     end
 
     class << self
-      attr_reader :name, :description, :arguments
+      attr_reader :name, :description, :defined_arguments
 
       def with_metadata(&block)
-        @arguments ||= []
+        @defined_arguments ||= []
 
         metadata_dsl = MetadataDSL.new
         metadata_dsl.instance_eval(&block)
@@ -54,12 +56,12 @@ module ModelContextProtocol
       end
 
       def with_argument(&block)
-        @arguments ||= []
+        @defined_arguments ||= []
 
         argument_dsl = ArgumentDSL.new
         argument_dsl.instance_eval(&block)
 
-        @arguments << {
+        @defined_arguments << {
           name: argument_dsl.name,
           description: argument_dsl.description,
           required: argument_dsl.required,
@@ -70,21 +72,21 @@ module ModelContextProtocol
       def inherited(subclass)
         subclass.instance_variable_set(:@name, @name)
         subclass.instance_variable_set(:@description, @description)
-        subclass.instance_variable_set(:@arguments, @arguments&.dup)
+        subclass.instance_variable_set(:@defined_arguments, @defined_arguments&.dup)
       end
 
-      def call(params)
-        new(params).call
+      def call(arguments, logger, context = {})
+        new(arguments, logger, context).call
       rescue ArgumentError => error
         raise ModelContextProtocol::Server::ParameterValidationError, error.message
       end
 
       def metadata
-        {name: @name, description: @description, arguments: @arguments}
+        {name: @name, description: @description, arguments: @defined_arguments}
       end
 
       def complete_for(arg_name, value)
-        arg = @arguments&.find { |a| a[:name] == arg_name.to_s }
+        arg = @defined_arguments&.find { |a| a[:name] == arg_name.to_s }
         completion = (arg && arg[:completion]) ? arg[:completion] : ModelContextProtocol::Server::NullCompletion
         completion.call(arg_name.to_s, value)
       end

data/lib/model_context_protocol/server/resource.rb

@@ -1,10 +1,12 @@
 module ModelContextProtocol
   class Server::Resource
-    attr_reader :mime_type, :uri
+    attr_reader :mime_type, :uri, :context, :logger
 
-    def initialize
+    def initialize(logger, context = {})
       @mime_type = self.class.mime_type
       @uri = self.class.uri
+      @context = context
+      @logger = logger
     end
 
     def call
@@ -56,8 +58,8 @@ module ModelContextProtocol
         subclass.instance_variable_set(:@uri, @uri)
       end
 
-      def call
-        new.call
+      def call(logger, context = {})
+        new(logger, context).call
       end
 
       def metadata

data/lib/model_context_protocol/server/session_store.rb

@@ -0,0 +1,108 @@
+# frozen_string_literal: true
+
+require "json"
+require "securerandom"
+
+module ModelContextProtocol
+  class Server
+    class SessionStore
+      def initialize(redis_client, ttl: 3600)
+        @redis = redis_client
+        @ttl = ttl
+      end
+
+      def create_session(session_id, data)
+        session_data = {
+          id: session_id,
+          server_instance: data[:server_instance],
+          context: data[:context] || {},
+          created_at: data[:created_at] || Time.now.to_f,
+          last_activity: Time.now.to_f,
+          active_stream: false
+        }
+
+        @redis.hset("session:#{session_id}", session_data.transform_values(&:to_json))
+        @redis.expire("session:#{session_id}", @ttl)
+        session_id
+      end
+
+      def mark_stream_active(session_id, server_instance)
+        @redis.multi do |multi|
+          multi.hset("session:#{session_id}",
+            "active_stream", true.to_json,
+            "stream_server", server_instance.to_json,
+            "last_activity", Time.now.to_f.to_json)
+          multi.expire("session:#{session_id}", @ttl)
+        end
+      end
+
+      def mark_stream_inactive(session_id)
+        @redis.multi do |multi|
+          multi.hset("session:#{session_id}",
+            "active_stream", false.to_json,
+            "stream_server", nil.to_json,
+            "last_activity", Time.now.to_f.to_json)
+          multi.expire("session:#{session_id}", @ttl)
+        end
+      end
+
+      def get_session_server(session_id)
+        server_data = @redis.hget("session:#{session_id}", "stream_server")
+        server_data ? JSON.parse(server_data) : nil
+      end
+
+      def session_exists?(session_id)
+        @redis.exists("session:#{session_id}") == 1
+      end
+
+      def session_has_active_stream?(session_id)
+        stream_data = @redis.hget("session:#{session_id}", "active_stream")
+        stream_data ? JSON.parse(stream_data) : false
+      end
+
+      def get_session_context(session_id)
+        context_data = @redis.hget("session:#{session_id}", "context")
+        context_data ? JSON.parse(context_data) : {}
+      end
+
+      def cleanup_session(session_id)
+        @redis.del("session:#{session_id}")
+      end
+
+      def route_message_to_session(session_id, message)
+        server_instance = get_session_server(session_id)
+        return false unless server_instance
+
+        # Publish to server-specific channel
+        @redis.publish("server:#{server_instance}:messages", {
+          session_id: session_id,
+          message: message
+        }.to_json)
+        true
+      end
+
+      def subscribe_to_server(server_instance, &block)
+        @redis.subscribe("server:#{server_instance}:messages") do |on|
+          on.message do |channel, message|
+            data = JSON.parse(message)
+            yield(data)
+          end
+        end
+      end
+
+      def get_all_active_sessions
+        keys = @redis.keys("session:*")
+        active_sessions = []
+
+        keys.each do |key|
+          session_id = key.sub("session:", "")
+          if session_has_active_stream?(session_id)
+            active_sessions << session_id
+          end
+        end
+
+        active_sessions
+      end
+    end
+  end
+end

data/lib/model_context_protocol/server/stdio_transport.rb

@@ -12,16 +12,19 @@ module ModelContextProtocol
       end
     end
 
-    attr_reader :logger, :router
+    attr_reader :router, :configuration
 
-    def initialize(logger:, router:)
-      @logger = logger
+    def initialize(router:, configuration:)
       @router = router
+      @configuration = configuration
     end
 
-    def begin
+    def handle
+      # Connect logger to transport
+      @configuration.logger.connect_transport(self)
+
       loop do
-        line = $stdin.gets
+        line = receive_message
         break unless line
 
         begin
@@ -31,18 +34,17 @@ module ModelContextProtocol
           result = router.route(message)
           send_message(Response[id: message["id"], result: result.serialized])
         rescue ModelContextProtocol::Server::ParameterValidationError => validation_error
-          log("Validation error: #{validation_error.message}")
+          @configuration.logger.error("Validation error", error: validation_error.message)
           send_message(
             ErrorResponse[id: message["id"], error: {code: -32602, message: validation_error.message}]
           )
         rescue JSON::ParserError => parser_error
-          log("Parser error: #{parser_error.message}")
+          @configuration.logger.error("Parser error", error: parser_error.message)
           send_message(
             ErrorResponse[id: "", error: {code: -32700, message: parser_error.message}]
           )
         rescue => error
-          log("Internal error: #{error.message}")
-          log(error.backtrace)
+          @configuration.logger.error("Internal error", error: error.message, backtrace: error.backtrace.first(5))
           send_message(
             ErrorResponse[id: message["id"], error: {code: -32603, message: error.message}]
           )
@@ -50,10 +52,23 @@ module ModelContextProtocol
       end
     end
 
+    def send_notification(method, params)
+      notification = {
+        jsonrpc: "2.0",
+        method: method,
+        params: params
+      }
+      $stdout.puts(JSON.generate(notification))
+      $stdout.flush
+    rescue IOError => e
+      # Handle broken pipe gracefully
+      @configuration.logger.debug("Failed to send notification", error: e.message) if @configuration.logging_enabled?
+    end
+
     private
 
-    def log(output, level = :error)
-      logger.send(level.to_sym, output)
+    def receive_message
+      $stdin.gets
     end
 
     def send_message(message)

data/lib/model_context_protocol/server/streamable_http_transport.rb

@@ -0,0 +1,291 @@
+require "json"
+require "securerandom"
+
+module ModelContextProtocol
+  class Server::StreamableHttpTransport
+    Response = Data.define(:id, :result) do
+      def serialized
+        {jsonrpc: "2.0", id:, result:}
+      end
+    end
+
+    ErrorResponse = Data.define(:id, :error) do
+      def serialized
+        {jsonrpc: "2.0", id:, error:}
+      end
+    end
+    def initialize(router:, configuration:)
+      @router = router
+      @configuration = configuration
+
+      transport_options = @configuration.transport_options
+      @redis = transport_options[:redis_client]
+
+      @session_store = ModelContextProtocol::Server::SessionStore.new(
+        @redis,
+        ttl: transport_options[:session_ttl] || 3600
+      )
+
+      @server_instance = "#{Socket.gethostname}-#{Process.pid}-#{SecureRandom.hex(4)}"
+      @local_streams = {}
+      @notification_queue = []
+
+      setup_redis_subscriber
+    end
+
+    def handle
+      @configuration.logger.connect_transport(self)
+
+      request = @configuration.transport_options[:request]
+      response = @configuration.transport_options[:response]
+
+      unless request && response
+        raise ArgumentError, "StreamableHTTP transport requires request and response objects in transport_options"
+      end
+
+      case request.method
+      when "POST"
+        handle_post_request(request)
+      when "GET"
+        handle_sse_request(request, response)
+      when "DELETE"
+        handle_delete_request(request)
+      else
+        error_response = ErrorResponse[id: nil, error: {code: -32601, message: "Method not allowed"}]
+        {json: error_response.serialized, status: 405}
+      end
+    end
+
+    def send_notification(method, params)
+      notification = {
+        jsonrpc: "2.0",
+        method: method,
+        params: params
+      }
+
+      if has_active_streams?
+        deliver_to_active_streams(notification)
+      else
+        @notification_queue << notification
+      end
+    end
+
+    private
+
+    def handle_post_request(request)
+      body_string = request.body.read
+      body = JSON.parse(body_string)
+      session_id = request.headers["Mcp-Session-Id"]
+
+      case body["method"]
+      when "initialize"
+        handle_initialization(body)
+      else
+        handle_regular_request(body, session_id)
+      end
+    rescue JSON::ParserError
+      error_response = ErrorResponse[id: "", error: {code: -32700, message: "Parse error"}]
+      {json: error_response.serialized, status: 400}
+    rescue ModelContextProtocol::Server::ParameterValidationError => validation_error
+      @configuration.logger.error("Validation error", error: validation_error.message)
+      error_response = ErrorResponse[id: body&.dig("id"), error: {code: -32602, message: validation_error.message}]
+      {json: error_response.serialized, status: 400}
+    rescue => e
+      @configuration.logger.error("Error handling POST request", error: e.message, backtrace: e.backtrace.first(5))
+      error_response = ErrorResponse[id: body&.dig("id"), error: {code: -32603, message: "Internal error"}]
+      {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
+      })
+
+      result = @router.route(body)
+      response = Response[id: body["id"], result: result.serialized]
+
+      {
+        json: response.serialized,
+        status: 200,
+        headers: {"Mcp-Session-Id" => session_id}
+      }
+    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}
+      end
+
+      result = @router.route(body)
+      response = Response[id: body["id"], result: result.serialized]
+
+      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}
+      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"}]
+        return {json: error_response.serialized, status: 400}
+      end
+
+      @session_store.mark_stream_active(session_id, @server_instance)
+
+      {
+        stream: true,
+        headers: {
+          "Content-Type" => "text/event-stream",
+          "Cache-Control" => "no-cache",
+          "Connection" => "keep-alive"
+        },
+        stream_proc: create_sse_stream_proc(session_id)
+      }
+    end
+
+    def handle_delete_request(request)
+      session_id = request.headers["Mcp-Session-Id"]
+
+      if session_id
+        cleanup_session(session_id)
+      end
+
+      {json: {success: true}, status: 200}
+    end
+
+    def create_sse_stream_proc(session_id)
+      proc do |stream|
+        register_local_stream(session_id, stream)
+
+        flush_notifications_to_stream(stream)
+
+        start_keepalive_thread(session_id, stream)
+
+        loop do
+          break unless stream_connected?(stream)
+          sleep 0.1
+        end
+      ensure
+        cleanup_local_stream(session_id)
+      end
+    end
+
+    def register_local_stream(session_id, stream)
+      @local_streams[session_id] = stream
+    end
+
+    def cleanup_local_stream(session_id)
+      @local_streams.delete(session_id)
+      @session_store.mark_stream_inactive(session_id)
+    end
+
+    def stream_connected?(stream)
+      return false unless stream
+
+      begin
+        stream.write(": ping\n\n")
+        stream.flush if stream.respond_to?(:flush)
+        true
+      rescue IOError, Errno::EPIPE, Errno::ECONNRESET
+        false
+      end
+    end
+
+    def start_keepalive_thread(session_id, stream)
+      Thread.new do
+        loop do
+          sleep 30
+          break unless stream_connected?(stream)
+
+          begin
+            send_ping_to_stream(stream)
+          rescue IOError, Errno::EPIPE, Errno::ECONNRESET
+            break
+          end
+        end
+      rescue => e
+        @configuration.logger.error("Keepalive thread error", error: e.message)
+      ensure
+        cleanup_local_stream(session_id)
+      end
+    end
+
+    def send_ping_to_stream(stream)
+      stream.write(": ping #{Time.now.iso8601}\n\n")
+      stream.flush if stream.respond_to?(:flush)
+    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)
+    end
+
+    def deliver_to_session_stream(session_id, data)
+      if @local_streams[session_id]
+        begin
+          send_to_stream(@local_streams[session_id], data)
+          return true
+        rescue IOError, Errno::EPIPE, Errno::ECONNRESET
+          cleanup_local_stream(session_id)
+        end
+      end
+
+      @session_store.route_message_to_session(session_id, data)
+    end
+
+    def cleanup_session(session_id)
+      cleanup_local_stream(session_id)
+      @session_store.cleanup_session(session_id)
+    end
+
+    def setup_redis_subscriber
+      Thread.new do
+        @session_store.subscribe_to_server(@server_instance) do |data|
+          session_id = data["session_id"]
+          message = data["message"]
+
+          if @local_streams[session_id]
+            begin
+              send_to_stream(@local_streams[session_id], message)
+            rescue IOError, Errno::EPIPE, Errno::ECONNRESET
+              cleanup_local_stream(session_id)
+            end
+          end
+        end
+      rescue => e
+        @configuration.logger.error("Redis subscriber error", error: e.message, backtrace: e.backtrace.first(5))
+        sleep 5
+        retry
+      end
+    end
+
+    def has_active_streams?
+      @local_streams.any?
+    end
+
+    def deliver_to_active_streams(notification)
+      @local_streams.each do |session_id, stream|
+        send_to_stream(stream, notification)
+      rescue IOError, Errno::EPIPE, Errno::ECONNRESET
+        cleanup_local_stream(session_id)
+      end
+    end
+
+    def flush_notifications_to_stream(stream)
+      while (notification = @notification_queue.shift)
+        send_to_stream(stream, notification)
+      end
+    end
+  end
+end

data/lib/model_context_protocol/server/tool.rb

@@ -2,11 +2,13 @@ require "json-schema"
 
 module ModelContextProtocol
   class Server::Tool
-    attr_reader :params
+    attr_reader :arguments, :context, :logger
 
-    def initialize(params)
-      validate!(params)
-      @params = params
+    def initialize(arguments, logger, context = {})
+      validate!(arguments)
+      @arguments = arguments
+      @context = context
+      @logger = logger
     end
 
     def call
@@ -68,8 +70,8 @@ module ModelContextProtocol
       end
     end
 
-    private def validate!(params)
-      JSON::Validator.validate!(self.class.input_schema, params)
+    private def validate!(arguments)
+      JSON::Validator.validate!(self.class.input_schema, arguments)
     end
 
     class << self
@@ -90,8 +92,8 @@ module ModelContextProtocol
         subclass.instance_variable_set(:@input_schema, @input_schema)
       end
 
-      def call(params)
-        new(params).call
+      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