From a78939654d77bee5c700d5a430f46a7ecdb73478 Mon Sep 17 00:00:00 2001 From: Jack Amadeo Date: Fri, 15 Aug 2025 14:01:35 -0400 Subject: [PATCH] feat(mcp): Persist OAuth credentials to keyring (#4007) --- Cargo.lock | 1 + crates/goose/Cargo.toml | 1 + crates/goose/src/agents/extension_manager.rs | 4 +- crates/goose/src/{oauth.rs => oauth/mod.rs} | 26 ++++++- .../goose/src/{ => oauth}/oauth_callback.html | 0 crates/goose/src/oauth/persist.rs | 72 +++++++++++++++++++ 6 files changed, 100 insertions(+), 4 deletions(-) rename crates/goose/src/{oauth.rs => oauth/mod.rs} (77%) rename crates/goose/src/{ => oauth}/oauth_callback.html (100%) create mode 100644 crates/goose/src/oauth/persist.rs diff --git a/Cargo.lock b/Cargo.lock index 23ce244e91..eeab7b562a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3317,6 +3317,7 @@ dependencies = [ "minijinja", "mockall", "nanoid", + "oauth2", "once_cell", "opentelemetry", "opentelemetry-otlp", diff --git a/crates/goose/Cargo.toml b/crates/goose/Cargo.toml index b553cbcc28..473b6320dc 100644 --- a/crates/goose/Cargo.toml +++ b/crates/goose/Cargo.toml @@ -100,6 +100,7 @@ unicode-normalization = "0.1" # Vector database for tool selection lancedb = "0.13" arrow = "52.2" +oauth2 = "5.0.0" [target.'cfg(target_os = "windows")'.dependencies] winapi = { version = "0.3", features = ["wincred"] } diff --git a/crates/goose/src/agents/extension_manager.rs b/crates/goose/src/agents/extension_manager.rs index bdbf9a316c..c8a2b671a8 100644 --- a/crates/goose/src/agents/extension_manager.rs +++ b/crates/goose/src/agents/extension_manager.rs @@ -285,7 +285,9 @@ impl ExtensionManager { .await; let client = if let Err(e) = client_res { // make an attempt at oauth, but failing that, return the original error, - // because this might not have been an auth error at all + // because this might not have been an auth error at all. + // TODO: when rmcp supports it, we should trigger this flow on 401s with + // WWW-Authenticate headers, not just any init error let am = match oauth_flow(uri, name).await { Ok(am) => am, Err(_) => return Err(e.into()), diff --git a/crates/goose/src/oauth.rs b/crates/goose/src/oauth/mod.rs similarity index 77% rename from crates/goose/src/oauth.rs rename to crates/goose/src/oauth/mod.rs index f8f458c550..1dcb954fa5 100644 --- a/crates/goose/src/oauth.rs +++ b/crates/goose/src/oauth/mod.rs @@ -9,6 +9,11 @@ use serde::Deserialize; use std::net::SocketAddr; use std::sync::Arc; use tokio::sync::{oneshot, Mutex}; +use tracing::warn; + +use crate::oauth::persist::{clear_credentials, load_cached_state, save_credentials}; + +mod persist; const CALLBACK_TEMPLATE: &str = include_str!("oauth_callback.html"); @@ -28,6 +33,18 @@ pub async fn oauth_flow( mcp_server_url: &String, name: &String, ) -> Result { + if let Ok(oauth_state) = load_cached_state(mcp_server_url, name).await { + if let Some(authorization_manager) = oauth_state.into_authorization_manager() { + if authorization_manager.refresh_token().await.is_ok() { + return Ok(authorization_manager); + } + } + + if let Err(e) = clear_credentials(name) { + warn!("error clearing bad credentials: {}", e); + } + } + let (code_sender, code_receiver) = oneshot::channel::(); let app_state = AppState { code_receiver: Arc::new(Mutex::new(Some(code_sender))), @@ -52,7 +69,6 @@ pub async fn oauth_flow( let used_addr = listener.local_addr()?; tokio::spawn(async move { let result = axum::serve(listener, app).await; - if let Err(e) = result { eprintln!("Callback server error: {}", e); } @@ -73,9 +89,13 @@ pub async fn oauth_flow( let auth_code = code_receiver.await?; oauth_state.handle_callback(&auth_code).await?; - let am = oauth_state + if let Err(e) = save_credentials(name, &oauth_state).await { + warn!("Failed to save credentials: {}", e); + } + + let auth_manager = oauth_state .into_authorization_manager() .ok_or_else(|| anyhow::anyhow!("Failed to get authorization manager"))?; - Ok(am) + Ok(auth_manager) } diff --git a/crates/goose/src/oauth_callback.html b/crates/goose/src/oauth/oauth_callback.html similarity index 100% rename from crates/goose/src/oauth_callback.html rename to crates/goose/src/oauth/oauth_callback.html diff --git a/crates/goose/src/oauth/persist.rs b/crates/goose/src/oauth/persist.rs new file mode 100644 index 0000000000..cf6fccaa14 --- /dev/null +++ b/crates/goose/src/oauth/persist.rs @@ -0,0 +1,72 @@ +use oauth2::{basic::BasicTokenType, EmptyExtraTokenFields, StandardTokenResponse}; +use reqwest::IntoUrl; +use rmcp::transport::{auth::OAuthState, AuthError}; +use serde::{Deserialize, Serialize}; + +use crate::config::Config; + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct SerializableCredentials { + pub client_id: String, + pub token_response: Option>, +} + +fn secret_key(name: &str) -> String { + format!("oauth_creds_{name}") +} + +pub async fn save_credentials( + name: &str, + oauth_state: &OAuthState, +) -> Result<(), Box> { + let config = Config::global(); + let (client_id, token_response) = oauth_state.get_credentials().await?; + + let credentials = SerializableCredentials { + client_id, + token_response, + }; + + let value = serde_json::to_value(&credentials)?; + let key = secret_key(name); + config.set_secret(&key, value)?; + + Ok(()) +} + +async fn load_credentials( + name: &str, +) -> Result> { + 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> { + let config = Config::global(); + + Ok(config.delete_secret(&secret_key(name))?) +} + +pub async fn load_cached_state( + base_url: U, + name: &str, +) -> Result { + 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(), + )) + } +}