From 45e887e72fcb42ea672ec7ceb46951af03d9be16 Mon Sep 17 00:00:00 2001 From: Casey Chow Date: Wed, 11 Feb 2026 13:49:29 -0500 Subject: [PATCH] fix: persist and reuse OAuth discovery base URL --- codex-rs/rmcp-client/src/auth_status.rs | 84 ++++++++++++------- .../rmcp-client/src/perform_oauth_login.rs | 69 +++++++++++++-- codex-rs/rmcp-client/src/rmcp_client.rs | 37 +++++--- codex-rs/rmcp-client/src/utils.rs | 46 ++++++++++ 4 files changed, 185 insertions(+), 51 deletions(-) diff --git a/codex-rs/rmcp-client/src/auth_status.rs b/codex-rs/rmcp-client/src/auth_status.rs index c752aec788..aac4a8198c 100644 --- a/codex-rs/rmcp-client/src/auth_status.rs +++ b/codex-rs/rmcp-client/src/auth_status.rs @@ -15,6 +15,7 @@ use crate::OAuthCredentialsStoreMode; use crate::oauth::has_oauth_tokens; use crate::utils::apply_default_headers; use crate::utils::build_default_headers; +use crate::utils::oauth_auth_url_candidates; const DISCOVERY_TIMEOUT: Duration = Duration::from_secs(5); const OAUTH_DISCOVERY_HEADER: &str = "MCP-Protocol-Version"; @@ -57,50 +58,69 @@ pub async fn supports_oauth_login(url: &str) -> Result { } async fn supports_oauth_login_with_headers(url: &str, default_headers: &HeaderMap) -> Result { - let base_url = Url::parse(url)?; + let mut supports = false; - // Use no_proxy to avoid a bug in the system-configuration crate that - // can result in a panic. See #8912. - let builder = Client::builder().timeout(DISCOVERY_TIMEOUT).no_proxy(); - let client = apply_default_headers(builder, default_headers).build()?; - - let mut last_error: Option = None; - for candidate_path in discovery_paths(base_url.path()) { - let mut discovery_url = base_url.clone(); - discovery_url.set_path(&candidate_path); - - let response = match client - .get(discovery_url.clone()) - .header(OAUTH_DISCOVERY_HEADER, OAUTH_DISCOVERY_VERSION) - .send() - .await - { - Ok(response) => response, + for candidate_url in oauth_auth_url_candidates(url) { + let base_url = match Url::parse(&candidate_url) { + Ok(base_url) => base_url, Err(err) => { - last_error = Some(err.into()); + debug!("Skipping OAuth discovery candidate `{candidate_url}`: {err:?}"); continue; } }; - if response.status() != StatusCode::OK { - continue; + // Use no_proxy to avoid a bug in the system-configuration crate that + // can result in a panic. See #8912. + let builder = Client::builder().timeout(DISCOVERY_TIMEOUT).no_proxy(); + let client = apply_default_headers(builder, default_headers).build()?; + + let mut last_error: Option = None; + for candidate_path in discovery_paths(base_url.path()) { + let mut discovery_url = base_url.clone(); + discovery_url.set_path(&candidate_path); + + let response = match client + .get(discovery_url.clone()) + .header(OAUTH_DISCOVERY_HEADER, OAUTH_DISCOVERY_VERSION) + .send() + .await + { + Ok(response) => response, + Err(err) => { + last_error = Some(err.into()); + continue; + } + }; + + if response.status() != StatusCode::OK { + continue; + } + + let metadata = match response.json::().await { + Ok(metadata) => metadata, + Err(err) => { + last_error = Some(err.into()); + continue; + } + }; + + if metadata.authorization_endpoint.is_some() && metadata.token_endpoint.is_some() { + supports = true; + break; + } } - let metadata = match response.json::().await { - Ok(metadata) => metadata, - Err(err) => { - last_error = Some(err.into()); - continue; - } - }; + if let Some(err) = last_error { + debug!("OAuth discovery requests failed for {candidate_url}: {err:?}"); + } - if metadata.authorization_endpoint.is_some() && metadata.token_endpoint.is_some() { - return Ok(true); + if supports { + break; } } - if let Some(err) = last_error { - debug!("OAuth discovery requests failed for {url}: {err:?}"); + if supports { + return Ok(true); } Ok(false) diff --git a/codex-rs/rmcp-client/src/perform_oauth_login.rs b/codex-rs/rmcp-client/src/perform_oauth_login.rs index 09b746837e..b2c33cbb2f 100644 --- a/codex-rs/rmcp-client/src/perform_oauth_login.rs +++ b/codex-rs/rmcp-client/src/perform_oauth_login.rs @@ -13,6 +13,7 @@ use tiny_http::Response; use tiny_http::Server; use tokio::sync::oneshot; use tokio::time::timeout; +use tracing::debug; use urlencoding::decode; use crate::OAuthCredentialsStoreMode; @@ -22,6 +23,7 @@ use crate::oauth::compute_expires_at_millis; use crate::save_oauth_tokens; use crate::utils::apply_default_headers; use crate::utils::build_default_headers; +use crate::utils::oauth_auth_url_candidates; struct OauthHeaders { http_headers: Option>, @@ -219,7 +221,7 @@ struct OauthLoginFlow { rx: oneshot::Receiver<(String, String)>, guard: CallbackServerGuard, server_name: String, - server_url: String, + oauth_server_url: String, store_mode: OAuthCredentialsStoreMode, launch_browser: bool, timeout: Duration, @@ -288,12 +290,61 @@ impl OauthLoginFlow { let default_headers = build_default_headers(http_headers, env_http_headers)?; let http_client = apply_default_headers(ClientBuilder::new(), &default_headers).build()?; - let mut oauth_state = OAuthState::new(server_url, Some(http_client)).await?; - let scope_refs: Vec<&str> = scopes.iter().map(String::as_str).collect(); - oauth_state - .start_authorization(&scope_refs, &redirect_uri, Some("Codex")) - .await?; - let auth_url = oauth_state.get_authorization_url().await?; + let mut last_error: Option = None; + let mut auth_url = None; + let mut oauth_state = None; + let mut oauth_server_url = None; + + for candidate_url in oauth_auth_url_candidates(server_url) { + let mut state = match OAuthState::new(&candidate_url, Some(http_client.clone())).await { + Ok(state) => state, + Err(err) => { + last_error = Some(err.into()); + continue; + } + }; + let scope_refs: Vec<&str> = scopes.iter().map(String::as_str).collect(); + if let Err(err) = state + .start_authorization(&scope_refs, &redirect_uri, Some("Codex")) + .await + { + last_error = Some(err.into()); + continue; + } + + auth_url = match state.get_authorization_url().await { + Ok(candidate_auth_url) => { + oauth_state = Some(state); + oauth_server_url = Some(candidate_url); + Some(candidate_auth_url) + } + Err(err) => { + let err_msg = format!("{err:?}"); + last_error = Some(anyhow!(err_msg.clone())); + debug!("OAuth state did not provide authorization URL: {err_msg}"); + continue; + } + }; + break; + } + + let oauth_state = match oauth_state { + Some(state) => state, + None => { + return Err(last_error.unwrap_or_else(|| { + anyhow!("No usable OAuth auth endpoint found for MCP server") + })); + } + }; + let auth_url = match auth_url { + Some(url) => url, + None => { + return Err(anyhow!( + "Internal OAuth login error: authorization URL not initialized after state setup" + )); + } + }; + let oauth_server_url = oauth_server_url.unwrap_or_else(|| server_url.to_string()); let timeout_secs = timeout_secs.unwrap_or(DEFAULT_OAUTH_TIMEOUT_SECS).max(1); let timeout = Duration::from_secs(timeout_secs as u64); @@ -303,7 +354,7 @@ impl OauthLoginFlow { rx, guard, server_name: server_name.to_string(), - server_url: server_url.to_string(), + oauth_server_url, store_mode, launch_browser, timeout, @@ -349,7 +400,7 @@ impl OauthLoginFlow { let expires_at = compute_expires_at_millis(&credentials); let stored = StoredOAuthTokens { server_name: self.server_name.clone(), - url: self.server_url.clone(), + url: self.oauth_server_url.clone(), client_id, token_response: WrappedOAuthTokenResponse(credentials), expires_at, diff --git a/codex-rs/rmcp-client/src/rmcp_client.rs b/codex-rs/rmcp-client/src/rmcp_client.rs index 0deccf3919..21a8ac0956 100644 --- a/codex-rs/rmcp-client/src/rmcp_client.rs +++ b/codex-rs/rmcp-client/src/rmcp_client.rs @@ -59,6 +59,7 @@ use crate::program_resolver; use crate::utils::apply_default_headers; use crate::utils::build_default_headers; use crate::utils::create_env_for_mcp_server; +use crate::utils::oauth_auth_url_candidates; use crate::utils::run_with_timeout; enum PendingTransport { @@ -244,21 +245,36 @@ impl RmcpClient { ) -> Result { let default_headers = build_default_headers(http_headers, env_http_headers)?; - let initial_oauth_tokens = match bearer_token { - Some(_) => None, - None => match load_oauth_tokens(server_name, url, store_mode) { - Ok(tokens) => tokens, - Err(err) => { - warn!("failed to read tokens for server `{server_name}`: {err}"); - None + let (initial_oauth_tokens, oauth_server_url) = match bearer_token { + Some(_) => (None, None), + None => { + let mut token_match = None; + for candidate_url in oauth_auth_url_candidates(url) { + match load_oauth_tokens(server_name, &candidate_url, store_mode) { + Ok(Some(tokens)) => { + token_match = Some((candidate_url, tokens)); + break; + } + Ok(None) => {} + Err(err) => { + warn!( + "failed to read tokens for server `{server_name}` at `{candidate_url}`: {err}" + ); + } + } } - }, + token_match.map_or((None, None), |(candidate_url, tokens)| { + (Some(tokens), Some(candidate_url)) + }) + } }; let transport = if let Some(initial_tokens) = initial_oauth_tokens.clone() { + let oauth_url = oauth_server_url.as_deref().unwrap_or(url); match create_oauth_transport_and_runtime( server_name, url, + oauth_url, initial_tokens.clone(), store_mode, default_headers.clone(), @@ -583,6 +599,7 @@ impl RmcpClient { async fn create_oauth_transport_and_runtime( server_name: &str, url: &str, + oauth_url: &str, initial_tokens: StoredOAuthTokens, credentials_store: OAuthCredentialsStoreMode, default_headers: HeaderMap, @@ -592,7 +609,7 @@ async fn create_oauth_transport_and_runtime( )> { let http_client = apply_default_headers(reqwest::Client::builder(), &default_headers).build()?; - let mut oauth_state = OAuthState::new(url.to_string(), Some(http_client.clone())).await?; + let mut oauth_state = OAuthState::new(oauth_url.to_string(), Some(http_client.clone())).await?; oauth_state .set_credentials( @@ -619,7 +636,7 @@ async fn create_oauth_transport_and_runtime( let runtime = OAuthPersistor::new( server_name.to_string(), - url.to_string(), + oauth_url.to_string(), auth_manager, credentials_store, Some(initial_tokens), diff --git a/codex-rs/rmcp-client/src/utils.rs b/codex-rs/rmcp-client/src/utils.rs index e47c1d14b6..91003aaeeb 100644 --- a/codex-rs/rmcp-client/src/utils.rs +++ b/codex-rs/rmcp-client/src/utils.rs @@ -6,6 +6,7 @@ use anyhow::Context; use anyhow::Result; use anyhow::anyhow; use reqwest::ClientBuilder; +use reqwest::Url; use reqwest::header::HeaderMap; use reqwest::header::HeaderName; use reqwest::header::HeaderValue; @@ -112,6 +113,33 @@ pub(crate) fn apply_default_headers( } } +pub(crate) fn oauth_auth_url_candidates(server_url: &str) -> Vec { + let mut candidates = vec![server_url.to_string()]; + + let Ok(mut base_url) = Url::parse(server_url) else { + return candidates; + }; + + let mut trimmed_path = base_url.path().trim_end_matches('/').to_string(); + if let Some(truncated) = trimmed_path.strip_suffix("/mcp") { + trimmed_path = truncated.to_string(); + if trimmed_path.is_empty() { + base_url.set_path("/"); + } else { + base_url.set_path(&trimmed_path); + } + base_url.set_query(None); + base_url.set_fragment(None); + + let fallback = base_url.to_string().trim_end_matches('/').to_string(); + if !candidates.contains(&fallback) { + candidates.push(fallback); + } + } + + candidates +} + #[cfg(unix)] pub(crate) const DEFAULT_ENV_VARS: &[&str] = &[ "HOME", @@ -215,4 +243,22 @@ mod tests { let env = create_env_for_mcp_server(None, &[custom_var.to_string()]); assert_eq!(env.get(custom_var), Some(&value.to_string())); } + + #[test] + fn oauth_auth_url_candidates_returns_input_when_not_mcp_path() { + let candidates = oauth_auth_url_candidates("https://example.com/api"); + assert_eq!(candidates, vec!["https://example.com/api".to_string()]); + } + + #[test] + fn oauth_auth_url_candidates_strips_trailing_mcp_segment() { + let candidates = oauth_auth_url_candidates("https://example.com/mcp"); + assert_eq!( + candidates, + vec![ + "https://example.com/mcp".to_string(), + "https://example.com".to_string() + ] + ); + } }