model-context-protocol-rb 0.4.0 → 0.5.0

data/lib/model_context_protocol/server/redis_config.rb

@@ -0,0 +1,108 @@
+require "singleton"
+
+module ModelContextProtocol
+  class Server::RedisConfig
+    include Singleton
+
+    class NotConfiguredError < StandardError
+      def initialize
+        super("Redis not configured. Call ModelContextProtocol::Server.configure_redis first")
+      end
+    end
+
+    attr_reader :manager
+
+    def self.configure(&block)
+      instance.configure(&block)
+    end
+
+    def self.configured?
+      instance.configured?
+    end
+
+    def self.pool
+      instance.pool
+    end
+
+    def self.shutdown!
+      instance.shutdown!
+    end
+
+    def self.reset!
+      instance.reset!
+    end
+
+    def self.stats
+      instance.stats
+    end
+
+    def self.pool_manager
+      instance.manager
+    end
+
+    def initialize
+      reset!
+    end
+
+    def configure(&block)
+      shutdown! if configured?
+
+      config = Configuration.new
+      yield(config) if block_given?
+
+      @manager = Server::RedisPoolManager.new(
+        redis_url: config.redis_url,
+        pool_size: config.pool_size,
+        pool_timeout: config.pool_timeout
+      )
+
+      if config.enable_reaper
+        @manager.configure_reaper(
+          enabled: true,
+          interval: config.reaper_interval,
+          idle_timeout: config.idle_timeout
+        )
+      end
+
+      @manager.start
+    end
+
+    def configured?
+      !@manager.nil? && !@manager.pool.nil?
+    end
+
+    def pool
+      raise NotConfiguredError unless configured?
+      @manager.pool
+    end
+
+    def shutdown!
+      @manager&.shutdown
+      @manager = nil
+    end
+
+    def reset!
+      shutdown!
+      @manager = nil
+    end
+
+    def stats
+      return {} unless configured?
+      @manager.stats
+    end
+
+    class Configuration
+      attr_accessor :redis_url, :pool_size, :pool_timeout,
+        :enable_reaper, :reaper_interval, :idle_timeout
+
+      def initialize
+        @redis_url = nil
+        @pool_size = 20
+        @pool_timeout = 5
+        @enable_reaper = true
+        @reaper_interval = 60
+        @idle_timeout = 300
+      end
+    end
+  end
+end

data/lib/model_context_protocol/server/redis_pool_manager.rb

@@ -0,0 +1,110 @@
+module ModelContextProtocol
+  class Server::RedisPoolManager
+    attr_reader :pool, :reaper_thread
+
+    def initialize(redis_url:, pool_size: 20, pool_timeout: 5)
+      @redis_url = redis_url
+      @pool_size = pool_size
+      @pool_timeout = pool_timeout
+      @pool = nil
+      @reaper_thread = nil
+      @reaper_config = {
+        enabled: false,
+        interval: 60,
+        idle_timeout: 300
+      }
+    end
+
+    def configure_reaper(enabled:, interval: 60, idle_timeout: 300)
+      @reaper_config = {
+        enabled: enabled,
+        interval: interval,
+        idle_timeout: idle_timeout
+      }
+    end
+
+    def start
+      validate!
+      create_pool
+      start_reaper if @reaper_config[:enabled]
+      true
+    end
+
+    def shutdown
+      stop_reaper
+      close_pool
+    end
+
+    def healthy?
+      return false unless @pool
+
+      @pool.with do |conn|
+        conn.ping == "PONG"
+      end
+    rescue
+      false
+    end
+
+    def reap_now
+      return unless @pool
+
+      @pool.reap(@reaper_config[:idle_timeout]) do |conn|
+        conn.close
+      end
+    end
+
+    def stats
+      return {} unless @pool
+
+      {
+        size: @pool.size,
+        available: @pool.available,
+        idle: @pool.instance_variable_get(:@idle_since)&.size || 0
+      }
+    end
+
+    private
+
+    def validate!
+      raise ArgumentError, "redis_url is required" if @redis_url.nil? || @redis_url.empty?
+      raise ArgumentError, "pool_size must be positive" if @pool_size <= 0
+      raise ArgumentError, "pool_timeout must be positive" if @pool_timeout <= 0
+    end
+
+    def create_pool
+      @pool = ConnectionPool.new(size: @pool_size, timeout: @pool_timeout) do
+        Redis.new(url: @redis_url)
+      end
+    end
+
+    def close_pool
+      @pool&.shutdown { |conn| conn.close }
+      @pool = nil
+    end
+
+    def start_reaper
+      return if @reaper_thread&.alive?
+
+      @reaper_thread = Thread.new do
+        loop do
+          sleep @reaper_config[:interval]
+          begin
+            reap_now
+          rescue => e
+            warn "Redis reaper error: #{e.message}"
+          end
+        end
+      end
+
+      @reaper_thread.name = "MCP-Redis-Reaper"
+    end
+
+    def stop_reaper
+      return unless @reaper_thread&.alive?
+
+      @reaper_thread.kill
+      @reaper_thread.join(5)
+      @reaper_thread = nil
+    end
+  end
+end

data/lib/model_context_protocol/server/resource.rb

@@ -1,5 +1,8 @@
 module ModelContextProtocol
   class Server::Resource
+    include ModelContextProtocol::Server::Cancellable
+    include ModelContextProtocol::Server::Progressable
+
     attr_reader :mime_type, :uri
 
     def initialize

data/lib/model_context_protocol/server/router.rb

@@ -1,3 +1,5 @@
+require_relative "cancellable"
+
 module ModelContextProtocol
   class Server::Router
     # Raised when an invalid method is provided.
@@ -12,14 +14,45 @@ module ModelContextProtocol
       @handlers[method] = handler
     end
 
-    def route(message)
+    # Route a message to its handler with request tracking support
+    #
+    # @param message [Hash] the JSON-RPC message
+    # @param request_store [Object] the request store for tracking cancellation
+    # @param session_id [String, nil] the session ID for HTTP transport
+    # @param transport [Object, nil] the transport for sending notifications
+    # @return [Object] the handler result, or nil if cancelled
+    def route(message, request_store: nil, session_id: nil, transport: nil)
       method = message["method"]
       handler = @handlers[method]
       raise MethodNotFoundError, "Method not found: #{method}" unless handler
 
-      with_environment(@configuration&.environment_variables) do
-        handler.call(message)
+      request_id = message["id"]
+      progress_token = message.dig("params", "_meta", "progressToken")
+
+      if request_id && request_store
+        request_store.register_request(request_id, session_id)
+      end
+
+      result = nil
+      begin
+        with_environment(@configuration&.environment_variables) do
+          context = {request_id:, request_store:, session_id:, progress_token:, transport:}
+
+          Thread.current[:mcp_context] = context
+
+          result = handler.call(message)
+        end
+      rescue Server::Cancellable::CancellationError
+        return nil
+      ensure
+        if request_id && request_store
+          request_store.unregister_request(request_id)
+        end
+
+        Thread.current[:mcp_context] = nil
       end
+
+      result
     end
 
     private

data/lib/model_context_protocol/server/stdio_transport/request_store.rb

@@ -0,0 +1,102 @@
+module ModelContextProtocol
+  class Server::StdioTransport
+    # Thread-safe in-memory storage for tracking active requests and their cancellation status.
+    # This store is used by StdioTransport to manage request lifecycle and handle cancellation.
+    class RequestStore
+      def initialize
+        @mutex = Mutex.new
+        @requests = {}
+      end
+
+      # Register a new request with its associated thread
+      #
+      # @param request_id [String] the unique request identifier
+      # @param thread [Thread] the thread processing this request (defaults to current thread)
+      # @return [void]
+      def register_request(request_id, thread = Thread.current)
+        @mutex.synchronize do
+          @requests[request_id] = {
+            thread:,
+            cancelled: false,
+            started_at: Time.now
+          }
+        end
+      end
+
+      # Mark a request as cancelled
+      #
+      # @param request_id [String] the unique request identifier
+      # @return [Boolean] true if request was found and marked cancelled, false otherwise
+      def mark_cancelled(request_id)
+        @mutex.synchronize do
+          if (request = @requests[request_id])
+            request[:cancelled] = true
+            return true
+          end
+          false
+        end
+      end
+
+      # Check if a request has been cancelled
+      #
+      # @param request_id [String] the unique request identifier
+      # @return [Boolean] true if the request is cancelled, false otherwise
+      def cancelled?(request_id)
+        @mutex.synchronize do
+          @requests[request_id]&.fetch(:cancelled, false) || false
+        end
+      end
+
+      # Unregister a request (typically called when request completes)
+      #
+      # @param request_id [String] the unique request identifier
+      # @return [Hash, nil] the removed request data, or nil if not found
+      def unregister_request(request_id)
+        @mutex.synchronize do
+          @requests.delete(request_id)
+        end
+      end
+
+      # Get information about a specific request
+      #
+      # @param request_id [String] the unique request identifier
+      # @return [Hash, nil] request information or nil if not found
+      def get_request(request_id)
+        @mutex.synchronize do
+          @requests[request_id]&.dup
+        end
+      end
+
+      # Get all active request IDs
+      #
+      # @return [Array<String>] list of active request IDs
+      def active_requests
+        @mutex.synchronize do
+          @requests.keys.dup
+        end
+      end
+
+      # Clean up old requests (useful for preventing memory leaks)
+      #
+      # @param max_age_seconds [Integer] maximum age of requests to keep
+      # @return [Array<String>] list of cleaned up request IDs
+      def cleanup_old_requests(max_age_seconds = 300)
+        cutoff_time = Time.now - max_age_seconds
+        removed_ids = []
+
+        @mutex.synchronize do
+          @requests.delete_if do |request_id, data|
+            if data[:started_at] < cutoff_time
+              removed_ids << request_id
+              true
+            else
+              false
+            end
+          end
+        end
+
+        removed_ids
+      end
+    end
+  end
+end

data/lib/model_context_protocol/server/stdio_transport.rb

@@ -1,3 +1,5 @@
+require_relative "stdio_transport/request_store"
+
 module ModelContextProtocol
   class Server::StdioTransport
     Response = Data.define(:id, :result) do
@@ -12,15 +14,15 @@ module ModelContextProtocol
       end
     end
 
-    attr_reader :router, :configuration
+    attr_reader :router, :configuration, :request_store
 
     def initialize(router:, configuration:)
       @router = router
       @configuration = configuration
+      @request_store = RequestStore.new
     end
 
     def handle
-      # Connect logger to transport
       @configuration.logger.connect_transport(self)
 
       loop do
@@ -29,10 +31,19 @@ module ModelContextProtocol
 
         begin
           message = JSON.parse(line.chomp)
-          next if message["method"].start_with?("notifications")
 
-          result = router.route(message)
-          send_message(Response[id: message["id"], result: result.serialized])
+          if message["method"] == "notifications/cancelled"
+            handle_cancellation(message)
+            next
+          end
+
+          next if message["method"]&.start_with?("notifications/")
+
+          result = router.route(message, request_store: @request_store, transport: self)
+
+          if result
+            send_message(Response[id: message["id"], result: result.serialized])
+          end
         rescue ModelContextProtocol::Server::ParameterValidationError => validation_error
           @configuration.logger.error("Validation error", error: validation_error.message)
           send_message(
@@ -61,12 +72,26 @@ module ModelContextProtocol
       $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
 
+    # Handle a cancellation notification from the client
+    #
+    # @param message [Hash] the cancellation notification message
+    def handle_cancellation(message)
+      params = message["params"]
+      return unless params
+
+      request_id = params["requestId"]
+      return unless request_id
+
+      @request_store.mark_cancelled(request_id)
+    rescue
+      nil
+    end
+
     def receive_message
       $stdin.gets
     end

data/lib/model_context_protocol/server/streamable_http_transport/event_counter.rb

@@ -0,0 +1,35 @@
+module ModelContextProtocol
+  class Server::StreamableHttpTransport
+    class EventCounter
+      COUNTER_KEY_PREFIX = "event_counter:"
+
+      def initialize(redis_client, server_instance)
+        @redis = redis_client
+        @server_instance = server_instance
+        @counter_key = "#{COUNTER_KEY_PREFIX}#{server_instance}"
+
+        if @redis.exists(@counter_key) == 0
+          @redis.set(@counter_key, 0)
+        end
+      end
+
+      def next_event_id
+        count = @redis.incr(@counter_key)
+        "#{@server_instance}-#{count}"
+      end
+
+      def current_count
+        count = @redis.get(@counter_key)
+        count ? count.to_i : 0
+      end
+
+      def reset
+        @redis.set(@counter_key, 0)
+      end
+
+      def set_count(value)
+        @redis.set(@counter_key, value.to_i)
+      end
+    end
+  end
+end

data/lib/model_context_protocol/server/streamable_http_transport/message_poller.rb

@@ -0,0 +1,101 @@
+require_relative "session_message_queue"
+
+module ModelContextProtocol
+  class Server::StreamableHttpTransport
+    class MessagePoller
+      POLL_INTERVAL = 0.1  # 100ms
+      BATCH_SIZE = 100
+
+      def initialize(redis_client, stream_registry, logger, &message_delivery_block)
+        @redis = redis_client
+        @stream_registry = stream_registry
+        @logger = logger
+        @message_delivery_block = message_delivery_block
+        @running = false
+        @poll_thread = nil
+      end
+
+      def start
+        return if @running
+
+        @running = true
+        @poll_thread = Thread.new do
+          poll_loop
+        rescue => e
+          @logger.error("Message poller thread error", error: e.message, backtrace: e.backtrace.first(5))
+          sleep 1
+          retry if @running
+        end
+
+        @poll_thread.name = "MCP-MessagePoller" if @poll_thread.respond_to?(:name=)
+
+        @logger.debug("Message poller started")
+      end
+
+      def stop
+        @running = false
+
+        if @poll_thread&.alive?
+          @poll_thread.kill
+          @poll_thread.join(timeout: 5)
+        end
+
+        @poll_thread = nil
+        @logger.debug("Message poller stopped")
+      end
+
+      def running?
+        @running && @poll_thread&.alive?
+      end
+
+      private
+
+      def poll_loop
+        while @running
+          begin
+            poll_and_deliver_messages
+          rescue => e
+            @logger.error("Error in message polling", error: e.message)
+          end
+
+          sleep POLL_INTERVAL
+        end
+      end
+
+      def poll_and_deliver_messages
+        local_sessions = @stream_registry.get_all_local_streams.keys
+        return if local_sessions.empty?
+
+        local_sessions.each_slice(BATCH_SIZE) do |session_batch|
+          poll_sessions_batch(session_batch)
+        end
+      end
+
+      def poll_sessions_batch(session_ids)
+        session_ids.each do |session_id|
+          queue = SessionMessageQueue.new(@redis, session_id)
+          messages = queue.poll_messages
+
+          next if messages.empty?
+
+          stream = @stream_registry.get_local_stream(session_id)
+          next unless stream
+
+          messages.each do |message|
+            deliver_message_to_stream(stream, message, session_id)
+          end
+        end
+      end
+
+      def deliver_message_to_stream(stream, message, session_id)
+        @message_delivery_block&.call(stream, message)
+      rescue IOError, Errno::EPIPE, Errno::ECONNRESET
+        @stream_registry.unregister_stream(session_id)
+        @logger.debug("Unregistered disconnected stream", session_id: session_id)
+      rescue => e
+        @logger.error("Error delivering message to stream",
+          session_id: session_id, error: e.message)
+      end
+    end
+  end
+end

data/lib/model_context_protocol/server/streamable_http_transport/notification_queue.rb

@@ -0,0 +1,80 @@
+require "json"
+
+module ModelContextProtocol
+  class Server::StreamableHttpTransport
+    class NotificationQueue
+      QUEUE_KEY_PREFIX = "notifications:"
+      DEFAULT_MAX_SIZE = 1000
+
+      def initialize(redis_client, server_instance, max_size: DEFAULT_MAX_SIZE)
+        @redis = redis_client
+        @server_instance = server_instance
+        @queue_key = "#{QUEUE_KEY_PREFIX}#{server_instance}"
+        @max_size = max_size
+      end
+
+      def push(notification)
+        notification_json = notification.to_json
+
+        @redis.multi do |multi|
+          multi.lpush(@queue_key, notification_json)
+          multi.ltrim(@queue_key, 0, @max_size - 1)
+        end
+      end
+
+      def pop
+        notification_json = @redis.rpop(@queue_key)
+        return nil unless notification_json
+
+        JSON.parse(notification_json)
+      end
+
+      def pop_all
+        notification_jsons = @redis.multi do |multi|
+          multi.lrange(@queue_key, 0, -1)
+          multi.del(@queue_key)
+        end.first
+
+        return [] if notification_jsons.empty?
+
+        notification_jsons.reverse.map do |notification_json|
+          JSON.parse(notification_json)
+        end
+      end
+
+      def peek_all
+        notification_jsons = @redis.lrange(@queue_key, 0, -1)
+        return [] if notification_jsons.empty?
+
+        notification_jsons.reverse.map do |notification_json|
+          JSON.parse(notification_json)
+        end
+      end
+
+      def size
+        @redis.llen(@queue_key)
+      end
+
+      def empty?
+        size == 0
+      end
+
+      def clear
+        @redis.del(@queue_key)
+      end
+
+      def push_bulk(notifications)
+        return if notifications.empty?
+
+        notification_jsons = notifications.map(&:to_json)
+
+        @redis.multi do |multi|
+          notification_jsons.each do |json|
+            multi.lpush(@queue_key, json)
+          end
+          multi.ltrim(@queue_key, 0, @max_size - 1)
+        end
+      end
+    end
+  end
+end