ruby_llm_swarm-mcp 0.8.0

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

@@ -0,0 +1,54 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module MCP
+    module Auth
+      # Service for managing OAuth session state (PKCE and CSRF state)
+      # Handles creation, validation, and cleanup of temporary session data
+      class SessionManager
+        attr_reader :storage
+
+        def initialize(storage)
+          @storage = storage
+        end
+
+        # Create a new OAuth session with PKCE and CSRF state
+        # @param server_url [String] MCP server URL
+        # @return [Hash] session data with :pkce and :state
+        def create_session(server_url)
+          pkce = PKCE.new
+          state = SecureRandom.urlsafe_base64(CSRF_STATE_SIZE)
+
+          storage.set_pkce(server_url, pkce)
+          storage.set_state(server_url, state)
+
+          { pkce: pkce, state: state }
+        end
+
+        # Validate state parameter and retrieve session data
+        # @param server_url [String] MCP server URL
+        # @param state [String] state parameter from callback
+        # @return [Hash] session data with :pkce and :client_info
+        # @raise [ArgumentError] if state is invalid
+        def validate_and_retrieve_session(server_url, state)
+          stored_state = storage.get_state(server_url)
+          unless stored_state && Security.secure_compare(stored_state, state)
+            raise ArgumentError, "Invalid state parameter"
+          end
+
+          {
+            pkce: storage.get_pkce(server_url),
+            client_info: storage.get_client_info(server_url)
+          }
+        end
+
+        # Clean up temporary session data
+        # @param server_url [String] MCP server URL
+        def cleanup_session(server_url)
+          storage.delete_pkce(server_url)
+          storage.delete_state(server_url)
+        end
+      end
+    end
+  end
+end

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

@@ -0,0 +1,236 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module MCP
+    module Auth
+      # Service for managing OAuth token operations
+      # Handles token exchange, refresh, and client credentials flows
+      class TokenManager
+        attr_reader :http_client, :logger
+
+        def initialize(http_client, logger)
+          @http_client = http_client
+          @logger = logger
+        end
+
+        # Exchange authorization code for access token
+        # @param server_metadata [ServerMetadata] server metadata
+        # @param client_info [ClientInfo] client info
+        # @param code [String] authorization code
+        # @param pkce [PKCE] PKCE parameters
+        # @param server_url [String] MCP server URL
+        # @return [Token] access token
+        def exchange_authorization_code(server_metadata, client_info, code, pkce, server_url)
+          logger.debug("Exchanging authorization code for access token")
+
+          registered_redirect_uri = client_info.metadata.redirect_uris.first
+          params = build_auth_code_params(client_info, code, pkce, registered_redirect_uri, server_url)
+
+          response = post_token_exchange(server_metadata, params)
+          response = retry_if_redirect_mismatch(response, server_metadata, params, registered_redirect_uri)
+
+          validate_token_response!(response, "Token exchange")
+          parse_token_response(response)
+        end
+
+        # Exchange client credentials for access token
+        # @param server_metadata [ServerMetadata] server metadata
+        # @param client_info [ClientInfo] client info with secret
+        # @param scope [String, nil] requested scope
+        # @param server_url [String] MCP server URL
+        # @return [Token] access token
+        def exchange_client_credentials(server_metadata, client_info, scope, server_url)
+          logger.debug("Exchanging client credentials for access token")
+
+          params = {
+            grant_type: "client_credentials",
+            client_id: client_info.client_id,
+            client_secret: client_info.client_secret,
+            scope: scope,
+            resource: server_url
+          }.compact
+
+          response = post_token_exchange(server_metadata, params)
+          validate_token_response!(response, "Token exchange")
+          parse_token_response(response)
+        end
+
+        # Refresh access token using refresh token
+        # @param server_metadata [ServerMetadata] server metadata
+        # @param client_info [ClientInfo] client info
+        # @param token [Token] current token with refresh_token
+        # @param server_url [String] MCP server URL
+        # @return [Token, nil] new token or nil if refresh failed
+        def refresh_token(server_metadata, client_info, token, server_url)
+          return nil unless token.refresh_token
+
+          logger.debug("Refreshing access token")
+
+          params = build_refresh_params(client_info, token, server_url)
+          response = post_token_refresh(server_metadata, params)
+
+          # Return nil on error responses
+          return nil if response.is_a?(HTTPX::ErrorResponse)
+          return nil unless response.status == 200
+
+          parse_refresh_response(response, token)
+        rescue JSON::ParserError => e
+          logger.warn("Invalid token refresh response: #{e.message}")
+          nil
+        rescue HTTPX::Error => e
+          logger.warn("Network error during token refresh: #{e.message}")
+          nil
+        end
+
+        private
+
+        # Build parameters for authorization code exchange
+        # @param client_info [ClientInfo] client info
+        # @param code [String] authorization code
+        # @param pkce [PKCE] PKCE parameters
+        # @param redirect_uri [String] redirect URI
+        # @param server_url [String] MCP server URL
+        # @return [Hash] token exchange parameters
+        def build_auth_code_params(client_info, code, pkce, redirect_uri, server_url)
+          params = {
+            grant_type: "authorization_code",
+            code: code,
+            redirect_uri: redirect_uri,
+            client_id: client_info.client_id,
+            code_verifier: pkce.code_verifier,
+            resource: server_url
+          }
+
+          add_client_secret_if_needed(params, client_info)
+          params
+        end
+
+        # Build parameters for token refresh
+        # @param client_info [ClientInfo] client info
+        # @param token [Token] current token
+        # @param server_url [String] MCP server URL
+        # @return [Hash] refresh parameters
+        def build_refresh_params(client_info, token, server_url)
+          params = {
+            grant_type: "refresh_token",
+            refresh_token: token.refresh_token,
+            client_id: client_info.client_id,
+            resource: server_url
+          }
+
+          add_client_secret_if_needed(params, client_info)
+          params
+        end
+
+        # Add client secret to params if needed
+        # @param params [Hash] token request parameters
+        # @param client_info [ClientInfo] client info
+        def add_client_secret_if_needed(params, client_info)
+          return unless client_info.client_secret
+          return unless client_info.metadata.token_endpoint_auth_method == "client_secret_post"
+
+          params[:client_secret] = client_info.client_secret
+        end
+
+        # Post token exchange request
+        # @param server_metadata [ServerMetadata] server metadata
+        # @param params [Hash] form parameters
+        # @return [HTTPX::Response] HTTP response
+        def post_token_exchange(server_metadata, params)
+          http_client.post(
+            server_metadata.token_endpoint,
+            headers: { "Content-Type" => "application/x-www-form-urlencoded" },
+            form: params
+          )
+        end
+
+        # Post token refresh request
+        # @param server_metadata [ServerMetadata] server metadata
+        # @param params [Hash] form parameters
+        # @return [HTTPX::Response] HTTP response
+        def post_token_refresh(server_metadata, params)
+          response = http_client.post(
+            server_metadata.token_endpoint,
+            headers: { "Content-Type" => "application/x-www-form-urlencoded" },
+            form: params
+          )
+
+          if response.is_a?(HTTPX::ErrorResponse)
+            logger.warn("Token refresh failed: #{response.error&.message || 'Request failed'}")
+          elsif response.status != 200
+            logger.warn("Token refresh failed: HTTP #{response.status}")
+          end
+          response
+        end
+
+        # Retry token exchange if redirect URI mismatch detected
+        # @param response [HTTPX::Response] initial response
+        # @param server_metadata [ServerMetadata] server metadata
+        # @param params [Hash] exchange parameters
+        # @param registered_redirect_uri [String] registered redirect URI
+        # @return [HTTPX::Response] response (possibly retried)
+        def retry_if_redirect_mismatch(response, server_metadata, params, registered_redirect_uri)
+          # Don't retry on error responses
+          return response if response.is_a?(HTTPX::ErrorResponse)
+          return response if response.status == 200
+
+          redirect_hint = HttpResponseHandler.extract_redirect_mismatch(response.body.to_s)
+          return response unless redirect_hint
+          return response if redirect_hint[:expected] == registered_redirect_uri
+
+          logger.warn("Redirect URI mismatch, retrying with: #{redirect_hint[:expected]}")
+          params[:redirect_uri] = redirect_hint[:expected]
+          post_token_exchange(server_metadata, params)
+        end
+
+        # Validate token response
+        # @param response [HTTPX::Response, HTTPX::ErrorResponse] HTTP response
+        # @param context [String] context for error messages
+        # @raise [Errors::TransportError] if response is invalid
+        def validate_token_response!(response, context)
+          # Handle HTTPX ErrorResponse
+          if response.is_a?(HTTPX::ErrorResponse)
+            error_message = response.error&.message || "Request failed"
+            raise Errors::TransportError.new(message: "#{context} failed: #{error_message}")
+          end
+
+          return if response.status == 200
+
+          raise Errors::TransportError.new(
+            message: "#{context} failed: HTTP #{response.status}",
+            code: response.status
+          )
+        end
+
+        # Parse token response
+        # @param response [HTTPX::Response] HTTP response
+        # @return [Token] parsed token
+        def parse_token_response(response)
+          data = JSON.parse(response.body.to_s)
+          Token.new(
+            access_token: data["access_token"],
+            token_type: data["token_type"] || "Bearer",
+            expires_in: data["expires_in"],
+            scope: data["scope"],
+            refresh_token: data["refresh_token"]
+          )
+        end
+
+        # Parse refresh response, preserving old refresh token if not provided
+        # @param response [HTTPX::Response] HTTP response
+        # @param old_token [Token] previous token
+        # @return [Token] new token
+        def parse_refresh_response(response, old_token)
+          data = JSON.parse(response.body.to_s)
+          Token.new(
+            access_token: data["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
+      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

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

@@ -0,0 +1,76 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module MCP
+    module Auth
+      # Utility class for building OAuth URLs
+      # Handles discovery URLs, authorization URLs, and URL normalization
+      class UrlBuilder
+        # Build discovery URL for OAuth server metadata
+        # @param server_url [String] MCP server URL
+        # @param discovery_type [Symbol] :authorization_server or :protected_resource
+        # @return [String] discovery URL
+        def self.build_discovery_url(server_url, discovery_type = :authorization_server)
+          uri = URI.parse(server_url)
+
+          # Extract ONLY origin (scheme + host + port)
+          origin = "#{uri.scheme}://#{uri.host}"
+          origin += ":#{uri.port}" if uri.port && !default_port?(uri)
+
+          # Two discovery endpoints supported
+          endpoint = if discovery_type == :authorization_server
+                       "oauth-authorization-server"
+                     else
+                       "oauth-protected-resource"
+                     end
+
+          "#{origin}/.well-known/#{endpoint}"
+        end
+
+        # Build OAuth authorization URL
+        # @param authorization_endpoint [String] auth server endpoint
+        # @param client_id [String] client ID
+        # @param redirect_uri [String] redirect URI
+        # @param scope [String, nil] requested scope
+        # @param state [String] CSRF state
+        # @param pkce [PKCE] PKCE parameters
+        # @param resource [String] resource indicator (RFC 8707)
+        # @return [String] authorization URL
+        def self.build_authorization_url(authorization_endpoint, client_id, redirect_uri, scope, state, pkce, resource) # rubocop:disable Metrics/ParameterLists
+          params = {
+            response_type: "code",
+            client_id: client_id,
+            redirect_uri: redirect_uri,
+            scope: scope,
+            state: state, # CSRF protection
+            code_challenge: pkce.code_challenge,
+            code_challenge_method: pkce.code_challenge_method, # S256
+            resource: resource # RFC 8707 - Resource Indicators
+          }.compact
+
+          uri = URI.parse(authorization_endpoint)
+          uri.query = URI.encode_www_form(params)
+          uri.to_s
+        end
+
+        # Get authorization base URL from server URL
+        # @param server_url [String] MCP server URL
+        # @return [String] authorization base URL (scheme + host + port)
+        def self.get_authorization_base_url(server_url)
+          uri = URI.parse(server_url)
+          origin = "#{uri.scheme}://#{uri.host}"
+          origin += ":#{uri.port}" if uri.port && !default_port?(uri)
+          origin
+        end
+
+        # Check if port is default for scheme
+        # @param uri [URI] parsed URI
+        # @return [Boolean] true if default port
+        def self.default_port?(uri)
+          (uri.scheme == "http" && uri.port == 80) ||
+            (uri.scheme == "https" && uri.port == 443)
+        end
+      end
+    end
+  end
+end