vector_mcp 0.3.4 → 0.5.0

data/lib/vector_mcp/transport/http_stream.rb

@@ -52,6 +52,18 @@ module VectorMCP
       DEFAULT_SESSION_TIMEOUT = 300 # 5 minutes
       DEFAULT_EVENT_RETENTION = 100 # Keep last 100 events for resumability
       DEFAULT_REQUEST_TIMEOUT = 30 # Default timeout for server-initiated requests
+      DEFAULT_MIN_THREADS = 4
+      DEFAULT_MAX_THREADS = 32
+
+      # Default allowed origins — restrict to localhost by default for security.
+      DEFAULT_ALLOWED_ORIGINS = %w[
+        http://localhost
+        https://localhost
+        http://127.0.0.1
+        https://127.0.0.1
+        http://[::1]
+        https://[::1]
+      ].freeze
 
       # Initializes a new HTTP Stream transport.
       #
@@ -62,7 +74,10 @@ module VectorMCP
       # @option options [String] :path_prefix ("/mcp") The base path for HTTP endpoints
       # @option options [Integer] :session_timeout (300) Session timeout in seconds
       # @option options [Integer] :event_retention (100) Number of events to retain for resumability
-      # @option options [Array<String>] :allowed_origins (["*"]) List of allowed origins for CORS. Use ["*"] to allow all origins.
+      # @option options [Integer] :min_threads (4) Minimum Puma thread pool size
+      # @option options [Integer] :max_threads (32) Maximum Puma thread pool size
+      # @option options [Array<String>] :allowed_origins Allowed origins for CORS validation.
+      #   Defaults to localhost origins only. Pass ["*"] to allow all origins (NOT recommended for production).
       def initialize(server, options = {})
         @server = server
         @logger = server.logger
@@ -95,13 +110,29 @@ module VectorMCP
         start_time = Time.now
         path = env["PATH_INFO"]
         method = env["REQUEST_METHOD"]
+        transport_context = build_transport_context(env, method, path, start_time)
 
-        # Processing HTTP request
+        logger.debug { "Processing HTTP request #{method} #{path}" }
+
+        transport_context = execute_transport_hooks(:before_request, transport_context)
+        raise transport_context.error if transport_context.error?
+        return transport_context.result if transport_context.result
 
         response = route_request(path, method, env)
+        transport_context.result = response
+        transport_context = execute_transport_hooks(:after_response, transport_context)
+        raise transport_context.error if transport_context.error?
+
+        response = transport_context.result || response
         log_request_completion(method, path, start_time, response[0])
         response
       rescue StandardError => e
+        if transport_context
+          transport_context.error = e
+          transport_context = execute_transport_hooks(:on_transport_error, transport_context)
+          return transport_context.result if transport_context.result
+        end
+
         handle_request_error(method, path, e)
       end
 
@@ -133,16 +164,6 @@ module VectorMCP
         @stream_handler.send_message_to_session(session, message)
       end
 
-      # Broadcasts a notification to all active sessions.
-      #
-      # @param method [String] The notification method name
-      # @param params [Hash, Array, nil] The notification parameters
-      # @return [Integer] Number of sessions the notification was sent to
-      def broadcast_notification(method, params = nil)
-        message = build_notification(method, params)
-        @session_manager.broadcast_message(message)
-      end
-
       # Sends a server-initiated JSON-RPC request compatible with Session expectations.
       # This method will block until a response is received or the timeout is reached.
       # For HTTP transport, this requires finding an appropriate session with streaming connection.
@@ -247,7 +268,7 @@ module VectorMCP
       #
       # @return [void]
       def start_puma_server
-        @puma_server = Puma::Server.new(self)
+        @puma_server = Puma::Server.new(self, nil, min_threads: @min_threads, max_threads: @max_threads)
         @puma_server.add_tcp_listener(@host, @port)
 
         @running = true
@@ -307,11 +328,28 @@ module VectorMCP
       # @param env [Hash] The Rack environment
       # @return [Array] Rack response triplet
       def route_request(path, method, env)
+        return route_mounted_request(path, method, env) if @mounted
+
+        # Standalone mode: unchanged behavior
         return handle_health_check if path == "/"
         return not_found_response unless path == @path_prefix
 
-        # Validate origin for security (MCP specification requirement)
+        validate_and_dispatch(method, env)
+      end
+
+      # Routes requests when mounted inside another Rack app (e.g., Rails).
+      # PATH_INFO is relative to the mount point: "/" = MCP endpoint, "/health" = health check.
+      def route_mounted_request(path, method, env)
+        return handle_health_check if path == "/health"
+        return not_found_response unless ["", "/"].include?(path)
+
+        validate_and_dispatch(method, env)
+      end
+
+      # Validates origin and dispatches to the appropriate handler by HTTP method.
+      def validate_and_dispatch(method, env)
         return forbidden_response("Origin not allowed") unless valid_origin?(env)
+        return unauthorized_oauth_response(env) if oauth_gate_should_reject?(env)
 
         case method
         when "POST"
@@ -325,6 +363,73 @@ module VectorMCP
         end
       end
 
+      # True when OAuth 2.1 resource server mode is enabled and the incoming
+      # request has not successfully authenticated. Opt-in: only activates when the
+      # server was configured with a +resource_metadata_url+ via +enable_authentication!+.
+      #
+      # @param env [Hash] The Rack environment
+      # @return [Boolean]
+      def oauth_gate_should_reject?(env)
+        return false unless oauth_resource_server_enabled?
+
+        !authenticate_transport_request(env).authenticated?
+      end
+
+      # @return [Boolean] true when the server is configured to act as an OAuth 2.1 resource server.
+      def oauth_resource_server_enabled?
+        return false unless @server.respond_to?(:oauth_resource_metadata_url)
+        return false if @server.oauth_resource_metadata_url.nil?
+
+        @server.auth_manager.required?
+      end
+
+      # Runs the configured authentication strategy against the Rack env and returns
+      # the resulting SessionContext. The request is normalized into the
+      # +{ headers:, params:, method:, path:, rack_env: }+ hash shape that the rest
+      # of the codebase's authentication pipeline uses (see
+      # +VectorMCP::Handlers::Core.extract_request_from_session+), so +:custom+
+      # strategy handlers see the same contract here as they do on the in-handler
+      # auth path. Errors in the strategy are logged and treated as unauthenticated
+      # rather than propagated, so a malformed token can never crash the request
+      # pipeline.
+      #
+      # @param env [Hash] The Rack environment
+      # @return [VectorMCP::Security::SessionContext]
+      def authenticate_transport_request(env)
+        normalized_request = @server.security_middleware.normalize_request(env)
+        @server.security_middleware.authenticate_request(normalized_request)
+      rescue StandardError => e
+        VectorMCP.logger_for("security").warn do
+          "OAuth transport auth strategy raised #{e.class}: #{e.message}"
+        end
+        VectorMCP::Security::SessionContext.anonymous
+      end
+
+      # Returns a 401 Rack response carrying a WWW-Authenticate header that points
+      # Claude Desktop (and other RFC 9728 clients) at the configured OAuth 2.1
+      # protected resource metadata document. The JSON-RPC error envelope in the
+      # body is for clients that parse bodies regardless of status code; the header
+      # and status are the parts that drive the discovery flow.
+      #
+      # @param env [Hash] The Rack environment
+      # @return [Array] Rack response triplet
+      def unauthorized_oauth_response(env)
+        VectorMCP.logger_for("security").info do
+          "OAuth 401 challenge issued for #{env["REQUEST_METHOD"]} #{env["PATH_INFO"]}"
+        end
+
+        header_value = %(Bearer realm="mcp", resource_metadata="#{@server.oauth_resource_metadata_url}")
+        body = {
+          jsonrpc: "2.0",
+          id: nil,
+          error: { code: -32_401, message: "Authentication required" }
+        }.to_json
+
+        [401,
+         { "Content-Type" => "application/json", "WWW-Authenticate" => header_value },
+         [body]]
+      end
+
       # Handles POST requests (client-to-server JSON-RPC)
       #
       # @param env [Hash] The Rack environment
@@ -339,14 +444,23 @@ module VectorMCP
         request_body = read_request_body(env)
         parsed = parse_json_message(request_body)
 
+        # MCP spec: POST body MUST be a single JSON-RPC message, not a batch array
+        if parsed.is_a?(Array)
+          return json_error_response(nil, -32_600, "Invalid Request",
+                                     { details: "Batch requests are not supported. Send a single JSON-RPC message per POST." })
+        end
+
         session = resolve_session_for_post(session_id, parsed, env)
         return session if session.is_a?(Array) # Rack error response
 
-        if parsed.is_a?(Array)
-          handle_batch_request(parsed, session)
-        else
-          handle_single_request(parsed, session, env)
+        # Validate MCP-Protocol-Version header (skip for initialize requests)
+        is_initialize = parsed.is_a?(Hash) && parsed["method"] == "initialize"
+        unless is_initialize
+          version_error = validate_protocol_version_header(env)
+          return version_error if version_error
         end
+
+        handle_single_request(parsed, session, env)
       rescue JSON::ParserError => e
         json_error_response(nil, -32_700, "Parse error", { details: e.message })
       end
@@ -363,59 +477,16 @@ module VectorMCP
           return [202, { "Mcp-Session-Id" => session.id }, []]
         end
 
-        result = @server.handle_message(message, session.context, session.id)
-        build_rpc_response(env, result, message["id"], session.id)
-      rescue VectorMCP::ProtocolError => e
-        build_protocol_error_response(env, e, session_id: session.id)
-      end
-
-      # Handles a batch of JSON-RPC messages per JSON-RPC 2.0 spec.
-      #
-      # @param messages [Array] Array of parsed JSON-RPC messages
-      # @param session [Session] The resolved session
-      # @return [Array] Rack response triplet
-      def handle_batch_request(messages, session)
-        return json_error_response(nil, -32_600, "Invalid Request", { details: "Empty batch" }) if messages.empty?
-
-        responses = messages.filter_map do |message|
-          next batch_invalid_item_error unless message.is_a?(Hash)
-
-          process_batch_item(message, session)
-        end
-
-        return [204, { "Mcp-Session-Id" => session.id }, []] if responses.empty?
-
-        headers = { "Content-Type" => "application/json", "Mcp-Session-Id" => session.id }
-        [200, headers, [responses.to_json]]
-      end
-
-      # Processes a single item within a batch request.
-      #
-      # @param message [Hash] A single JSON-RPC message
-      # @param session [Session] The resolved session
-      # @return [Hash, nil] Response hash or nil for notifications/outgoing responses
-      def process_batch_item(message, session)
-        if outgoing_response?(message)
-          handle_outgoing_response(message)
-          return nil
+        # Notifications: has method, no id -> 202 Accepted with no body (MCP spec requirement)
+        if message["method"] && !message.key?("id")
+          @server.handle_message(message, session.context, session.id)
+          return [202, { "Mcp-Session-Id" => session.id }, []]
         end
 
         result = @server.handle_message(message, session.context, session.id)
-        return nil if result.nil? && message["id"].nil?
-
-        { jsonrpc: "2.0", id: message["id"], result: result }
+        build_rpc_response(env, result, message["id"], session.id)
       rescue VectorMCP::ProtocolError => e
-        { jsonrpc: "2.0", id: e.request_id, error: { code: e.code, message: e.message, data: e.details } }
-      rescue StandardError => e
-        { jsonrpc: "2.0", id: message["id"],
-          error: { code: -32_603, message: "Internal error", data: { details: e.message } } }
-      end
-
-      # Returns an error object for non-Hash items in a batch.
-      #
-      # @return [Hash] JSON-RPC error object
-      def batch_invalid_item_error
-        { jsonrpc: "2.0", id: nil, error: { code: -32_600, message: "Invalid Request" } }
+        build_protocol_error_response(env, e, session_id: session.id)
       end
 
       # Resolves or creates the session for a POST request following MCP spec rules:
@@ -425,12 +496,11 @@ module VectorMCP
       # - no session_id + other request → 400 Bad Request
       #
       # @param session_id [String, nil] Client-supplied Mcp-Session-Id header value
-      # @param message [Hash, Array] Parsed JSON-RPC message or batch array
+      # @param message [Hash] Parsed JSON-RPC message
       # @param env [Hash] Rack environment
       # @return [Session, Array] Session object or Rack error response triplet
       def resolve_session_for_post(session_id, message, env)
-        first_message = message.is_a?(Array) ? message.first : message
-        is_initialize = first_message.is_a?(Hash) && first_message["method"] == "initialize"
+        is_initialize = message.is_a?(Hash) && message["method"] == "initialize"
 
         if session_id
           session = @session_manager.get_session(session_id)
@@ -464,6 +534,9 @@ module VectorMCP
         session = @session_manager.get_or_create_session(session_id, env)
         return not_found_response unless session
 
+        version_error = validate_protocol_version_header(env)
+        return version_error if version_error
+
         @stream_handler.handle_streaming_request(env, session)
       end
 
@@ -475,6 +548,9 @@ module VectorMCP
         session_id = extract_session_id(env)
         return bad_request_response("Missing Mcp-Session-Id header") unless session_id
 
+        version_error = validate_protocol_version_header(env)
+        return version_error if version_error
+
         success = @session_manager.terminate_session(session_id)
         if success
           [204, {}, []]
@@ -590,10 +666,11 @@ module VectorMCP
         accept.include?("text/event-stream")
       end
 
-      def format_sse_event(data, type, event_id)
+      def format_sse_event(data, type, event_id, retry_ms: nil)
         lines = []
-        lines << "id: #{event_id}"
+        lines << "id: #{event_id}" if event_id
         lines << "event: #{type}" if type
+        lines << "retry: #{retry_ms}" if retry_ms
         lines << "data: #{data}"
         lines << ""
         "#{lines.join("\n")}\n"
@@ -602,8 +679,14 @@ module VectorMCP
       def sse_rpc_response(result, request_id, headers = {}, session_id: nil)
         response = { jsonrpc: "2.0", id: request_id, result: result }
         event_data = response.to_json
+        stream_id = generate_sse_stream_id(session_id, :post)
+
+        # Priming event per MCP spec: event ID + empty data field
+        prime_event_id = @event_store.store_event("", nil, session_id: session_id, stream_id: stream_id)
+        prime_event = "id: #{prime_event_id}\ndata:\n\n"
 
-        event_id = @event_store.store_event(event_data, "message", session_id: session_id)
+        # Actual response event
+        event_id = @event_store.store_event(event_data, "message", session_id: session_id, stream_id: stream_id)
         sse_event = format_sse_event(event_data, "message", event_id)
 
         response_headers = {
@@ -613,7 +696,7 @@ module VectorMCP
           "X-Accel-Buffering" => "no"
         }.merge(headers)
 
-        [200, response_headers, [sse_event]]
+        [200, response_headers, [prime_event, sse_event]]
       end
 
       def sse_error_response(id, code, err_message, data = nil, session_id: nil)
@@ -621,8 +704,13 @@ module VectorMCP
         error_obj[:data] = data if data
         response = { jsonrpc: "2.0", id: id, error: error_obj }
         event_data = response.to_json
+        stream_id = generate_sse_stream_id(session_id, :post)
+
+        # Priming event per MCP spec
+        prime_event_id = @event_store.store_event("", nil, session_id: session_id, stream_id: stream_id)
+        prime_event = "id: #{prime_event_id}\ndata:\n\n"
 
-        event_id = @event_store.store_event(event_data, "message", session_id: session_id)
+        event_id = @event_store.store_event(event_data, "message", session_id: session_id, stream_id: stream_id)
         sse_event = format_sse_event(event_data, "message", event_id)
 
         response_headers = {
@@ -630,7 +718,7 @@ module VectorMCP
           "Cache-Control" => "no-cache"
         }
 
-        [200, response_headers, [sse_event]]
+        [200, response_headers, [prime_event, sse_event]]
       end
 
       def not_found_response(message = "Not Found")
@@ -642,7 +730,8 @@ module VectorMCP
       end
 
       def forbidden_response(message = "Forbidden")
-        [403, { "Content-Type" => "text/plain" }, [message]]
+        error = { jsonrpc: "2.0", error: { code: -32_600, message: message } }
+        [403, { "Content-Type" => "application/json" }, [error.to_json]]
       end
 
       def method_not_allowed_response(allowed_methods)
@@ -654,11 +743,26 @@ module VectorMCP
         [406, { "Content-Type" => "text/plain" }, [message]]
       end
 
+      # Validates the MCP-Protocol-Version header per spec.
+      # Returns nil if valid, or a 400 Rack response if unsupported.
+      def validate_protocol_version_header(env)
+        version = env["HTTP_MCP_PROTOCOL_VERSION"]
+        return nil if version.nil? # Backwards compatibility: assume 2025-03-26
+
+        unless VectorMCP::Server::SUPPORTED_PROTOCOL_VERSIONS.include?(version)
+          return bad_request_response("Unsupported MCP-Protocol-Version: #{version}")
+        end
+
+        nil
+      end
+
       def valid_post_accept?(env)
         accept = env["HTTP_ACCEPT"]
         return true if accept.nil? || accept.strip.empty?
+        return true if accept.include?("*/*")
 
-        accept.include?("application/json") || accept.include?("*/*")
+        # MCP spec: client MUST include both application/json AND text/event-stream
+        accept.include?("application/json") && accept.include?("text/event-stream")
       end
 
       def valid_get_accept?(env)
@@ -668,7 +772,13 @@ module VectorMCP
         accept.include?("text/event-stream") || accept.include?("*/*")
       end
 
-      # Validates the Origin header for security
+      # Validates the Origin header for security.
+      #
+      # Matches are checked both exactly and as prefix (so that
+      # +http://localhost+ in the allowed list matches +http://localhost:3000+).
+      #
+      # Requests without an Origin header are allowed through because they
+      # originate from non-browser contexts (curl, server-to-server, etc.).
       #
       # @param env [Hash] The Rack environment
       # @return [Boolean] True if origin is allowed, false otherwise
@@ -678,7 +788,9 @@ module VectorMCP
         origin = env["HTTP_ORIGIN"]
         return true if origin.nil? # Allow requests without Origin header (e.g., server-to-server)
 
-        @allowed_origins.include?(origin)
+        @allowed_origins.any? do |allowed|
+          origin == allowed || origin.start_with?("#{allowed}:")
+        end
       end
 
       # Logging and error handling
@@ -692,6 +804,29 @@ module VectorMCP
         [500, { "Content-Type" => "text/plain" }, ["Internal Server Error"]]
       end
 
+      def build_transport_context(env, method, path, start_time)
+        request_context = VectorMCP::RequestContext.from_rack_env(env, "http_stream")
+
+        VectorMCP::Middleware::Context.new(
+          operation_type: :transport,
+          operation_name: "#{method} #{path}",
+          params: request_context.to_h,
+          session: nil,
+          server: @server,
+          metadata: {
+            start_time: start_time,
+            path: path,
+            method: method
+          }
+        )
+      end
+
+      def execute_transport_hooks(hook_type, context)
+        return context unless @server.respond_to?(:middleware_manager) && @server.middleware_manager
+
+        @server.middleware_manager.execute_hooks(hook_type, context)
+      end
+
       def handle_fatal_error(error)
         logger.fatal { "Fatal error in HttpStream transport: #{error.message}" }
         exit(1)
@@ -699,15 +834,12 @@ module VectorMCP
 
       # Request tracking helpers for server-initiated requests
 
-      # Sets up tracking for an outgoing request using pooled condition variables.
+      # Sets up tracking for an outgoing request.
       #
       # @param request_id [String] The request ID to track
       # @return [void]
       def setup_request_tracking(request_id)
-        @request_mutex.synchronize do
-          # Create IVar for thread-safe request tracking (no race conditions)
-          @outgoing_request_ivars[request_id] = Concurrent::IVar.new
-        end
+        @outgoing_request_ivars[request_id] = Concurrent::IVar.new
       end
 
       # Waits for a response to an outgoing request.
@@ -718,11 +850,7 @@ module VectorMCP
       # @return [Hash] The response data
       # @raise [VectorMCP::SamplingTimeoutError] if timeout occurs
       def wait_for_response(request_id, method, timeout)
-        ivar = nil
-        @request_mutex.synchronize do
-          ivar = @outgoing_request_ivars[request_id]
-        end
-
+        ivar = @outgoing_request_ivars[request_id]
         return nil unless ivar
 
         begin
@@ -764,24 +892,11 @@ module VectorMCP
         response[:result]
       end
 
-      # Cleans up tracking for a request and returns condition variable to pool.
+      # Cleans up tracking for a request.
       #
       # @param request_id [String] The request ID to clean up
       # @return [void]
       def cleanup_request_tracking(request_id)
-        @request_mutex.synchronize do
-          cleanup_request_tracking_unsafe(request_id)
-        end
-      end
-
-      # Internal cleanup method that assumes mutex is already held.
-      # This prevents recursive locking when called from within synchronized blocks.
-      #
-      # @param request_id [String] The request ID to clean up
-      # @return [void]
-      # @api private
-      def cleanup_request_tracking_unsafe(request_id)
-        # Remove IVar for this request (no condition variable cleanup needed)
         @outgoing_request_ivars.delete(request_id)
       end
 
@@ -810,10 +925,7 @@ module VectorMCP
       def handle_outgoing_response(message)
         request_id = message["id"]
 
-        ivar = nil
-        @request_mutex.synchronize do
-          ivar = @outgoing_request_ivars[request_id]
-        end
+        ivar = @outgoing_request_ivars[request_id]
 
         unless ivar
           logger.debug { "Received response for request ID #{request_id} but no thread is waiting (likely timed out)" }
@@ -822,11 +934,6 @@ module VectorMCP
 
         # Convert keys to symbols for consistency and put response in IVar
         response_data = deep_transform_keys(message, &:to_sym)
-
-        # Store in both places for compatibility with tests
-        @outgoing_request_responses[request_id] = response_data
-
-        # IVar handles thread-safe response delivery - no race conditions possible
         if ivar.try_set(response_data)
           logger.debug { "Response delivered to waiting thread for request ID #{request_id}" }
         else
@@ -839,8 +946,8 @@ module VectorMCP
       #
       # @param obj [Object] The object to transform (Hash, Array, or other)
       # @return [Object] The transformed object
-      def deep_transform_keys(obj, &block)
-        transform_object_keys(obj, &block)
+      def deep_transform_keys(obj, &)
+        transform_object_keys(obj, &)
       end
 
       # Core transformation logic extracted for better maintainability
@@ -874,7 +981,20 @@ module VectorMCP
         @path_prefix = normalize_path_prefix(options[:path_prefix] || DEFAULT_PATH_PREFIX)
         @session_timeout = options[:session_timeout] || DEFAULT_SESSION_TIMEOUT
         @event_retention = options[:event_retention] || DEFAULT_EVENT_RETENTION
-        @allowed_origins = options[:allowed_origins] || ["*"]
+        @min_threads = options[:min_threads] || DEFAULT_MIN_THREADS
+        @max_threads = options[:max_threads] || DEFAULT_MAX_THREADS
+        @allowed_origins = options[:allowed_origins] || DEFAULT_ALLOWED_ORIGINS
+        @mounted = options.fetch(:mounted, false)
+
+        warn_on_permissive_origins if @allowed_origins.include?("*")
+      end
+
+      # Logs a security warning when wildcard origin is configured.
+      def warn_on_permissive_origins
+        logger.warn do
+          "[SECURITY] allowed_origins includes '*', which permits cross-origin requests from any website. " \
+            "This is not recommended for production. Specify explicit origins instead."
+        end
       end
 
       # Initialize core HTTP stream components
@@ -886,11 +1006,7 @@ module VectorMCP
 
       # Initialize request tracking system and ID generation for server-initiated requests
       def initialize_request_tracking
-        # Use IVars for thread-safe request/response handling (eliminates condition variable races)
         @outgoing_request_ivars = Concurrent::Hash.new
-        # Keep compatibility with tests that expect @outgoing_request_responses
-        @outgoing_request_responses = Concurrent::Hash.new
-        @request_mutex = Mutex.new
         initialize_request_id_generation
       end
 
@@ -901,6 +1017,11 @@ module VectorMCP
         @request_id_counter = Concurrent::AtomicFixnum.new(0)
       end
 
+      def generate_sse_stream_id(session_id, origin)
+        session_label = session_id || "anonymous"
+        "#{session_label}-#{origin}-#{SecureRandom.hex(4)}"
+      end
+
       # Generate a unique, thread-safe request ID for server-initiated requests
       #
       # @return [String] A unique request ID in format: vecmcp_http_{pid}_{random}_{counter}
@@ -929,10 +1050,7 @@ module VectorMCP
 
         logger.debug { "Cleaning up #{@outgoing_request_ivars.size} pending requests" }
 
-        @request_mutex.synchronize do
-          # IVars will timeout naturally, just clear the tracking
-          @outgoing_request_ivars.clear
-        end
+        @outgoing_request_ivars.clear
       end
 
       # Finds the first session with an active streaming connection.