[RFC] feat(oauth): proactive token refresh to avoid re-auth on every session (#8386)

Signed-off-by: Vincenzo Palazzo <vincenzopalazzodev@gmail.com>
Signed-off-by: Douwe Osinga <douwe@squareup.com>
Co-authored-by: Douwe Osinga <douwe@squareup.com>

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<String, String>
 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<Box<dyn McpClientTrait>> {
+    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?)),
         }

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);
         }
     }