legate 0.1.0

data/lib/legate/auth/token_manager.rb

@@ -0,0 +1,362 @@
+# frozen_string_literal: true
+
+require_relative 'token_store'
+require_relative 'exchanged_credential'
+require_relative 'error'
+
+module Legate
+  module Auth
+    # TokenManager is responsible for managing the lifecycle of authentication tokens.
+    # It provides a centralized system for token acquisition, refresh, and invalidation.
+    # This class works with the TokenStore for persistence and the various authentication
+    # schemes for token operations.
+    class TokenManager
+      # Default configuration values
+      DEFAULT_CONFIG = {
+        refresh_buffer: 60,       # Seconds before expiration to trigger refresh
+        retry_max_attempts: 3,    # Maximum number of refresh retry attempts
+        retry_delay: 2,           # Initial delay between retries (seconds)
+        retry_backoff: 1.5,       # Backoff multiplier for subsequent retries
+        auto_refresh: true,       # Whether to automatically refresh tokens
+        background_refresh: false # Whether to refresh tokens in background
+      }.freeze
+
+      # Initialize a new TokenManager
+      # @param token_store [Legate::Auth::TokenStore] The token store for persistence
+      # @param config [Hash] Configuration options
+      def initialize(token_store, config = {})
+        @token_store = token_store
+        @config = DEFAULT_CONFIG.merge(config)
+        @callbacks = {
+          before_expiry: [],
+          refresh_success: [],
+          refresh_failure: [],
+          invalidated: []
+        }
+        @lock = Mutex.new
+      end
+
+      # Get a token for the given scheme and credential
+      # @param scheme [Legate::Auth::Scheme] The authentication scheme
+      # @param credential [Legate::Auth::Credential] The credential
+      # @param force_refresh [Boolean] Whether to force a token refresh
+      # @return [Legate::Auth::ExchangedCredential, nil] The token or nil if not available
+      def get_token(scheme, credential, force_refresh: false)
+        raise ArgumentError, 'Scheme must be an Legate::Auth::Scheme' unless scheme.is_a?(Legate::Auth::Scheme)
+
+        cache_key = generate_cache_key(scheme, credential)
+
+        # Use a mutex to prevent race conditions during token retrieval/refresh
+        @lock.synchronize do
+          # Try to get the token from the store
+          token = @token_store.get(cache_key)
+
+          # If no token, refresh it
+          return refresh_token(scheme, credential, nil, cache_key) if token.nil?
+
+          # Check if a scheme supports refresh
+          supports_refresh = scheme.respond_to?(:supports_refresh?) && scheme.supports_refresh?
+          is_refreshable = token.respond_to?(:refreshable?) && token.refreshable?
+
+          # If force refresh is requested, refresh the token regardless of expiration
+          if force_refresh
+            # Check if the token is refreshable before attempting to refresh
+            return refresh_token(scheme, credential, token, cache_key) if supports_refresh && is_refreshable
+
+            # For tokens that aren't refreshable, create a new one
+            invalidate_token(cache_key)
+            new_token = exchange_token(scheme, credential)
+            if new_token
+              @token_store.store(cache_key, new_token)
+              trigger_callback(:refresh_success, new_token, scheme, credential)
+              return new_token
+            end
+
+          end
+
+          # Check if token needs refresh based on expiration
+          if needs_refresh?(token)
+            # Only try to refresh if the scheme supports it and the token is refreshable
+            return refresh_token(scheme, credential, token, cache_key) if supports_refresh && is_refreshable
+
+            # If not refreshable, invalidate and return nil
+            invalidate_token(cache_key)
+            return nil
+
+          end
+
+          # Check if token is approaching expiration and trigger callback
+          if approaching_expiration?(token)
+            trigger_callback(:before_expiry, token, scheme, credential)
+
+            # If auto_refresh is enabled and we can refresh, do it
+            return refresh_token(scheme, credential, token, cache_key) if @config[:auto_refresh] && supports_refresh && is_refreshable
+          end
+
+          token
+        end
+      end
+
+      # Explicitly refresh a token
+      # @param scheme [Legate::Auth::Scheme] The authentication scheme
+      # @param credential [Legate::Auth::Credential] The credential
+      # @param token [Legate::Auth::ExchangedCredential, nil] The current token, if available
+      # @return [Legate::Auth::ExchangedCredential, nil] The refreshed token or nil on failure
+      def refresh_token(scheme, credential, token = nil, cache_key = nil)
+        raise ArgumentError, 'Scheme must be an Legate::Auth::Scheme' unless scheme.is_a?(Legate::Auth::Scheme)
+
+        cache_key ||= generate_cache_key(scheme, credential)
+
+        # If we don't have a token and it's an oauth or service account scheme,
+        # we need to authenticate from scratch
+        if token.nil?
+          if %i[oauth2 oidc service_account].include?(scheme.scheme_type)
+            # For these schemes, we need a complete authentication flow
+            # which can't be handled here - return nil to indicate need for full auth
+            return nil
+          end
+
+          # For other schemes, we can simply apply the credential
+          begin
+            # Basic auth, API key, etc. - create a new token directly
+            token = exchange_token(scheme, credential)
+            if token
+              @token_store.store(cache_key, token)
+              trigger_callback(:refresh_success, token, scheme, credential)
+            end
+            return token
+          rescue Legate::Auth::Error => e
+            Legate.logger.error("Failed to create token: #{e.message}")
+            trigger_callback(:refresh_failure, nil, scheme, credential, error: e)
+            return nil
+          end
+        end
+
+        # Token exists - attempt to refresh it if scheme supports refresh
+        supports_refresh = scheme.respond_to?(:supports_refresh?) && scheme.supports_refresh?
+        is_refreshable = token.respond_to?(:refreshable?) && token.refreshable?
+
+        if supports_refresh && is_refreshable
+          begin
+            refreshed = scheme.refresh_token(token, credential)
+            if refreshed
+              @token_store.store(cache_key, refreshed)
+              trigger_callback(:refresh_success, refreshed, scheme, credential)
+              return refreshed
+            else
+              # Handle the case where refresh_token returns nil but doesn't raise an error
+              Legate.logger.error('Failed to refresh token: refresh_token returned nil')
+              trigger_callback(:refresh_failure, token, scheme, credential, error: nil)
+              return nil
+            end
+          rescue Legate::Auth::TokenRefreshError => e
+            Legate.logger.error("Failed to refresh token: #{e.message}")
+            trigger_callback(:refresh_failure, token, scheme, credential, error: e)
+            return nil
+          end
+        end
+
+        # Scheme doesn't support refresh or token isn't refreshable
+        # Return existing token if it's not expired
+        return token unless token.expired?
+
+        # Otherwise, invalidate it
+        invalidate_token(cache_key)
+        nil
+      end
+
+      # Invalidate a token, removing it from the store
+      # @param cache_key [String] The cache key for the token
+      # @return [Boolean] True if the token was invalidated
+      def invalidate_token(cache_key)
+        result = @token_store.clear(cache_key)
+        trigger_callback(:invalidated, nil, nil, nil, cache_key: cache_key) if result
+        result
+      end
+
+      # Revoke a token with the authentication provider
+      # @param scheme [Legate::Auth::Scheme] The authentication scheme
+      # @param credential [Legate::Auth::Credential] The credential
+      # @param token [Legate::Auth::ExchangedCredential] The token to revoke
+      # @return [Boolean] True if the token was revoked
+      def revoke_token(scheme, credential, token)
+        raise ArgumentError, 'Scheme must be an Legate::Auth::Scheme' unless scheme.is_a?(Legate::Auth::Scheme)
+        raise ArgumentError, 'Token must be an ExchangedCredential' unless token.is_a?(Legate::Auth::ExchangedCredential)
+
+        # Check if scheme supports revocation
+        unless scheme.respond_to?(:revoke_token)
+          Legate.logger.warn("Scheme #{scheme.scheme_type} does not support token revocation")
+          return false
+        end
+
+        begin
+          # Attempt to revoke the token
+          result = scheme.revoke_token(token, credential)
+
+          # Invalidate the token in our store if revocation succeeded
+          if result
+            cache_key = generate_cache_key(scheme, credential)
+            invalidate_token(cache_key)
+          end
+
+          result
+        rescue Legate::Auth::Error => e
+          Legate.logger.error("Failed to revoke token: #{e.message}")
+          false
+        rescue NotImplementedError => e
+          Legate.logger.error("#{e.message}")
+          false
+        end
+      end
+
+      # Register a callback for token lifecycle events
+      # @param event [Symbol] The event to register for (:before_expiry, :refresh_success, :refresh_failure, :invalidated)
+      # @param callback [Proc] The callback to execute
+      # @return [self]
+      def on(event, &callback)
+        raise ArgumentError, "Unknown event: #{event}. Valid events: #{@callbacks.keys.join(', ')}" unless @callbacks.key?(event)
+
+        @callbacks[event] << callback
+        self
+      end
+
+      private
+
+      # Check if a token needs to be refreshed
+      # @param token [Legate::Auth::ExchangedCredential] The token to check
+      # @return [Boolean] True if the token needs to be refreshed
+      def needs_refresh?(token)
+        token.expired?(@config[:refresh_buffer])
+      end
+
+      # Check if a token is approaching expiration
+      # @param token [Legate::Auth::ExchangedCredential] The token to check
+      # @return [Boolean] True if the token is approaching expiration
+      def approaching_expiration?(token)
+        return false unless token.expires_at
+
+        # Consider a token approaching expiration if it's within 2x the refresh buffer
+        buffer = @config[:refresh_buffer] * 2
+        (token.expires_at - Time.now) <= buffer
+      end
+
+      # Generate a cache key for the 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
+
+      # Perform token refresh with retry logic
+      # @param scheme [Legate::Auth::Scheme] The authentication scheme
+      # @param credential [Legate::Auth::Credential] The credential
+      # @param token [Legate::Auth::ExchangedCredential] The current token
+      # @param cache_key [String] The cache key
+      # @return [Legate::Auth::ExchangedCredential, nil] The refreshed token or nil on failure
+      def perform_token_refresh(scheme, credential, token, cache_key)
+        attempts = 0
+        delay = @config[:retry_delay]
+
+        loop do
+          refreshed = scheme.refresh_token(token, credential)
+          if refreshed
+            @token_store.store(cache_key, refreshed)
+            trigger_callback(:refresh_success, refreshed, scheme, credential)
+            return refreshed
+          else
+            # Handle the case where refresh_token returns nil but doesn't raise an error
+            Legate.logger.error('Failed to refresh token: refresh_token returned nil')
+            trigger_callback(:refresh_failure, token, scheme, credential, error: nil)
+            return nil
+          end
+        rescue Legate::Auth::TokenRefreshError => e
+          attempts += 1
+
+          # Check if we've exceeded max attempts
+          if attempts >= @config[:retry_max_attempts]
+            Legate.logger.error("Failed to refresh token after #{attempts} attempts: #{e.message}")
+            trigger_callback(:refresh_failure, token, scheme, credential, error: e)
+            return nil
+          end
+
+          # Exponential backoff for retry
+          sleep delay
+          delay *= @config[:retry_backoff]
+        end
+      end
+
+      # Exchange a credential for a token
+      # @param scheme [Legate::Auth::Scheme] The authentication scheme
+      # @param credential [Legate::Auth::Credential] The credential
+      # @return [Legate::Auth::ExchangedCredential, nil] The exchanged token or nil on failure
+      def exchange_token(scheme, credential)
+        case scheme.scheme_type
+        when :api_key
+          # Create a simple exchanged credential for API key
+          Legate::Auth::ExchangedCredential.new(
+            auth_type: :api_key,
+            access_token: credential[:api_key],
+            token_type: 'ApiKey'
+          )
+        when :http_bearer
+          # Create a simple exchanged credential for Bearer token
+          Legate::Auth::ExchangedCredential.new(
+            auth_type: :http_bearer,
+            access_token: credential[:bearer_token],
+            token_type: 'Bearer'
+          )
+        else
+          # Other types like OAuth2 require a more complex flow
+          # and cannot be handled directly here
+          nil
+        end
+      end
+
+      # Trigger a callback for the given event
+      # @param event [Symbol] The event that occurred
+      # @param token [Legate::Auth::ExchangedCredential, nil] The token involved in the event
+      # @param scheme [Legate::Auth::Scheme, nil] The authentication scheme
+      # @param credential [Legate::Auth::Credential, nil] The credential
+      # @param extras [Hash] Extra information to pass to the callback
+      def trigger_callback(event, token, scheme, credential, extras = {})
+        return unless @callbacks.key?(event)
+
+        @callbacks[event].each do |callback|
+          # Create a hash with all the data
+          data = {
+            event: event,
+            token: token,
+            scheme: scheme,
+            credential: credential
+          }.merge(extras)
+
+          callback.call(data)
+        rescue StandardError => e
+          Legate.logger.error("Error in #{event} callback: #{e.message}")
+        end
+      end
+    end
+  end
+end

data/lib/legate/auth/token_store.rb

@@ -0,0 +1,86 @@
+# File: lib/legate/auth/token_store.rb
+# frozen_string_literal: true
+
+require_relative 'exchanged_credential'
+
+module Legate
+  module Auth
+    # Provides a token store for caching authentication tokens
+    # This class wraps a session service and provides methods
+    # for storing, retrieving, and clearing tokens
+    class TokenStore
+      # Initialize a new token store
+      # @param session_service [Legate::SessionService::Base] The session service to use
+      def initialize(session_service)
+        @session_service = session_service
+        @scope = 'auth' # Scoped state namespace for authentication
+      end
+
+      # Store a token in the cache
+      # @param key [String] The cache key
+      # @param token [Legate::Auth::ExchangedCredential] The token to store
+      # @return [Boolean] True if the token was stored
+      def store(key, token)
+        return false unless token.is_a?(Legate::Auth::ExchangedCredential)
+
+        begin
+          # Serialize token to hash
+          token_data = token.to_h
+
+          # Store in scoped state
+          @session_service.save_scoped_state(@scope, key, token_data)
+          true
+        rescue StandardError => e
+          Legate.logger.error("Failed to store token: #{e.message}")
+          false
+        end
+      end
+
+      # Get a token from the cache
+      # @param key [String] The cache key
+      # @return [Legate::Auth::ExchangedCredential, nil] The token or nil if not found or expired
+      def get(key)
+        # Retrieve from scoped state
+        token_data = @session_service.load_scoped_state(@scope, key)
+        return nil unless token_data
+
+        # Deserialize to token object
+        token = Legate::Auth::ExchangedCredential.from_h(token_data)
+
+        # Check expiration
+        if token.expired?
+          Legate.logger.debug("Retrieved expired token from cache (key: #{key})")
+          # Clear expired token
+          clear(key)
+          return nil
+        end
+
+        token
+      rescue StandardError => e
+        Legate.logger.error("Failed to retrieve token: #{e.message}")
+        nil
+      end
+
+      # Clear a token from the cache
+      # @param key [String] The cache key to clear
+      # @return [Boolean] True if the token was cleared
+      def clear(key)
+        @session_service.clear_scoped_state(@scope, key)
+        true
+      rescue StandardError => e
+        Legate.logger.error("Failed to clear token: #{e.message}")
+        false
+      end
+
+      # Clear all tokens from the cache
+      # @return [Boolean] True if all tokens were cleared
+      def clear_all
+        @session_service.clear_scoped_state(@scope, '*')
+        true
+      rescue StandardError => e
+        Legate.logger.error("Failed to clear all tokens: #{e.message}")
+        false
+      end
+    end
+  end
+end

data/lib/legate/auth/tool_context_extension.rb

@@ -0,0 +1,97 @@
+# frozen_string_literal: true
+
+require_relative 'runner'
+require_relative 'token_manager'
+require_relative 'token_store'
+
+module Legate
+  module Auth
+    # Extension for Legate::ToolContext that adds fiber-based authentication support
+    # This module is meant to be included in the Legate::ToolContext class to add
+    # authentication-related methods for tools.
+    module ToolContextExtension
+      # Get or create an authentication runner for this context
+      # @return [Legate::Auth::Runner] The authentication runner
+      def auth_runner
+        @auth_runner ||= begin
+          # Create the token store
+          token_store = get_token_store
+
+          # Create a token manager
+          token_manager = Legate::Auth::TokenManager.new(token_store)
+
+          # Create the runner
+          Legate::Auth::Runner.new(
+            session_service: session_service,
+            token_store: token_store,
+            token_manager: token_manager
+          )
+        end
+      end
+
+      # Get a token store for this context
+      # @return [Legate::Auth::TokenStore] The token store
+      def get_token_store
+        @token_store ||= if session_service.respond_to?(:scoped_state_container)
+                           Legate::Auth::TokenStore.new(session_service)
+                         elsif defined?(Legate::Auth) && Legate::Auth.respond_to?(:token_store)
+                           Legate::Auth.token_store
+                         else
+                           Legate::Auth::TokenStore.new
+                         end
+      end
+
+      # Run a block with authentication support
+      # @param handler [Proc, nil] Optional handler for authentication requests
+      # @yield The block to run
+      # @return [Object] The result of the block
+      def with_authentication(&block)
+        raise ArgumentError, 'Block is required' unless block_given?
+
+        # Get or create the authentication runner
+        runner = auth_runner
+
+        # Run the block with authentication support
+        runner.run(block, self) do |auth_request|
+          # Here we return nil to indicate that the auth request should be yielded
+          # to the tool's caller for handling. In a real implementation, this could
+          # handle authentication UI or delegate to another component.
+          nil
+        end
+      end
+
+      # Start an authentication session
+      # @param scheme [Legate::Auth::Scheme] The authentication scheme to use
+      # @param credential [Legate::Auth::Credential] The credential to use
+      # @param options [Hash] Additional options for the authentication session
+      # @return [Legate::Auth::ExchangedCredential] The authenticated credential
+      def auth_session(scheme, credential, **options)
+        # This method will be dynamically replaced by the auth_runner when
+        # running in a fiber context. This implementation is just for fallback
+        # when not running in a fiber.
+        raise NotImplementedError, 'Authentication session not available outside of with_authentication block'
+      end
+
+      # Handle an authentication response (for tools that handle responses)
+      # @param request_id [String] The request ID
+      # @param response [Hash] The response
+      # @return [Hash] The result of handling the response
+      def handle_auth_response(request_id, response)
+        runner = auth_runner
+        runner.handle_auth_response(request_id, response)
+      end
+
+      # Cancel an authentication flow (for tools that handle responses)
+      # @param request_id [String] The request ID
+      # @param reason [String, nil] Optional reason for cancellation
+      # @return [Boolean] True if the flow was successfully cancelled
+      def cancel_auth_flow(request_id, reason = nil)
+        runner = auth_runner
+        runner.cancel_auth_flow(request_id, reason)
+      end
+    end
+  end
+end
+
+# Extend the ToolContext class if it's already defined
+Legate::ToolContext.include(Legate::Auth::ToolContextExtension) if defined?(Legate::ToolContext)