ruby_llm-mcp 0.8.0 → 1.0.0

data/lib/ruby_llm/mcp/auth/flows/authorization_code_flow.rb

@@ -22,15 +22,18 @@ module RubyLLM
           # @param server_url [String] MCP server URL
           # @param redirect_uri [String] redirect URI for callback
           # @param scope [String, nil] requested scope
+          # @param resource_metadata [String, nil] explicit resource metadata URL hint
           # @param https_validator [Proc] callback to validate HTTPS usage
           # @return [String] authorization URL for user to visit
-          def start(server_url, redirect_uri, scope, https_validator: nil)
+          def start(server_url, redirect_uri, scope, resource_metadata: nil, https_validator: nil)
             logger.debug("Starting OAuth authorization flow for #{server_url}")
 
             # 1. Discover authorization server
-            server_metadata = discoverer.discover(server_url)
+            server_metadata = discoverer.discover(server_url, resource_metadata_url: resource_metadata)
             raise Errors::TransportError.new(message: "OAuth server discovery failed") unless server_metadata
 
+            validate_pkce_support!(server_metadata)
+
             # 2. Register client (or get cached client)
             client_info = client_registrar.get_or_register(
               server_url,
@@ -98,6 +101,20 @@ module RubyLLM
             logger.info("OAuth authorization completed successfully")
             token
           end
+
+          private
+
+          # If the server advertises PKCE methods, S256 must be supported per MCP authorization requirements.
+          def validate_pkce_support!(server_metadata)
+            methods = server_metadata.code_challenge_methods_supported
+            normalized_methods = Array(methods)
+            return if normalized_methods.empty? || normalized_methods.include?("S256")
+
+            raise Errors::TransportError.new(
+              message: "Authorization server does not support required PKCE method S256 " \
+                       "(advertised: #{normalized_methods.join(', ')})"
+            )
+          end
         end
       end
     end

data/lib/ruby_llm/mcp/auth/flows/client_credentials_flow.rb

@@ -21,12 +21,13 @@ module RubyLLM
           # @param server_url [String] MCP server URL
           # @param redirect_uri [String] redirect URI (used for registration only)
           # @param scope [String, nil] requested scope
+          # @param resource_metadata [String, nil] explicit resource metadata URL hint
           # @return [Token] access token
-          def execute(server_url, redirect_uri, scope)
+          def execute(server_url, redirect_uri, scope, resource_metadata: nil)
             logger.debug("Starting OAuth client credentials flow")
 
             # 1. Discover authorization server
-            server_metadata = discoverer.discover(server_url)
+            server_metadata = discoverer.discover(server_url, resource_metadata_url: resource_metadata)
             raise Errors::TransportError.new(message: "OAuth server discovery failed") unless server_metadata
 
             # 2. Register client (or get cached client) with client credentials grant

data/lib/ruby_llm/mcp/auth/http_response_handler.rb

@@ -1,7 +1,5 @@
 # frozen_string_literal: true
 
-require "json"
-
 module RubyLLM
   module MCP
     module Auth

data/lib/ruby_llm/mcp/auth/memory_storage.rb

@@ -7,64 +7,83 @@ module RubyLLM
       # Stores tokens, client registrations, server metadata, and temporary session data
       class MemoryStorage
         def initialize
+          @mutex = Mutex.new
           @tokens = {}
           @client_infos = {}
           @server_metadata = {}
           @pkce_data = {}
           @state_data = {}
+          @resource_metadata = {}
         end
 
         # Token storage
         def get_token(server_url)
-          @tokens[server_url]
+          @mutex.synchronize { @tokens[server_url] }
         end
 
         def set_token(server_url, token)
-          @tokens[server_url] = token
+          @mutex.synchronize { @tokens[server_url] = token }
+        end
+
+        def delete_token(server_url)
+          @mutex.synchronize { @tokens.delete(server_url) }
         end
 
         # Client registration storage
         def get_client_info(server_url)
-          @client_infos[server_url]
+          @mutex.synchronize { @client_infos[server_url] }
         end
 
         def set_client_info(server_url, client_info)
-          @client_infos[server_url] = client_info
+          @mutex.synchronize { @client_infos[server_url] = client_info }
         end
 
         # Server metadata caching
         def get_server_metadata(server_url)
-          @server_metadata[server_url]
+          @mutex.synchronize { @server_metadata[server_url] }
         end
 
         def set_server_metadata(server_url, metadata)
-          @server_metadata[server_url] = metadata
+          @mutex.synchronize { @server_metadata[server_url] = metadata }
         end
 
         # PKCE state management (temporary)
         def get_pkce(server_url)
-          @pkce_data[server_url]
+          @mutex.synchronize { @pkce_data[server_url] }
         end
 
         def set_pkce(server_url, pkce)
-          @pkce_data[server_url] = pkce
+          @mutex.synchronize { @pkce_data[server_url] = pkce }
         end
 
         def delete_pkce(server_url)
-          @pkce_data.delete(server_url)
+          @mutex.synchronize { @pkce_data.delete(server_url) }
         end
 
         # State parameter management (temporary)
         def get_state(server_url)
-          @state_data[server_url]
+          @mutex.synchronize { @state_data[server_url] }
         end
 
         def set_state(server_url, state)
-          @state_data[server_url] = state
+          @mutex.synchronize { @state_data[server_url] = state }
         end
 
         def delete_state(server_url)
-          @state_data.delete(server_url)
+          @mutex.synchronize { @state_data.delete(server_url) }
+        end
+
+        # Resource metadata management
+        def get_resource_metadata(server_url)
+          @mutex.synchronize { @resource_metadata[server_url] }
+        end
+
+        def set_resource_metadata(server_url, metadata)
+          @mutex.synchronize { @resource_metadata[server_url] = metadata }
+        end
+
+        def delete_resource_metadata(server_url)
+          @mutex.synchronize { @resource_metadata.delete(server_url) }
         end
       end
     end

data/lib/ruby_llm/mcp/auth/oauth_provider.rb

@@ -1,8 +1,5 @@
 # frozen_string_literal: true
 
-require "httpx"
-require "uri"
-
 module RubyLLM
   module MCP
     module Auth
@@ -117,21 +114,31 @@ module RubyLLM
         end
 
         # Start OAuth authorization flow (authorization code grant)
+        # @param resource_metadata [String, nil] explicit resource metadata URL hint
         # @return [String] authorization URL for user to visit
-        def start_authorization_flow
+        def start_authorization_flow(resource_metadata: nil)
+          hint = resource_metadata || @resource_metadata_hint
           @auth_code_flow.start(
             server_url,
             redirect_uri,
             scope,
+            resource_metadata: hint,
             https_validator: method(:validate_https_endpoint)
           )
         end
 
         # Perform client credentials flow (application authentication without user)
         # @param scope [String] optional scope override
+        # @param resource_metadata [String, nil] explicit resource metadata URL hint
         # @return [Token] access token
-        def client_credentials_flow(scope: nil)
-          @client_creds_flow.execute(server_url, redirect_uri, scope || self.scope)
+        def client_credentials_flow(scope: nil, resource_metadata: nil)
+          hint = resource_metadata || @resource_metadata_hint
+          @client_creds_flow.execute(
+            server_url,
+            redirect_uri,
+            scope || self.scope,
+            resource_metadata: hint
+          )
         end
 
         # Complete OAuth authorization flow after callback
@@ -153,6 +160,91 @@ module RubyLLM
           request.headers["Authorization"] = token.to_header
         end
 
+        # Handle authentication challenge from server (401 response)
+        # Attempts to refresh token or raises error if interactive auth required
+        # @param www_authenticate [String, nil] WWW-Authenticate header value
+        # @param resource_metadata [String, nil] Resource metadata URL from response/challenge
+        # @param resource_metadata_url [String, nil] Legacy alias for resource_metadata
+        # @param requested_scope [String, nil] Scope from WWW-Authenticate challenge
+        # @return [Boolean] true if authentication was refreshed successfully
+        # @raise [Errors::AuthenticationRequiredError] if interactive auth is required
+        def handle_authentication_challenge(www_authenticate: nil, resource_metadata: nil, resource_metadata_url: nil,
+                                            requested_scope: nil)
+          resolved_resource_metadata = resource_metadata || resource_metadata_url
+          logger.debug("Handling authentication challenge")
+          logger.debug("  WWW-Authenticate: #{www_authenticate}") if www_authenticate
+          logger.debug("  Resource metadata URL: #{resolved_resource_metadata}") if resolved_resource_metadata
+          logger.debug("  Requested scope: #{requested_scope}") if requested_scope
+
+          final_requested_scope, final_resource_metadata = resolve_challenge_context(
+            www_authenticate,
+            resource_metadata,
+            resource_metadata_url,
+            requested_scope
+          )
+
+          # Persist discovery hint for browser-based fallback flows.
+          @resource_metadata_hint = final_resource_metadata if final_resource_metadata
+
+          update_scope_if_needed(final_requested_scope)
+
+          # Try to refresh existing token
+          token = storage.get_token(server_url)
+          if token&.refresh_token
+            logger.debug("Attempting token refresh with existing refresh token")
+            refreshed_token = refresh_token(token, resource_metadata: final_resource_metadata)
+            return true if refreshed_token
+          end
+
+          # If we have client credentials, try that flow
+          if grant_type == :client_credentials
+            logger.debug("Attempting client credentials flow")
+            begin
+              new_token = client_credentials_flow(
+                scope: final_requested_scope,
+                resource_metadata: final_resource_metadata
+              )
+              return true if new_token
+            rescue StandardError => e
+              logger.warn("Client credentials flow failed: #{e.message}")
+            end
+          end
+
+          # Cannot automatically authenticate - interactive auth required
+          logger.warn("Cannot automatically authenticate - interactive authorization required")
+          raise Errors::AuthenticationRequiredError.new(
+            message: "OAuth authentication required. Token refresh failed and interactive authorization is needed."
+          )
+        end
+
+        # Parse WWW-Authenticate header to extract challenge parameters
+        # @param header [String] WWW-Authenticate header value
+        # @return [Hash] parsed challenge information
+        def parse_www_authenticate(header)
+          result = {}
+
+          # Example: Bearer realm="example", scope="mcp:read mcp:write", resource_metadata="https://..."
+          if header =~ /Bearer\s+(.+)/i
+            params = ::Regexp.last_match(1)
+            parsed_params = {}
+            params.scan(/([a-zA-Z_][a-zA-Z0-9_-]*)="([^"]*)"/) do |key, value|
+              parsed_params[key.downcase] = value
+            end
+
+            # Extract scope
+            result[:scope] = parsed_params["scope"] if parsed_params["scope"]
+
+            # Extract resource metadata URL (spec + legacy alias)
+            result[:resource_metadata] = parsed_params["resource_metadata"] || parsed_params["resource_metadata_url"]
+            result[:resource_metadata_url] = result[:resource_metadata] if result[:resource_metadata]
+
+            # Extract realm
+            result[:realm] = parsed_params["realm"] if parsed_params["realm"]
+          end
+
+          result
+        end
+
         private
 
         # Create HTTP client for OAuth requests
@@ -165,7 +257,7 @@ module RubyLLM
           headers["MCP-Protocol-Version"] = RubyLLM::MCP.config.protocol_version
 
           HTTPX.plugin(:follow_redirects).with(
-            timeout: { total: DEFAULT_OAUTH_TIMEOUT },
+            timeout: { request_timeout: DEFAULT_OAUTH_TIMEOUT },
             headers: headers
           )
         end
@@ -207,11 +299,13 @@ module RubyLLM
 
         # Refresh access token using refresh token
         # @param token [Token] current token with refresh_token
+        # @param resource_metadata [String, nil] explicit resource metadata URL hint
         # @return [Token, nil] new token or nil if refresh failed
-        def refresh_token(token)
+        def refresh_token(token, resource_metadata: nil)
           return nil unless token.refresh_token
 
-          server_metadata = @discoverer.discover(server_url)
+          hint = resource_metadata || @resource_metadata_hint
+          server_metadata = @discoverer.discover(server_url, resource_metadata_url: hint)
           client_info = storage.get_client_info(server_url)
           return nil unless server_metadata && client_info
 
@@ -220,6 +314,27 @@ module RubyLLM
           logger.debug("Token refreshed successfully") if new_token
           new_token
         end
+
+        # Resolve requested scope and resource metadata from inputs and WWW-Authenticate header.
+        # @return [Array(String, String)] [requested_scope, resource_metadata]
+        def resolve_challenge_context(www_authenticate, resource_metadata, resource_metadata_url, requested_scope)
+          final_resource_metadata = resource_metadata || resource_metadata_url
+          final_requested_scope = requested_scope
+          return [final_requested_scope, final_resource_metadata] unless www_authenticate
+
+          challenge_info = parse_www_authenticate(www_authenticate)
+          final_requested_scope ||= challenge_info[:scope]
+          final_resource_metadata ||= challenge_info[:resource_metadata]
+          [final_requested_scope, final_resource_metadata]
+        end
+
+        # Update provider scope only when a different challenge scope is provided.
+        def update_scope_if_needed(new_scope)
+          return unless new_scope && new_scope != scope
+
+          logger.debug("Updating scope from '#{scope}' to '#{new_scope}'")
+          self.scope = new_scope
+        end
       end
     end
   end

data/lib/ruby_llm/mcp/auth/session_manager.rb

@@ -1,7 +1,5 @@
 # frozen_string_literal: true
 
-require "securerandom"
-
 module RubyLLM
   module MCP
     module Auth

data/lib/ruby_llm/mcp/auth/token_manager.rb

@@ -71,9 +71,17 @@ module RubyLLM
 
           # Return nil on error responses
           return nil if response.is_a?(HTTPX::ErrorResponse)
-          return nil unless response.status == 200
+
+          if response.status != 200
+            oauth_error = extract_oauth_error(response.body.to_s)
+            raise_oauth_error!("Token refresh", oauth_error, response.status) if oauth_error
+            return nil
+          end
 
           parse_refresh_response(response, token)
+        rescue Errors::TransportError => e
+          logger.warn(e.message)
+          nil
         rescue JSON::ParserError => e
           logger.warn("Invalid token refresh response: #{e.message}")
           nil
@@ -194,6 +202,9 @@ module RubyLLM
             raise Errors::TransportError.new(message: "#{context} failed: #{error_message}")
           end
 
+          oauth_error = extract_oauth_error(response.body.to_s)
+          raise_oauth_error!(context, oauth_error, response.status) if oauth_error
+
           return if response.status == 200
 
           raise Errors::TransportError.new(
@@ -207,8 +218,18 @@ module RubyLLM
         # @return [Token] parsed token
         def parse_token_response(response)
           data = JSON.parse(response.body.to_s)
+          raise_oauth_error!("Token exchange", extract_oauth_error(data), response.status)
+
+          access_token = data["access_token"]
+          if access_token.nil? || access_token.empty?
+            raise Errors::TransportError.new(
+              message: "Token exchange failed: invalid token response (missing access_token)",
+              code: response.status
+            )
+          end
+
           Token.new(
-            access_token: data["access_token"],
+            access_token: access_token,
             token_type: data["token_type"] || "Bearer",
             expires_in: data["expires_in"],
             scope: data["scope"],
@@ -222,14 +243,64 @@ module RubyLLM
         # @return [Token] new token
         def parse_refresh_response(response, old_token)
           data = JSON.parse(response.body.to_s)
+          raise_oauth_error!("Token refresh", extract_oauth_error(data), response.status)
+
+          access_token = data["access_token"]
+          if access_token.nil? || access_token.empty?
+            raise Errors::TransportError.new(
+              message: "Token refresh failed: invalid token response (missing access_token)",
+              code: response.status
+            )
+          end
+
           Token.new(
-            access_token: data["access_token"],
+            access_token: access_token,
             token_type: data["token_type"] || "Bearer",
             expires_in: data["expires_in"],
             scope: data["scope"],
             refresh_token: data["refresh_token"] || old_token.refresh_token
           )
         end
+
+        # Extract OAuth error fields from JSON response data
+        # @param source [String, Hash] response body string or parsed JSON hash
+        # @return [Hash, nil] OAuth error fields or nil
+        def extract_oauth_error(source)
+          data = source.is_a?(Hash) ? source : JSON.parse(source)
+          error = data["error"] || data[:error]
+          return nil unless error
+
+          {
+            error: error,
+            error_description: data["error_description"] || data[:error_description],
+            error_uri: data["error_uri"] || data[:error_uri]
+          }
+        rescue JSON::ParserError
+          nil
+        end
+
+        # Raise TransportError for OAuth error responses
+        # @param context [String] context for the error
+        # @param oauth_error [Hash, nil] OAuth error fields
+        # @param status_code [Integer, nil] HTTP response status code
+        # @raise [Errors::TransportError] when oauth_error is present
+        def raise_oauth_error!(context, oauth_error, status_code)
+          return unless oauth_error
+
+          error = oauth_error[:error]
+          description = oauth_error[:error_description]
+          error_uri = oauth_error[:error_uri]
+
+          message = "#{context} failed: OAuth error '#{error}'"
+          message += ": #{description}" if description
+          message += " (#{error_uri})" if error_uri
+
+          raise Errors::TransportError.new(
+            message: message,
+            code: status_code,
+            error: error
+          )
+        end
       end
     end
   end

data/lib/ruby_llm/mcp/auth/transport_oauth_helper.rb

@@ -0,0 +1,107 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module MCP
+    module Auth
+      # Helper module for preparing OAuth providers for transports
+      # This keeps OAuth logic out of the Native module while making it reusable
+      module TransportOauthHelper
+        module_function
+
+        # Check if OAuth configuration is present
+        # @param config [Hash] transport configuration hash
+        # @return [Boolean] true if OAuth config is present
+        def oauth_config_present?(config)
+          oauth_config = config[:oauth] || config["oauth"]
+          return false if oauth_config.nil?
+
+          # If it's an OAuth provider instance, it's present
+          return true if oauth_config.respond_to?(:access_token)
+
+          # If it's a hash, check if it's not empty
+          !oauth_config.empty?
+        end
+
+        # Create OAuth provider from configuration
+        # Accepts either a provider instance or a configuration hash
+        # @param config [Hash] transport configuration hash (will be modified)
+        # @return [OAuthProvider, BrowserOAuthProvider, nil] OAuth provider or nil
+        def create_oauth_provider(config)
+          oauth_config = config.delete(:oauth) || config.delete("oauth")
+          return nil unless oauth_config
+
+          # If provider key exists with an instance, use it
+          if oauth_config.is_a?(Hash) && (oauth_config[:provider] || oauth_config["provider"])
+            return oauth_config[:provider] || oauth_config["provider"]
+          end
+
+          # If oauth_config itself is a provider instance, use it directly
+          if oauth_config.respond_to?(:access_token) && oauth_config.respond_to?(:start_authorization_flow)
+            return oauth_config
+          end
+
+          # Otherwise create new provider from config hash
+          server_url = determine_server_url(config)
+          return nil unless server_url
+
+          redirect_uri = oauth_config[:redirect_uri] || oauth_config["redirect_uri"] || "http://localhost:8080/callback"
+          scope = oauth_config[:scope] || oauth_config["scope"]
+          storage = oauth_config[:storage] || oauth_config["storage"]
+          grant_type = oauth_config[:grant_type] || oauth_config["grant_type"] || :authorization_code
+
+          RubyLLM::MCP::Auth::OAuthProvider.new(
+            server_url: server_url,
+            redirect_uri: redirect_uri,
+            scope: scope,
+            logger: MCP.logger,
+            storage: storage,
+            grant_type: grant_type
+          )
+        end
+
+        # Determine server URL from transport config
+        # @param config [Hash] transport configuration hash
+        # @return [String, nil] server URL or nil
+        def determine_server_url(config)
+          config[:url] || config["url"]
+        end
+
+        # Prepare HTTP transport configuration with OAuth provider
+        # @param config [Hash] transport configuration hash (will be modified)
+        # @param oauth_provider [OAuthProvider, nil] OAuth provider instance
+        # @return [Hash] prepared configuration
+        def prepare_http_transport_config(config, oauth_provider)
+          options = {
+            version: config.delete(:version) || config.delete("version"),
+            headers: config.delete(:headers) || config.delete("headers"),
+            oauth_provider: oauth_provider,
+            reconnection: config.delete(:reconnection) || config.delete("reconnection"),
+            reconnection_options: config.delete(:reconnection_options) || config.delete("reconnection_options"),
+            rate_limit: config.delete(:rate_limit) || config.delete("rate_limit"),
+            session_id: config.delete(:session_id) || config.delete("session_id")
+          }.compact
+
+          config[:options] = options
+          config
+        end
+
+        # Prepare stdio transport configuration
+        # @param config [Hash] transport configuration hash (will be modified)
+        # @return [Hash] prepared configuration
+        def prepare_stdio_transport_config(config)
+          # Remove OAuth config from stdio transport (not supported)
+          config.delete(:oauth)
+          config.delete("oauth")
+
+          options = {
+            args: config.delete(:args) || config.delete("args"),
+            env: config.delete(:env) || config.delete("env")
+          }.compact
+
+          config[:options] = options unless options.empty?
+          config
+        end
+      end
+    end
+  end
+end