vector_mcp 0.3.1 → 0.3.3

data/lib/vector_mcp/session.rb

@@ -3,6 +3,7 @@
 require_relative "sampling/request"
 require_relative "sampling/result"
 require_relative "errors"
+require_relative "request_context"
 
 module VectorMCP
   # Represents the state of a single client-server connection session in MCP.
@@ -13,8 +14,9 @@ module VectorMCP
   # @attr_reader protocol_version [String] The MCP protocol version used by the server.
   # @attr_reader client_info [Hash, nil] Information about the client, received during initialization.
   # @attr_reader client_capabilities [Hash, nil] Capabilities supported by the client, received during initialization.
+  # @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
+    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
 
     # Initializes a new session.
@@ -22,7 +24,8 @@ module VectorMCP
     # @param server [VectorMCP::Server] The server instance managing this session.
     # @param transport [VectorMCP::Transport::Base, nil] The transport handling this session. Required for sampling.
     # @param id [String] A unique identifier for this session (e.g., from transport layer).
-    def initialize(server, transport = nil, id: SecureRandom.uuid)
+    # @param request_context [RequestContext, Hash, nil] The request context for this session.
+    def initialize(server, transport = nil, id: SecureRandom.uuid, request_context: nil)
       @server = server
       @transport = transport # Store the transport for sending requests
       @id = id
@@ -31,6 +34,16 @@ module VectorMCP
       @client_capabilities = nil
       @data = {} # Initialize user data hash
       @logger = server.logger
+
+      # Initialize request context
+      @request_context = case request_context
+                         when RequestContext
+                           request_context
+                         when Hash
+                           RequestContext.new(**request_context)
+                         else
+                           RequestContext.new
+                         end
     end
 
     # Marks the session as initialized using parameters from the client's `initialize` request.
@@ -75,6 +88,75 @@ module VectorMCP
       @initialized_state == :succeeded
     end
 
+    # Sets the request context for this session.
+    # This method should be called by transport layers to populate request-specific data.
+    #
+    # @param context [RequestContext, Hash] The request context to set.
+    #   Can be a RequestContext object or a hash of attributes.
+    # @return [RequestContext] The newly set request context.
+    # @raise [ArgumentError] If the context is not a RequestContext or Hash.
+    def request_context=(context)
+      @request_context = case context
+                         when RequestContext
+                           context
+                         when Hash
+                           RequestContext.new(**context)
+                         else
+                           raise ArgumentError, "Request context must be a RequestContext or Hash, got #{context.class}"
+                         end
+    end
+
+    # Updates the request context with new data.
+    # This merges the provided attributes with the existing context.
+    #
+    # @param attributes [Hash] The attributes to merge into the request context.
+    # @return [RequestContext] The updated request context.
+    def update_request_context(**attributes)
+      current_attrs = @request_context.to_h
+
+      # Deep merge nested hashes like headers and params
+      merged_attrs = current_attrs.dup
+      attributes.each do |key, value|
+        merged_attrs[key] = if value.is_a?(Hash) && current_attrs[key].is_a?(Hash)
+                              current_attrs[key].merge(value)
+                            else
+                              value
+                            end
+      end
+
+      @request_context = RequestContext.new(**merged_attrs)
+    end
+
+    # Convenience method to check if the session has request headers.
+    #
+    # @return [Boolean] True if the request context has headers, false otherwise.
+    def request_headers?
+      @request_context.headers?
+    end
+
+    # Convenience method to check if the session has request parameters.
+    #
+    # @return [Boolean] True if the request context has parameters, false otherwise.
+    def request_params?
+      @request_context.params?
+    end
+
+    # Convenience method to get a request header value.
+    #
+    # @param name [String] The header name.
+    # @return [String, nil] The header value or nil if not found.
+    def request_header(name)
+      @request_context.header(name)
+    end
+
+    # Convenience method to get a request parameter value.
+    #
+    # @param name [String] The parameter name.
+    # @return [String, nil] The parameter value or nil if not found.
+    def request_param(name)
+      @request_context.param(name)
+    end
+
     # Helper to check client capabilities later if needed
     # def supports?(capability_key)
     #   @client_capabilities.key?(capability_key.to_s)
@@ -95,10 +177,43 @@ module VectorMCP
     def sample(request_params, timeout: nil)
       validate_sampling_preconditions
 
-      sampling_req_obj = VectorMCP::Sampling::Request.new(request_params)
-      @logger.info("[Session #{@id}] Sending sampling/createMessage request to client.")
+      # Create middleware context for sampling
+      context = VectorMCP::Middleware::Context.new(
+        operation_type: :sampling,
+        operation_name: "createMessage",
+        params: request_params,
+        session: self,
+        server: @server,
+        metadata: { start_time: Time.now, timeout: timeout }
+      )
+
+      # Execute before_sampling_request hooks
+      context = @server.middleware_manager.execute_hooks(:before_sampling_request, context)
+      raise context.error if context.error?
+
+      begin
+        sampling_req_obj = VectorMCP::Sampling::Request.new(request_params)
+        @logger.debug("[Session #{@id}] Sending sampling/createMessage request to client.")
 
-      send_sampling_request(sampling_req_obj, timeout)
+        result = send_sampling_request(sampling_req_obj, timeout)
+
+        # Set result in context
+        context.result = result
+
+        # Execute after_sampling_response hooks
+        context = @server.middleware_manager.execute_hooks(:after_sampling_response, context)
+
+        context.result
+      rescue StandardError => e
+        # Set error in context and execute error hooks
+        context.error = e
+        context = @server.middleware_manager.execute_hooks(:on_sampling_error, context)
+
+        # Re-raise unless middleware handled the error
+        raise e unless context.result
+
+        context.result
+      end
     end
 
     private
@@ -128,17 +243,25 @@ module VectorMCP
       send_request_kwargs = {}
       send_request_kwargs[:timeout] = timeout if timeout
 
-      raw_result = @transport.send_request(*send_request_args, **send_request_kwargs)
+      # For HTTP transport, we need to use send_request_to_session to target this specific session
+      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
+                     @transport.send_request(*send_request_args, **send_request_kwargs)
+                   end
+
       VectorMCP::Sampling::Result.new(raw_result)
     rescue ArgumentError => e
       @logger.error("[Session #{@id}] Invalid parameters for sampling request or result: #{e.message}")
-      raise VectorMCP::SamplingError, "Invalid sampling parameters or malformed client response: #{e.message}", details: { original_error: e.to_s }
+      raise VectorMCP::SamplingError.new("Invalid sampling parameters or malformed client response: #{e.message}",
+                                         details: { original_error: e.to_s })
     rescue VectorMCP::SamplingError => e
       @logger.warn("[Session #{@id}] Sampling request failed: #{e.message}")
       raise e
     rescue StandardError => e
       @logger.error("[Session #{@id}] Unexpected error during sampling: #{e.class.name}: #{e.message}")
-      raise VectorMCP::SamplingError, "An unexpected error occurred during sampling: #{e.message}", details: { original_error: e.to_s }
+      raise VectorMCP::SamplingError.new("An unexpected error occurred during sampling: #{e.message}", details: { original_error: e.to_s })
     end
   end
 end

data/lib/vector_mcp/transport/base_session_manager.rb

@@ -0,0 +1,320 @@
+# frozen_string_literal: true
+
+require "securerandom"
+require "concurrent-ruby"
+
+module VectorMCP
+  module Transport
+    # Base session manager providing unified session lifecycle management across all transports.
+    # This abstract base class defines the standard interface that all transport session managers
+    # should implement, ensuring consistent session handling regardless of transport type.
+    #
+    # @abstract Subclass and implement transport-specific methods
+    class BaseSessionManager
+      # Session data structure for unified session management
+      Session = Struct.new(:id, :context, :created_at, :last_accessed_at, :metadata) do
+        def touch!
+          self.last_accessed_at = Time.now
+        end
+
+        def expired?(timeout)
+          Time.now - last_accessed_at > timeout
+        end
+
+        def age
+          Time.now - created_at
+        end
+      end
+
+      attr_reader :transport, :session_timeout, :logger
+
+      # Initializes a new session manager.
+      #
+      # @param transport [Object] The parent transport instance
+      # @param session_timeout [Integer] Session timeout in seconds (default: 300)
+      def initialize(transport, session_timeout = 300)
+        @transport = transport
+        @session_timeout = session_timeout
+        @logger = transport.logger
+        @sessions = Concurrent::Hash.new
+        @cleanup_timer = nil
+
+        start_cleanup_timer if auto_cleanup_enabled?
+        logger.debug { "#{self.class.name} initialized with session_timeout: #{session_timeout}" }
+      end
+
+      # Gets an existing session by ID.
+      #
+      # @param session_id [String] The session ID
+      # @return [Session, nil] The session if found and valid
+      def get_session(session_id)
+        return nil unless session_id
+
+        session = @sessions[session_id]
+        return nil unless session && !session.expired?(@session_timeout)
+
+        session.touch!
+        session
+      end
+
+      # Gets an existing session or creates a new one.
+      #
+      # @param session_id [String, nil] The session ID (optional)
+      # @return [Session] The existing or newly created session
+      def get_or_create_session(session_id = nil)
+        if session_id
+          session = get_session(session_id)
+          return session if session
+
+          # If session_id was provided but not found, create with that ID
+          return create_session(session_id)
+        end
+
+        create_session
+      end
+
+      # Creates a new session.
+      #
+      # @param session_id [String, nil] Optional specific session ID to use
+      # @return [Session] The newly created session
+      def create_session(session_id = nil)
+        session_id ||= generate_session_id
+        now = Time.now
+
+        # Create VectorMCP session context
+        session_context = VectorMCP::Session.new(@transport.server, @transport, id: session_id)
+
+        # Create internal session record with transport-specific metadata
+        session = Session.new(
+          session_id,
+          session_context,
+          now,
+          now,
+          create_session_metadata
+        )
+
+        @sessions[session_id] = session
+
+        logger.info { "Session created: #{session_id}" }
+        on_session_created(session)
+        session
+      end
+
+      # Terminates a session by ID.
+      #
+      # @param session_id [String] The session ID to terminate
+      # @return [Boolean] True if session was found and terminated
+      def session_terminated?(session_id)
+        session = @sessions.delete(session_id)
+        return false unless session
+
+        on_session_terminated(session)
+        logger.info { "Session terminated: #{session_id}" }
+        true
+      end
+
+      # Gets the current number of active sessions.
+      #
+      # @return [Integer] Number of active sessions
+      def session_count
+        @sessions.size
+      end
+
+      # Gets all active session IDs.
+      #
+      # @return [Array<String>] Array of session IDs
+      def active_session_ids
+        @sessions.keys
+      end
+
+      # Checks if any sessions exist.
+      #
+      # @return [Boolean] True if at least one session exists
+      def sessions?
+        !@sessions.empty?
+      end
+
+      # Cleans up all sessions and stops the cleanup timer.
+      #
+      # @return [void]
+      def cleanup_all_sessions
+        logger.info { "Cleaning up all sessions: #{@sessions.size}" }
+
+        @sessions.each_value do |session|
+          on_session_terminated(session)
+        end
+
+        @sessions.clear
+        stop_cleanup_timer
+      end
+
+      # Updates session metadata.
+      #
+      # @param session_id [String] The session ID
+      # @param metadata [Hash] Metadata to merge
+      # @return [Boolean] True if session was found and updated
+      def session_metadata_updated?(session_id, metadata)
+        session = @sessions[session_id]
+        return false unless session
+
+        session.metadata.merge!(metadata)
+        session.touch!
+        true
+      end
+
+      # Gets session metadata.
+      #
+      # @param session_id [String] The session ID
+      # @return [Hash, nil] Session metadata or nil if session not found
+      def get_session_metadata(session_id)
+        session = @sessions[session_id]
+        session&.metadata
+      end
+
+      # Finds sessions matching criteria.
+      #
+      # @param criteria [Hash] Search criteria
+      # @option criteria [Symbol] :created_after Time to search after
+      # @option criteria [Symbol] :metadata Hash of metadata to match
+      # @return [Array<Session>] Matching sessions
+      def find_sessions(criteria = {})
+        @sessions.values.select do |session|
+          matches_criteria?(session, criteria)
+        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.
+      #
+      # @param session [Session] The newly created session
+      # @return [void]
+      def on_session_created(session)
+        # Override in subclasses
+      end
+
+      # Hook called when a session is terminated. Override in subclasses for transport-specific cleanup.
+      #
+      # @param session [Session] The session being terminated
+      # @return [void]
+      def on_session_terminated(session)
+        # Override in subclasses
+      end
+
+      # Creates transport-specific session metadata. Override in subclasses.
+      #
+      # @return [Hash] Initial metadata for the session
+      def create_session_metadata
+        {}
+      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).
+      #
+      # @return [Boolean] True if auto-cleanup should be enabled
+      def auto_cleanup_enabled?
+        true
+      end
+
+      # Checks if a message can be sent to the given session.
+      # Override in subclasses based on transport capabilities.
+      #
+      # @param session [Session] The session to check
+      # @return [Boolean] True if messaging is supported for this session
+      def can_send_message_to_session?(_session)
+        false # Override in subclasses
+      end
+
+      # Sends a message to a specific session.
+      # Override in subclasses based on transport messaging mechanism.
+      #
+      # @param session [Session] The target session
+      # @param message [Hash] The message to send
+      # @return [Boolean] True if message was sent successfully
+      def message_sent_to_session?(_session, _message)
+        false # Override in subclasses
+      end
+
+      private
+
+      # Generates a cryptographically secure session ID.
+      #
+      # @return [String] A unique session ID
+      def generate_session_id
+        SecureRandom.uuid
+      end
+
+      # Starts the automatic cleanup timer if auto-cleanup is enabled.
+      #
+      # @return [void]
+      def start_cleanup_timer
+        return unless auto_cleanup_enabled?
+
+        # Run cleanup every 60 seconds
+        @cleanup_timer = Concurrent::TimerTask.new(execution_interval: 60) do
+          cleanup_expired_sessions
+        end
+        @cleanup_timer.execute
+      end
+
+      # Stops the automatic cleanup timer.
+      #
+      # @return [void]
+      def stop_cleanup_timer
+        @cleanup_timer&.shutdown
+        @cleanup_timer = nil
+      end
+
+      # Cleans up expired sessions.
+      #
+      # @return [void]
+      def cleanup_expired_sessions
+        expired_sessions = []
+
+        @sessions.each do |session_id, session|
+          expired_sessions << session_id if session.expired?(@session_timeout)
+        end
+
+        expired_sessions.each do |session_id|
+          session = @sessions.delete(session_id)
+          on_session_terminated(session) if session
+        end
+
+        return unless expired_sessions.any?
+
+        logger.debug { "Cleaned up expired sessions: #{expired_sessions.size}" }
+      end
+
+      # Checks if a session matches the given criteria.
+      #
+      # @param session [Session] The session to check
+      # @param criteria [Hash] The search criteria
+      # @return [Boolean] True if session matches all criteria
+      def matches_criteria?(session, criteria)
+        return false if criteria[:created_after] && session.created_at <= criteria[:created_after]
+
+        criteria[:metadata]&.each do |key, value|
+          return false unless session.metadata[key] == value
+        end
+
+        true
+      end
+    end
+  end
+end

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

@@ -0,0 +1,151 @@
+# frozen_string_literal: true
+
+require "concurrent-ruby"
+require "securerandom"
+
+module VectorMCP
+  module Transport
+    class HttpStream
+      # Manages Server-Sent Events storage for resumable connections.
+      #
+      # Handles:
+      # - Event storage with unique IDs
+      # - Event replay from a specific Last-Event-ID
+      # - Circular buffer for memory efficiency
+      # - Thread-safe operations
+      #
+      # @api private
+      class EventStore
+        # Event data structure
+        Event = Struct.new(:id, :data, :type, :timestamp) do
+          def to_sse_format
+            lines = []
+            lines << "id: #{id}"
+            lines << "event: #{type}" if type
+            lines << "data: #{data}"
+            lines << ""
+            lines.join("\n")
+          end
+        end
+
+        attr_reader :max_events, :logger
+
+        # Initializes a new event store.
+        #
+        # @param max_events [Integer] Maximum number of events to retain
+        def initialize(max_events)
+          @max_events = max_events
+          @events = Concurrent::Array.new
+          @event_index = Concurrent::Hash.new # event_id -> index for fast lookup
+          @current_sequence = Concurrent::AtomicFixnum.new(0)
+        end
+
+        # Stores a new event and returns its ID.
+        #
+        # @param data [String] The event data
+        # @param type [String] The event type (optional)
+        # @return [String] The generated event ID
+        def store_event(data, type = nil)
+          event_id = generate_event_id
+          timestamp = Time.now
+
+          event = Event.new(event_id, data, type, timestamp)
+
+          # Add to events array
+          @events.push(event)
+
+          # Update index
+          @event_index[event_id] = @events.length - 1
+
+          # Maintain circular buffer
+          if @events.length > @max_events
+            removed_event = @events.shift
+            @event_index.delete(removed_event.id)
+
+            # Update all indices after removal
+            @event_index.transform_values! { |index| index - 1 }
+          end
+
+          event_id
+        end
+
+        # Retrieves events starting from a specific event ID.
+        #
+        # @param last_event_id [String] The last event ID received by client
+        # @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?
+
+          # Return events after the last_event_id
+          start_index = last_index + 1
+          return [] if start_index >= @events.length
+
+          @events[start_index..]
+        end
+
+        # Gets the total number of stored events.
+        #
+        # @return [Integer] Number of events currently stored
+        def event_count
+          @events.length
+        end
+
+        # Gets the oldest event ID (for debugging/monitoring).
+        #
+        # @return [String, nil] The oldest event ID or nil if no events
+        def oldest_event_id
+          @events.first&.id
+        end
+
+        # Gets the newest event ID (for debugging/monitoring).
+        #
+        # @return [String, nil] The newest event ID or nil if no events
+        def newest_event_id
+          @events.last&.id
+        end
+
+        # Checks if an event ID exists in the store.
+        #
+        # @param event_id [String] The event ID to check
+        # @return [Boolean] True if event exists
+        def event_exists?(event_id)
+          @event_index.key?(event_id)
+        end
+
+        # Clears all stored events.
+        #
+        # @return [void]
+        def clear
+          @events.clear
+          @event_index.clear
+        end
+
+        # Gets statistics about the event store.
+        #
+        # @return [Hash] Statistics hash
+        def stats
+          {
+            total_events: event_count,
+            max_events: @max_events,
+            oldest_event_id: oldest_event_id,
+            newest_event_id: newest_event_id,
+            memory_usage_ratio: event_count.to_f / @max_events
+          }
+        end
+
+        private
+
+        # Generates a unique event ID.
+        #
+        # @return [String] A unique event ID
+        def generate_event_id
+          sequence = @current_sequence.increment
+          "#{Time.now.to_i}-#{sequence}-#{SecureRandom.hex(4)}"
+        end
+      end
+    end
+  end
+end