model-context-protocol-rb 0.5.0 → 0.6.0

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

@@ -6,10 +6,10 @@ module ModelContextProtocol
       POLL_INTERVAL = 0.1  # 100ms
       BATCH_SIZE = 100
 
-      def initialize(redis_client, stream_registry, logger, &message_delivery_block)
+      def initialize(redis_client, stream_registry, client_logger, &message_delivery_block)
         @redis = redis_client
         @stream_registry = stream_registry
-        @logger = logger
+        @client_logger = client_logger
         @message_delivery_block = message_delivery_block
         @running = false
         @poll_thread = nil
@@ -22,14 +22,14 @@ module ModelContextProtocol
         @poll_thread = Thread.new do
           poll_loop
         rescue => e
-          @logger.error("Message poller thread error", error: e.message, backtrace: e.backtrace.first(5))
+          @client_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")
+        @client_logger.debug("Message poller started")
       end
 
       def stop
@@ -37,11 +37,11 @@ module ModelContextProtocol
 
         if @poll_thread&.alive?
           @poll_thread.kill
-          @poll_thread.join(timeout: 5)
+          @poll_thread.join(5)
         end
 
         @poll_thread = nil
-        @logger.debug("Message poller stopped")
+        @client_logger.debug("Message poller stopped")
       end
 
       def running?
@@ -55,7 +55,7 @@ module ModelContextProtocol
           begin
             poll_and_deliver_messages
           rescue => e
-            @logger.error("Error in message polling", error: e.message)
+            @client_logger.error("Error in message polling", error: e.message)
           end
 
           sleep POLL_INTERVAL
@@ -91,9 +91,9 @@ module ModelContextProtocol
         @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)
+        @client_logger.debug("Unregistered disconnected stream", session_id: session_id)
       rescue => e
-        @logger.error("Error delivering message to stream",
+        @client_logger.error("Error delivering message to stream",
           session_id: session_id, error: e.message)
       end
     end

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

@@ -19,10 +19,10 @@ module ModelContextProtocol
 
       # Register a new request with its associated session
       #
-      # @param request_id [String] the unique request identifier
+      # @param jsonrpc_request_id [String] the unique JSON-RPC 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)
+      def register_request(jsonrpc_request_id, session_id = nil)
         request_data = {
           session_id: session_id,
           server_instance: @server_instance,
@@ -30,11 +30,11 @@ module ModelContextProtocol
         }
 
         @redis.multi do |multi|
-          multi.set("#{REQUEST_KEY_PREFIX}#{request_id}",
+          multi.set("#{REQUEST_KEY_PREFIX}#{jsonrpc_request_id}",
             request_data.to_json, ex: @ttl)
 
           if session_id
-            multi.set("#{SESSION_KEY_PREFIX}#{session_id}:#{request_id}",
+            multi.set("#{SESSION_KEY_PREFIX}#{session_id}:#{jsonrpc_request_id}",
               true, ex: @ttl)
           end
         end
@@ -42,34 +42,34 @@ module ModelContextProtocol
 
       # Mark a request as cancelled
       #
-      # @param request_id [String] the unique request identifier
+      # @param jsonrpc_request_id [String] the unique JSON-RPC request identifier
       # @param reason [String] optional reason for cancellation
       # @return [Boolean] true if cancellation was recorded
-      def mark_cancelled(request_id, reason = nil)
+      def mark_cancelled(jsonrpc_request_id, reason = nil)
         cancellation_data = {
           cancelled_at: Time.now.to_f,
           reason: reason
         }
 
-        result = @redis.set("#{CANCELLED_KEY_PREFIX}#{request_id}",
+        result = @redis.set("#{CANCELLED_KEY_PREFIX}#{jsonrpc_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
+      # @param jsonrpc_request_id [String] the unique JSON-RPC request identifier
       # @return [Boolean] true if the request is cancelled, false otherwise
-      def cancelled?(request_id)
-        @redis.exists("#{CANCELLED_KEY_PREFIX}#{request_id}") == 1
+      def cancelled?(jsonrpc_request_id)
+        @redis.exists("#{CANCELLED_KEY_PREFIX}#{jsonrpc_request_id}") == 1
       end
 
       # Get cancellation information for a request
       #
-      # @param request_id [String] the unique request identifier
+      # @param jsonrpc_request_id [String] the unique JSON-RPC 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}")
+      def get_cancellation_info(jsonrpc_request_id)
+        data = @redis.get("#{CANCELLED_KEY_PREFIX}#{jsonrpc_request_id}")
         data ? JSON.parse(data) : nil
       rescue JSON::ParserError
         nil
@@ -77,13 +77,13 @@ module ModelContextProtocol
 
       # Unregister a request (typically called when request completes)
       #
-      # @param request_id [String] the unique request identifier
+      # @param jsonrpc_request_id [String] the unique JSON-RPC request identifier
       # @return [void]
-      def unregister_request(request_id)
-        request_data = @redis.get("#{REQUEST_KEY_PREFIX}#{request_id}")
+      def unregister_request(jsonrpc_request_id)
+        request_data = @redis.get("#{REQUEST_KEY_PREFIX}#{jsonrpc_request_id}")
 
-        keys_to_delete = ["#{REQUEST_KEY_PREFIX}#{request_id}",
-          "#{CANCELLED_KEY_PREFIX}#{request_id}"]
+        keys_to_delete = ["#{REQUEST_KEY_PREFIX}#{jsonrpc_request_id}",
+          "#{CANCELLED_KEY_PREFIX}#{jsonrpc_request_id}"]
 
         if request_data
           begin
@@ -91,7 +91,7 @@ module ModelContextProtocol
             session_id = data["session_id"]
 
             if session_id
-              keys_to_delete << "#{SESSION_KEY_PREFIX}#{session_id}:#{request_id}"
+              keys_to_delete << "#{SESSION_KEY_PREFIX}#{session_id}:#{jsonrpc_request_id}"
             end
           rescue JSON::ParserError
             nil
@@ -103,10 +103,10 @@ module ModelContextProtocol
 
       # Get information about a specific request
       #
-      # @param request_id [String] the unique request identifier
+      # @param jsonrpc_request_id [String] the unique JSON-RPC 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}")
+      def get_request(jsonrpc_request_id)
+        data = @redis.get("#{REQUEST_KEY_PREFIX}#{jsonrpc_request_id}")
         data ? JSON.parse(data) : nil
       rescue JSON::ParserError
         nil
@@ -114,10 +114,10 @@ module ModelContextProtocol
 
       # Check if a request is currently active
       #
-      # @param request_id [String] the unique request identifier
+      # @param jsonrpc_request_id [String] the unique JSON-RPC request identifier
       # @return [Boolean] true if the request is active, false otherwise
-      def active?(request_id)
-        @redis.exists("#{REQUEST_KEY_PREFIX}#{request_id}") == 1
+      def active?(jsonrpc_request_id)
+        @redis.exists("#{REQUEST_KEY_PREFIX}#{jsonrpc_request_id}") == 1
       end
 
       # Clean up all requests associated with a session
@@ -131,20 +131,20 @@ module ModelContextProtocol
         return [] if request_keys.empty?
 
         # Extract request IDs from the keys
-        request_ids = request_keys.map do |key|
+        jsonrpc_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}"
+        jsonrpc_request_ids.each do |jsonrpc_request_id|
+          all_keys << "#{REQUEST_KEY_PREFIX}#{jsonrpc_request_id}"
+          all_keys << "#{CANCELLED_KEY_PREFIX}#{jsonrpc_request_id}"
         end
         all_keys.concat(request_keys)
 
         @redis.del(*all_keys) unless all_keys.empty?
-        request_ids
+        jsonrpc_request_ids
       end
 
       # Get all active request IDs for a specific session
@@ -196,21 +196,21 @@ module ModelContextProtocol
 
       # Refresh the TTL for an active request
       #
-      # @param request_id [String] the unique request identifier
+      # @param jsonrpc_request_id [String] the unique JSON-RPC 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}")
+      def refresh_request_ttl(jsonrpc_request_id)
+        request_data = @redis.get("#{REQUEST_KEY_PREFIX}#{jsonrpc_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)
+          multi.expire("#{REQUEST_KEY_PREFIX}#{jsonrpc_request_id}", @ttl)
+          multi.expire("#{CANCELLED_KEY_PREFIX}#{jsonrpc_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)
+              multi.expire("#{SESSION_KEY_PREFIX}#{session_id}:#{jsonrpc_request_id}", @ttl)
             end
           rescue JSON::ParserError
             nil

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

@@ -0,0 +1,231 @@
+require "json"
+
+module ModelContextProtocol
+  class Server::StreamableHttpTransport
+    # Redis-based distributed storage for tracking server-initiated requests and their response status.
+    # This store is used by StreamableHttpTransport to manage outgoing request lifecycle (like pings)
+    # across multiple server instances and handle timeouts in a distributed environment.
+    class ServerRequestStore
+      REQUEST_KEY_PREFIX = "server_request:pending:"
+      SESSION_KEY_PREFIX = "server_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 server-initiated request with its associated session
+      #
+      # @param request_id [String] the unique JSON-RPC request identifier
+      # @param session_id [String] the session identifier (can be nil for sessionless requests)
+      # @param type [Symbol] the type of request (e.g., :ping)
+      # @return [void]
+      def register_request(request_id, session_id = nil, type: :ping)
+        request_data = {
+          session_id: session_id,
+          server_instance: @server_instance,
+          type: type.to_s,
+          created_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 server-initiated request as completed (response received)
+      #
+      # @param request_id [String] the unique JSON-RPC request identifier
+      # @return [Boolean] true if request was pending, false if not found
+      def mark_completed(request_id)
+        request_data = @redis.get("#{REQUEST_KEY_PREFIX}#{request_id}")
+        return false unless request_data
+
+        unregister_request(request_id)
+        true
+      end
+
+      # Check if a server-initiated request is still pending
+      #
+      # @param request_id [String] the unique JSON-RPC request identifier
+      # @return [Boolean] true if the request is pending, false otherwise
+      def pending?(request_id)
+        @redis.exists("#{REQUEST_KEY_PREFIX}#{request_id}") == 1
+      end
+
+      # Get information about a specific pending request
+      #
+      # @param request_id [String] the unique JSON-RPC 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
+
+      # Find requests that have exceeded the specified timeout
+      #
+      # @param timeout_seconds [Integer] timeout in seconds
+      # @return [Array<Hash>] array of expired request info with request_id and session_id
+      def get_expired_requests(timeout_seconds)
+        current_time = Time.now.to_f
+        expired_requests = []
+
+        # Get all pending request keys
+        request_keys = @redis.keys("#{REQUEST_KEY_PREFIX}*")
+        return expired_requests if request_keys.empty?
+
+        # Get all request data in batch
+        request_values = @redis.mget(request_keys)
+
+        request_keys.each_with_index do |key, index|
+          next unless request_values[index]
+
+          begin
+            request_data = JSON.parse(request_values[index])
+            created_at = request_data["created_at"]
+
+            if created_at && (current_time - created_at) > timeout_seconds
+              request_id = key.sub(REQUEST_KEY_PREFIX, "")
+              expired_requests << {
+                request_id: request_id,
+                session_id: request_data["session_id"],
+                type: request_data["type"],
+                age: current_time - created_at
+              }
+            end
+          rescue JSON::ParserError
+            # Skip malformed entries
+            next
+          end
+        end
+
+        expired_requests
+      end
+
+      # Clean up expired requests based on timeout
+      #
+      # @param timeout_seconds [Integer] timeout in seconds
+      # @return [Array<String>] list of cleaned up request IDs
+      def cleanup_expired_requests(timeout_seconds)
+        expired_requests = get_expired_requests(timeout_seconds)
+
+        expired_requests.each do |request_info|
+          unregister_request(request_info[:request_id])
+        end
+
+        expired_requests.map { |r| r[:request_id] }
+      end
+
+      # Unregister a request (typically called when request completes or times out)
+      #
+      # @param request_id [String] the unique JSON-RPC 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}"]
+
+        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
+
+      # Clean up all server 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}"
+        end
+        all_keys.concat(request_keys)
+
+        @redis.del(*all_keys) unless all_keys.empty?
+        request_ids
+      end
+
+      # Get all pending request IDs for a specific session
+      #
+      # @param session_id [String] the session identifier
+      # @return [Array<String>] list of pending 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 pending request IDs across all sessions
+      #
+      # @return [Array<String>] list of all pending request IDs
+      def get_all_pending_requests
+        pattern = "#{REQUEST_KEY_PREFIX}*"
+        request_keys = @redis.keys(pattern)
+
+        request_keys.map do |key|
+          key.sub(REQUEST_KEY_PREFIX, "")
+        end
+      end
+
+      # Refresh the TTL for a pending request
+      #
+      # @param request_id [String] the unique JSON-RPC 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)
+
+          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