model-context-protocol-rb 0.3.1 → 0.3.3

data/lib/model_context_protocol/server/configuration.rb

@@ -1,3 +1,5 @@
+require_relative "mcp_logger"
+
 module ModelContextProtocol
   class Server::Configuration
     # Raised when configured with invalid name.
@@ -12,16 +14,74 @@ module ModelContextProtocol
     # Raised when a required environment variable is not set
     class MissingRequiredEnvironmentVariable < StandardError; end
 
-    attr_accessor :enable_log, :name, :registry, :version
+    # Raised when transport configuration is invalid
+    class InvalidTransportError < StandardError; end
+
+    # Raised when an invalid log level is provided
+    class InvalidLogLevelError < StandardError; end
+
+    # Valid MCP log levels per the specification
+    VALID_LOG_LEVELS = %w[debug info notice warning error critical alert emergency].freeze
+
+    attr_accessor :name, :registry, :version, :transport
+    attr_reader :logger
+
+    def initialize
+      # Always create a logger - enabled by default
+      @logging_enabled = true
+      @default_log_level = "info"
+      @logger = ModelContextProtocol::Server::MCPLogger.new(
+        logger_name: "server",
+        level: @default_log_level,
+        enabled: @logging_enabled
+      )
+    end
 
     def logging_enabled?
-      enable_log || false
+      @logging_enabled
+    end
+
+    def logging_enabled=(value)
+      @logging_enabled = value
+      @logger = ModelContextProtocol::Server::MCPLogger.new(
+        logger_name: "server",
+        level: @default_log_level,
+        enabled: value
+      )
+    end
+
+    def default_log_level=(level)
+      unless VALID_LOG_LEVELS.include?(level.to_s)
+        raise InvalidLogLevelError, "Invalid log level: #{level}. Valid levels are: #{VALID_LOG_LEVELS.join(", ")}"
+      end
+
+      @default_log_level = level.to_s
+      @logger.set_mcp_level(@default_log_level)
+    end
+
+    def transport_type
+      case transport
+      when Hash
+        transport[:type] || transport["type"]
+      when Symbol, String
+        transport.to_sym
+      end
+    end
+
+    def transport_options
+      case transport
+      when Hash
+        transport.except(:type, "type").transform_keys(&:to_sym)
+      else
+        {}
+      end
     end
 
     def validate!
       raise InvalidServerNameError unless valid_name?
       raise InvalidRegistryError unless valid_registry?
       raise InvalidServerVersionError unless valid_version?
+      validate_transport!
 
       validate_environment_variables!
     end
@@ -51,6 +111,14 @@ module ModelContextProtocol
       environment_variables[key.to_s.upcase] = value
     end
 
+    def context
+      @context ||= {}
+    end
+
+    def context=(context_hash = {})
+      @context = context_hash
+    end
+
     private
 
     def required_environment_variables
@@ -74,5 +142,29 @@ module ModelContextProtocol
     def valid_version?
       version&.is_a?(String)
     end
+
+    def validate_transport!
+      case transport_type
+      when :streamable_http
+        validate_streamable_http_transport!
+      when :stdio, nil
+        # stdio transport has no required options
+      else
+        raise InvalidTransportError, "Unknown transport type: #{transport_type}" if transport_type
+      end
+    end
+
+    def validate_streamable_http_transport!
+      options = transport_options
+
+      unless options[:redis_client]
+        raise InvalidTransportError, "streamable_http transport requires redis_client option"
+      end
+
+      redis_client = options[:redis_client]
+      unless redis_client.respond_to?(:hset) && redis_client.respond_to?(:expire)
+        raise InvalidTransportError, "redis_client must be a Redis-compatible client"
+      end
+    end
   end
 end

data/lib/model_context_protocol/server/mcp_logger.rb

@@ -0,0 +1,109 @@
+require "logger"
+require "forwardable"
+require "json"
+
+module ModelContextProtocol
+  class Server::MCPLogger
+    extend Forwardable
+
+    def_delegators :@internal_logger, :datetime_format=, :formatter=, :progname, :progname=
+
+    LEVEL_MAP = {
+      "debug" => Logger::DEBUG,
+      "info" => Logger::INFO,
+      "notice" => Logger::INFO,
+      "warning" => Logger::WARN,
+      "error" => Logger::ERROR,
+      "critical" => Logger::FATAL,
+      "alert" => Logger::FATAL,
+      "emergency" => Logger::UNKNOWN
+    }.freeze
+
+    REVERSE_LEVEL_MAP = {
+      Logger::DEBUG => "debug",
+      Logger::INFO => "info",
+      Logger::WARN => "warning",
+      Logger::ERROR => "error",
+      Logger::FATAL => "critical",
+      Logger::UNKNOWN => "emergency"
+    }.freeze
+
+    attr_accessor :transport
+    attr_reader :logger_name, :enabled
+
+    def initialize(logger_name: "server", level: "info", enabled: true)
+      @logger_name = logger_name
+      @enabled = enabled
+      @internal_logger = Logger.new(nil)
+      @internal_logger.level = LEVEL_MAP[level] || Logger::INFO
+      @transport = nil
+      @queued_messages = []
+    end
+
+    %i[debug info warn error fatal unknown].each do |severity|
+      define_method(severity) do |message = nil, **data, &block|
+        return true unless @enabled
+        add(Logger.const_get(severity.to_s.upcase), message, data, &block)
+      end
+    end
+
+    def add(severity, message = nil, data = {}, &block)
+      return true unless @enabled
+      return true if severity < @internal_logger.level
+
+      message = block.call if message.nil? && block_given?
+      send_notification(severity, message, data)
+      true
+    end
+
+    def level=(value)
+      @internal_logger.level = value
+    end
+
+    def level
+      @internal_logger.level
+    end
+
+    def set_mcp_level(mcp_level)
+      self.level = LEVEL_MAP[mcp_level] || Logger::INFO
+    end
+
+    def connect_transport(transport)
+      @transport = transport
+      flush_queued_messages if @enabled
+    end
+
+    private
+
+    def send_notification(severity, message, data)
+      return unless @enabled
+
+      notification_params = {
+        level: REVERSE_LEVEL_MAP[severity] || "info",
+        logger: @logger_name,
+        data: format_data(message, data)
+      }
+
+      if @transport
+        @transport.send_notification("notifications/message", notification_params)
+      else
+        @queued_messages << notification_params
+      end
+    end
+
+    def format_data(message, additional_data)
+      data = {}
+      data[:message] = message.to_s if message
+      data.merge!(additional_data) unless additional_data.empty?
+      data
+    end
+
+    def flush_queued_messages
+      return unless @transport && @enabled
+      @queued_messages.each do |params|
+        @transport.send_notification("notifications/message", params)
+      end
+      @queued_messages.clear
+    end
+  end
+end

data/lib/model_context_protocol/server/prompt.rb

@@ -1,32 +1,33 @@
 module ModelContextProtocol
   class Server::Prompt
-    attr_reader :params, :description
+    attr_reader :params, :context, :logger
 
-    def initialize(params)
+    def initialize(params, logger, context = {})
       validate!(params)
-      @description = self.class.description
       @params = params
+      @context = context
+      @logger = logger
     end
 
     def call
       raise NotImplementedError, "Subclasses must implement the call method"
     end
 
-    Response = Data.define(:messages, :prompt) do
+    Response = Data.define(:messages, :description) do
       def serialized
-        {description: prompt.description, messages:}
+        {description:, messages:}
       end
     end
     private_constant :Response
 
     private def respond_with(messages:)
-      Response[messages:, prompt: self]
+      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] }
+      required_args = arguments.select { |arg| arg[:required] }.map { |arg| arg[:name].to_sym }
+      valid_arg_names = arguments.map { |arg| arg[:name].to_sym }
 
       missing_args = required_args - params.keys
       unless missing_args.empty?
@@ -45,21 +46,37 @@ module ModelContextProtocol
       attr_reader :name, :description, :arguments
 
       def with_metadata(&block)
-        metadata = instance_eval(&block)
+        @arguments ||= []
 
-        @name = metadata[:name]
-        @description = metadata[:description]
-        @arguments = metadata[:arguments]
+        metadata_dsl = MetadataDSL.new
+        metadata_dsl.instance_eval(&block)
+
+        @name = metadata_dsl.name
+        @description = metadata_dsl.description
+      end
+
+      def with_argument(&block)
+        @arguments ||= []
+
+        argument_dsl = ArgumentDSL.new
+        argument_dsl.instance_eval(&block)
+
+        @arguments << {
+          name: argument_dsl.name,
+          description: argument_dsl.description,
+          required: argument_dsl.required,
+          completion: argument_dsl.completion
+        }
       end
 
       def inherited(subclass)
         subclass.instance_variable_set(:@name, @name)
         subclass.instance_variable_set(:@description, @description)
-        subclass.instance_variable_set(:@arguments, @arguments)
+        subclass.instance_variable_set(:@arguments, @arguments&.dup)
       end
 
-      def call(params)
-        new(params).call
+      def call(params, logger, context = {})
+        new(params, logger, context).call
       rescue ArgumentError => error
         raise ModelContextProtocol::Server::ParameterValidationError, error.message
       end
@@ -67,6 +84,46 @@ module ModelContextProtocol
       def metadata
         {name: @name, description: @description, arguments: @arguments}
       end
+
+      def complete_for(arg_name, value)
+        arg = @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
+    end
+
+    class MetadataDSL
+      def name(value = nil)
+        @name = value if value
+        @name
+      end
+
+      def description(value = nil)
+        @description = value if value
+        @description
+      end
+    end
+
+    class ArgumentDSL
+      def name(value = nil)
+        @name = value if value
+        @name
+      end
+
+      def description(value = nil)
+        @description = value if value
+        @description
+      end
+
+      def required(value = nil)
+        @required = value unless value.nil?
+        @required
+      end
+
+      def completion(klass = nil)
+        @completion = klass unless klass.nil?
+        @completion
+      end
     end
   end
 end

data/lib/model_context_protocol/server/registry.rb

@@ -12,6 +12,7 @@ module ModelContextProtocol
     def initialize
       @prompts = []
       @resources = []
+      @resource_templates = []
       @tools = []
       @prompts_options = {}
       @resources_options = {}
@@ -28,6 +29,10 @@ module ModelContextProtocol
       instance_eval(&block) if block
     end
 
+    def resource_templates(&block)
+      instance_eval(&block) if block
+    end
+
     def tools(options = {}, &block)
       @tools_options = options
       instance_eval(&block) if block
@@ -42,6 +47,8 @@ module ModelContextProtocol
         @prompts << entry
       when ->(ancestors) { ancestors.include?(ModelContextProtocol::Server::Resource) }
         @resources << entry
+      when ->(ancestors) { ancestors.include?(ModelContextProtocol::Server::ResourceTemplate) }
+        @resource_templates << entry
       when ->(ancestors) { ancestors.include?(ModelContextProtocol::Server::Tool) }
         @tools << entry
       else
@@ -58,6 +65,11 @@ module ModelContextProtocol
       entry ? entry[:klass] : nil
     end
 
+    def find_resource_template(uri)
+      entry = @resource_templates.find { |r| uri == r[:uriTemplate] }
+      entry ? entry[:klass] : nil
+    end
+
     def find_tool(name)
       find_by_name(@tools, name)
     end
@@ -70,6 +82,10 @@ module ModelContextProtocol
       ResourcesData[resources: @resources.map { |entry| entry.except(:klass) }]
     end
 
+    def resource_templates_data
+      ResourceTemplatesData[resource_templates: @resource_templates.map { |entry| entry.except(:klass, :completions) }]
+    end
+
     def tools_data
       ToolsData[tools: @tools.map { |entry| entry.except(:klass) }]
     end
@@ -88,6 +104,12 @@ module ModelContextProtocol
       end
     end
 
+    ResourceTemplatesData = Data.define(:resource_templates) do
+      def serialized
+        {resourceTemplates: resource_templates}
+      end
+    end
+
     ToolsData = Data.define(:tools) do
       def serialized
         {tools:}

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
@@ -40,12 +42,13 @@ module ModelContextProtocol
       attr_reader :name, :description, :mime_type, :uri
 
       def with_metadata(&block)
-        metadata = instance_eval(&block)
+        metadata_dsl = MetadataDSL.new
+        metadata_dsl.instance_eval(&block)
 
-        @name = metadata[:name]
-        @description = metadata[:description]
-        @mime_type = metadata[:mime_type]
-        @uri = metadata[:uri]
+        @name = metadata_dsl.name
+        @description = metadata_dsl.description
+        @mime_type = metadata_dsl.mime_type
+        @uri = metadata_dsl.uri
       end
 
       def inherited(subclass)
@@ -55,12 +58,34 @@ module ModelContextProtocol
         subclass.instance_variable_set(:@uri, @uri)
       end
 
-      def call
-        new.call
+      def call(logger, context = {})
+        new(logger, context).call
       end
 
       def metadata
-        {name: @name, description: @description, mime_type: @mime_type, uri: @uri}
+        {name: @name, description: @description, mimeType: @mime_type, uri: @uri}
+      end
+    end
+
+    class MetadataDSL
+      def name(value = nil)
+        @name = value if value
+        @name
+      end
+
+      def description(value = nil)
+        @description = value if value
+        @description
+      end
+
+      def mime_type(value = nil)
+        @mime_type = value if value
+        @mime_type
+      end
+
+      def uri(value = nil)
+        @uri = value if value
+        @uri
       end
     end
   end

data/lib/model_context_protocol/server/resource_template.rb

@@ -0,0 +1,93 @@
+module ModelContextProtocol
+  class Server::ResourceTemplate
+    class << self
+      attr_reader :name, :description, :mime_type, :uri_template, :completions
+
+      def with_metadata(&block)
+        metadata_dsl = MetadataDSL.new
+        metadata_dsl.instance_eval(&block)
+
+        @name = metadata_dsl.name
+        @description = metadata_dsl.description
+        @mime_type = metadata_dsl.mime_type
+        @uri_template = metadata_dsl.uri_template
+        @completions = metadata_dsl.completions
+      end
+
+      def inherited(subclass)
+        subclass.instance_variable_set(:@name, @name)
+        subclass.instance_variable_set(:@description, @description)
+        subclass.instance_variable_set(:@mime_type, @mime_type)
+        subclass.instance_variable_set(:@uri_template, @uri_template)
+        subclass.instance_variable_set(:@completions, @completions&.dup)
+      end
+
+      def complete_for(param_name, value)
+        completion = if @completions && @completions[param_name.to_s]
+          @completions[param_name.to_s]
+        else
+          ModelContextProtocol::Server::NullCompletion
+        end
+
+        completion.call(param_name.to_s, value)
+      end
+
+      def metadata
+        {
+          name: @name,
+          description: @description,
+          mimeType: @mime_type,
+          uriTemplate: @uri_template,
+          completions: @completions&.transform_keys(&:to_s)
+        }
+      end
+    end
+
+    class MetadataDSL
+      attr_reader :completions
+
+      def initialize
+        @completions = {}
+      end
+
+      def name(value = nil)
+        @name = value if value
+        @name
+      end
+
+      def description(value = nil)
+        @description = value if value
+        @description
+      end
+
+      def mime_type(value = nil)
+        @mime_type = value if value
+        @mime_type
+      end
+
+      def uri_template(value = nil, &block)
+        @uri_template = value if value
+
+        if block_given?
+          completion_dsl = CompletionDSL.new
+          completion_dsl.instance_eval(&block)
+          @completions = completion_dsl.completions
+        end
+
+        @uri_template
+      end
+    end
+
+    class CompletionDSL
+      attr_reader :completions
+
+      def initialize
+        @completions = {}
+      end
+
+      def completion(param_name, completion_class)
+        @completions[param_name.to_s] = completion_class
+      end
+    end
+  end
+end

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