model-context-protocol-rb 0.4.0 → 0.5.1

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

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

@@ -0,0 +1,224 @@
+require "json"
+
+module ModelContextProtocol
+  class Server::StreamableHttpTransport
+    # Redis-based distributed storage for tracking active requests and their cancellation status.
+    # This store is used by StreamableHttpTransport to manage request lifecycle across multiple
+    # server instances and handle cancellation in a distributed environment.
+    class RequestStore
+      REQUEST_KEY_PREFIX = "request:active:"
+      CANCELLED_KEY_PREFIX = "request:cancelled:"
+      SESSION_KEY_PREFIX = "request:session:"
+      DEFAULT_TTL = 60  # 1 minute TTL for request entries
+
+      def initialize(redis_client, server_instance, ttl: DEFAULT_TTL)
+        @redis = redis_client
+        @server_instance = server_instance
+        @ttl = ttl
+      end
+
+      # Register a new request with its associated session
+      #
+      # @param request_id [String] the unique request identifier
+      # @param session_id [String] the session identifier (can be nil for sessionless requests)
+      # @return [void]
+      def register_request(request_id, session_id = nil)
+        request_data = {
+          session_id: session_id,
+          server_instance: @server_instance,
+          started_at: Time.now.to_f
+        }
+
+        @redis.multi do |multi|
+          multi.set("#{REQUEST_KEY_PREFIX}#{request_id}",
+            request_data.to_json, ex: @ttl)
+
+          if session_id
+            multi.set("#{SESSION_KEY_PREFIX}#{session_id}:#{request_id}",
+              true, ex: @ttl)
+          end
+        end
+      end
+
+      # Mark a request as cancelled
+      #
+      # @param request_id [String] the unique request identifier
+      # @param reason [String] optional reason for cancellation
+      # @return [Boolean] true if cancellation was recorded
+      def mark_cancelled(request_id, reason = nil)
+        cancellation_data = {
+          cancelled_at: Time.now.to_f,
+          reason: reason
+        }
+
+        result = @redis.set("#{CANCELLED_KEY_PREFIX}#{request_id}",
+          cancellation_data.to_json, ex: @ttl)
+        result == "OK"
+      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)
+        @redis.exists("#{CANCELLED_KEY_PREFIX}#{request_id}") == 1
+      end
+
+      # Get cancellation information for a request
+      #
+      # @param request_id [String] the unique request identifier
+      # @return [Hash, nil] cancellation data or nil if not cancelled
+      def get_cancellation_info(request_id)
+        data = @redis.get("#{CANCELLED_KEY_PREFIX}#{request_id}")
+        data ? JSON.parse(data) : nil
+      rescue JSON::ParserError
+        nil
+      end
+
+      # Unregister a request (typically called when request completes)
+      #
+      # @param request_id [String] the unique request identifier
+      # @return [void]
+      def unregister_request(request_id)
+        request_data = @redis.get("#{REQUEST_KEY_PREFIX}#{request_id}")
+
+        keys_to_delete = ["#{REQUEST_KEY_PREFIX}#{request_id}",
+          "#{CANCELLED_KEY_PREFIX}#{request_id}"]
+
+        if request_data
+          begin
+            data = JSON.parse(request_data)
+            session_id = data["session_id"]
+
+            if session_id
+              keys_to_delete << "#{SESSION_KEY_PREFIX}#{session_id}:#{request_id}"
+            end
+          rescue JSON::ParserError
+            nil
+          end
+        end
+
+        @redis.del(*keys_to_delete) unless keys_to_delete.empty?
+      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)
+        data = @redis.get("#{REQUEST_KEY_PREFIX}#{request_id}")
+        data ? JSON.parse(data) : nil
+      rescue JSON::ParserError
+        nil
+      end
+
+      # Check if a request is currently active
+      #
+      # @param request_id [String] the unique request identifier
+      # @return [Boolean] true if the request is active, false otherwise
+      def active?(request_id)
+        @redis.exists("#{REQUEST_KEY_PREFIX}#{request_id}") == 1
+      end
+
+      # Clean up all requests associated with a session
+      # This is typically called when a session is terminated
+      #
+      # @param session_id [String] the session identifier
+      # @return [Array<String>] list of cleaned up request IDs
+      def cleanup_session_requests(session_id)
+        pattern = "#{SESSION_KEY_PREFIX}#{session_id}:*"
+        request_keys = @redis.keys(pattern)
+        return [] if request_keys.empty?
+
+        # Extract request IDs from the keys
+        request_ids = request_keys.map do |key|
+          key.sub("#{SESSION_KEY_PREFIX}#{session_id}:", "")
+        end
+
+        # Delete all related keys
+        all_keys = []
+        request_ids.each do |request_id|
+          all_keys << "#{REQUEST_KEY_PREFIX}#{request_id}"
+          all_keys << "#{CANCELLED_KEY_PREFIX}#{request_id}"
+        end
+        all_keys.concat(request_keys)
+
+        @redis.del(*all_keys) unless all_keys.empty?
+        request_ids
+      end
+
+      # Get all active request IDs for a specific session
+      #
+      # @param session_id [String] the session identifier
+      # @return [Array<String>] list of active request IDs for the session
+      def get_session_requests(session_id)
+        pattern = "#{SESSION_KEY_PREFIX}#{session_id}:*"
+        request_keys = @redis.keys(pattern)
+
+        request_keys.map do |key|
+          key.sub("#{SESSION_KEY_PREFIX}#{session_id}:", "")
+        end
+      end
+
+      # Get all active request IDs across all sessions
+      #
+      # @return [Array<String>] list of all active request IDs
+      def get_all_active_requests
+        pattern = "#{REQUEST_KEY_PREFIX}*"
+        request_keys = @redis.keys(pattern)
+
+        request_keys.map do |key|
+          key.sub(REQUEST_KEY_PREFIX, "")
+        end
+      end
+
+      # Clean up expired requests based on TTL
+      # This method can be called periodically to ensure cleanup
+      #
+      # @return [Integer] number of expired requests cleaned up
+      def cleanup_expired_requests
+        active_keys = @redis.keys("#{REQUEST_KEY_PREFIX}*")
+        expired_count = 0
+        key_exists_without_expiration = -1
+        key_does_not_exist = -2
+
+        active_keys.each do |key|
+          ttl = @redis.ttl(key)
+          if ttl == key_exists_without_expiration
+            @redis.expire(key, @ttl)
+          elsif ttl == key_does_not_exist
+            expired_count += 1
+          end
+        end
+
+        expired_count
+      end
+
+      # Refresh the TTL for an active request
+      #
+      # @param request_id [String] the unique request identifier
+      # @return [Boolean] true if TTL was refreshed, false if request doesn't exist
+      def refresh_request_ttl(request_id)
+        request_data = @redis.get("#{REQUEST_KEY_PREFIX}#{request_id}")
+        return false unless request_data
+
+        @redis.multi do |multi|
+          multi.expire("#{REQUEST_KEY_PREFIX}#{request_id}", @ttl)
+          multi.expire("#{CANCELLED_KEY_PREFIX}#{request_id}", @ttl)
+
+          begin
+            data = JSON.parse(request_data)
+            session_id = data["session_id"]
+            if session_id
+              multi.expire("#{SESSION_KEY_PREFIX}#{session_id}:#{request_id}", @ttl)
+            end
+          rescue JSON::ParserError
+            nil
+          end
+        end
+
+        true
+      end
+    end
+  end
+end

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

@@ -0,0 +1,120 @@
+require "json"
+require "securerandom"
+
+module ModelContextProtocol
+  class Server::StreamableHttpTransport
+    class SessionMessageQueue
+      QUEUE_KEY_PREFIX = "session_messages:"
+      LOCK_KEY_PREFIX = "session_lock:"
+      DEFAULT_TTL = 3600  # 1 hour
+      MAX_MESSAGES = 1000
+      LOCK_TIMEOUT = 5  # seconds
+
+      def initialize(redis_client, session_id, ttl: DEFAULT_TTL)
+        @redis = redis_client
+        @session_id = session_id
+        @queue_key = "#{QUEUE_KEY_PREFIX}#{session_id}"
+        @lock_key = "#{LOCK_KEY_PREFIX}#{session_id}"
+        @ttl = ttl
+      end
+
+      def push_message(message)
+        message_json = serialize_message(message)
+
+        @redis.multi do |multi|
+          multi.lpush(@queue_key, message_json)
+          multi.expire(@queue_key, @ttl)
+          multi.ltrim(@queue_key, 0, MAX_MESSAGES - 1)
+        end
+      end
+
+      def push_messages(messages)
+        return if messages.empty?
+
+        message_jsons = messages.map { |msg| serialize_message(msg) }
+
+        @redis.multi do |multi|
+          message_jsons.each do |json|
+            multi.lpush(@queue_key, json)
+          end
+          multi.expire(@queue_key, @ttl)
+          multi.ltrim(@queue_key, 0, MAX_MESSAGES - 1)
+        end
+      end
+
+      def poll_messages
+        lua_script = <<~LUA
+          local messages = redis.call('lrange', KEYS[1], 0, -1)
+          if #messages > 0 then
+            redis.call('del', KEYS[1])
+          end
+          return messages
+        LUA
+
+        messages = @redis.eval(lua_script, keys: [@queue_key])
+        return [] unless messages && !messages.empty?
+        messages.reverse.map { |json| deserialize_message(json) }
+      rescue
+        []
+      end
+
+      def peek_messages
+        messages = @redis.lrange(@queue_key, 0, -1)
+        messages.reverse.map { |json| deserialize_message(json) }
+      rescue
+        []
+      end
+
+      def has_messages?
+        @redis.exists(@queue_key) > 0
+      rescue
+        false
+      end
+
+      def message_count
+        @redis.llen(@queue_key)
+      rescue
+        0
+      end
+
+      def clear
+        @redis.del(@queue_key)
+      rescue
+      end
+
+      def with_lock(timeout: LOCK_TIMEOUT, &block)
+        lock_id = SecureRandom.hex(16)
+
+        acquired = @redis.set(@lock_key, lock_id, nx: true, ex: timeout)
+        return false unless acquired
+
+        begin
+          yield
+        ensure
+          lua_script = <<~LUA
+            if redis.call("get", KEYS[1]) == ARGV[1] then
+              return redis.call("del", KEYS[1])
+            else
+              return 0
+            end
+          LUA
+          @redis.eval(lua_script, keys: [@lock_key], argv: [lock_id])
+        end
+
+        true
+      end
+
+      private
+
+      def serialize_message(message)
+        message.is_a?(String) ? message : message.to_json
+      end
+
+      def deserialize_message(json)
+        JSON.parse(json)
+      rescue JSON::ParserError
+        json
+      end
+    end
+  end
+end

data/lib/model_context_protocol/server/{session_store.rb → streamable_http_transport/session_store.rb}

@@ -1,10 +1,9 @@
-# frozen_string_literal: true
-
 require "json"
 require "securerandom"
+require_relative "session_message_queue"
 
 module ModelContextProtocol
-  class Server
+  class Server::StreamableHttpTransport
     class SessionStore
       def initialize(redis_client, ttl: 3600)
         @redis = redis_client
@@ -69,25 +68,40 @@ module ModelContextProtocol
         @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
+      def queue_message_for_session(session_id, message)
+        return false unless session_exists?(session_id)
 
-        # Publish to server-specific channel
-        @redis.publish("server:#{server_instance}:messages", {
-          session_id: session_id,
-          message: message
-        }.to_json)
+        queue = SessionMessageQueue.new(@redis, session_id, ttl: @ttl)
+        queue.push_message(message)
         true
+      rescue
+        false
+      end
+
+      def poll_messages_for_session(session_id)
+        return [] unless session_exists?(session_id)
+
+        queue = SessionMessageQueue.new(@redis, session_id, ttl: @ttl)
+        queue.poll_messages
+      rescue
+        []
       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)
+      def get_sessions_with_messages
+        session_keys = @redis.keys("session:*")
+        sessions_with_messages = []
+
+        session_keys.each do |key|
+          session_id = key.sub("session:", "")
+          queue = SessionMessageQueue.new(@redis, session_id, ttl: @ttl)
+          if queue.has_messages?
+            sessions_with_messages << session_id
           end
         end
+
+        sessions_with_messages
+      rescue
+        []
       end
 
       def get_all_active_sessions