fix: Implement a CredentialStore for auth (#5741)

crates/goose/src/oauth/persist.rs

@@ -1,71 +1,54 @@
-use oauth2::{basic::BasicTokenType, EmptyExtraTokenFields, StandardTokenResponse};
-use reqwest::IntoUrl;
-use rmcp::transport::{auth::OAuthState, AuthError};
-use serde::{Deserialize, Serialize};
+use rmcp::transport::auth::{AuthError, CredentialStore, StoredCredentials};
 
 use crate::config::Config;
 
-#[derive(Debug, Clone, Serialize, Deserialize)]
-pub struct SerializableCredentials {
-    pub client_id: String,
-    pub token_response: Option<StandardTokenResponse<EmptyExtraTokenFields, BasicTokenType>>,
+/// Goose-specific credential store that uses the Config system
+///
+/// This implementation stores OAuth credentials in the goose configuration
+/// system, which handles secure storage (e.g., keychain integration).
+
+#[derive(Clone)]
+pub struct GooseCredentialStore {
+    name: String,
 }
 
-fn secret_key(name: &str) -> String {
-    format!("oauth_creds_{name}")
-}
+impl GooseCredentialStore {
+    pub fn new(name: String) -> Self {
+        Self { name }
+    }
 
-pub async fn save_credentials(
-    name: &str,
-    oauth_state: &OAuthState,
-) -> Result<(), Box<dyn std::error::Error>> {
-    let config = Config::global();
-    let (client_id, token_response) = oauth_state.get_credentials().await?;
-
-    let credentials = SerializableCredentials {
-        client_id,
-        token_response,
-    };
-
-    let key = secret_key(name);
-    config.set_secret(&key, &credentials)?;
-
-    Ok(())
-}
-
-async fn load_credentials(
-    name: &str,
-) -> Result<SerializableCredentials, Box<dyn std::error::Error>> {
-    let config = Config::global();
-    let key = secret_key(name);
-    let credentials: SerializableCredentials = config.get_secret(&key)?;
-
-    Ok(credentials)
-}
-
-pub fn clear_credentials(name: &str) -> Result<(), Box<dyn std::error::Error>> {
-    let config = Config::global();
-
-    Ok(config.delete_secret(&secret_key(name))?)
-}
-
-pub async fn load_cached_state<U: IntoUrl>(
-    base_url: U,
-    name: &str,
-) -> Result<OAuthState, AuthError> {
-    let credentials = load_credentials(name)
-        .await
-        .map_err(|e| AuthError::InternalError(format!("Failed to load credentials: {}", e)))?;
-
-    if let Some(token_response) = credentials.token_response {
-        let mut oauth_state = OAuthState::new(base_url, None).await?;
-        oauth_state
-            .set_credentials(&credentials.client_id, token_response)
-            .await?;
-        Ok(oauth_state)
-    } else {
-        Err(AuthError::InternalError(
-            "No token response in cached credentials".to_string(),
-        ))
+    fn secret_key(&self) -> String {
+        format!("oauth_creds_{}", self.name)
+    }
+}
+
+#[async_trait::async_trait]
+impl CredentialStore for GooseCredentialStore {
+    async fn load(&self) -> Result<Option<StoredCredentials>, AuthError> {
+        let config = Config::global();
+        let key = self.secret_key();
+
+        match config.get_secret::<StoredCredentials>(&key) {
+            Ok(credentials) => Ok(Some(credentials)),
+            Err(_) => Ok(None), // No credentials found
+        }
+    }
+
+    async fn save(&self, credentials: StoredCredentials) -> Result<(), AuthError> {
+        let config = Config::global();
+        let key = self.secret_key();
+
+        config
+            .set_secret(&key, &credentials)
+            .map_err(|e| AuthError::InternalError(format!("Failed to save credentials: {}", e)))
+    }
+
+    async fn clear(&self) -> Result<(), AuthError> {
+        let config = Config::global();
+        let key = self.secret_key();
+
+        config
+            .delete_secret(&key)
+            .map_err(|e| AuthError::InternalError(format!("Failed to clear credentials: {}", e)))
     }
 }