owlen/crates/platform/auth/src/refresh.rs

//! Background Token Refresh
//!
//! Provides automatic background token refresh for OAuth credentials before they expire.
//! This prevents authentication failures during long-running sessions.
//!
//! # Architecture
//!
//! The `TokenRefresher` spawns a background tokio task that:
//! 1. Runs every `check_interval` (default: 5 minutes)
//! 2. Scans all stored OAuth tokens
//! 3. Refreshes any tokens expiring within `refresh_threshold` (default: 10 minutes)
//! 4. Updates both credential store and in-memory cache on success
//!
//! # Usage
//!
//! ```rust,ignore
//! use auth_manager::{AuthManager, TokenRefresher, RefreshConfig};
//!
//! let auth_manager = AuthManager::new()?;
//! let refresher = TokenRefresher::start(auth_manager.clone(), RefreshConfig::default());
//!
//! // The refresher runs in the background
//! // Stop it when shutting down
//! refresher.stop();
//! ```

use crate::{AuthError, AuthManager, Result};
use llm_core::{AuthMethod, OAuthProvider, ProviderType};
use llm_anthropic::AnthropicAuth;
use llm_openai::OpenAIAuth;
use std::sync::Arc;
use std::time::{Duration, SystemTime, UNIX_EPOCH};
use tokio::sync::watch;
use tracing::{debug, error, info, warn};

/// Configuration for background token refresh
#[derive(Debug, Clone)]
pub struct RefreshConfig {
    /// How often to check for expiring tokens (default: 5 minutes)
    pub check_interval: Duration,

    /// Refresh tokens expiring within this threshold (default: 10 minutes)
    pub refresh_threshold: Duration,

    /// Whether to enable refresh (can be disabled for testing)
    pub enabled: bool,
}

impl Default for RefreshConfig {
    fn default() -> Self {
        Self {
            check_interval: Duration::from_secs(5 * 60),     // 5 minutes
            refresh_threshold: Duration::from_secs(10 * 60), // 10 minutes
            enabled: true,
        }
    }
}

impl RefreshConfig {
    /// Create a config for testing with shorter intervals
    pub fn for_testing() -> Self {
        Self {
            check_interval: Duration::from_secs(1),
            refresh_threshold: Duration::from_secs(5),
            enabled: true,
        }
    }

    /// Create a disabled config
    pub fn disabled() -> Self {
        Self {
            enabled: false,
            ..Default::default()
        }
    }
}

/// Handle to a running background token refresher
pub struct TokenRefresher {
    /// Channel to signal shutdown
    shutdown_sender: watch::Sender<bool>,

    /// Join handle for the background task
    task_handle: tokio::task::JoinHandle<()>,
}

impl TokenRefresher {
    /// Start the background token refresh task
    ///
    /// Returns a handle that can be used to stop the refresher.
    pub fn start(auth_manager: Arc<AuthManager>, config: RefreshConfig) -> Self {
        let (shutdown_sender, shutdown_receiver) = watch::channel(false);

        let task_handle = tokio::spawn(async move {
            run_refresh_loop(auth_manager, config, shutdown_receiver).await;
        });

        Self {
            shutdown_sender,
            task_handle,
        }
    }

    /// Stop the background refresh task
    pub fn stop(self) {
        // Signal shutdown
        let _ = self.shutdown_sender.send(true);
        // The task will exit on next check interval
    }

    /// Stop and wait for the task to complete
    pub async fn stop_and_wait(self) {
        let _ = self.shutdown_sender.send(true);
        let _ = self.task_handle.await;
    }

    /// Check if the refresh task is still running
    pub fn is_running(&self) -> bool {
        !self.task_handle.is_finished()
    }
}

/// Main refresh loop running in the background
async fn run_refresh_loop(
    auth_manager: Arc<AuthManager>,
    config: RefreshConfig,
    mut shutdown: watch::Receiver<bool>,
) {
    if !config.enabled {
        info!("Token refresh disabled");
        return;
    }

    info!(
        "Starting token refresh task (check_interval: {:?}, threshold: {:?})",
        config.check_interval, config.refresh_threshold
    );

    loop {
        tokio::select! {
            _ = tokio::time::sleep(config.check_interval) => {
                if let Err(error) = refresh_expiring_tokens(&auth_manager, &config).await {
                    error!(?error, "Error during token refresh cycle");
                }
            }
            _ = shutdown.changed() => {
                if *shutdown.borrow() {
                    info!("Token refresh task shutting down");
                    break;
                }
            }
        }
    }
}

/// Check all providers and refresh any tokens that are close to expiring
async fn refresh_expiring_tokens(auth_manager: &AuthManager, config: &RefreshConfig) -> Result<()> {
    let providers = [ProviderType::Anthropic, ProviderType::OpenAI];
    let threshold_secs = config.refresh_threshold.as_secs();

    for provider in providers {
        match check_and_refresh_provider(auth_manager, provider, threshold_secs).await {
            Ok(refreshed) => {
                if refreshed {
                    info!(?provider, "Successfully refreshed token");
                }
            }
            Err(error) => {
                warn!(?provider, ?error, "Failed to refresh token");
            }
        }
    }

    Ok(())
}

/// Check a single provider and refresh if needed
///
/// Returns `Ok(true)` if token was refreshed, `Ok(false)` if no refresh needed.
async fn check_and_refresh_provider(
    auth_manager: &AuthManager,
    provider: ProviderType,
    threshold_secs: u64,
) -> Result<bool> {
    // Get current auth
    let current_auth = match auth_manager.get_auth(provider) {
        Ok(auth) => auth,
        Err(_) => return Ok(false), // Not authenticated, nothing to refresh
    };

    // Only OAuth tokens can be refreshed
    let (_current_access_token, refresh_token_str, expires_at) = match &current_auth {
        AuthMethod::OAuth {
            access_token,
            refresh_token: Some(refresh),
            expires_at: Some(exp),
        } => (access_token.clone(), refresh.clone(), *exp),
        _ => return Ok(false), // Not OAuth or no refresh token
    };

    // Check if expiring soon
    let now_secs = SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .map(|d| d.as_secs())
        .unwrap_or(0);

    if expires_at > now_secs + threshold_secs {
        debug!(
            ?provider,
            expires_in_secs = expires_at - now_secs,
            "Token not expiring soon, skipping refresh"
        );
        return Ok(false);
    }

    info!(
        ?provider,
        expires_in_secs = expires_at.saturating_sub(now_secs),
        "Token expiring soon, refreshing"
    );

    // Perform the refresh using provider-specific OAuth implementation
    let new_auth = perform_token_refresh(provider, &refresh_token_str).await?;

    // Store the new credentials
    match &new_auth {
        AuthMethod::OAuth {
            access_token,
            refresh_token,
            expires_at: new_expires_at,
        } => {
            // Calculate expires_in from expires_at for store_oauth
            let expires_in = new_expires_at.map(|exp| {
                let now = SystemTime::now()
                    .duration_since(UNIX_EPOCH)
                    .map(|d| d.as_secs())
                    .unwrap_or(0);
                exp.saturating_sub(now)
            });

            auth_manager.store_oauth(
                provider,
                access_token,
                refresh_token.as_deref(),
                expires_in,
            )?;
        }
        _ => {
            return Err(AuthError::OAuth(
                "Refresh returned non-OAuth auth method".to_string(),
            ));
        }
    }

    Ok(true)
}

/// Perform token refresh for a specific provider
async fn perform_token_refresh(
    provider: ProviderType,
    refresh_token: &str,
) -> Result<AuthMethod> {
    match provider {
        ProviderType::Anthropic => {
            let auth_client = AnthropicAuth::new();
            auth_client
                .refresh_token(refresh_token)
                .await
                .map_err(|e| AuthError::OAuth(format!("Anthropic refresh failed: {}", e)))
        }
        ProviderType::OpenAI => {
            let auth_client = OpenAIAuth::new();
            auth_client
                .refresh_token(refresh_token)
                .await
                .map_err(|e| AuthError::OAuth(format!("OpenAI refresh failed: {}", e)))
        }
        ProviderType::Ollama => {
            // Ollama doesn't use OAuth tokens
            Err(AuthError::NotSupported(
                "Ollama does not support OAuth refresh".to_string(),
            ))
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_default_config() {
        let config = RefreshConfig::default();
        assert_eq!(config.check_interval, Duration::from_secs(5 * 60));
        assert_eq!(config.refresh_threshold, Duration::from_secs(10 * 60));
        assert!(config.enabled);
    }

    #[test]
    fn test_testing_config() {
        let config = RefreshConfig::for_testing();
        assert_eq!(config.check_interval, Duration::from_secs(1));
        assert_eq!(config.refresh_threshold, Duration::from_secs(5));
        assert!(config.enabled);
    }

    #[test]
    fn test_disabled_config() {
        let config = RefreshConfig::disabled();
        assert!(!config.enabled);
    }

    #[tokio::test]
    async fn test_refresher_starts_and_stops() {
        // Create auth manager
        let auth_manager = Arc::new(AuthManager::new().unwrap());

        // Start refresher with disabled config (won't actually run checks)
        let refresher = TokenRefresher::start(
            auth_manager,
            RefreshConfig::disabled(),
        );

        // Give it a moment to start
        tokio::time::sleep(Duration::from_millis(50)).await;

        // Should not be running since disabled
        // Note: task finishes almost immediately when disabled
        tokio::time::sleep(Duration::from_millis(50)).await;
        assert!(!refresher.is_running());

        // Stop is safe even if not running
        refresher.stop();
    }

    #[tokio::test]
    async fn test_refresher_with_short_interval() {
        let auth_manager = Arc::new(AuthManager::new().unwrap());

        // Start with very short interval
        let config = RefreshConfig {
            check_interval: Duration::from_millis(100),
            refresh_threshold: Duration::from_secs(1),
            enabled: true,
        };

        let refresher = TokenRefresher::start(auth_manager, config);

        // Let it run a couple cycles
        tokio::time::sleep(Duration::from_millis(250)).await;

        // Should still be running
        assert!(refresher.is_running());

        // Stop and wait
        refresher.stop_and_wait().await;
    }
}