vector_mcp 0.4.0 → 0.6.0

data/lib/vector_mcp/security/strategies/api_key.rb

@@ -31,18 +31,14 @@ module VectorMCP
         end
 
         # Authenticate a request using API key
-        # @param request [Hash] the request object
+        # @param request [VectorMCP::RequestContext, Hash] the request
         # @return [Hash, false] user info hash or false if authentication failed
         def authenticate(request)
-          api_key = extract_api_key(request)
+          api_key = extract_api_key(VectorMCP::RequestContext.coerce(request))
           return false unless api_key&.length&.positive?
 
           if secure_key_match?(api_key)
-            {
-              api_key: api_key,
-              strategy: "api_key",
-              authenticated_at: Time.now
-            }
+            { api_key: api_key }
           else
             false
           end
@@ -76,57 +72,16 @@ module VectorMCP
           matched
         end
 
-        # Extract API key from various request formats
-        # @param request [Hash] the request object
+        # Extract API key from the request context
+        # @param request [VectorMCP::RequestContext] the request context
         # @return [String, nil] the extracted API key
         def extract_api_key(request)
-          headers = normalize_headers(request)
-
-          from_headers = extract_from_headers(headers)
+          from_headers = extract_from_headers(request.headers)
           return from_headers if from_headers
 
           return nil unless @allow_query_params
 
-          params = normalize_params(request)
-          extract_from_params(params)
-        end
-
-        # Normalize headers to handle different formats
-        # @param request [Hash] the request object
-        # @return [Hash] normalized headers
-        def normalize_headers(request)
-          # Check if it's a Rack environment (has REQUEST_METHOD)
-          if request["REQUEST_METHOD"]
-            extract_headers_from_rack_env(request)
-          else
-            request[:headers] || request["headers"] || {}
-          end
-        end
-
-        # Normalize params to handle different formats
-        # @param request [Hash] the request object
-        # @return [Hash] normalized params
-        def normalize_params(request)
-          # Check if it's a Rack environment (has REQUEST_METHOD)
-          if request["REQUEST_METHOD"]
-            extract_params_from_rack_env(request)
-          else
-            request[:params] || request["params"] || {}
-          end
-        end
-
-        # Extract headers from Rack environment
-        # @param env [Hash] the Rack environment
-        # @return [Hash] normalized headers
-        def extract_headers_from_rack_env(env)
-          VectorMCP::Util.extract_headers_from_rack_env(env)
-        end
-
-        # Extract params from Rack environment
-        # @param env [Hash] the Rack environment
-        # @return [Hash] normalized params
-        def extract_params_from_rack_env(env)
-          VectorMCP::Util.extract_params_from_rack_env(env)
+          extract_from_params(request.params)
         end
 
         # Extract API key from headers

data/lib/vector_mcp/security/strategies/custom.rb

@@ -16,20 +16,23 @@ module VectorMCP
           @handler = handler
         end
 
-        # Authenticate a request using the custom handler
-        # @param request [Hash] the request object
-        # @return [Object, false] result from custom handler or false if authentication failed
+        # Authenticate a request using the custom handler.
+        # The handler receives a {VectorMCP::RequestContext}, which supports
+        # hash-style access (`request[:headers]`, `request[:params]`) for
+        # compatibility with handlers written against the legacy request hash.
+        # If the handler returns a Hash with a :user key, the value is extracted
+        # so that AuthManager receives the user data directly.
+        # @param request [VectorMCP::RequestContext, Hash] the request
+        # @return [Object, nil, false] user data or false if authentication failed.
+        #   A return of nil (from { user: nil }) signals "authenticated, no user object."
         def authenticate(request)
-          result = @handler.call(request)
+          result = @handler.call(VectorMCP::RequestContext.coerce(request))
+          return false unless result && result != false
 
-          # Ensure result includes strategy info if it's successful
-          if result && result != false
-            format_successful_result(result)
-          else
-            false
-          end
-        rescue NoMemoryError, StandardError
-          # Log error but return false for security
+          return result unless result.is_a?(Hash) && result.key?(:user)
+
+          result[:user]
+        rescue StandardError, NoMemoryError
           false
         end
 
@@ -38,33 +41,6 @@ module VectorMCP
         def configured?
           !@handler.nil?
         end
-
-        private
-
-        # Format successful authentication result with strategy metadata
-        # @param result [Object] the result from the custom handler
-        # @return [Object] formatted result with strategy metadata
-        def format_successful_result(result)
-          case result
-          when Hash
-            # If result has a :user key, extract it and use as main user data
-            if result.key?(:user)
-              user_data = result[:user]
-              # For nil user, return a marker that will become nil in session context
-              return :authenticated_nil_user if user_data.nil?
-
-              user_data
-            else
-              result.merge(strategy: "custom", authenticated_at: Time.now)
-            end
-          else
-            {
-              user: result,
-              strategy: "custom",
-              authenticated_at: Time.now
-            }
-          end
-        end
       end
     end
   end

data/lib/vector_mcp/security/strategies/jwt_token.rb

@@ -35,24 +35,15 @@ module VectorMCP
         end
 
         # Authenticate a request using JWT token
-        # @param request [Hash] the request object
+        # @param request [VectorMCP::RequestContext, Hash] the request
         # @return [Hash, false] decoded JWT payload or false if authentication failed
         def authenticate(request)
-          token = extract_token(request)
+          token = extract_token(VectorMCP::RequestContext.coerce(request))
           return false unless token
 
           begin
             decoded = JWT.decode(token, @secret, true, @options)
-            payload = decoded[0] # First element is the payload
-            headers = decoded[1] # Second element is the headers
-
-            # Return user info from JWT payload
-            {
-              **payload,
-              strategy: "jwt",
-              authenticated_at: Time.now,
-              jwt_headers: headers
-            }
+            decoded[0]
           rescue JWT::ExpiredSignature, JWT::InvalidIssuerError, JWT::InvalidAudienceError,
                  JWT::VerificationError, JWT::DecodeError, StandardError
             false # Token validation failed
@@ -79,19 +70,18 @@ module VectorMCP
 
         private
 
-        # Extract JWT token from request
-        # @param request [Hash] the request object
+        # Extract JWT token from the request context
+        # @param request [VectorMCP::RequestContext] the request context
         # @return [String, nil] the extracted token
         def extract_token(request)
-          headers = request[:headers] || request["headers"] || {}
+          headers = request.headers
 
           from_headers = extract_from_auth_header(headers) || extract_from_jwt_header(headers)
           return from_headers if from_headers
 
           return nil unless @allow_query_params
 
-          params = request[:params] || request["params"] || {}
-          extract_from_params(params)
+          extract_from_params(request.params)
         end
 
         # Extract token from Authorization header

data/lib/vector_mcp/server/capabilities.rb

@@ -39,19 +39,7 @@ module VectorMCP
       # Notifies connected clients that the list of available prompts has changed.
       # @return [void]
       def notify_prompts_list_changed
-        return unless transport && @prompts_list_changed
-
-        notification_method = "notifications/prompts/list_changed"
-        begin
-          if transport.respond_to?(:send_notification)
-            logger.debug("Sending prompts list changed notification.")
-            transport.send_notification(notification_method)
-          else
-            logger.warn("Transport does not support sending notifications/prompts/list_changed.")
-          end
-        rescue StandardError => e
-          logger.error("Failed to send prompts list changed notification: #{e.class.name}: #{e.message}")
-        end
+        send_list_changed_notification("prompts") if @prompts_list_changed
       end
 
       # Resets the `roots_list_changed` flag to false.
@@ -63,19 +51,7 @@ module VectorMCP
       # Notifies connected clients that the list of available roots has changed.
       # @return [void]
       def notify_roots_list_changed
-        return unless transport && @roots_list_changed
-
-        notification_method = "notifications/roots/list_changed"
-        begin
-          if transport.respond_to?(:send_notification)
-            logger.debug("Sending roots list changed notification.")
-            transport.send_notification(notification_method)
-          else
-            logger.warn("Transport does not support sending notifications/roots/list_changed.")
-          end
-        rescue StandardError => e
-          logger.error("Failed to send roots list changed notification: #{e.class.name}: #{e.message}")
-        end
+        send_list_changed_notification("roots") if @roots_list_changed
       end
 
       # Registers a session as a subscriber to prompt list changes.
@@ -87,6 +63,26 @@ module VectorMCP
 
       private
 
+      # Sends a `notifications/<kind>/list_changed` notification to the transport.
+      # No-op if no transport is attached. Logs a warning if the transport does not
+      # implement `send_notification` (intentional extension point for alternate
+      # transports).
+      # @api private
+      # @param kind [String] One of "prompts" or "roots".
+      def send_list_changed_notification(kind)
+        return unless transport
+
+        notification_method = "notifications/#{kind}/list_changed"
+        if transport.respond_to?(:send_notification)
+          logger.debug("Sending #{kind} list changed notification.")
+          transport.send_notification(notification_method)
+        else
+          logger.warn("Transport does not support sending #{notification_method}.")
+        end
+      rescue StandardError => e
+        logger.error("Failed to send #{kind} list changed notification: #{e.class.name}: #{e.message}")
+      end
+
       # Configures sampling capabilities based on provided configuration.
       # @api private
       def configure_sampling_capabilities(config)

data/lib/vector_mcp/server/message_handling.rb

@@ -10,27 +10,30 @@ module VectorMCP
       # This is the main dispatch point for messages received by a transport.
       #
       # @param message [Hash] The parsed JSON-RPC message object.
-      # @param session [VectorMCP::Session] The client session associated with this message.
-      # @param session_id [String] A unique identifier for the underlying transport connection.
+      # @param session [VectorMCP::Session, VectorMCP::Invocation] The session (or per-request
+      #   invocation view of it) this message belongs to.
+      # @param session_id [String, nil] Deprecated: defaults to session.id; used only for log prefixes.
       # @return [Object, nil] For requests, returns the result data to be sent in the JSON-RPC response.
       #   For notifications, returns `nil`.
       # @raise [VectorMCP::ProtocolError] if the message is invalid or an error occurs during handling.
-      def handle_message(message, session, session_id)
+      def handle_message(message, session, session_id = nil)
+        session_id ||= session.id
         id = message["id"]
         method = message["method"]
         params = message["params"] || {}
 
-        if id && method # Request
+        case classify_message(id, method)
+        when :request
           logger.debug("[#{session_id}] Request [#{id}]: #{method} with params: #{VectorMCP::LogFilter.filter_hash(params).inspect}")
           handle_request(id, method, params, session)
-        elsif method # Notification
+        when :notification
           logger.debug("[#{session_id}] Notification: #{method} with params: #{VectorMCP::LogFilter.filter_hash(params).inspect}")
           handle_notification(method, params, session)
-          nil # Notifications do not have a return value to send back to client
-        elsif id # Invalid: Has ID but no method
+          nil
+        when :invalid_missing_method
           logger.warn("[#{session_id}] Invalid message: Has ID [#{id}] but no method. #{message.inspect}")
           raise VectorMCP::InvalidRequestError.new("Request object must include a 'method' member.", request_id: id)
-        else # Invalid: No ID and no method
+        when :invalid_missing_both
           logger.warn("[#{session_id}] Invalid message: Missing both 'id' and 'method'. #{message.inspect}")
           raise VectorMCP::InvalidRequestError.new("Invalid message format", request_id: nil)
         end
@@ -60,6 +63,16 @@ module VectorMCP
 
       private
 
+      # Classify a JSON-RPC message based on presence of id and method fields.
+      # @api private
+      def classify_message(id, method)
+        return :request if id && method
+        return :notification if method
+        return :invalid_missing_method if id
+
+        :invalid_missing_both
+      end
+
       # Internal handler for JSON-RPC requests.
       # @api private
       def handle_request(id, method, params, session)
@@ -74,9 +87,7 @@ module VectorMCP
       # Validates that the session is properly initialized for the given request.
       # @api private
       def validate_session_initialization(id, method, _params, session)
-        # Handle both direct VectorMCP::Session and BaseSessionManager::Session wrapper
-        actual_session = session.respond_to?(:context) ? session.context : session
-        return if actual_session.initialized?
+        return if session.initialized?
 
         # Allow "initialize" even if not marked initialized yet by server
         return if method == "initialize"
@@ -115,9 +126,7 @@ module VectorMCP
       # Internal handler for JSON-RPC notifications.
       # @api private
       def handle_notification(method, params, session)
-        # Handle both direct VectorMCP::Session and BaseSessionManager::Session wrapper
-        actual_session = session.respond_to?(:context) ? session.context : session
-        unless actual_session.initialized? || method == "initialized"
+        unless session.initialized? || method == "initialized"
           logger.warn("Ignoring notification '#{method}' before session is initialized. Params: #{params.inspect}")
           return
         end
@@ -162,9 +171,7 @@ module VectorMCP
       # @api private
       def session_method(method_name)
         lambda do |params, session, _server|
-          # Handle both direct VectorMCP::Session and BaseSessionManager::Session wrapper
-          actual_session = session.respond_to?(:context) ? session.context : session
-          actual_session.public_send(method_name, params)
+          session.public_send(method_name, params)
         end
       end
     end

data/lib/vector_mcp/server/registry.rb

@@ -131,114 +131,74 @@ module VectorMCP
       end
 
       # Helper method to register an image resource from a file path.
+      # Thin wrapper: delegates schema-building to Definitions::Resource.from_image_file,
+      # then stores the result via register_resource.
       #
       # @param uri [String] Unique URI for the resource.
       # @param file_path [String] Path to the image file.
       # @param name [String, nil] Human-readable name (auto-generated if nil).
       # @param description [String, nil] Description (auto-generated if nil).
-      # @return [VectorMCP::Definitions::Resource] The registered resource.
+      # @return [self]
       # @raise [ArgumentError] If the file doesn't exist or isn't a valid image.
       def register_image_resource(uri:, file_path:, name: nil, description: nil)
         resource = VectorMCP::Definitions::Resource.from_image_file(
-          uri: uri,
-          file_path: file_path,
-          name: name,
-          description: description
-        )
-
-        register_resource(
-          uri: resource.uri,
-          name: resource.name,
-          description: resource.description,
-          mime_type: resource.mime_type,
-          &resource.handler
+          uri: uri, file_path: file_path, name: name, description: description
         )
+        register_resource(uri: resource.uri, name: resource.name,
+                          description: resource.description, mime_type: resource.mime_type, &resource.handler)
       end
 
       # Helper method to register an image resource from binary data.
+      # Thin wrapper: delegates to Definitions::Resource.from_image_data.
       #
       # @param uri [String] Unique URI for the resource.
       # @param image_data [String] Binary image data.
       # @param name [String] Human-readable name.
       # @param description [String, nil] Description (auto-generated if nil).
       # @param mime_type [String, nil] MIME type (auto-detected if nil).
-      # @return [VectorMCP::Definitions::Resource] The registered resource.
-      # @raise [ArgumentError] If the data isn't valid image data.
+      # @return [self]
       def register_image_resource_from_data(uri:, image_data:, name:, description: nil, mime_type: nil)
         resource = VectorMCP::Definitions::Resource.from_image_data(
-          uri: uri,
-          image_data: image_data,
-          name: name,
-          description: description,
-          mime_type: mime_type
-        )
-
-        register_resource(
-          uri: resource.uri,
-          name: resource.name,
-          description: resource.description,
-          mime_type: resource.mime_type,
-          &resource.handler
+          uri: uri, image_data: image_data, name: name, description: description, mime_type: mime_type
         )
+        register_resource(uri: resource.uri, name: resource.name,
+                          description: resource.description, mime_type: resource.mime_type, &resource.handler)
       end
 
       # Helper method to register a tool that accepts image inputs.
+      # Thin wrapper: delegates schema-building to Definitions::Tool.with_image_support.
       #
       # @param name [String] Unique name for the tool.
       # @param description [String] Human-readable description.
       # @param image_parameter [String] Name of the image parameter (default: "image").
       # @param additional_parameters [Hash] Additional JSON Schema properties.
       # @param required_parameters [Array<String>] List of required parameter names.
-      # @param block [Proc] The tool handler block.
-      # @return [VectorMCP::Definitions::Tool] The registered tool.
-      def register_image_tool(name:, description:, image_parameter: "image", additional_parameters: {}, required_parameters: [], &)
-        # Build the input schema with image support
-        image_property = {
-          type: "string",
-          description: "Base64 encoded image data or file path to image",
-          contentEncoding: "base64",
-          contentMediaType: "image/*"
-        }
-
-        properties = { image_parameter => image_property }.merge(additional_parameters)
-
-        input_schema = {
-          type: "object",
-          properties: properties,
-          required: required_parameters
-        }
-
-        register_tool(
-          name: name,
-          description: description,
-          input_schema: input_schema,
-          &
+      # @return [self]
+      def register_image_tool(name:, description:, image_parameter: "image",
+                              additional_parameters: {}, required_parameters: [], &handler)
+        tool = VectorMCP::Definitions::Tool.with_image_support(
+          name: name, description: description, image_parameter: image_parameter,
+          additional_parameters: additional_parameters, required_parameters: required_parameters, &handler
         )
+        register_tool(name: tool.name, description: tool.description,
+                      input_schema: tool.input_schema, &tool.handler)
       end
 
       # Helper method to register a prompt that supports image arguments.
+      # Thin wrapper: delegates to Definitions::Prompt.with_image_support.
       #
       # @param name [String] Unique name for the prompt.
       # @param description [String] Human-readable description.
       # @param image_argument [String] Name of the image argument (default: "image").
       # @param additional_arguments [Array<Hash>] Additional prompt arguments.
-      # @param block [Proc] The prompt handler block.
-      # @return [VectorMCP::Definitions::Prompt] The registered prompt.
-      def register_image_prompt(name:, description:, image_argument: "image", additional_arguments: [], &)
+      # @return [self]
+      def register_image_prompt(name:, description:, image_argument: "image", additional_arguments: [], &handler)
         prompt = VectorMCP::Definitions::Prompt.with_image_support(
-          name: name,
-          description: description,
-          image_argument_name: image_argument,
-          additional_arguments: additional_arguments,
-          &
-        )
-
-        register_prompt(
-          name: prompt.name,
-          description: prompt.description,
-          arguments: prompt.arguments,
-          &prompt.handler
+          name: name, description: description, image_argument_name: image_argument,
+          additional_arguments: additional_arguments, &handler
         )
+        register_prompt(name: prompt.name, description: prompt.description,
+                        arguments: prompt.arguments, &prompt.handler)
       end
 
       private
@@ -262,6 +222,34 @@ module VectorMCP
         raise ArgumentError, "Invalid input_schema structure: #{e.message}"
       end
 
+      # Schema for a single prompt argument definition. Each entry names the
+      # required/optional key, whether it is required, and the rule that validates
+      # its value. Rules return nil on success or an error message fragment.
+      PROMPT_ARG_SCHEMA = {
+        "name" => {
+          required: true,
+          missing_message: "missing :name",
+          rule: lambda { |v|
+            next "must be a String or Symbol. Found: #{v.class}" unless v.is_a?(String) || v.is_a?(Symbol)
+
+            "cannot be empty." if v.to_s.strip.empty?
+          }
+        },
+        "description" => {
+          required: false,
+          rule: ->(v) { "must be a String if provided. Found: #{v.class}" unless v.nil? || v.is_a?(String) }
+        },
+        "required" => {
+          required: false,
+          rule: ->(v) { "must be true or false if provided. Found: #{v.inspect}" unless [true, false].include?(v) }
+        },
+        "type" => {
+          required: false,
+          rule: ->(v) { "must be a String if provided (e.g., JSON schema type). Found: #{v.class}" unless v.nil? || v.is_a?(String) }
+        }
+      }.freeze
+      private_constant :PROMPT_ARG_SCHEMA
+
       # Validates the structure of the `arguments` array provided to {#register_prompt}.
       # @api private
       def validate_prompt_arguments(argument_defs)
@@ -270,75 +258,37 @@ module VectorMCP
         argument_defs.each_with_index { |arg, idx| validate_single_prompt_argument(arg, idx) }
       end
 
-      # Defines the keys allowed in a prompt argument definition hash.
-      ALLOWED_PROMPT_ARG_KEYS = %w[name description required type].freeze
-      private_constant :ALLOWED_PROMPT_ARG_KEYS
-
-      # Validates a single prompt argument definition hash.
+      # Validates a single prompt argument definition hash against PROMPT_ARG_SCHEMA.
       # @api private
       def validate_single_prompt_argument(arg, idx)
         raise ArgumentError, "Prompt argument definition at index #{idx} must be a Hash. Found: #{arg.class}" unless arg.is_a?(Hash)
 
-        validate_prompt_arg_name!(arg, idx)
-        validate_prompt_arg_description!(arg, idx)
-        validate_prompt_arg_required_flag!(arg, idx)
-        validate_prompt_arg_type!(arg, idx)
-        validate_prompt_arg_unknown_keys!(arg, idx)
-      end
-
-      # Validates the :name key of a prompt argument definition.
-      # @api private
-      def validate_prompt_arg_name!(arg, idx)
-        name_val = arg[:name] || arg["name"]
-        raise ArgumentError, "Prompt argument at index #{idx} missing :name" if name_val.nil?
-        unless name_val.is_a?(String) || name_val.is_a?(Symbol)
-          raise ArgumentError, "Prompt argument :name at index #{idx} must be a String or Symbol. Found: #{name_val.class}"
-        end
-        raise ArgumentError, "Prompt argument :name at index #{idx} cannot be empty." if name_val.to_s.strip.empty?
-      end
-
-      # Validates the :description key of a prompt argument definition.
-      # @api private
-      def validate_prompt_arg_description!(arg, idx)
-        return unless arg.key?(:description) || arg.key?("description")
-
-        desc_val = arg[:description] || arg["description"]
-        return if desc_val.nil? || desc_val.is_a?(String)
-
-        raise ArgumentError, "Prompt argument :description at index #{idx} must be a String if provided. Found: #{desc_val.class}"
-      end
-
-      # Validates the :required key of a prompt argument definition.
-      # @api private
-      def validate_prompt_arg_required_flag!(arg, idx)
-        return unless arg.key?(:required) || arg.key?("required")
-
-        req_val = arg[:required] || arg["required"]
-        return if [true, false].include?(req_val)
-
-        raise ArgumentError, "Prompt argument :required at index #{idx} must be true or false if provided. Found: #{req_val.inspect}"
+        PROMPT_ARG_SCHEMA.each { |key, spec| validate_prompt_arg_field(arg, idx, key, spec) }
+        validate_prompt_arg_unknown_keys(arg, idx)
       end
 
-      # Validates the :type key of a prompt argument definition.
+      # Validates a single field of a prompt argument hash against its schema spec.
       # @api private
-      def validate_prompt_arg_type!(arg, idx)
-        return unless arg.key?(:type) || arg.key?("type")
+      def validate_prompt_arg_field(arg, idx, key, spec)
+        present = arg.key?(key.to_sym) || arg.key?(key)
+        value = arg[key.to_sym] || arg[key]
 
-        type_val = arg[:type] || arg["type"]
-        return if type_val.nil? || type_val.is_a?(String)
+        raise ArgumentError, "Prompt argument at index #{idx} #{spec[:missing_message]}" if spec[:required] && value.nil?
+        return unless present
 
-        raise ArgumentError, "Prompt argument :type at index #{idx} must be a String if provided (e.g., JSON schema type). Found: #{type_val.class}"
+        error_fragment = spec[:rule].call(value)
+        raise ArgumentError, "Prompt argument :#{key} at index #{idx} #{error_fragment}" if error_fragment
       end
 
-      # Checks for any unknown keys in a prompt argument definition.
+      # Checks a prompt argument hash for keys outside PROMPT_ARG_SCHEMA.
       # @api private
-      def validate_prompt_arg_unknown_keys!(arg, idx)
-        unknown_keys = arg.transform_keys(&:to_s).keys - ALLOWED_PROMPT_ARG_KEYS
+      def validate_prompt_arg_unknown_keys(arg, idx)
+        unknown_keys = arg.transform_keys(&:to_s).keys - PROMPT_ARG_SCHEMA.keys
         return if unknown_keys.empty?
 
         raise ArgumentError,
               "Prompt argument definition at index #{idx} contains unknown keys: #{unknown_keys.join(", ")}. " \
-              "Allowed: #{ALLOWED_PROMPT_ARG_KEYS.join(", ")}."
+              "Allowed: #{PROMPT_ARG_SCHEMA.keys.join(", ")}."
       end
     end
   end