From edb5b84a4820aaee0014af4a32569191d05548f1 Mon Sep 17 00:00:00 2001 From: Vincenzo Palazzo Date: Thu, 14 May 2026 19:26:07 +0200 Subject: [PATCH] [RFC] feat(oauth): proactive token refresh to avoid re-auth on every session (#8386) Signed-off-by: Vincenzo Palazzo Signed-off-by: Douwe Osinga Co-authored-by: Douwe Osinga --- crates/goose/src/agents/extension_manager.rs | 98 ++++++++++++++------ crates/goose/src/oauth/mod.rs | 18 +++- 2 files changed, 85 insertions(+), 31 deletions(-) diff --git a/crates/goose/src/agents/extension_manager.rs b/crates/goose/src/agents/extension_manager.rs index ecbca90a..fe03598b 100644 --- a/crates/goose/src/agents/extension_manager.rs +++ b/crates/goose/src/agents/extension_manager.rs @@ -41,14 +41,14 @@ use crate::builtin_extension::get_builtin_extension; use crate::config::extensions::name_to_key; use crate::config::search_path::SearchPaths; use crate::config::{get_all_extensions, Config}; -use crate::oauth::oauth_flow; +use crate::oauth::{oauth_flow, GooseCredentialStore}; use crate::prompt_template; use crate::subprocess::configure_subprocess; use rmcp::model::{ CallToolRequestParams, CallToolResult, Content, ErrorCode, ErrorData, GetPromptResult, Meta, Prompt, Resource, ResourceContents, ServerInfo, Tool, }; -use rmcp::transport::auth::AuthClient; +use rmcp::transport::auth::{AuthClient, CredentialStore}; use schemars::_private::NoSerialize; use serde_json::Value; @@ -507,6 +507,40 @@ pub(crate) fn substitute_env_vars(value: &str, env_map: &HashMap const GOOSE_USER_AGENT: reqwest::header::HeaderValue = reqwest::header::HeaderValue::from_static(concat!("goose/", env!("CARGO_PKG_VERSION"))); +#[allow(clippy::too_many_arguments)] +async fn connect_with_auth( + auth_manager: rmcp::transport::AuthorizationManager, + uri: &str, + timeout: Duration, + provider: SharedProvider, + client_name: String, + capabilities: GooseMcpClientCapabilities, + roots_dir: &std::path::Path, +) -> ExtensionResult> { + let mut auth_headers = HeaderMap::new(); + auth_headers.insert(reqwest::header::USER_AGENT, GOOSE_USER_AGENT); + let auth_http_client = reqwest::Client::builder() + .default_headers(auth_headers) + .build() + .map_err(|_| ExtensionError::ConfigError("could not construct http client".to_string()))?; + let auth_client = AuthClient::new(auth_http_client, auth_manager); + let transport = StreamableHttpClientTransport::with_client( + auth_client, + StreamableHttpClientTransportConfig::with_uri(uri), + ); + Ok(Box::new( + McpClient::connect( + transport, + timeout, + provider, + client_name, + capabilities, + roots_dir.to_path_buf(), + ) + .await?, + )) +} + #[allow(clippy::too_many_arguments)] async fn create_streamable_http_client( uri: &str, @@ -567,6 +601,32 @@ async fn create_streamable_http_client( let timeout_duration = Duration::from_secs(resolve_timeout(timeout)); + // If we have stored OAuth credentials, try refreshing and connecting directly. + // This avoids the unnecessary 401 → browser re-auth cycle on every new session. + let credential_store = GooseCredentialStore::new(name.to_string()); + if credential_store.load().await.is_ok_and(|c| c.is_some()) { + match oauth_flow(&uri.to_string(), &name.to_string()).await { + Ok(auth_manager) => { + return connect_with_auth( + auth_manager, + uri, + timeout_duration, + provider, + client_name, + capabilities, + roots_dir, + ) + .await; + } + Err(e) => { + warn!( + "[OAuth:{}] Proactive refresh failed: {}, falling back to unauthenticated attempt", + name, e + ); + } + } + } + let client_res = McpClient::connect( transport, timeout_duration, @@ -580,30 +640,16 @@ async fn create_streamable_http_client( if should_attempt_oauth_fallback(&client_res) { match oauth_flow(&uri.to_string(), &name.to_string()).await { Ok(auth_manager) => { - let mut auth_headers = HeaderMap::new(); - auth_headers.insert(reqwest::header::USER_AGENT, GOOSE_USER_AGENT); - let auth_http_client = reqwest::Client::builder() - .default_headers(auth_headers) - .build() - .map_err(|_| { - ExtensionError::ConfigError("could not construct http client".to_string()) - })?; - let auth_client = AuthClient::new(auth_http_client, auth_manager); - let transport = StreamableHttpClientTransport::with_client( - auth_client, - StreamableHttpClientTransportConfig::with_uri(uri), - ); - Ok(Box::new( - McpClient::connect( - transport, - timeout_duration, - provider, - client_name, - capabilities, - roots_dir.to_path_buf(), - ) - .await?, - )) + connect_with_auth( + auth_manager, + uri, + timeout_duration, + provider, + client_name, + capabilities, + roots_dir, + ) + .await } Err(_) => Ok(Box::new(client_res?)), } diff --git a/crates/goose/src/oauth/mod.rs b/crates/goose/src/oauth/mod.rs index ef0aefa0..aae1dc20 100644 --- a/crates/goose/src/oauth/mod.rs +++ b/crates/goose/src/oauth/mod.rs @@ -1,5 +1,7 @@ mod persist; +pub use persist::GooseCredentialStore; + use axum::extract::{Query, State}; use axum::response::Html; use axum::routing::get; @@ -14,8 +16,6 @@ use std::sync::Arc; use tokio::sync::{oneshot, Mutex}; use tracing::warn; -use crate::oauth::persist::GooseCredentialStore; - const CALLBACK_TEMPLATE: &str = include_str!("oauth_callback.html"); const CLIENT_METADATA_URL: &str = "https://goose-docs.ai/oauth/client-metadata.json"; @@ -39,12 +39,20 @@ pub async fn oauth_flow( auth_manager.set_credential_store(credential_store.clone()); if auth_manager.initialize_from_store().await? { - if auth_manager.refresh_token().await.is_ok() { - return Ok(auth_manager); + match auth_manager.refresh_token().await { + Ok(_) => { + return Ok(auth_manager); + } + Err(e) => { + warn!( + "[OAuth:{}] Token refresh failed: {} - clearing stored credentials and falling back to browser auth", + name, e + ); + } } if let Err(e) = credential_store.clear().await { - warn!("error clearing bad credentials: {}", e); + warn!("[OAuth:{}] error clearing bad credentials: {}", name, e); } }