Add OAuth token caching and automatic refresh

- Add TokenCache for persisting OAuth tokens to disk
- Add TokenProvider trait for dynamic token management
- Implement OAuthTokenProvider with automatic token refresh
- Update HttpTransport to use dynamic token provider
- Add token refresh support in OAuthDiscovery
- Fix control flow for OAuth re-authentication when refresh fails
- Upgrade all dependencies to latest versions

Author: dreygur

Cargo.lock

@@ -1,6 +1,6 @@
 [package]
 name = "mcp-client"
-version = "0.1.0"
+version = "0.2.0"
 edition = "2021"
 authors = ["Rakibul Yeasin <ryeasin03@gmail.com>"]
 

crates/mcp-client/Cargo.toml

@@ -2,10 +2,11 @@ use crate::error::{ClientError, Result};
 use reqwest::Client;
 use serde::{Deserialize, Serialize};
 use sha2::{Digest, Sha256};
+use std::path::PathBuf;
 use std::sync::Arc;
 use std::time::{SystemTime, UNIX_EPOCH};
 use tokio::sync::RwLock;
-use tracing::{debug, info};
+use tracing::{debug, info, warn};
 
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct OAuthClientConfig {
@@ -384,6 +385,289 @@ impl OAuthClient {
     }
 }
 
+/// Trait for providing OAuth tokens dynamically
+#[async_trait::async_trait]
+pub trait TokenProvider: Send + Sync {
+    /// Get a valid access token, refreshing if necessary
+    async fn get_token(&self) -> Result<String>;
+}
+
+/// Static token provider - just returns a fixed token string
+pub struct StaticTokenProvider {
+    token: String,
+}
+
+impl StaticTokenProvider {
+    pub fn new(token: String) -> Self {
+        Self { token }
+    }
+}
+
+#[async_trait::async_trait]
+impl TokenProvider for StaticTokenProvider {
+    async fn get_token(&self) -> Result<String> {
+        Ok(self.token.clone())
+    }
+}
+
+/// OAuth token provider - manages token lifecycle with automatic refresh
+pub struct OAuthTokenProvider {
+    server_name: String,
+    oauth_client: Arc<OAuthClient>,
+    token_cache: Arc<TokenCache>,
+    current_token: Arc<RwLock<Option<ClientToken>>>,
+}
+
+impl OAuthTokenProvider {
+    pub fn new(
+        server_name: String,
+        oauth_client: Arc<OAuthClient>,
+        token_cache: Arc<TokenCache>,
+        initial_token: Option<ClientToken>,
+    ) -> Self {
+        Self {
+            server_name,
+            oauth_client,
+            token_cache,
+            current_token: Arc::new(RwLock::new(initial_token)),
+        }
+    }
+
+    /// Check if token needs refresh (expired or expiring soon)
+    fn needs_refresh(token: &ClientToken) -> bool {
+        if let Some(expires_at) = token.expires_at {
+            let now = SystemTime::now()
+                .duration_since(UNIX_EPOCH)
+                .map(|d| d.as_secs())
+                .unwrap_or(0);
+            // Refresh 60 seconds before expiry
+            now + 60 >= expires_at
+        } else {
+            false
+        }
+    }
+}
+
+#[async_trait::async_trait]
+impl TokenProvider for OAuthTokenProvider {
+    async fn get_token(&self) -> Result<String> {
+        // Check current token
+        {
+            let token_guard = self.current_token.read().await;
+            if let Some(ref token) = *token_guard {
+                if !Self::needs_refresh(token) {
+                    return Ok(token.access_token.clone());
+                }
+            }
+        }
+
+        // Need to refresh - acquire write lock
+        let mut token_guard = self.current_token.write().await;
+
+        // Double-check after acquiring write lock
+        if let Some(ref token) = *token_guard {
+            if !Self::needs_refresh(token) {
+                return Ok(token.access_token.clone());
+            }
+
+            // Try to refresh
+            if token.refresh_token.is_some() {
+                debug!("Refreshing expired token for: {}", self.server_name);
+                self.oauth_client.set_token(token.clone()).await;
+
+                match self.oauth_client.refresh_token().await {
+                    Ok(new_token) => {
+                        info!("Token refreshed for: {}", self.server_name);
+                        // Update cache
+                        if let Err(e) = self.token_cache.save(&self.server_name, &new_token) {
+                            warn!("Failed to update token cache: {}", e);
+                        }
+                        let access_token = new_token.access_token.clone();
+                        *token_guard = Some(new_token);
+                        return Ok(access_token);
+                    }
+                    Err(e) => {
+                        warn!("Token refresh failed for '{}': {}", self.server_name, e);
+                    }
+                }
+            }
+        }
+
+        // No valid token available
+        Err(ClientError::OAuthError(
+            "No valid token available and refresh failed".to_string(),
+        ))
+    }
+}
+
+/// Cached token entry with metadata
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct CachedToken {
+    pub token: ClientToken,
+    pub server_name: String,
+    pub created_at: u64,
+}
+
+/// Token cache for persisting OAuth tokens to disk
+pub struct TokenCache {
+    cache_dir: PathBuf,
+}
+
+impl TokenCache {
+    /// Create a new token cache using the default cache directory
+    pub fn new() -> Result<Self> {
+        let cache_dir = Self::default_cache_dir()?;
+        Self::with_dir(cache_dir)
+    }
+
+    /// Create a new token cache with a custom directory
+    pub fn with_dir(cache_dir: PathBuf) -> Result<Self> {
+        std::fs::create_dir_all(&cache_dir)
+            .map_err(|e| ClientError::OAuthError(format!("Failed to create cache directory: {}", e)))?;
+        Ok(Self { cache_dir })
+    }
+
+    /// Get the default cache directory
+    fn default_cache_dir() -> Result<PathBuf> {
+        let base = dirs::cache_dir()
+            .or_else(dirs::home_dir)
+            .ok_or_else(|| ClientError::OAuthError("Cannot determine cache directory".to_string()))?;
+        Ok(base.join("mcp-connect").join("tokens"))
+    }
+
+    /// Generate a cache key from server name
+    fn cache_key(server_name: &str) -> String {
+        use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine};
+        let mut hasher = Sha256::new();
+        hasher.update(server_name.as_bytes());
+        let hash = hasher.finalize();
+        URL_SAFE_NO_PAD.encode(&hash[..16])
+    }
+
+    /// Get the cache file path for a server
+    fn cache_path(&self, server_name: &str) -> PathBuf {
+        let key = Self::cache_key(server_name);
+        self.cache_dir.join(format!("{}.json", key))
+    }
+
+    /// Load a cached token for a server
+    pub fn load(&self, server_name: &str) -> Option<CachedToken> {
+        let path = self.cache_path(server_name);
+        match std::fs::read_to_string(&path) {
+            Ok(content) => {
+                match serde_json::from_str::<CachedToken>(&content) {
+                    Ok(cached) => {
+                        debug!("Loaded cached token for server: {}", server_name);
+                        Some(cached)
+                    }
+                    Err(e) => {
+                        warn!("Failed to parse cached token: {}", e);
+                        None
+                    }
+                }
+            }
+            Err(_) => None,
+        }
+    }
+
+    /// Save a token to the cache
+    pub fn save(&self, server_name: &str, token: &ClientToken) -> Result<()> {
+        let now = SystemTime::now()
+            .duration_since(UNIX_EPOCH)
+            .map(|d| d.as_secs())
+            .unwrap_or(0);
+
+        let cached = CachedToken {
+            token: token.clone(),
+            server_name: server_name.to_string(),
+            created_at: now,
+        };
+
+        let path = self.cache_path(server_name);
+        let content = serde_json::to_string_pretty(&cached)
+            .map_err(|e| ClientError::OAuthError(format!("Failed to serialize token: {}", e)))?;
+
+        std::fs::write(&path, content)
+            .map_err(|e| ClientError::OAuthError(format!("Failed to write token cache: {}", e)))?;
+
+        info!("Saved token to cache for server: {}", server_name);
+        Ok(())
+    }
+
+    /// Remove a cached token
+    pub fn remove(&self, server_name: &str) -> Result<()> {
+        let path = self.cache_path(server_name);
+        if path.exists() {
+            std::fs::remove_file(&path)
+                .map_err(|e| ClientError::OAuthError(format!("Failed to remove cached token: {}", e)))?;
+            debug!("Removed cached token for server: {}", server_name);
+        }
+        Ok(())
+    }
+
+    /// Check if a cached token is still valid (not expired)
+    pub fn is_token_valid(token: &ClientToken) -> bool {
+        if let Some(expires_at) = token.expires_at {
+            let now = SystemTime::now()
+                .duration_since(UNIX_EPOCH)
+                .map(|d| d.as_secs())
+                .unwrap_or(0);
+            // Consider token expired 60 seconds before actual expiry for safety
+            now + 60 < expires_at
+        } else {
+            true // No expiration means valid
+        }
+    }
+
+    /// Load a valid token, or return None if expired/missing
+    pub fn load_valid(&self, server_name: &str) -> Option<ClientToken> {
+        self.load(server_name).and_then(|cached| {
+            if Self::is_token_valid(&cached.token) {
+                Some(cached.token)
+            } else {
+                debug!("Cached token for '{}' is expired", server_name);
+                None
+            }
+        })
+    }
+
+    /// Load token and refresh if expired (requires OAuthClient)
+    pub async fn load_or_refresh(
+        &self,
+        server_name: &str,
+        oauth_client: &OAuthClient,
+    ) -> Result<ClientToken> {
+        if let Some(cached) = self.load(server_name) {
+            if Self::is_token_valid(&cached.token) {
+                debug!("Using valid cached token for: {}", server_name);
+                return Ok(cached.token);
+            }
+
+            // Token expired, try to refresh
+            if cached.token.refresh_token.is_some() {
+                debug!("Attempting to refresh expired token for: {}", server_name);
+                oauth_client.set_token(cached.token).await;
+                match oauth_client.refresh_token().await {
+                    Ok(new_token) => {
+                        self.save(server_name, &new_token)?;
+                        return Ok(new_token);
+                    }
+                    Err(e) => {
+                        warn!("Failed to refresh token for '{}': {}", server_name, e);
+                        // Remove invalid cached token
+                        let _ = self.remove(server_name);
+                    }
+                }
+            } else {
+                debug!("No refresh token available for: {}", server_name);
+                let _ = self.remove(server_name);
+            }
+        }
+
+        Err(ClientError::OAuthError("No valid cached token available".to_string()))
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

crates/mcp-client/src/auth.rs

@@ -7,5 +7,8 @@ pub mod oauth_discovery;
 pub use client::McpRemoteClient;
 pub use error::ClientError;
 pub use transport::{HttpTransport, StdioTransport, TcpTransport};
-pub use auth::{OAuthClient, OAuthClientConfig, ClientToken};
+pub use auth::{
+    OAuthClient, OAuthClientConfig, ClientToken, TokenCache, CachedToken,
+    TokenProvider, StaticTokenProvider, OAuthTokenProvider,
+};
 pub use oauth_discovery::{OAuthDiscovery, OAuthMetadata, OAuthRequirement};