vector_mcp 0.3.3 → 0.4.0

data/lib/vector_mcp/server.rb

@@ -5,8 +5,6 @@ require "logger"
 require_relative "definitions"
 require_relative "session"
 require_relative "errors"
-require_relative "transport/stdio" # Default transport
-# require_relative "transport/sse" # Load on demand to avoid async dependencies
 require_relative "handlers/core" # Default handlers
 require_relative "util" # Needed if not using Handlers::Core
 require_relative "server/registry"
@@ -26,7 +24,7 @@ module VectorMCP
   # It manages tools, resources, prompts, and handles the MCP message lifecycle.
   #
   # A server instance is typically initialized, configured with capabilities (tools,
-  # resources, prompts), and then run with a chosen transport mechanism (e.g., Stdio, SSE).
+  # resources, prompts), and then run with a chosen transport mechanism (e.g., HttpStream).
   #
   # @example Creating and running a simple server
   #   server = VectorMCP::Server.new(name: "MySimpleServer", version: "1.0")
@@ -39,7 +37,7 @@ module VectorMCP
   #     args["message"]
   #   end
   #
-  #   server.run(transport: :stdio) # Runs with Stdio transport by default
+  #   server.run # Runs with HttpStream transport by default
   #
   # @!attribute [r] logger
   #   @return [Logger] The logger instance for this server.
@@ -68,7 +66,10 @@ module VectorMCP
     include MessageHandling
 
     # The specific version of the Model Context Protocol this server implements.
-    PROTOCOL_VERSION = "2024-11-05"
+    PROTOCOL_VERSION = "2025-11-25"
+
+    # All protocol versions this server accepts via the MCP-Protocol-Version header.
+    SUPPORTED_PROTOCOL_VERSIONS = %w[2025-11-25 2025-03-26 2024-11-05].freeze
 
     attr_reader :logger, :name, :version, :protocol_version, :tools, :resources, :prompts, :roots, :in_flight_requests,
                 :auth_manager, :authorization, :security_middleware, :middleware_manager
@@ -134,33 +135,19 @@ module VectorMCP
 
     # Runs the server using the specified transport mechanism.
     #
-    # @param transport [:stdio, :sse, :http_stream, VectorMCP::Transport::Base] The transport to use.
-    #   Can be a symbol (`:stdio`, `:sse`, `:http_stream`) or an initialized transport instance.
-    #   If a symbol is provided, the method will instantiate the corresponding transport class.
-    #   If `:sse` is chosen, it uses Puma as the HTTP server (deprecated).
-    #   If `:http_stream` is chosen, it uses the MCP-compliant streamable HTTP transport.
-    # @param options [Hash] Transport-specific options (e.g., `:host`, `:port` for HTTP transports).
+    # @param transport [:http_stream, VectorMCP::Transport::Base] The transport to use.
+    #   Can be the symbol `:http_stream` or an initialized transport instance.
+    #   If `:http_stream` is provided, the method will instantiate the MCP-compliant streamable HTTP transport.
+    # @param options [Hash] Transport-specific options (e.g., `:host`, `:port`).
     #   These are passed to the transport's constructor if a symbol is provided for `transport`.
     # @return [void]
     # @raise [ArgumentError] if an unsupported transport symbol is given.
-    # @raise [NotImplementedError] if `:sse` transport is specified (currently a placeholder).
-    def run(transport: :stdio, **options)
+    def run(transport: :http_stream, **)
       active_transport = case transport
-                         when :stdio
-                           VectorMCP::Transport::Stdio.new(self, **options)
-                         when :sse
-                           begin
-                             require_relative "transport/sse"
-                             logger.warn("SSE transport is deprecated. Please use :http_stream instead.")
-                             VectorMCP::Transport::SSE.new(self, **options)
-                           rescue LoadError => e
-                             logger.fatal("SSE transport requires additional dependencies.")
-                             raise NotImplementedError, "SSE transport dependencies not available: #{e.message}"
-                           end
                          when :http_stream
                            begin
                              require_relative "transport/http_stream"
-                             VectorMCP::Transport::HttpStream.new(self, **options)
+                             VectorMCP::Transport::HttpStream.new(self, **)
                            rescue LoadError => e
                              logger.fatal("HttpStream transport requires additional dependencies.")
                              raise NotImplementedError, "HttpStream transport dependencies not available: #{e.message}"
@@ -176,6 +163,26 @@ module VectorMCP
       active_transport.run
     end
 
+    # Returns the MCP server as a Rack application suitable for mounting inside
+    # another Rack-based framework (e.g., Rails, Sinatra).
+    #
+    # Unlike {#run}, this method does NOT start its own HTTP server or block.
+    # The returned object responds to `#call(env)` and can be mounted directly:
+    #
+    #   # config/routes.rb (Rails)
+    #   mount MCP_APP => "/mcp"
+    #
+    # Call `server.transport.stop` on application shutdown to clean up resources.
+    #
+    # @param options [Hash] Transport options (e.g., :session_timeout, :event_retention, :allowed_origins)
+    # @return [VectorMCP::Transport::HttpStream] A Rack-compatible app
+    def rack_app(**)
+      require_relative "transport/http_stream"
+      active_transport = VectorMCP::Transport::HttpStream.new(self, mounted: true, **)
+      self.transport = active_transport
+      active_transport
+    end
+
     # --- Security Configuration ---
 
     # Enable authentication with specified strategy and configuration
@@ -190,7 +197,7 @@ module VectorMCP
 
       case strategy
       when :api_key
-        add_api_key_auth(options[:keys] || [])
+        add_api_key_auth(options[:keys] || [], allow_query_params: options[:allow_query_params] || false)
       when :jwt
         add_jwt_auth(options)
       when :custom
@@ -216,9 +223,9 @@ module VectorMCP
     # Enable authorization with optional policy configuration block
     # @param block [Proc] optional block for configuring authorization policies
     # @return [void]
-    def enable_authorization!(&block)
+    def enable_authorization!(&)
       @authorization.enable!
-      instance_eval(&block) if block_given?
+      instance_eval(&) if block_given?
       @logger.info("Authorization enabled")
     end
 
@@ -232,29 +239,29 @@ module VectorMCP
     # Add authorization policy for tools
     # @param block [Proc] policy block that receives (user, action, tool)
     # @return [void]
-    def authorize_tools(&block)
-      @authorization.add_policy(:tool, &block)
+    def authorize_tools(&)
+      @authorization.add_policy(:tool, &)
     end
 
     # Add authorization policy for resources
     # @param block [Proc] policy block that receives (user, action, resource)
     # @return [void]
-    def authorize_resources(&block)
-      @authorization.add_policy(:resource, &block)
+    def authorize_resources(&)
+      @authorization.add_policy(:resource, &)
     end
 
     # Add authorization policy for prompts
     # @param block [Proc] policy block that receives (user, action, prompt)
     # @return [void]
-    def authorize_prompts(&block)
-      @authorization.add_policy(:prompt, &block)
+    def authorize_prompts(&)
+      @authorization.add_policy(:prompt, &)
     end
 
     # Add authorization policy for roots
     # @param block [Proc] policy block that receives (user, action, root)
     # @return [void]
-    def authorize_roots(&block)
-      @authorization.add_policy(:root, &block)
+    def authorize_roots(&)
+      @authorization.add_policy(:root, &)
     end
 
     # Check if security features are enabled
@@ -313,9 +320,10 @@ module VectorMCP
 
     # Add API key authentication strategy
     # @param keys [Array<String>] array of valid API keys
+    # @param allow_query_params [Boolean] whether to accept API keys from query parameters
     # @return [void]
-    def add_api_key_auth(keys)
-      strategy = Security::Strategies::ApiKey.new(keys: keys)
+    def add_api_key_auth(keys, allow_query_params: false)
+      strategy = Security::Strategies::ApiKey.new(keys: keys, allow_query_params: allow_query_params)
       @auth_manager.add_strategy(:api_key, strategy)
     end
 
@@ -330,8 +338,8 @@ module VectorMCP
     # Add custom authentication strategy
     # @param handler [Proc] custom authentication handler block
     # @return [void]
-    def add_custom_auth(&block)
-      strategy = Security::Strategies::Custom.new(&block)
+    def add_custom_auth(&)
+      strategy = Security::Strategies::Custom.new(&)
       @auth_manager.add_strategy(:custom, strategy)
     end
 
@@ -346,7 +354,7 @@ module VectorMCP
 
   module Transport
     # Dummy base class placeholder used only for argument validation in tests.
-    # Real transport classes (e.g., Stdio, SSE) are separate concrete classes.
+    # Real transport classes (e.g., HttpStream) are separate concrete classes.
     class Base # :nodoc:
     end
   end

data/lib/vector_mcp/session.rb

@@ -4,6 +4,7 @@ require_relative "sampling/request"
 require_relative "sampling/result"
 require_relative "errors"
 require_relative "request_context"
+require_relative "security/session_context"
 
 module VectorMCP
   # Represents the state of a single client-server connection session in MCP.
@@ -17,7 +18,7 @@ module VectorMCP
   # @attr_reader request_context [RequestContext] The request context for this session.
   class Session
     attr_reader :server_info, :server_capabilities, :protocol_version, :client_info, :client_capabilities, :server, :transport, :id, :request_context
-    attr_accessor :data # For user-defined session-specific storage
+    attr_accessor :data, :security_context # For user-defined session-specific storage and resolved auth context
 
     # Initializes a new session.
     #
@@ -33,6 +34,7 @@ module VectorMCP
       @client_info = nil
       @client_capabilities = nil
       @data = {} # Initialize user data hash
+      @security_context = Security::SessionContext.anonymous
       @logger = server.logger
 
       # Initialize request context
@@ -243,11 +245,11 @@ module VectorMCP
       send_request_kwargs = {}
       send_request_kwargs[:timeout] = timeout if timeout
 
-      # For HTTP transport, we need to use send_request_to_session to target this specific session
+      # Prefer session-targeted request sending when available
       raw_result = if @transport.respond_to?(:send_request_to_session)
                      @transport.send_request_to_session(@id, *send_request_args, **send_request_kwargs)
                    else
-                     # Fallback to generic send_request for other transports
+                     # Fallback to generic send_request for transports without session targeting
                      @transport.send_request(*send_request_args, **send_request_kwargs)
                    end
 

data/lib/vector_mcp/tool.rb

@@ -0,0 +1,221 @@
+# frozen_string_literal: true
+
+require "date"
+require "time"
+require_relative "errors"
+
+module VectorMCP
+  # Abstract base class for declarative tool definitions.
+  #
+  # Subclass this to define tools using a class-level DSL instead of
+  # the block-based +register_tool+ API. The two styles are fully
+  # interchangeable -- both produce the same +VectorMCP::Definitions::Tool+
+  # struct that the rest of the system consumes.
+  #
+  # @example
+  #   class ListProviders < VectorMCP::Tool
+  #     tool_name   "list_providers"
+  #     description "List providers filtered by category or status"
+  #
+  #     param :category, type: :string, desc: "Filter by category slug"
+  #     param :active,   type: :boolean, default: true
+  #
+  #     def call(args, session)
+  #       Provider.where(active: args.fetch("active", true))
+  #     end
+  #   end
+  #
+  #   server.register(ListProviders)
+  #
+  class Tool
+    # Maps Ruby symbol types to JSON Schema property fragments.
+    # Each value is merged into the generated property hash, so it may carry
+    # both +type+ and +format+ (or any other JSON Schema keyword).
+    TYPE_MAP = {
+      string: { "type" => "string" },
+      integer: { "type" => "integer" },
+      number: { "type" => "number" },
+      boolean: { "type" => "boolean" },
+      array: { "type" => "array" },
+      object: { "type" => "object" },
+      date: { "type" => "string", "format" => "date" },
+      datetime: { "type" => "string", "format" => "date-time" }
+    }.freeze
+
+    # Maps Ruby symbol types to a coercer lambda. Types not listed here
+    # pass their values through unchanged. Coercers receive the raw value
+    # (or nil) and return the coerced value. They must be total over the
+    # values JSON Schema validation would accept.
+    COERCERS = {
+      date: ->(v) { v.nil? || v.is_a?(Date) ? v : Date.parse(v.to_s) },
+      datetime: ->(v) { v.nil? || v.is_a?(Time) ? v : Time.parse(v.to_s) }
+    }.freeze
+
+    # Ensures each subclass gets its own +@params+ array so sibling
+    # classes do not share mutable state.
+    def self.inherited(subclass)
+      super
+      subclass.instance_variable_set(:@params, [])
+    end
+
+    # Sets or retrieves the tool name.
+    #
+    # When called with an argument, stores the name.
+    # When called without, returns the stored name or derives one from the class name.
+    #
+    # @param name [String, Symbol, nil] The tool name to set.
+    # @return [String] The tool name.
+    def self.tool_name(name = nil)
+      if name
+        @tool_name = name.to_s
+      else
+        @tool_name || derive_tool_name
+      end
+    end
+
+    # Sets or retrieves the tool description.
+    #
+    # @param text [String, nil] The description to set.
+    # @return [String, nil] The description.
+    def self.description(text = nil)
+      if text
+        @description = text
+      else
+        @description
+      end
+    end
+
+    # Declares a parameter for the tool's input schema.
+    #
+    # @param name [Symbol, String] The parameter name.
+    # @param type [Symbol] The parameter type (:string, :integer, :number, :boolean, :array, :object).
+    # @param desc [String, nil] A human-readable description.
+    # @param required [Boolean] Whether the parameter is required (default: false).
+    # @param options [Hash] Additional JSON Schema keywords (enum:, default:, format:, items:, etc.).
+    def self.param(name, type: :string, desc: nil, required: false, **options)
+      @params << {
+        name: name.to_s,
+        type: type,
+        desc: desc,
+        required: required,
+        options: options
+      }
+    end
+
+    # Builds a +VectorMCP::Definitions::Tool+ struct from the DSL metadata.
+    #
+    # @return [VectorMCP::Definitions::Tool]
+    # @raise [ArgumentError] If the subclass is missing a description or +#call+ method.
+    def self.to_definition
+      validate_tool_class!
+
+      VectorMCP::Definitions::Tool.new(
+        tool_name,
+        description,
+        build_input_schema,
+        build_handler
+      )
+    end
+
+    # The handler method that subclasses must implement.
+    #
+    # @param _args [Hash] The tool arguments (string keys).
+    # @param _session [VectorMCP::Session] The current session.
+    # @return [Object] The tool result.
+    def call(_args, _session)
+      raise NotImplementedError, "#{self.class.name} must implement #call(args, session)"
+    end
+
+    # Derives a snake_case tool name from the class name.
+    def self.derive_tool_name
+      base = name&.split("::")&.last
+      return "unnamed_tool" unless base
+
+      base.gsub(/([a-z\d])([A-Z])/, '\1_\2').downcase
+    end
+    private_class_method :derive_tool_name
+
+    # Validates that the subclass is properly configured.
+    def self.validate_tool_class!
+      raise ArgumentError, "#{name || self} must declare a description" unless description
+      return if method_defined?(:call) && instance_method(:call).owner != VectorMCP::Tool
+
+      raise ArgumentError, "#{name || self} must implement #call"
+    end
+    private_class_method :validate_tool_class!
+
+    # Builds a 2-arity handler lambda. A new instance is created per invocation.
+    # Arguments are coerced based on the declared param types before the
+    # handler sees them (e.g. :date param strings become Date objects).
+    def self.build_handler
+      klass = self
+      params = @params
+      ->(args, session) { klass.new.call(klass.coerce_args(args, params), session) }
+    end
+    private_class_method :build_handler
+
+    # Applies coercers to the raw argument hash. Returns a new hash; does
+    # not mutate the original. Keys without a coercible type pass through.
+    # Keys that are absent from +args+ stay absent — coercion only fires
+    # for keys actually present.
+    #
+    # A parse failure on a client-supplied value is translated into
+    # +VectorMCP::InvalidParamsError+ (JSON-RPC -32602) so the client sees
+    # a "bad request" response instead of a generic internal error.
+    # This is needed because the +json-schema+ gem does not enforce
+    # +format: date+ upstream (it does enforce +format: date-time+), so
+    # malformed +:date+ values would otherwise crash inside +Date.parse+.
+    def self.coerce_args(args, params)
+      coerced = args.dup
+      params.each do |param|
+        name = param[:name]
+        next unless coerced.key?(name)
+
+        coercer = COERCERS[param[:type]]
+        next unless coercer
+
+        begin
+          coerced[name] = coercer.call(coerced[name])
+        rescue ArgumentError, TypeError => e
+          # Date::Error < ArgumentError in Ruby 3.2+, so ArgumentError alone covers Date.parse failures.
+          raise VectorMCP::InvalidParamsError.new(
+            "Invalid #{param[:type]} value for param '#{name}': #{e.message}",
+            details: { param: name, type: param[:type], message: e.message }
+          )
+        end
+      end
+      coerced
+    end
+
+    # Builds a JSON Schema hash from the declared params.
+    def self.build_input_schema
+      properties = {}
+      required = []
+
+      @params.each do |param|
+        type_fragment = TYPE_MAP.fetch(param[:type]) do
+          raise ArgumentError, "Unknown param type :#{param[:type]} for param '#{param[:name]}' in #{name}. " \
+                               "Valid types: #{TYPE_MAP.keys.join(", ")}"
+        end
+
+        prop = type_fragment.dup
+        prop["description"] = param[:desc] if param[:desc]
+
+        param[:options].each do |key, value|
+          prop[key.to_s] = value
+        end
+
+        properties[param[:name]] = prop
+        required << param[:name] if param[:required]
+      end
+
+      schema = {
+        "type" => "object",
+        "properties" => properties
+      }
+      schema["required"] = required unless required.empty?
+      schema
+    end
+    private_class_method :build_input_schema
+  end
+end

data/lib/vector_mcp/transport/base_session_manager.rb

@@ -183,22 +183,6 @@ module VectorMCP
         end
       end
 
-      # Broadcasts a message to all sessions that support messaging.
-      #
-      # @param message [Hash] The message to broadcast
-      # @return [Integer] Number of sessions the message was sent to
-      def broadcast_message(message)
-        count = 0
-        @sessions.each_value do |session|
-          next unless can_send_message_to_session?(session)
-
-          count += 1 if message_sent_to_session?(session, message)
-        end
-
-        # Message broadcasted to recipients
-        count
-      end
-
       protected
 
       # Hook called when a session is created. Override in subclasses for transport-specific logic.
@@ -225,7 +209,7 @@ module VectorMCP
       end
 
       # Determines if this session manager should enable automatic cleanup.
-      # Override in subclasses that don't need automatic cleanup (e.g., stdio with single session).
+      # Override in subclasses that don't need automatic cleanup (e.g., transports with single persistent session).
       #
       # @return [Boolean] True if auto-cleanup should be enabled
       def auto_cleanup_enabled?

data/lib/vector_mcp/transport/http_stream/event_store.rb

@@ -17,7 +17,7 @@ module VectorMCP
       # @api private
       class EventStore
         # Event data structure
-        Event = Struct.new(:id, :data, :type, :timestamp) do
+        Event = Struct.new(:id, :data, :type, :timestamp, :session_id, :stream_id) do
           def to_sse_format
             lines = []
             lines << "id: #{id}"
@@ -44,12 +44,14 @@ module VectorMCP
         #
         # @param data [String] The event data
         # @param type [String] The event type (optional)
+        # @param session_id [String, nil] The session ID to scope this event to
+        # @param stream_id [String, nil] The stream ID to scope this event to
         # @return [String] The generated event ID
-        def store_event(data, type = nil)
+        def store_event(data, type = nil, session_id: nil, stream_id: nil)
           event_id = generate_event_id
           timestamp = Time.now
 
-          event = Event.new(event_id, data, type, timestamp)
+          event = Event.new(event_id, data, type, timestamp, session_id, stream_id)
 
           # Add to events array
           @events.push(event)
@@ -69,21 +71,39 @@ module VectorMCP
           event_id
         end
 
-        # Retrieves events starting from a specific event ID.
+        # Retrieves events starting from a specific event ID, optionally filtered by session.
         #
         # @param last_event_id [String] The last event ID received by client
+        # @param session_id [String, nil] Filter events to this session only
+        # @param stream_id [String, nil] Filter events to this stream only
         # @return [Array<Event>] Array of events after the specified ID
-        def get_events_after(last_event_id)
-          return @events.to_a if last_event_id.nil?
-
-          last_index = @event_index[last_event_id]
-          return [] if last_index.nil?
+        def get_events_after(last_event_id, session_id: nil, stream_id: nil)
+          events = if last_event_id.nil?
+                     @events.to_a
+                   else
+                     last_index = @event_index[last_event_id]
+                     return [] if last_index.nil?
+
+                     start_index = last_index + 1
+                     return [] if start_index >= @events.length
+
+                     @events[start_index..]
+                   end
+
+          events = events.select { |e| e.session_id == session_id } if session_id
+          events = events.select { |e| e.stream_id == stream_id } if stream_id
+          events
+        end
 
-          # Return events after the last_event_id
-          start_index = last_index + 1
-          return [] if start_index >= @events.length
+        # Retrieves a specific event by ID.
+        #
+        # @param event_id [String] The event ID to look up
+        # @return [Event, nil] The stored event, or nil if it is no longer retained
+        def get_event(event_id)
+          index = @event_index[event_id]
+          return nil if index.nil?
 
-          @events[start_index..]
+          @events[index]
         end
 
         # Gets the total number of stored events.

data/lib/vector_mcp/transport/http_stream/session_manager.rb

@@ -34,7 +34,7 @@ module VectorMCP
           end
 
           def streaming?
-            metadata[:streaming_connection] && !metadata[:streaming_connection].nil?
+            !streaming_connection.nil? || !streaming_connections.empty?
           end
 
           def streaming_connection
@@ -44,6 +44,25 @@ module VectorMCP
           def streaming_connection=(connection)
             metadata[:streaming_connection] = connection
           end
+
+          def streaming_connections
+            metadata[:streaming_connections] ||= Concurrent::Hash.new
+          end
+
+          def add_streaming_connection(connection)
+            streaming_connections[connection.stream_id] = connection
+            self.streaming_connection = connection
+          end
+
+          def remove_streaming_connection(connection = nil)
+            if connection
+              streaming_connections.delete(connection.stream_id)
+              self.streaming_connection = streaming_connections.values.first if streaming_connection == connection
+            else
+              streaming_connections.clear
+              self.streaming_connection = nil
+            end
+          end
         end
 
         # Initializes a new HTTP stream session manager.
@@ -64,7 +83,7 @@ module VectorMCP
                             end
 
           # Pre-allocate metadata hash for better performance
-          metadata = { streaming_connection: nil }
+          metadata = create_session_metadata
 
           # Create internal session record with streaming connection metadata
           session = Session.new(session_id, session_context, now, now, metadata)
@@ -75,7 +94,9 @@ module VectorMCP
           session
         end
 
-        # Override to add rack_env support
+        # Override to add rack_env support.
+        # Returns nil when a session_id is provided but not found (expired or unknown).
+        # Callers are responsible for returning 404 in that case.
         def get_or_create_session(session_id = nil, rack_env = nil)
           if session_id
             session = get_session(session_id)
@@ -88,8 +109,8 @@ module VectorMCP
               return session
             end
 
-            # If session_id was provided but not found, create with that ID
-            return create_session(session_id, rack_env)
+            # Session ID provided but not found — signal 404 to caller
+            return nil
           end
 
           create_session(nil, rack_env)
@@ -129,17 +150,18 @@ module VectorMCP
         # @param connection [Object] The streaming connection object
         # @return [void]
         def set_streaming_connection(session, connection)
-          session.streaming_connection = connection
+          session.add_streaming_connection(connection)
           session.touch!
-          logger.debug { "Streaming connection associated: #{session.id}" }
+          logger.debug { "Streaming connection associated: #{session.id} (stream #{connection.stream_id})" }
         end
 
         # Removes streaming connection from a session.
         #
         # @param session [Session] The session to remove streaming from
+        # @param connection [Object, nil] The specific connection to remove, or nil to clear all
         # @return [void]
-        def remove_streaming_connection(session)
-          session.streaming_connection = nil
+        def remove_streaming_connection(session, connection = nil)
+          session.remove_streaming_connection(connection)
           session.touch!
           logger.debug { "Streaming connection removed: #{session.id}" }
         end
@@ -153,7 +175,7 @@ module VectorMCP
 
         # Override: Returns metadata for new HTTP stream sessions.
         def create_session_metadata
-          { streaming_connection: nil }
+          { streaming_connection: nil, streaming_connections: Concurrent::Hash.new }
         end
 
         # Override: Checks if a session can receive messages (has streaming connection).
@@ -173,15 +195,18 @@ module VectorMCP
         # @param session [Session] The session whose connection to close
         # @return [void]
         def close_streaming_connection(session)
-          return unless session&.streaming_connection
+          return unless session&.streaming?
+
+          connections = session.streaming_connections.values
+          connections = [session.streaming_connection] if connections.empty? && session.streaming_connection
 
-          begin
-            session.streaming_connection.close
+          connections.each do |connection|
+            connection.close
           rescue StandardError => e
             logger.warn { "Error closing streaming connection for #{session.id}: #{e.message}" }
           end
 
-          session.streaming_connection = nil
+          session.remove_streaming_connection
         end
       end
     end