model-context-protocol-rb 0.4.0 → 0.5.1

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

data/lib/model_context_protocol/server/streamable_http_transport.rb

@@ -14,29 +14,62 @@ module ModelContextProtocol
         {jsonrpc: "2.0", id:, error:}
       end
     end
+
     def initialize(router:, configuration:)
       @router = router
       @configuration = configuration
 
       transport_options = @configuration.transport_options
-      @redis = transport_options[:redis_client]
+      @redis_pool = ModelContextProtocol::Server::RedisConfig.pool
       @require_sessions = transport_options.fetch(:require_sessions, false)
       @default_protocol_version = transport_options.fetch(:default_protocol_version, "2025-03-26")
-      @session_protocol_versions = {}  # Track protocol versions per session
+      @session_protocol_versions = {}
       @validate_origin = transport_options.fetch(:validate_origin, true)
       @allowed_origins = transport_options.fetch(:allowed_origins, ["http://localhost", "https://localhost", "http://127.0.0.1", "https://127.0.0.1"])
+      @redis = ModelContextProtocol::Server::RedisClientProxy.new(@redis_pool)
 
-      @session_store = ModelContextProtocol::Server::SessionStore.new(
+      @session_store = SessionStore.new(
         @redis,
         ttl: transport_options[:session_ttl] || 3600
       )
 
       @server_instance = "#{Socket.gethostname}-#{Process.pid}-#{SecureRandom.hex(4)}"
-      @local_streams = {}
-      @notification_queue = []
-      @sse_event_counter = 0
+      @stream_registry = StreamRegistry.new(@redis, @server_instance)
+      @notification_queue = NotificationQueue.new(@redis, @server_instance)
+      @event_counter = EventCounter.new(@redis, @server_instance)
+      @request_store = RequestStore.new(@redis, @server_instance)
+      @stream_monitor_thread = nil
+      @message_poller = MessagePoller.new(@redis, @stream_registry, @configuration.logger) do |stream, message|
+        send_to_stream(stream, message)
+      end
 
-      setup_redis_subscriber
+      start_message_poller
+      start_stream_monitor
+    end
+
+    def shutdown
+      @configuration.logger.info("Shutting down StreamableHttpTransport")
+
+      # Stop the message poller
+      @message_poller&.stop
+
+      # Stop the stream monitor thread
+      if @stream_monitor_thread&.alive?
+        @stream_monitor_thread.kill
+        @stream_monitor_thread.join(timeout: 5)
+      end
+
+      # Unregister all local streams
+      @stream_registry.get_all_local_streams.each do |session_id, stream|
+        @stream_registry.unregister_stream(session_id)
+        @session_store.mark_stream_inactive(session_id)
+      rescue => e
+        @configuration.logger.error("Error during stream cleanup", session_id: session_id, error: e.message)
+      end
+
+      @redis_pool.checkin(@redis) if @redis_pool && @redis
+
+      @configuration.logger.info("StreamableHttpTransport shutdown complete")
     end
 
     def handle
@@ -68,10 +101,10 @@ module ModelContextProtocol
         params: params
       }
 
-      if has_active_streams?
+      if @stream_registry.has_any_local_streams?
         deliver_to_active_streams(notification)
       else
-        @notification_queue << notification
+        @notification_queue.push(notification)
       end
     end
 
@@ -96,7 +129,6 @@ module ModelContextProtocol
 
       protocol_version = env["HTTP_MCP_PROTOCOL_VERSION"]
       if protocol_version
-        # Check if this matches a known negotiated version
         valid_versions = @session_protocol_versions.values.compact.uniq
         unless valid_versions.empty? || valid_versions.include?(protocol_version)
           error_response = ErrorResponse[id: nil, error: {code: -32600, message: "Invalid MCP protocol version: #{protocol_version}. Expected one of: #{valid_versions.join(", ")}"}]
@@ -133,9 +165,37 @@ module ModelContextProtocol
       end
     end
 
+    def create_progressive_request_sse_stream_proc(request_body, session_id)
+      proc do |stream|
+        temp_stream_id = session_id || "temp-#{SecureRandom.hex(8)}"
+        @stream_registry.register_stream(temp_stream_id, stream)
+
+        begin
+          result = @router.route(request_body, request_store: @request_store, session_id: session_id, transport: self)
+
+          if result
+            response = Response[id: request_body["id"], result: result.serialized]
+
+            event_id = next_event_id
+            send_sse_event(stream, response.serialized, event_id)
+          else
+            event_id = next_event_id
+            send_sse_event(stream, {}, event_id)
+          end
+
+          # Close stream immediately when work is complete
+          close_stream(temp_stream_id, reason: "request_completed")
+        rescue IOError, Errno::EPIPE, Errno::ECONNRESET
+          # Client disconnected during processing
+        ensure
+          # Fallback cleanup
+          @stream_registry.unregister_stream(temp_stream_id)
+        end
+      end
+    end
+
     def next_event_id
-      @sse_event_counter += 1
-      "#{@server_instance}-#{@sse_event_counter}"
+      @event_counter.next_event_id
     end
 
     def send_sse_event(stream, data, event_id = nil)
@@ -147,6 +207,20 @@ module ModelContextProtocol
       stream.flush if stream.respond_to?(:flush)
     end
 
+    def close_stream(session_id, reason: "completed")
+      if (stream = @stream_registry.get_local_stream(session_id))
+        begin
+          send_sse_event(stream, {type: "stream_complete", reason: reason})
+          stream.close
+        rescue IOError, Errno::EPIPE, Errno::ECONNRESET, Errno::ENOTCONN, Errno::EBADF
+          nil
+        end
+
+        @stream_registry.unregister_stream(session_id)
+        @session_store.mark_stream_inactive(session_id) if @require_sessions
+      end
+    end
+
     def handle_post_request(env)
       validation_error = validate_headers(env)
       return validation_error if validation_error
@@ -176,7 +250,7 @@ module ModelContextProtocol
     end
 
     def handle_initialization(body, accept_header)
-      result = @router.route(body)
+      result = @router.route(body, transport: self)
       response = Response[id: body["id"], result: result.serialized]
       response_headers = {}
 
@@ -235,21 +309,19 @@ module ModelContextProtocol
 
       case message_type
       when :notification, :response
-        if session_id && @session_store.session_has_active_stream?(session_id)
+        if body["method"] == "notifications/cancelled"
+          handle_cancellation(body, session_id)
+        elsif session_id && @session_store.session_has_active_stream?(session_id)
           deliver_to_session_stream(session_id, body)
         end
         {json: {}, status: 202}
 
       when :request
-        result = @router.route(body)
-        response = Response[id: body["id"], result: result.serialized]
-
-        if session_id && @session_store.session_has_active_stream?(session_id)
-          deliver_to_session_stream(session_id, response.serialized)
-          return {json: {accepted: true}, status: 200}
-        end
+        has_progress_token = body.dig("params", "_meta", "progressToken")
+        should_stream = (accept_header.include?("text/event-stream") && !accept_header.include?("application/json")) ||
+          has_progress_token
 
-        if accept_header.include?("text/event-stream") && !accept_header.include?("application/json")
+        if should_stream
           {
             stream: true,
             headers: {
@@ -257,14 +329,27 @@ module ModelContextProtocol
               "Cache-Control" => "no-cache",
               "Connection" => "keep-alive"
             },
-            stream_proc: create_request_sse_stream_proc(response.serialized)
+            stream_proc: create_progressive_request_sse_stream_proc(body, session_id)
           }
         else
-          {
-            json: response.serialized,
-            status: 200,
-            headers: {"Content-Type" => "application/json"}
-          }
+          result = @router.route(body, request_store: @request_store, session_id: session_id, transport: self)
+
+          if result
+            response = Response[id: body["id"], result: result.serialized]
+
+            if session_id && @session_store.session_has_active_stream?(session_id)
+              deliver_to_session_stream(session_id, response.serialized)
+              return {json: {accepted: true}, status: 200}
+            end
+
+            {
+              json: response.serialized,
+              status: 200,
+              headers: {"Content-Type" => "application/json"}
+            }
+          else
+            {json: {}, status: 204}
+          end
         end
       end
     end
@@ -315,7 +400,7 @@ module ModelContextProtocol
 
     def create_sse_stream_proc(session_id, last_event_id = nil)
       proc do |stream|
-        register_local_stream(session_id, stream) if session_id
+        @stream_registry.register_stream(session_id, stream) if session_id
 
         if last_event_id
           replay_messages_after_event_id(stream, session_id, last_event_id)
@@ -323,26 +408,15 @@ module ModelContextProtocol
           flush_notifications_to_stream(stream)
         end
 
-        start_keepalive_thread(session_id, stream)
-
         loop do
           break unless stream_connected?(stream)
           sleep 0.1
         end
       ensure
-        cleanup_local_stream(session_id) if session_id
+        @stream_registry.unregister_stream(session_id) if session_id
       end
     end
 
-    def register_local_stream(session_id, stream)
-      @local_streams[session_id] = stream
-    end
-
-    def cleanup_local_stream(session_id)
-      @local_streams.delete(session_id)
-      @session_store.mark_stream_inactive(session_id)
-    end
-
     def stream_connected?(stream)
       return false unless stream
 
@@ -350,27 +424,44 @@ module ModelContextProtocol
         stream.write(": ping\n\n")
         stream.flush if stream.respond_to?(:flush)
         true
-      rescue IOError, Errno::EPIPE, Errno::ECONNRESET
+      rescue IOError, Errno::EPIPE, Errno::ECONNRESET, Errno::ENOTCONN, Errno::EBADF
         false
       end
     end
 
-    def start_keepalive_thread(session_id, stream)
-      Thread.new do
+    def start_stream_monitor
+      @stream_monitor_thread = Thread.new do
         loop do
-          sleep 30
-          break unless stream_connected?(stream)
+          sleep 30 # Check every 30 seconds
 
           begin
-            send_ping_to_stream(stream)
-          rescue IOError, Errno::EPIPE, Errno::ECONNRESET
-            break
+            monitor_streams
+          rescue => e
+            @configuration.logger.error("Stream monitor error", error: e.message)
           end
         end
       rescue => e
-        @configuration.logger.error("Keepalive thread error", error: e.message)
-      ensure
-        cleanup_local_stream(session_id)
+        @configuration.logger.error("Stream monitor thread error", error: e.message)
+        sleep 5
+        retry
+      end
+    end
+
+    def monitor_streams
+      expired_sessions = @stream_registry.cleanup_expired_streams
+      expired_sessions.each do |session_id|
+        @session_store.mark_stream_inactive(session_id)
+      end
+
+      @stream_registry.get_all_local_streams.each do |session_id, stream|
+        if stream_connected?(stream)
+          send_ping_to_stream(stream)
+          @stream_registry.refresh_heartbeat(session_id)
+        else
+          close_stream(session_id, reason: "client_disconnected")
+        end
+      rescue IOError, Errno::EPIPE, Errno::ECONNRESET, Errno::ENOTCONN, Errno::EBADF
+        close_stream(session_id, reason: "network_error")
       end
     end
 
@@ -389,60 +480,70 @@ module ModelContextProtocol
     end
 
     def deliver_to_session_stream(session_id, data)
-      if @local_streams[session_id]
+      if @stream_registry.has_local_stream?(session_id)
+        stream = @stream_registry.get_local_stream(session_id)
         begin
-          send_to_stream(@local_streams[session_id], data)
+          send_to_stream(stream, data)
           return true
         rescue IOError, Errno::EPIPE, Errno::ECONNRESET
-          cleanup_local_stream(session_id)
+          close_stream(session_id, reason: "client_disconnected")
         end
       end
 
-      @session_store.route_message_to_session(session_id, data)
+      @session_store.queue_message_for_session(session_id, data)
     end
 
     def cleanup_session(session_id)
-      cleanup_local_stream(session_id)
+      @stream_registry.unregister_stream(session_id)
       @session_store.cleanup_session(session_id)
+      @request_store.cleanup_session_requests(session_id)
     end
 
-    def setup_redis_subscriber
-      Thread.new do
-        @session_store.subscribe_to_server(@server_instance) do |data|
-          session_id = data["session_id"]
-          message = data["message"]
-
-          if @local_streams[session_id]
-            begin
-              send_to_stream(@local_streams[session_id], message)
-            rescue IOError, Errno::EPIPE, Errno::ECONNRESET
-              cleanup_local_stream(session_id)
-            end
-          end
-        end
-      rescue => e
-        @configuration.logger.error("Redis subscriber error", error: e.message, backtrace: e.backtrace.first(5))
-        sleep 5
-        retry
-      end
+    def start_message_poller
+      @message_poller.start
     end
 
     def has_active_streams?
-      @local_streams.any?
+      @stream_registry.has_any_local_streams?
     end
 
     def deliver_to_active_streams(notification)
-      @local_streams.each do |session_id, stream|
+      @stream_registry.get_all_local_streams.each do |session_id, stream|
         send_to_stream(stream, notification)
       rescue IOError, Errno::EPIPE, Errno::ECONNRESET
-        cleanup_local_stream(session_id)
+        close_stream(session_id, reason: "client_disconnected")
       end
     end
 
     def flush_notifications_to_stream(stream)
-      while (notification = @notification_queue.shift)
+      notifications = @notification_queue.pop_all
+      notifications.each do |notification|
         send_to_stream(stream, notification)
       end
     end
+
+    # Handle a cancellation notification from the client
+    #
+    # @param message [Hash] the cancellation notification message
+    # @param session_id [String, nil] the session ID if available
+    def handle_cancellation(message, session_id = nil)
+      params = message["params"]
+      return unless params
+
+      request_id = params["requestId"]
+      reason = params["reason"]
+
+      return unless request_id
+
+      @request_store.mark_cancelled(request_id, reason)
+    rescue
+      nil
+    end
+
+    def cleanup
+      @message_poller&.stop
+      @stream_monitor_thread&.kill
+      @redis = nil
+    end
   end
 end

data/lib/model_context_protocol/server/tool.rb

@@ -5,7 +5,9 @@ module ModelContextProtocol
     # Raised when output schema validation fails.
     class OutputSchemaValidationError < StandardError; end
 
+    include ModelContextProtocol::Server::Cancellable
     include ModelContextProtocol::Server::ContentHelpers
+    include ModelContextProtocol::Server::Progressable
 
     attr_reader :arguments, :context, :logger
 
@@ -107,6 +109,8 @@ module ModelContextProtocol
         raise ModelContextProtocol::Server::ParameterValidationError, validation_error.message
       rescue OutputSchemaValidationError, ModelContextProtocol::Server::ResponseArgumentsError => tool_error
         raise tool_error, tool_error.message
+      rescue Server::Cancellable::CancellationError
+        raise
       rescue => error
         ErrorResponse[error: error.message]
       end

data/lib/model_context_protocol/server.rb

@@ -8,7 +8,7 @@ module ModelContextProtocol
     # Raised when invalid parameters are provided.
     class ParameterValidationError < StandardError; end
 
-    attr_reader :configuration, :router
+    attr_reader :configuration, :router, :transport
 
     def initialize
       @configuration = Configuration.new
@@ -20,7 +20,7 @@ module ModelContextProtocol
     def start
       configuration.validate!
 
-      transport = case configuration.transport_type
+      @transport = case configuration.transport_type
       when :stdio, nil
         StdioTransport.new(router: @router, configuration: @configuration)
       when :streamable_http
@@ -32,7 +32,7 @@ module ModelContextProtocol
         raise ArgumentError, "Unknown transport: #{configuration.transport_type}"
       end
 
-      transport.handle
+      @transport.handle
     end
 
     private
@@ -281,5 +281,11 @@ module ModelContextProtocol
         end
       end
     end
+
+    class << self
+      def configure_redis(&block)
+        RedisConfig.configure(&block)
+      end
+    end
   end
 end

data/lib/model_context_protocol/version.rb

@@ -1,3 +1,3 @@
 module ModelContextProtocol
-  VERSION = "0.4.0"
+  VERSION = "0.5.1"
 end