model-context-protocol-rb 0.4.0 → 0.5.0

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,31 @@ 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
+        ensure
+          @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)
@@ -176,7 +230,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 +289,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 +309,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 +380,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 +388,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
 
@@ -355,22 +409,41 @@ module ModelContextProtocol
       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
+          @stream_registry.unregister_stream(session_id)
+          @session_store.mark_stream_inactive(session_id)
+        end
+      rescue IOError, Errno::EPIPE, Errno::ECONNRESET
+        @stream_registry.unregister_stream(session_id)
+        @session_store.mark_stream_inactive(session_id)
       end
     end
 
@@ -389,60 +462,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)
+          @stream_registry.unregister_stream(session_id)
         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)
+        @stream_registry.unregister_stream(session_id)
       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.0"
 end

data/tasks/templates/dev-http.erb

@@ -4,6 +4,8 @@ require "bundler/setup"
 require "rack"
 require 'rackup/handler/webrick'
 require "webrick"
+require "webrick/https"
+require "openssl"
 require "securerandom"
 require "redis"
 require "logger"
@@ -12,6 +14,14 @@ require 'stringio'
 
 require_relative "../lib/model_context_protocol"
 
+ModelContextProtocol::Server.configure_redis do |config|
+  config.redis_url = "redis://localhost:6379/0"
+  config.pool_size = 10
+  config.enable_reaper = true
+  config.reaper_interval = 10
+  config.idle_timeout = 15
+end
+
 Dir[File.join(__dir__, "../spec/support/**/*.rb")].each { |file| require file }
 
 logger = Logger.new(STDOUT)
@@ -73,6 +83,11 @@ class MCPHttpApp
       @logger.info("  Request: #{body_content}") unless body_content.empty?
     end
 
+    if ModelContextProtocol::Server::RedisConfig.configured?
+      pool_stats = ModelContextProtocol::Server::RedisConfig.stats
+      @logger.info("  Redis Pool: #{pool_stats}")
+    end
+
     env['rack.input'] = StringIO.new(body_content)
     request = Rack::Request.new(env)
 
@@ -89,30 +104,26 @@ class MCPHttpApp
       }, [""]]
     end
 
-    begin
-      redis_client = Redis.new(url: ENV.fetch("REDIS_URL", "redis://localhost:6379/0"))
-      @logger.debug("Testing Redis connection...")
-      redis_client.ping
-      @logger.info("Redis connected successfully")
-    end
-
     transport_config = {
       type: :streamable_http,
-      env: env,
-      redis_client: redis_client,
-      require_sessions: false,  # Optional sessions for easier testing
+      env:,
+      require_sessions: false,
       session_ttl: 3600,
-      validate_origin: false,   # Disable for testing
       allowed_origins: ["*"]
     }
 
     @logger.debug("Creating MCP server with transport config")
     server = create_mcp_server(transport_config)
+    transport = nil
 
     begin
       @logger.debug("Starting MCP server")
       result = server.start
 
+      if server.respond_to?(:transport)
+        transport = server.transport
+      end
+
       case result
       when Hash
         if result[:stream]
@@ -134,6 +145,10 @@ class MCPHttpApp
               @logger.info("← ERROR RESPONSE (code: #{response_json[:error][:code]})")
             elsif status == 202
               @logger.info("← NOTIFICATION ACCEPTED [NO RESPONSE REQUIRED]")
+            elsif response_json[:accepted] == true && status == 200
+              method = request_json['method'] rescue 'unknown'
+              id = request_json['id'] rescue 'unknown'
+              @logger.info("← #{method} RESPONSE (id: #{id}) [DELIVERED VIA SSE STREAM]")
             elsif response_json[:result]
               method = request_json['method'] rescue 'unknown'
               @logger.info("← #{method} RESPONSE (id: #{response_json[:id]})")
@@ -167,6 +182,7 @@ class MCPHttpApp
       @logger.debug("Full backtrace:\n#{e.backtrace.join("\n")}")
       [500, {"Content-Type" => "application/json"}, [%Q({"error": "Internal server error: #{e.message}"})]]
     ensure
+      transport&.cleanup if transport&.respond_to?(:cleanup)
       Thread.current[:request_id] = nil
     end
   end
@@ -217,13 +233,18 @@ class MCPHttpApp
           register TestToolWithMixedContentResponse
           register TestToolWithResourceResponse
           register TestToolWithToolErrorResponse
+          register TestToolWithCancellableSleep
         end
       end
     end
   end
 end
 
-logger.info("Starting MCP HTTP Development Server on http://localhost:9292/mcp")
+use_ssl = ENV['SSL'] == 'true'
+port = use_ssl ? 9293 : 9292
+protocol = use_ssl ? 'https' : 'http'
+
+logger.info("Starting MCP #{protocol.upcase} Development Server on #{protocol}://localhost:#{port}/mcp")
 
 app = Rack::Builder.new do
   map '/mcp' do
@@ -231,7 +252,30 @@ app = Rack::Builder.new do
   end
 end
 
-server = WEBrick::HTTPServer.new(Port: 9292, Host: '0.0.0.0')
+server_options = {
+  Port: port,
+  Host: '0.0.0.0'
+}
+
+if use_ssl
+  cert_path = File.join(__dir__, '..', 'tmp', 'ssl', 'server.crt')
+  key_path = File.join(__dir__, '..', 'tmp', 'ssl', 'server.key')
+
+  unless File.exist?(cert_path) && File.exist?(key_path)
+    logger.error("SSL certificates not found at tmp/ssl/")
+    logger.error("Generate them with: openssl req -x509 -newkey rsa:4096 -keyout tmp/ssl/server.key -out tmp/ssl/server.crt -days 365 -nodes -subj \"/C=US/ST=Dev/L=Dev/O=Dev/CN=localhost\"")
+    exit(1)
+  end
+
+  server_options.merge!(
+    SSLEnable: true,
+    SSLCertificate: OpenSSL::X509::Certificate.new(File.read(cert_path)),
+    SSLPrivateKey: OpenSSL::PKey::RSA.new(File.read(key_path)),
+    SSLVerifyClient: OpenSSL::SSL::VERIFY_NONE
+  )
+end
+
+server = WEBrick::HTTPServer.new(server_options)
 server.mount '/', Rackup::Handler::WEBrick, app
 
 ['INT', 'TERM'].each do |signal|
@@ -241,4 +285,4 @@ server.mount '/', Rackup::Handler::WEBrick, app
   end
 end
 
-server.start
+server.start

metadata

@@ -1,14 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: model-context-protocol-rb
 version: !ruby/object:Gem::Version
-  version: 0.4.0
+  version: 0.5.0
 platform: ruby
 authors:
 - Dick Davis
 autorequire:
 bindir: bin
 cert_chain: []
-date: 2025-09-08 00:00:00.000000000 Z
+date: 2025-09-23 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: json-schema
@@ -38,6 +38,48 @@ dependencies:
     - - "~>"
       - !ruby/object:Gem::Version
         version: '2.8'
+- !ruby/object:Gem::Dependency
+  name: redis
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '5.0'
+  type: :runtime
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '5.0'
+- !ruby/object:Gem::Dependency
+  name: connection_pool
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '2.4'
+  type: :runtime
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '2.4'
+- !ruby/object:Gem::Dependency
+  name: concurrent-ruby
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '1.3'
+  type: :runtime
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '1.3'
 description:
 email:
 - dick@hey.com
@@ -56,20 +98,32 @@ files:
 - Rakefile
 - lib/model_context_protocol.rb
 - lib/model_context_protocol/server.rb
+- lib/model_context_protocol/server/cancellable.rb
 - lib/model_context_protocol/server/completion.rb
 - lib/model_context_protocol/server/configuration.rb
 - lib/model_context_protocol/server/content.rb
 - lib/model_context_protocol/server/content_helpers.rb
 - lib/model_context_protocol/server/mcp_logger.rb
 - lib/model_context_protocol/server/pagination.rb
+- lib/model_context_protocol/server/progressable.rb
 - lib/model_context_protocol/server/prompt.rb
+- lib/model_context_protocol/server/redis_client_proxy.rb
+- lib/model_context_protocol/server/redis_config.rb
+- lib/model_context_protocol/server/redis_pool_manager.rb
 - lib/model_context_protocol/server/registry.rb
 - lib/model_context_protocol/server/resource.rb
 - lib/model_context_protocol/server/resource_template.rb
 - lib/model_context_protocol/server/router.rb
-- lib/model_context_protocol/server/session_store.rb
 - lib/model_context_protocol/server/stdio_transport.rb
+- lib/model_context_protocol/server/stdio_transport/request_store.rb
 - lib/model_context_protocol/server/streamable_http_transport.rb
+- lib/model_context_protocol/server/streamable_http_transport/event_counter.rb
+- lib/model_context_protocol/server/streamable_http_transport/message_poller.rb
+- lib/model_context_protocol/server/streamable_http_transport/notification_queue.rb
+- lib/model_context_protocol/server/streamable_http_transport/request_store.rb
+- lib/model_context_protocol/server/streamable_http_transport/session_message_queue.rb
+- lib/model_context_protocol/server/streamable_http_transport/session_store.rb
+- lib/model_context_protocol/server/streamable_http_transport/stream_registry.rb
 - lib/model_context_protocol/server/tool.rb
 - lib/model_context_protocol/version.rb
 - tasks/mcp.rake