legate 0.1.0

data/lib/legate/agents/sequential_agent.rb

@@ -0,0 +1,190 @@
+# frozen_string_literal: true
+
+# File: lib/legate/agents/sequential_agent.rb
+require_relative '../agent'
+
+module Legate
+  module Agents
+    # SequentialAgent executes a series of sub-agents in a predefined order.
+    # Each sub-agent is executed one after another, with the same session and input.
+    class SequentialAgent < Legate::Agent
+      # Override run_task to execute sub-agents in sequence
+      # @param session_id [String] The session ID
+      # @param user_input [String] User input to process
+      # @param session_service [Legate::SessionService::Base] Session service for persistence
+      # @return [Legate::Event] The final agent event
+      def run_task(session_id:, user_input:, session_service:)
+        # Verify we have sequential sub-agents defined
+        unless @definition.sequential_sub_agent_names&.any?
+          err_msg = "SequentialAgent '#{name}' has no sequential_sub_agent_names defined."
+          Legate.logger.error(err_msg)
+          return Legate::Event.new(role: :agent, content: {
+                                     status: :error,
+                                     error_message: err_msg,
+                                     error_class: 'ConfigurationError'
+                                   })
+        end
+
+        # --- Pre-execution Checks --- #
+        unless running?
+          err_msg = "Agent '#{name}' is not running. Call agent.start before run_task, " \
+                    'or use agent.ask (which starts automatically).'
+          Legate.logger.error(err_msg)
+          return Legate::Event.new(role: :agent, content: { status: :error, error_message: err_msg })
+        end
+
+        session = session_service.get_session(session_id: session_id)
+        unless session
+          err_msg = "Session not found: #{session_id}"
+          Legate.logger.error(err_msg)
+          return Legate::Event.new(role: :agent, content: { status: :error, error_message: err_msg })
+        end
+        # --------------------------- #
+
+        # Log user input to the SequentialAgent itself
+        user_event = Legate::Event.new(role: :user, content: user_input)
+        session_service.append_event(session_id: session_id, event: user_event)
+
+        # Log the execution sequence start
+        Legate.logger.info("SequentialAgent '#{name}' starting execution of #{@definition.sequential_sub_agent_names.size} sub-agents in sequence.")
+
+        # Track results of all sub-agents
+        all_results = []
+        final_result = nil
+        current_input = user_input # Start with the original user input
+
+        # Execute each sub-agent in order
+        @definition.sequential_sub_agent_names.each_with_index do |sub_agent_name, index|
+          sub_agent = find_sub_agent(sub_agent_name)
+          unless sub_agent
+            err_msg = "Sub-agent '#{sub_agent_name}' not found for SequentialAgent '#{name}'."
+            Legate.logger.error(err_msg)
+            final_result = {
+              status: :error,
+              error_message: err_msg,
+              error_class: 'MissingSubAgentError',
+              step: index + 1,
+              total_steps: @definition.sequential_sub_agent_names.size,
+              previous_results: all_results.map.with_index { |r, i| { agent: @definition.sequential_sub_agent_names.to_a[i], result: r } }
+            }
+            break # Stop the sequence on error
+          end
+
+          # Start the sub-agent if it's not already running
+          sub_agent.start unless sub_agent.running?
+
+          # If this is not the first agent, try to augment its input with previous results
+          if index > 0 && all_results.any?
+            # Get previous agent name - handle both Array and Set types
+            previous_agent_name = nil
+            previous_agent_name = if @definition.sequential_sub_agent_names.is_a?(Set)
+                                    # For Set, convert to Array and get the element
+                                    @definition.sequential_sub_agent_names.to_a[index - 1]
+                                  else
+                                    # For Array, access directly
+                                    @definition.sequential_sub_agent_names[index - 1]
+                                  end
+
+            previous_output_key = find_sub_agent(previous_agent_name)&.definition&.output_key
+
+            if previous_output_key && session_service.respond_to?(:get_state)
+              # Get the previous agent's result from session state
+              previous_result = session_service.get_state(session_id: session_id, key: previous_output_key)
+
+              if previous_result && previous_result.is_a?(Hash) && previous_result[:result]
+                # Enhanced input with previous result
+                current_input = "#{current_input}\n\nHere is the result from the previous step (#{previous_agent_name}):\n#{previous_result[:result]}"
+                Legate.logger.info("Enhanced input for sub-agent '#{sub_agent_name}' with previous result from '#{previous_agent_name}'")
+              end
+            end
+          end
+
+          # Execute the sub-agent with the updated input
+          begin
+            Legate.logger.info("SequentialAgent '#{name}' executing sub-agent '#{sub_agent_name}' (step #{index + 1}/#{@definition.sequential_sub_agent_names.size}).")
+            sub_result = sub_agent.run_task(
+              session_id: session_id,
+              user_input: current_input,
+              session_service: session_service
+            )
+
+            # Record the result
+            all_results << sub_result.content
+
+            # Check for error to break sequence
+            if sub_result.content[:status] == :error
+              Legate.logger.warn("Sub-agent '#{sub_agent_name}' returned error, breaking sequence: #{sub_result.content[:error_message]}")
+              final_result = {
+                status: :error,
+                error_message: "Error in sub-agent '#{sub_agent_name}': #{sub_result.content[:error_message]}",
+                error_class: sub_result.content[:error_class] || 'SubAgentError',
+                step: index + 1,
+                total_steps: @definition.sequential_sub_agent_names.size,
+                sub_agent: sub_agent_name.to_s,
+                sub_result: sub_result.content,
+                previous_results: all_results.map.with_index { |r, i| { agent: @definition.sequential_sub_agent_names.to_a[i], result: r } }
+              }
+              break # Stop the sequence on error
+            elsif sub_result.content[:result]
+              # If the sub-agent succeeded, update the current_input to include its result for the next agent
+              # Store the raw result for potential state-based chaining in next iteration
+              current_input = sub_result.content[:result].to_s
+            end
+          rescue StandardError => e
+            Legate.logger.error("Error executing sub-agent '#{sub_agent_name}': #{e.class} - #{e.message}\n#{e.backtrace.join("\n")}")
+            final_result = {
+              status: :error,
+              error_message: "Exception in sub-agent '#{sub_agent_name}': #{e.message}",
+              error_class: e.class.name,
+              step: index + 1,
+              total_steps: @definition.sequential_sub_agent_names.size,
+              sub_agent: sub_agent_name.to_s,
+              previous_results: all_results.map.with_index { |r, i| { agent: @definition.sequential_sub_agent_names.to_a[i], result: r } }
+            }
+            break # Stop the sequence on error
+          end
+        end
+
+        # If we didn't set a final_result due to an error, create a success result with all sub-results
+        if final_result.nil?
+          final_result = {
+            status: :success,
+            result: "Completed sequential execution of #{@definition.sequential_sub_agent_names.size} sub-agents",
+            steps_completed: @definition.sequential_sub_agent_names.size,
+            sub_results: all_results.map.with_index { |r, i| { agent: @definition.sequential_sub_agent_names.to_a[i], result: r } }
+          }
+        end
+
+        # Create the final event
+        final_agent_event = Legate::Event.new(role: :agent, content: final_result)
+
+        # Log the final event to the session
+        session_service.append_event(session_id: session_id, event: final_agent_event)
+
+        # --- MAS: Store result in session state if output_key is defined --- #
+        if @definition.respond_to?(:output_key) && @definition.output_key && final_agent_event
+          output_value = final_agent_event.content # Store the entire content hash
+
+          # Make sure the output value is serializable
+          serialized_value = begin
+            # Convert to JSON and back to ensure it's serializable
+            JSON.parse(output_value.to_json)
+          rescue StandardError => e
+            Legate.logger.warn("SequentialAgent '#{@name}': Failed to serialize output value: #{e.message}. Using simplified value.")
+            { status: output_value[:status], result: output_value[:result].to_s }
+          end
+
+          Legate.logger.info("SequentialAgent '#{@name}' storing output to session state with key '#{@definition.output_key}' for session '#{session_id}'.")
+          if session_service.respond_to?(:set_state)
+            session_service.set_state(session_id: session_id, key: @definition.output_key, value: serialized_value)
+          else
+            Legate.logger.warn("SequentialAgent '#{@name}': Session service does not support :set_state. Cannot store output for key '#{@definition.output_key}'.")
+          end
+        end
+        # --- End MAS State Management --- #
+
+        final_agent_event
+      end
+    end
+  end
+end

data/lib/legate/agents.rb

@@ -0,0 +1,14 @@
+# frozen_string_literal: true
+
+# File: lib/legate/agents.rb
+# This manifest file loads all agent implementations for workflow agents
+
+require_relative 'agents/sequential_agent'
+require_relative 'agents/parallel_agent'
+require_relative 'agents/loop_agent'
+
+module Legate
+  module Agents
+    # This module contains specialized agent implementations for workflow composition
+  end
+end

data/lib/legate/auth/config.rb

@@ -0,0 +1,148 @@
+# File: lib/legate/auth/config.rb
+# frozen_string_literal: true
+
+require 'securerandom'
+
+module Legate
+  module Auth
+    # Configuration container used during the authentication flow.
+    # Holds the authentication scheme, credential, and request/response details
+    # needed for interactive authentication flows.
+    class Config
+      # @return [Legate::Auth::Scheme] The authentication scheme
+      attr_reader :scheme
+
+      # @return [Legate::Auth::Credential] The credential information
+      attr_reader :credential
+
+      # @return [String, nil] The unique ID for this authentication request
+      attr_reader :auth_request_id
+
+      # @return [String, nil] The authorization URI for interactive flows
+      attr_accessor :auth_uri
+
+      # @return [String, nil] The redirect URI for OAuth2/OIDC flows
+      attr_accessor :redirect_uri
+
+      # @return [String, nil] The state parameter for CSRF protection
+      attr_accessor :state
+
+      # @return [Hash, nil] The PKCE parameters (code_verifier, etc.)
+      attr_accessor :pkce
+
+      # @return [String, nil] The authorization response URI from the provider
+      attr_accessor :response_uri
+
+      # For backwards compatibility
+      alias auth_response_uri response_uri
+      alias auth_response_uri= response_uri=
+
+      # @return [Hash, nil] Additional options for the authentication process
+      attr_accessor :options
+
+      # Initialize a new authentication configuration
+      # @param scheme [Legate::Auth::Scheme] The authentication scheme
+      # @param credential [Legate::Auth::Credential] The credential information
+      # @param auth_request_id [String, nil] The unique ID for this authentication request
+      # @param options [Hash, nil] Additional options for the authentication process
+      def initialize(scheme:, credential:, auth_request_id: nil, options: {})
+        @scheme = scheme
+        @credential = credential
+        @auth_request_id = auth_request_id || Legate::Auth.generate_request_id
+        @options = options || {}
+        @auth_uri = nil
+        @redirect_uri = nil
+        @state = nil
+        @pkce = nil
+        @response_uri = nil
+      end
+
+      # Build the authorization URI for interactive flows
+      # @param redirect_uri [String, nil] The redirect URI for the authorization request
+      # @param state [String, nil] A state parameter for CSRF protection
+      # @return [String, Hash] The authorization URI or a hash with URI and additional parameters
+      def build_authorization_uri(redirect_uri = nil, state = nil)
+        @redirect_uri = redirect_uri
+        @state = state || @options[:state] || SecureRandom.hex(16)
+
+        # For OAuth2 schemes with detailed return values including PKCE
+        result = @scheme.build_authorization_uri(self, @redirect_uri, @state)
+
+        if result.is_a?(Hash) && result[:uri]
+          @auth_uri = result[:uri]
+          @state = result[:state] if result[:state]
+          @pkce = result[:pkce] if result[:pkce]
+          @auth_uri
+        else
+          # For backwards compatibility with simpler schemes
+          @auth_uri = result
+        end
+      end
+
+      # Convert to a hash for serialization
+      # @param include_credentials [Boolean] Whether to include credential details (use carefully)
+      # @return [Hash] A hash representation of the config
+      def to_h(include_credentials: false)
+        {
+          auth_request_id: @auth_request_id,
+          scheme_type: @scheme.scheme_type,
+          auth_uri: @auth_uri,
+          redirect_uri: @redirect_uri,
+          state: @state,
+          pkce: @pkce,
+          response_uri: @response_uri,
+          options: @options
+        }.tap do |h|
+          h[:credential] = @credential.to_h if include_credentials
+        end
+      end
+
+      # Creates a Config from a hash representation
+      # @param hash [Hash] The hash representation
+      # @param scheme [Legate::Auth::Scheme] The authentication scheme (required if not recreating from complete data)
+      # @param credential [Legate::Auth::Credential] The credential information (required if not recreating from complete data)
+      # @return [Legate::Auth::Config] A new Config instance
+      # @raise [Legate::Auth::ConfigurationError] If required parameters are missing
+      def self.from_h(hash, scheme: nil, credential: nil)
+        scheme ||= hash[:scheme]
+        credential ||= hash[:credential]
+
+        raise Legate::Auth::ConfigurationError, 'Scheme and credential must be provided' unless scheme && credential
+
+        config = new(
+          scheme: scheme,
+          credential: credential,
+          auth_request_id: hash[:auth_request_id],
+          options: hash[:options] || {}
+        )
+
+        config.auth_uri = hash[:auth_uri]
+        config.redirect_uri = hash[:redirect_uri]
+        config.state = hash[:state]
+        config.pkce = hash[:pkce]
+
+        # Handle both new and old response URI keys
+        config.response_uri = hash[:response_uri] || hash[:auth_response_uri]
+
+        config
+      end
+
+      # Validates a response against this configuration
+      # @param response_config [Legate::Auth::Config] The response configuration
+      # @return [Boolean] True if the response is valid for this request
+      # @raise [Legate::Auth::ConfigurationError] If the response is invalid
+      def validate_response!(response_config)
+        # Check request ID
+        raise Legate::Auth::ConfigurationError, 'Authentication response ID does not match request ID' unless response_config.auth_request_id == @auth_request_id
+
+        # Check that we have an auth response URI
+        raise Legate::Auth::ConfigurationError, 'Authentication response does not contain a response URI' unless response_config.response_uri
+
+        # Check state if we had one
+        raise Legate::Auth::ConfigurationError, 'Authentication response state does not match request state' if @state && response_config.state && response_config.state != @state
+
+        true
+      end
+    end
+  end
+end

data/lib/legate/auth/coordinator.rb

@@ -0,0 +1,218 @@
+# frozen_string_literal: true
+
+require_relative 'error'
+require_relative 'config'
+require_relative 'credential'
+require_relative 'exchanged_credential'
+require_relative 'token_store'
+
+module Legate
+  module Auth
+    # Base class for authentication coordinators that handle the fiber-based authentication flow.
+    # Coordinators are responsible for managing the state of an authentication flow,
+    # including pausing execution, waiting for user input/authentication, and resuming
+    # with appropriate credentials.
+    class Coordinator
+      # Default timeout for authentication flow (in seconds)
+      DEFAULT_TIMEOUT = 300
+
+      # Authentication status codes
+      module Status
+        PENDING = :pending
+        COMPLETED = :completed
+        FAILED = :failed
+        TIMEOUT = :timeout
+        CANCELLED = :cancelled
+      end
+
+      # Initialize a new authentication coordinator
+      # @param scheme [Legate::Auth::Scheme] The authentication scheme to use
+      # @param credential [Legate::Auth::Credential] The credential to use
+      # @param session_service [Legate::SessionService::Base] The session service for state persistence
+      # @param token_store [Legate::Auth::TokenStore, nil] Optional token store for caching tokens
+      # @param timeout [Integer, nil] Optional timeout in seconds (nil for no timeout)
+      def initialize(scheme:, credential:, session_service:, token_store: nil, timeout: DEFAULT_TIMEOUT)
+        @scheme = scheme
+        @credential = credential
+        @session_service = session_service
+        @token_store = token_store || TokenStore.new(session_service)
+        @timeout = timeout
+        @request_id = SecureRandom.uuid
+        @status = Status::PENDING
+        @start_time = nil
+        @auth_fiber = nil
+        @result = nil
+        @error = nil
+      end
+
+      # Start the authentication flow
+      # @return [Hash] The authentication request details to be sent to the client
+      def start
+        # Create a new fiber for this authentication flow
+        @auth_fiber = Fiber.new do
+          @start_time = Time.now
+          @result = authenticate
+          @status = Status::COMPLETED
+        rescue Legate::Auth::Error => e
+          @error = e
+          @status = Status::FAILED
+          nil
+        rescue StandardError => e
+          @error = Legate::Auth::Error.new("Unexpected error during authentication: #{e.message}", e)
+          @status = Status::FAILED
+          nil
+        end
+
+        # Start the fiber to get the initial authentication request
+        auth_request = @auth_fiber.resume
+
+        # Record start of authentication in session
+        save_auth_state
+
+        # Return the authentication request to be sent to the client
+        {
+          request_id: @request_id,
+          scheme_type: @scheme.scheme_type,
+          auth_request: auth_request
+        }
+      end
+
+      # Resume the authentication flow with a response from the client
+      # @param response [Hash] The response from the client
+      # @return [Legate::Auth::ExchangedCredential, nil] The resulting credential or nil on failure
+      def resume(response)
+        raise Legate::Auth::Error, "Authentication flow is not in progress (status: #{@status})" unless @status == Status::PENDING
+
+        # Check for timeout
+        if @timeout && Time.now - @start_time > @timeout
+          @status = Status::TIMEOUT
+          @error = Legate::Auth::Error.new("Authentication timed out after #{@timeout} seconds")
+          save_auth_state
+          return nil
+        end
+
+        # Resume the fiber with the client response
+        begin
+          result = @auth_fiber.resume(response)
+
+          # If the fiber yields again, we need more input from the client
+          if @auth_fiber.alive?
+            # Return the next authentication request
+            return {
+              request_id: @request_id,
+              scheme_type: @scheme.scheme_type,
+              auth_request: result
+            }
+          end
+
+          # Authentication completed
+          @status = Status::COMPLETED if @status == Status::PENDING
+          save_auth_state
+
+          # Store the token if we have one
+          if @result && @result.is_a?(Legate::Auth::ExchangedCredential) && @token_store
+            cache_key = generate_cache_key(@scheme, @credential)
+            @token_store.store(cache_key, @result)
+          end
+
+          @result
+        rescue StandardError => e
+          @error = Legate::Auth::Error.new("Error resuming authentication: #{e.message}", e)
+          @status = Status::FAILED
+          save_auth_state
+          nil
+        end
+      end
+
+      # Cancel the authentication flow
+      # @param reason [String, nil] Optional reason for cancellation
+      # @return [Boolean] True if the flow was successfully cancelled
+      def cancel(reason = nil)
+        return false unless @status == Status::PENDING
+
+        @status = Status::CANCELLED
+        @error = Legate::Auth::Error.new("Authentication cancelled#{reason ? ": #{reason}" : ''}")
+        save_auth_state
+        true
+      end
+
+      # Get the current status of the authentication flow
+      # @return [Symbol] The current status
+      attr_reader :status
+
+      # Get the error that occurred during authentication, if any
+      # @return [Legate::Auth::Error, nil] The error or nil if no error occurred
+      attr_reader :error
+
+      # Get the result of the authentication flow
+      # @return [Legate::Auth::ExchangedCredential, nil] The resulting credential or nil if not completed
+      attr_reader :result
+
+      # Check if the authentication flow is complete
+      # @return [Boolean] True if the flow is complete (success or failure)
+      def complete?
+        @status != Status::PENDING
+      end
+
+      # Check if the authentication flow is successful
+      # @return [Boolean] True if the flow completed successfully
+      def success?
+        @status == Status::COMPLETED && @result
+      end
+
+      protected
+
+      # Main authentication method to be implemented by subclasses
+      # This method should use Fiber.yield to pause execution and wait for client input
+      # @return [Legate::Auth::ExchangedCredential] The authenticated credential
+      # @raise [Legate::Auth::Error] If authentication fails
+      def authenticate
+        raise NotImplementedError, "Subclasses must implement the 'authenticate' method"
+      end
+
+      # Generate a cache key for the authentication token
+      # @param scheme [Legate::Auth::Scheme] The authentication scheme
+      # @param credential [Legate::Auth::Credential] The credential
+      # @return [String] The cache key
+      def generate_cache_key(scheme, credential)
+        require 'digest/sha2'
+
+        # Create a unique key based on scheme and credential
+        parts = [
+          scheme.scheme_type.to_s,
+          credential.auth_type.to_s
+        ]
+
+        # Add scheme-specific information
+        case scheme.scheme_type
+        when :api_key
+          parts << credential[:api_key, resolve_env: false].to_s
+        when :http_bearer
+          parts << credential[:bearer_token, resolve_env: false].to_s
+        when :oauth2, :oidc
+          parts << credential[:client_id, resolve_env: false].to_s
+          parts << (credential[:scope, resolve_env: false] || '').to_s
+        when :service_account
+          parts << credential[:client_email, resolve_env: false].to_s
+        end
+
+        "auth_#{Digest::SHA256.hexdigest(parts.join(':'))}"
+      end
+
+      private
+
+      # Save the current authentication state to the session
+      def save_auth_state
+        state = {
+          request_id: @request_id,
+          scheme_type: @scheme.scheme_type,
+          status: @status,
+          start_time: @start_time&.iso8601,
+          error: @error&.message
+        }
+
+        @session_service.save_scoped_state('auth_flow', @request_id, state)
+      end
+    end
+  end
+end

data/lib/legate/auth/coordinators/oauth2_coordinator.rb

@@ -0,0 +1,99 @@
+# frozen_string_literal: true
+
+require_relative '../coordinator'
+require_relative '../schemes/oauth2'
+require_relative '../config'
+
+module Legate
+  module Auth
+    module Coordinators
+      # OAuth2Coordinator handles the interactive OAuth2 authentication flow using fibers.
+      # It manages pausing execution to request user authorization, and resuming once
+      # the authorization code is received.
+      class OAuth2Coordinator < Coordinator
+        # Authentication steps for OAuth2
+        module Steps
+          AUTHORIZATION = :authorization
+          TOKEN_EXCHANGE = :token_exchange
+        end
+
+        # Initialize a new OAuth2 coordinator
+        # @param scheme [Legate::Auth::Schemes::OAuth2] The OAuth2 scheme
+        # @param credential [Legate::Auth::Credential] The credential with client information
+        # @param session_service [Legate::SessionService::Base] The session service
+        # @param token_store [Legate::Auth::TokenStore, nil] Optional token store
+        # @param timeout [Integer, nil] Optional timeout in seconds
+        # @param redirect_uri [String, nil] Optional redirect URI
+        def initialize(scheme:, credential:, session_service:, token_store: nil, timeout: DEFAULT_TIMEOUT, redirect_uri: nil)
+          super(scheme: scheme, credential: credential, session_service: session_service, token_store: token_store, timeout: timeout)
+
+          raise ArgumentError, "Expected an OAuth2 scheme, got #{scheme.class}" unless scheme.is_a?(Legate::Auth::Schemes::OAuth2)
+
+          raise ArgumentError, "Credential must have auth_type :oauth2, got #{credential.auth_type}" unless credential.auth_type == :oauth2
+
+          @redirect_uri = redirect_uri
+          @current_step = Steps::AUTHORIZATION
+          @auth_config = nil
+        end
+
+        protected
+
+        # Implement the OAuth2 authentication flow
+        # @return [Legate::Auth::ExchangedCredential] The authenticated credential
+        # @raise [Legate::Auth::Error] If authentication fails
+        def authenticate
+          # Step 1: Create authorization request
+          @current_step = Steps::AUTHORIZATION
+          authorization_response = request_authorization
+
+          # Step 2: Exchange code for tokens
+          @current_step = Steps::TOKEN_EXCHANGE
+          exchange_code_for_token(authorization_response)
+        end
+
+        private
+
+        # Request authorization from the user
+        # @return [Hash] The authorization response from the client
+        def request_authorization
+          # Create a config for the authorization request
+          @auth_config = Legate::Auth::Config.new(
+            scheme: @scheme,
+            credential: @credential,
+            options: { request_id: @request_id }
+          )
+
+          # Build the authorization URI
+          authorization_uri = @auth_config.build_authorization_uri(@redirect_uri)
+
+          # Yield to pause execution and wait for authorization response
+          response = Fiber.yield({
+                                   type: 'authorization_request',
+                                   authorization_url: authorization_uri[:uri],
+                                   state: authorization_uri[:state],
+                                   redirect_uri: @redirect_uri
+                                 })
+
+          # Validate the response
+          raise Legate::Auth::Error, "Invalid authorization response: expected Hash, got #{response.class}" unless response.is_a?(Hash)
+
+          raise Legate::Auth::Error, 'Missing response_uri in authorization response' unless response['response_uri']
+
+          response
+        end
+
+        # Exchange the authorization code for tokens
+        # @param authorization_response [Hash] The authorization response from the client
+        # @return [Legate::Auth::ExchangedCredential] The authenticated credential
+        # @raise [Legate::Auth::TokenExchangeError] If token exchange fails
+        def exchange_code_for_token(authorization_response)
+          # Update the config with the response URI
+          @auth_config.response_uri = authorization_response['response_uri']
+
+          # Exchange the code for tokens
+          @scheme.exchange_token(@auth_config, @credential)
+        end
+      end
+    end
+  end
+end