model-context-protocol-rb 0.4.0 → 0.5.0

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

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

@@ -0,0 +1,119 @@
+# frozen_string_literal: true
+
+require "json"
+
+module ModelContextProtocol
+  class Server::StreamableHttpTransport
+    class StreamRegistry
+      STREAM_KEY_PREFIX = "stream:active:"
+      HEARTBEAT_KEY_PREFIX = "stream:heartbeat:"
+      DEFAULT_TTL = 60 # 1 minute TTL for stream entries
+
+      def initialize(redis_client, server_instance, ttl: DEFAULT_TTL)
+        @redis = redis_client
+        @server_instance = server_instance
+        @ttl = ttl
+        @local_streams = {} # Keep local reference for direct stream access
+      end
+
+      def register_stream(session_id, stream)
+        @local_streams[session_id] = stream
+
+        # Store stream registration in Redis with TTL
+        @redis.multi do |multi|
+          multi.set("#{STREAM_KEY_PREFIX}#{session_id}", @server_instance, ex: @ttl)
+          multi.set("#{HEARTBEAT_KEY_PREFIX}#{session_id}", Time.now.to_f, ex: @ttl)
+        end
+      end
+
+      def unregister_stream(session_id)
+        @local_streams.delete(session_id)
+
+        @redis.multi do |multi|
+          multi.del("#{STREAM_KEY_PREFIX}#{session_id}")
+          multi.del("#{HEARTBEAT_KEY_PREFIX}#{session_id}")
+        end
+      end
+
+      def get_local_stream(session_id)
+        @local_streams[session_id]
+      end
+
+      def has_local_stream?(session_id)
+        @local_streams.key?(session_id)
+      end
+
+      def get_stream_server(session_id)
+        @redis.get("#{STREAM_KEY_PREFIX}#{session_id}")
+      end
+
+      def stream_active?(session_id)
+        @redis.exists("#{STREAM_KEY_PREFIX}#{session_id}") == 1
+      end
+
+      def refresh_heartbeat(session_id)
+        @redis.multi do |multi|
+          multi.set("#{HEARTBEAT_KEY_PREFIX}#{session_id}", Time.now.to_f, ex: @ttl)
+          multi.expire("#{STREAM_KEY_PREFIX}#{session_id}", @ttl)
+        end
+      end
+
+      def get_all_local_streams
+        @local_streams.dup
+      end
+
+      def has_any_local_streams?
+        !@local_streams.empty?
+      end
+
+      def cleanup_expired_streams
+        # Get all local stream session IDs
+        local_session_ids = @local_streams.keys
+
+        # Check which ones are still active in Redis
+        pipeline_results = @redis.pipelined do |pipeline|
+          local_session_ids.each do |session_id|
+            pipeline.exists("#{STREAM_KEY_PREFIX}#{session_id}")
+          end
+        end
+
+        # Remove expired streams from local storage
+        expired_sessions = []
+        local_session_ids.each_with_index do |session_id, index|
+          if pipeline_results[index] == 0 # Stream expired in Redis
+            @local_streams.delete(session_id)
+            expired_sessions << session_id
+          end
+        end
+
+        expired_sessions
+      end
+
+      def get_stale_streams(max_age_seconds = 90)
+        current_time = Time.now.to_f
+        stale_streams = []
+
+        # Get all heartbeat keys
+        heartbeat_keys = @redis.keys("#{HEARTBEAT_KEY_PREFIX}*")
+
+        return stale_streams if heartbeat_keys.empty?
+
+        # Get all heartbeat timestamps
+        heartbeat_values = @redis.mget(heartbeat_keys)
+
+        heartbeat_keys.each_with_index do |key, index|
+          next unless heartbeat_values[index]
+
+          session_id = key.sub(HEARTBEAT_KEY_PREFIX, "")
+          last_heartbeat = heartbeat_values[index].to_f
+
+          if current_time - last_heartbeat > max_age_seconds
+            stale_streams << session_id
+          end
+        end
+
+        stale_streams
+      end
+    end
+  end
+end