model-context-protocol-rb 0.4.0 → 0.5.1

data/lib/model_context_protocol/server/cancellable.rb

@@ -0,0 +1,54 @@
+require "concurrent-ruby"
+
+module ModelContextProtocol
+  module Server::Cancellable
+    # Raised when a request has been cancelled by the client
+    class CancellationError < StandardError; end
+
+    # Execute a block with automatic cancellation support for blocking I/O operations.
+    # This method uses Concurrent::TimerTask to poll for cancellation every 100ms
+    # and can interrupt even blocking operations like HTTP requests or database queries.
+    #
+    # @param interval [Float] polling interval in seconds (default: 0.1)
+    # @yield block to execute with cancellation support
+    # @return [Object] the result of the block
+    # @raise [CancellationError] if the request is cancelled during execution
+    #
+    # @example
+    #   cancellable do
+    #     response = Net::HTTP.get(URI('https://slow-api.example.com'))
+    #     process_response(response)
+    #   end
+    def cancellable(interval: 0.1, &block)
+      context = Thread.current[:mcp_context]
+      executing_thread = Concurrent::AtomicReference.new(nil)
+
+      timer_task = Concurrent::TimerTask.new(execution_interval: interval) do
+        if context && context[:request_store] && context[:request_id]
+          if context[:request_store].cancelled?(context[:request_id])
+            thread = executing_thread.get
+            thread&.raise(CancellationError, "Request was cancelled") if thread&.alive?
+          end
+        end
+      end
+
+      begin
+        executing_thread.set(Thread.current)
+
+        if context && context[:request_store] && context[:request_id]
+          if context[:request_store].cancelled?(context[:request_id])
+            raise CancellationError, "Request #{context[:request_id]} was cancelled"
+          end
+        end
+
+        timer_task.execute
+
+        result = block.call
+        result
+      ensure
+        executing_thread.set(nil)
+        timer_task&.shutdown if timer_task&.running?
+      end
+    end
+  end
+end

data/lib/model_context_protocol/server/configuration.rb

@@ -200,15 +200,10 @@ module ModelContextProtocol
     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"
+      unless ModelContextProtocol::Server::RedisConfig.configured?
+        raise InvalidTransportError,
+          "streamable_http transport requires Redis configuration. " \
+          "Call ModelContextProtocol::Server.configure_redis in an initializer."
       end
     end
 

data/lib/model_context_protocol/server/progressable.rb

@@ -0,0 +1,72 @@
+require "concurrent-ruby"
+
+module ModelContextProtocol
+  module Server::Progressable
+    # Execute a block with automatic time-based progress reporting.
+    # Uses Concurrent::TimerTask to send progress notifications at regular intervals.
+    #
+    # @param max_duration [Numeric] Expected duration in seconds
+    # @param message [String, nil] Optional custom progress message
+    # @yield block to execute with progress tracking
+    # @return [Object] the result of the block
+    #
+    # @example
+    #   progressable(max_duration: 30) do  # 30 seconds
+    #     perform_long_operation
+    #   end
+    def progressable(max_duration:, message: nil, &block)
+      context = Thread.current[:mcp_context]
+
+      return yield unless context && context[:progress_token] && context[:transport]
+
+      progress_token = context[:progress_token]
+      transport = context[:transport]
+      start_time = Time.now
+      update_interval = [1.0, max_duration * 0.05].max
+
+      timer_task = Concurrent::TimerTask.new(execution_interval: update_interval) do
+        elapsed_seconds = Time.now - start_time
+        progress_pct = [(elapsed_seconds / max_duration) * 100, 99].min
+
+        progress_message = if message
+          "#{message} (#{elapsed_seconds.round(1)}s / ~#{max_duration}s)"
+        else
+          "Processing... (#{elapsed_seconds.round(1)}s / ~#{max_duration}s)"
+        end
+
+        begin
+          transport.send_notification("notifications/progress", {
+            progressToken: progress_token,
+            progress: progress_pct.round(1),
+            total: 100,
+            message: progress_message
+          })
+        rescue
+          nil
+        end
+
+        timer_task.shutdown if elapsed_seconds >= max_duration
+      end
+
+      begin
+        timer_task.execute
+        result = yield
+
+        begin
+          transport.send_notification("notifications/progress", {
+            progressToken: progress_token,
+            progress: 100,
+            total: 100,
+            message: "Completed"
+          })
+        rescue
+          nil
+        end
+
+        result
+      ensure
+        timer_task&.shutdown if timer_task&.running?
+      end
+    end
+  end
+end

data/lib/model_context_protocol/server/prompt.rb

@@ -1,6 +1,8 @@
 module ModelContextProtocol
   class Server::Prompt
-    include Server::ContentHelpers
+    include ModelContextProtocol::Server::Cancellable
+    include ModelContextProtocol::Server::ContentHelpers
+    include ModelContextProtocol::Server::Progressable
 
     attr_reader :arguments, :context, :logger
 

data/lib/model_context_protocol/server/redis_client_proxy.rb

@@ -0,0 +1,134 @@
+# frozen_string_literal: true
+
+module ModelContextProtocol
+  class Server
+    class RedisClientProxy
+      def initialize(pool)
+        @pool = pool
+      end
+
+      def get(key)
+        with_connection { |redis| redis.get(key) }
+      end
+
+      def set(key, value, **options)
+        with_connection { |redis| redis.set(key, value, **options) }
+      end
+
+      def del(*keys)
+        with_connection { |redis| redis.del(*keys) }
+      end
+
+      def exists(*keys)
+        with_connection { |redis| redis.exists(*keys) }
+      end
+
+      def expire(key, seconds)
+        with_connection { |redis| redis.expire(key, seconds) }
+      end
+
+      def ttl(key)
+        with_connection { |redis| redis.ttl(key) }
+      end
+
+      def hget(key, field)
+        with_connection { |redis| redis.hget(key, field) }
+      end
+
+      def hset(key, *args)
+        with_connection { |redis| redis.hset(key, *args) }
+      end
+
+      def hgetall(key)
+        with_connection { |redis| redis.hgetall(key) }
+      end
+
+      def lpush(key, *values)
+        with_connection { |redis| redis.lpush(key, *values) }
+      end
+
+      def rpop(key)
+        with_connection { |redis| redis.rpop(key) }
+      end
+
+      def lrange(key, start, stop)
+        with_connection { |redis| redis.lrange(key, start, stop) }
+      end
+
+      def llen(key)
+        with_connection { |redis| redis.llen(key) }
+      end
+
+      def ltrim(key, start, stop)
+        with_connection { |redis| redis.ltrim(key, start, stop) }
+      end
+
+      def incr(key)
+        with_connection { |redis| redis.incr(key) }
+      end
+
+      def decr(key)
+        with_connection { |redis| redis.decr(key) }
+      end
+
+      def keys(pattern)
+        with_connection { |redis| redis.keys(pattern) }
+      end
+
+      def multi(&block)
+        with_connection do |redis|
+          redis.multi do |multi|
+            multi_wrapper = RedisMultiWrapper.new(multi)
+            block.call(multi_wrapper)
+          end
+        end
+      end
+
+      def pipelined(&block)
+        with_connection do |redis|
+          redis.pipelined do |pipeline|
+            pipeline_wrapper = RedisMultiWrapper.new(pipeline)
+            block.call(pipeline_wrapper)
+          end
+        end
+      end
+
+      def mget(*keys)
+        with_connection { |redis| redis.mget(*keys) }
+      end
+
+      def eval(script, keys: [], argv: [])
+        with_connection { |redis| redis.eval(script, keys: keys, argv: argv) }
+      end
+
+      def ping
+        with_connection { |redis| redis.ping }
+      end
+
+      def flushdb
+        with_connection { |redis| redis.flushdb }
+      end
+
+      private
+
+      def with_connection(&block)
+        @pool.with(&block)
+      end
+
+      # Wrapper for Redis multi/pipeline operations
+      class RedisMultiWrapper
+        def initialize(multi)
+          @multi = multi
+        end
+
+        def method_missing(method, *args, **kwargs, &block)
+          @multi.send(method, *args, **kwargs, &block)
+        end
+
+        def respond_to_missing?(method, include_private = false)
+          @multi.respond_to?(method, include_private)
+        end
+      end
+    end
+  end
+end

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