From 9d809e68cd606f0d6753733f16143bd032661bd4 Mon Sep 17 00:00:00 2001 From: GraciousGazelles <55840159+GraciousGazelles@users.noreply.github.com> Date: Sat, 20 Jun 2026 14:30:22 +1000 Subject: [PATCH 1/4] Serialize MCP OAuth refresh ownership --- codex-rs/cli/src/mcp_cmd.rs | 4 +- codex-rs/rmcp-client/src/lib.rs | 4 +- codex-rs/rmcp-client/src/oauth.rs | 668 ++++++++++++++++-- .../src/perform_oauth_device_login.rs | 4 +- .../rmcp-client/src/perform_oauth_login.rs | 4 +- codex-rs/rmcp-client/src/rmcp_client.rs | 209 +++--- .../rmcp-client/src/streamable_http_retry.rs | 121 +++- .../tests/streamable_http_oauth_startup.rs | 2 +- 8 files changed, 829 insertions(+), 187 deletions(-) diff --git a/codex-rs/cli/src/mcp_cmd.rs b/codex-rs/cli/src/mcp_cmd.rs index 4def1e0649aa..97c1d62cf5ac 100644 --- a/codex-rs/cli/src/mcp_cmd.rs +++ b/codex-rs/cli/src/mcp_cmd.rs @@ -27,7 +27,7 @@ use codex_mcp::resolve_oauth_scopes; use codex_mcp::should_retry_without_scopes; use codex_protocol::protocol::McpAuthStatus; use codex_rmcp_client::DeviceAuthorizationPrompt; -use codex_rmcp_client::delete_oauth_tokens; +use codex_rmcp_client::delete_oauth_tokens_locked; use codex_rmcp_client::perform_oauth_device_login; use codex_rmcp_client::perform_oauth_login; use codex_utils_cli::CliConfigOverrides; @@ -649,7 +649,7 @@ async fn run_logout(config_overrides: &CliConfigOverrides, logout_args: LogoutAr _ => bail!("OAuth logout is only supported for streamable_http transports."), }; - match delete_oauth_tokens(&name, &url, config.mcp_oauth_credentials_store_mode) { + match delete_oauth_tokens_locked(&name, &url, config.mcp_oauth_credentials_store_mode).await { Ok(true) => println!("Removed OAuth credentials for '{name}'."), Ok(false) => println!("No OAuth credentials stored for '{name}'."), Err(err) => return Err(anyhow!("failed to delete OAuth credentials: {err}")), diff --git a/codex-rs/rmcp-client/src/lib.rs b/codex-rs/rmcp-client/src/lib.rs index d18d0acce6bc..b564f01653ac 100644 --- a/codex-rs/rmcp-client/src/lib.rs +++ b/codex-rs/rmcp-client/src/lib.rs @@ -21,8 +21,10 @@ pub use in_process_transport::InProcessTransportFactory; pub use oauth::StoredOAuthTokens; pub use oauth::WrappedOAuthTokenResponse; pub use oauth::delete_oauth_tokens; -pub(crate) use oauth::load_oauth_tokens; +pub use oauth::delete_oauth_tokens_locked; +pub(crate) use oauth::load_oauth_tokens_with_source; pub use oauth::save_oauth_tokens; +pub use oauth::save_oauth_tokens_locked; pub use perform_oauth_device_login::DeviceAuthorizationPrompt; pub use perform_oauth_device_login::perform_oauth_device_login; pub use perform_oauth_login::OAuthProviderError; diff --git a/codex-rs/rmcp-client/src/oauth.rs b/codex-rs/rmcp-client/src/oauth.rs index 3ab905b83e2d..d007cec83fe3 100644 --- a/codex-rs/rmcp-client/src/oauth.rs +++ b/codex-rs/rmcp-client/src/oauth.rs @@ -35,11 +35,14 @@ use sha2::Digest; use sha2::Sha256; use std::collections::BTreeMap; use std::fs; +use std::fs::File; +use std::fs::OpenOptions; use std::io::ErrorKind; use std::io::Write; use std::path::PathBuf; use std::sync::Arc; use std::time::Duration; +use std::time::Instant; use std::time::SystemTime; use std::time::UNIX_EPOCH; use tracing::warn; @@ -47,12 +50,22 @@ use tracing::warn; use codex_keyring_store::DefaultKeyringStore; use codex_keyring_store::KeyringStore; use rmcp::transport::auth::AuthorizationManager; +use rmcp::transport::auth::CredentialStore as _; +use rmcp::transport::auth::InMemoryCredentialStore; +use rmcp::transport::auth::StoredCredentials; use tokio::sync::Mutex; +use tokio::time::sleep; +use tokio::time::timeout; use codex_utils_home_dir::find_codex_home; const KEYRING_SERVICE: &str = "Codex MCP Credentials"; -const REFRESH_SKEW_MILLIS: u64 = 30_000; +const REFRESH_SKEW_MILLIS: u64 = 60_000; +const REFRESH_LOCK_DIR: &str = "mcp-oauth-refresh-locks"; +const LOCK_ACQUIRE_TIMEOUT: Duration = Duration::from_secs(60); +const CREDENTIAL_LOCK_RETRY_SLEEP: Duration = Duration::from_millis(500); +const STORE_LOCK_RETRY_SLEEP: Duration = Duration::from_millis(50); +const REFRESH_TRANSACTION_TIMEOUT: Duration = Duration::from_secs(45); #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] pub struct StoredOAuthTokens { @@ -84,20 +97,67 @@ pub(crate) enum StoredOAuthTokenStatus { AuthorizationRequired, } +/// Concrete credential store resolved for one MCP OAuth client lifecycle. +/// +/// A client that loaded credentials from one backend must reread, refresh, and +/// persist through that same backend. Falling back mid-refresh can replay stale +/// rotating refresh tokens from another store. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum ResolvedOAuthCredentialStore { + File, + Keyring, +} + +#[derive(Debug)] +pub(crate) struct LoadedOAuthTokens { + pub(crate) tokens: StoredOAuthTokens, + pub(crate) store: ResolvedOAuthCredentialStore, +} + pub(crate) fn load_oauth_tokens( server_name: &str, url: &str, store_mode: OAuthCredentialsStoreMode, ) -> Result> { + Ok(load_oauth_tokens_with_source(server_name, url, store_mode)?.map(|loaded| loaded.tokens)) +} + +pub(crate) fn load_oauth_tokens_with_source( + server_name: &str, + url: &str, + store_mode: OAuthCredentialsStoreMode, +) -> Result> { let keyring_store = DefaultKeyringStore; + load_oauth_tokens_with_keyring_store(&keyring_store, server_name, url, store_mode) +} + +fn load_oauth_tokens_with_keyring_store( + keyring_store: &K, + server_name: &str, + url: &str, + store_mode: OAuthCredentialsStoreMode, +) -> Result> { match store_mode { OAuthCredentialsStoreMode::Auto => { - load_oauth_tokens_from_keyring_with_fallback_to_file(&keyring_store, server_name, url) + load_oauth_tokens_from_keyring_with_fallback_to_file_with_source( + keyring_store, + server_name, + url, + ) } - OAuthCredentialsStoreMode::File => load_oauth_tokens_from_file(server_name, url), + OAuthCredentialsStoreMode::File => Ok(load_oauth_tokens_from_file(server_name, url)?.map( + |tokens| LoadedOAuthTokens { + tokens, + store: ResolvedOAuthCredentialStore::File, + }, + )), OAuthCredentialsStoreMode::Keyring => { - load_oauth_tokens_from_keyring(&keyring_store, server_name, url) + Ok(load_oauth_tokens_from_keyring(keyring_store, server_name, url) .with_context(|| "failed to read OAuth tokens from keyring".to_string()) + .map(|tokens| LoadedOAuthTokens { + tokens, + store: ResolvedOAuthCredentialStore::Keyring, + })?) } } } @@ -158,22 +218,41 @@ fn load_oauth_tokens_from_keyring_with_fallback_to_file( server_name: &str, url: &str, ) -> Result> { + Ok(load_oauth_tokens_from_keyring_with_fallback_to_file_with_source( + keyring_store, + server_name, + url, + )? + .map(|loaded| loaded.tokens)) +} + +fn load_oauth_tokens_from_keyring_with_fallback_to_file_with_source( + keyring_store: &K, + server_name: &str, + url: &str, +) -> Result> { match load_oauth_tokens_from_keyring(keyring_store, server_name, url) { - Ok(Some(tokens)) => Ok(Some(tokens)), - Ok(None) => load_oauth_tokens_from_file_best_effort(server_name, url), + Ok(Some(tokens)) => Ok(Some(LoadedOAuthTokens { + tokens, + store: ResolvedOAuthCredentialStore::Keyring, + })), + Ok(None) => load_oauth_tokens_from_file_best_effort_with_source(server_name, url), Err(error) => { warn!("failed to read OAuth tokens from keyring: {error}"); - load_oauth_tokens_from_file_best_effort(server_name, url) + load_oauth_tokens_from_file_best_effort_with_source(server_name, url) } } } -fn load_oauth_tokens_from_file_best_effort( +fn load_oauth_tokens_from_file_best_effort_with_source( server_name: &str, url: &str, -) -> Result> { +) -> Result> { match load_oauth_tokens_from_file(server_name, url) { - Ok(tokens) => Ok(tokens), + Ok(tokens) => Ok(tokens.map(|tokens| LoadedOAuthTokens { + tokens, + store: ResolvedOAuthCredentialStore::File, + })), Err(error) => { warn!("failed to read OAuth tokens from fallback file: {error:?}"); Ok(None) @@ -218,6 +297,20 @@ pub fn save_oauth_tokens( } } +pub async fn save_oauth_tokens_locked( + server_name: &str, + tokens: &StoredOAuthTokens, + store_mode: OAuthCredentialsStoreMode, +) -> Result<()> { + let _lock = RefreshCredentialLock::acquire_for_server_with_timeout( + server_name, + &tokens.url, + LOCK_ACQUIRE_TIMEOUT, + ) + .await?; + save_oauth_tokens(server_name, tokens, store_mode) +} + fn save_oauth_tokens_with_keyring( keyring_store: &K, server_name: &str, @@ -269,6 +362,20 @@ pub fn delete_oauth_tokens( delete_oauth_tokens_from_keyring_and_file(&keyring_store, store_mode, server_name, url) } +pub async fn delete_oauth_tokens_locked( + server_name: &str, + url: &str, + store_mode: OAuthCredentialsStoreMode, +) -> Result { + let _lock = RefreshCredentialLock::acquire_for_server_with_timeout( + server_name, + url, + LOCK_ACQUIRE_TIMEOUT, + ) + .await?; + delete_oauth_tokens(server_name, url, store_mode) +} + fn delete_oauth_tokens_from_keyring_and_file( keyring_store: &K, store_mode: OAuthCredentialsStoreMode, @@ -305,7 +412,7 @@ struct OAuthPersistorInner { server_name: String, url: String, authorization_manager: Arc>, - store_mode: OAuthCredentialsStoreMode, + credential_store: ResolvedOAuthCredentialStore, last_credentials: Mutex>, } @@ -314,7 +421,7 @@ impl OAuthPersistor { server_name: String, url: String, authorization_manager: Arc>, - store_mode: OAuthCredentialsStoreMode, + credential_store: ResolvedOAuthCredentialStore, initial_credentials: Option, ) -> Self { Self { @@ -322,96 +429,423 @@ impl OAuthPersistor { server_name, url, authorization_manager, - store_mode, + credential_store, last_credentials: Mutex::new(initial_credentials), }), } } - /// Persists the latest stored credentials if they have changed. - /// Deletes the credentials if they are no longer present. - #[expect( - clippy::await_holding_invalid_type, - reason = "AuthorizationManager async access must be serialized through its mutex" - )] - pub(crate) async fn persist_if_needed(&self) -> Result<()> { - let (client_id, maybe_credentials) = { - let manager = self.inner.authorization_manager.clone(); - let guard = manager.lock().await; - guard.get_credentials().await - }?; - - match maybe_credentials { - Some(credentials) => { - let mut last_credentials = self.inner.last_credentials.lock().await; - let new_token_response = WrappedOAuthTokenResponse(credentials.clone()); - let same_token = last_credentials - .as_ref() - .map(|prev| prev.token_response == new_token_response) - .unwrap_or(false); - let expires_at = if same_token { - last_credentials.as_ref().and_then(|prev| prev.expires_at) - } else { - compute_expires_at_millis(&credentials) - }; - let stored = StoredOAuthTokens { - server_name: self.inner.server_name.clone(), - url: self.inner.url.clone(), - client_id, - token_response: new_token_response, - expires_at, - }; - if last_credentials.as_ref() != Some(&stored) { - save_oauth_tokens(&self.inner.server_name, &stored, self.inner.store_mode)?; - *last_credentials = Some(stored); - } + pub(crate) async fn refresh_if_needed(&self) -> Result<()> { + self.refresh_if_needed_with_timeout(REFRESH_TRANSACTION_TIMEOUT) + .await + } + + pub(crate) async fn refresh_if_needed_with_timeout( + &self, + overall_timeout: Duration, + ) -> Result<()> { + let expires_at = { + let guard = self.inner.last_credentials.lock().await; + guard.as_ref().and_then(|tokens| tokens.expires_at) + }; + + if !token_needs_refresh(expires_at) { + return Ok(()); + } + + self.refresh_transaction(RefreshReason::Expiry, overall_timeout) + .await + } + + pub(crate) async fn refresh_after_unauthorized(&self) -> Result<()> { + self.refresh_after_unauthorized_with_timeout(REFRESH_TRANSACTION_TIMEOUT) + .await + } + + pub(crate) async fn refresh_after_unauthorized_with_timeout( + &self, + overall_timeout: Duration, + ) -> Result<()> { + self.refresh_transaction(RefreshReason::Unauthorized, overall_timeout) + .await + } + + fn load_resolved_credentials(&self) -> Result> { + let keyring_store = DefaultKeyringStore; + match self.inner.credential_store { + ResolvedOAuthCredentialStore::File => { + load_oauth_tokens_from_file(&self.inner.server_name, &self.inner.url) + .context("failed to reread OAuth tokens from resolved file storage") } - None => { - let mut last_serialized = self.inner.last_credentials.lock().await; - if last_serialized.take().is_some() - && let Err(error) = delete_oauth_tokens( - &self.inner.server_name, - &self.inner.url, - self.inner.store_mode, - ) - { - warn!( - "failed to remove OAuth tokens for server {}: {error}", - self.inner.server_name - ); - } + ResolvedOAuthCredentialStore::Keyring => { + load_oauth_tokens_from_keyring( + &keyring_store, + &self.inner.server_name, + &self.inner.url, + ) + .context( + "failed to reread OAuth tokens from resolved keyring storage; refusing file fallback", + ) } } + } - Ok(()) + fn save_resolved_credentials(&self, tokens: &StoredOAuthTokens) -> Result<()> { + let keyring_store = DefaultKeyringStore; + match self.inner.credential_store { + ResolvedOAuthCredentialStore::File => save_oauth_tokens_to_file(tokens), + ResolvedOAuthCredentialStore::Keyring => { + save_oauth_tokens_with_keyring(&keyring_store, &self.inner.server_name, tokens) + } + } } #[expect( clippy::await_holding_invalid_type, reason = "AuthorizationManager async access must be serialized through its mutex" )] - pub(crate) async fn refresh_if_needed(&self) -> Result<()> { - let expires_at = { - let guard = self.inner.last_credentials.lock().await; - guard.as_ref().and_then(|tokens| tokens.expires_at) + async fn refresh_transaction( + &self, + reason: RefreshReason, + overall_timeout: Duration, + ) -> Result<()> { + let local_access_token = { + let last_credentials = self.inner.last_credentials.lock().await; + last_credentials + .as_ref() + .map(|tokens| tokens.token_response.0.access_token().secret().to_string()) }; - if !token_needs_refresh(expires_at) { + let deadline = Instant::now() + overall_timeout; + let lock_timeout = + remaining_refresh_timeout(deadline, overall_timeout, LOCK_ACQUIRE_TIMEOUT)?; + let _lock = RefreshCredentialLock::acquire_for_server_with_timeout( + &self.inner.server_name, + &self.inner.url, + lock_timeout, + ) + .await?; + + // Reread after acquiring the cross-process lock. If another Codex + // client refreshed the rotating token first, this process must adopt + // that newer state instead of replaying its stale refresh token. + let Some(latest) = self.load_resolved_credentials()? else { + self.clear_manager_credentials().await; + let mut last_credentials = self.inner.last_credentials.lock().await; + *last_credentials = None; + anyhow::bail!( + "OAuth tokens for server {} were removed before refresh; authorization required", + self.inner.server_name + ); + }; + + let latest_access_token = latest.token_response.0.access_token().secret(); + let should_adopt = !token_needs_refresh(latest.expires_at) + && match reason { + RefreshReason::Expiry => true, + RefreshReason::Unauthorized => { + local_access_token.as_deref() != Some(latest_access_token) + } + }; + if should_adopt { + self.adopt_credentials(latest).await?; return Ok(()); } + let manager = self.inner.authorization_manager.clone(); + let mut guard = manager.lock().await; + if let Err(error) = + install_tokens_in_manager_guard(&mut guard, &latest, CredentialExposure::Refresh).await { - let manager = self.inner.authorization_manager.clone(); - let guard = manager.lock().await; - guard.refresh_token().await.with_context(|| { + install_tokens_in_manager_guard(&mut guard, &latest, CredentialExposure::Request) + .await + .context("failed to restore request-only OAuth credentials")?; + return Err(error).context("failed to stage OAuth credentials for refresh"); + } + + let refresh_request_timeout = + match remaining_refresh_timeout(deadline, overall_timeout, REFRESH_TRANSACTION_TIMEOUT) + { + Ok(refresh_request_timeout) => refresh_request_timeout, + Err(error) => { + install_tokens_in_manager_guard( + &mut guard, + &latest, + CredentialExposure::Request, + ) + .await + .context("failed to restore request-only OAuth credentials")?; + return Err(error); + } + }; + + let refresh_result = match timeout(refresh_request_timeout, guard.refresh_token()).await { + Ok(Ok(token_response)) => Ok(refreshed_tokens(token_response, &latest, &self.inner)), + Ok(Err(error)) => Err(error).with_context(|| { format!( "failed to refresh OAuth tokens for server {}", self.inner.server_name ) + }), + Err(_) => Err(anyhow::anyhow!( + "timed out after {overall_timeout:?} refreshing OAuth tokens for server {}", + self.inner.server_name + )), + }; + + let request_tokens = refresh_result.as_ref().unwrap_or(&latest); + install_tokens_in_manager_guard(&mut guard, request_tokens, CredentialExposure::Request) + .await + .context("failed to restore request-only OAuth credentials")?; + drop(guard); + + let refreshed = refresh_result?; + self.save_resolved_credentials(&refreshed)?; + let mut last_credentials = self.inner.last_credentials.lock().await; + *last_credentials = Some(refreshed); + Ok(()) + } + + pub(crate) async fn adopt_credentials(&self, tokens: StoredOAuthTokens) -> Result<()> { + install_tokens_in_manager(&self.inner.authorization_manager, &tokens).await?; + let mut last_credentials = self.inner.last_credentials.lock().await; + *last_credentials = Some(tokens); + Ok(()) + } + + async fn clear_manager_credentials(&self) { + let manager = self.inner.authorization_manager.clone(); + let mut guard = manager.lock().await; + guard.set_credential_store(InMemoryCredentialStore::new()); + } +} + +#[derive(Clone, Copy)] +enum RefreshReason { + Expiry, + Unauthorized, +} + +fn remaining_refresh_timeout( + deadline: Instant, + overall_timeout: Duration, + maximum: Duration, +) -> Result { + let remaining = deadline.saturating_duration_since(Instant::now()); + if remaining.is_zero() { + anyhow::bail!("timed out after {overall_timeout:?} refreshing OAuth tokens"); + } + Ok(remaining.min(maximum)) +} + +struct RefreshCredentialLock { + _file: File, +} + +impl RefreshCredentialLock { + async fn acquire_for_server_with_timeout( + server_name: &str, + url: &str, + acquire_timeout: Duration, + ) -> Result { + let key = compute_store_key(server_name, url)?; + Self::acquire_with_timeout(&key, acquire_timeout) + .await + .with_context(|| format!("failed to acquire OAuth credential lock for {server_name}")) + } + + async fn acquire_with_timeout(store_key: &str, acquire_timeout: Duration) -> Result { + let path = refresh_lock_path(store_key)?; + if let Some(parent) = path.parent() { + fs::create_dir_all(parent)?; + } + + let file = OpenOptions::new() + .read(true) + .write(true) + .create(true) + .truncate(false) + .open(&path) + .with_context(|| format!("failed to open OAuth refresh lock {}", path.display()))?; + + match timeout(acquire_timeout, async { + loop { + match file.try_lock() { + Ok(()) => return Ok(()), + Err(std::fs::TryLockError::WouldBlock) => { + sleep(CREDENTIAL_LOCK_RETRY_SLEEP).await; + } + Err(error) => return Err(std::io::Error::from(error)), + } + } + }) + .await + { + Ok(Ok(())) => {} + Ok(Err(error)) => { + return Err(error).with_context(|| { + format!("failed to lock OAuth refresh lock {}", path.display()) + }); + } + Err(_) => anyhow::bail!( + "timed out after {acquire_timeout:?} waiting for OAuth refresh lock {}", + path.display() + ), + } + + Ok(Self { _file: file }) + } +} + +#[derive(Clone, Copy)] +enum OAuthStore { + File, +} + +impl OAuthStore { + fn lock_filename(self) -> &'static str { + match self { + Self::File => "file-store.lock", + } + } +} + +struct OAuthStoreLock { + _file: File, +} + +impl OAuthStoreLock { + fn acquire(store: OAuthStore) -> Result { + let path = oauth_store_lock_path(store)?; + if let Some(parent) = path.parent() { + fs::create_dir_all(parent)?; + } + + let file = OpenOptions::new() + .read(true) + .write(true) + .create(true) + .truncate(false) + .open(&path) + .with_context(|| { + format!("failed to open MCP OAuth store lock {}", path.display()) })?; + let started = Instant::now(); + + loop { + match file.try_lock() { + Ok(()) => return Ok(Self { _file: file }), + Err(std::fs::TryLockError::WouldBlock) + if started.elapsed() >= LOCK_ACQUIRE_TIMEOUT => + { + anyhow::bail!( + "timed out after {LOCK_ACQUIRE_TIMEOUT:?} waiting for MCP OAuth store lock {}", + path.display() + ); + } + Err(std::fs::TryLockError::WouldBlock) => { + std::thread::sleep(STORE_LOCK_RETRY_SLEEP); + } + Err(error) => { + return Err(std::io::Error::from(error)) + .with_context(|| format!("failed to lock MCP OAuth store lock {}", path.display())); + } + } } + } +} + +#[expect( + clippy::await_holding_invalid_type, + reason = "AuthorizationManager async access must be serialized through its mutex" +)] +async fn install_tokens_in_manager( + authorization_manager: &Arc>, + tokens: &StoredOAuthTokens, +) -> Result<()> { + let manager = authorization_manager.clone(); + let mut guard = manager.lock().await; + install_tokens_in_manager_guard(&mut guard, tokens, CredentialExposure::Request).await +} + +async fn install_tokens_in_manager_guard( + authorization_manager: &mut AuthorizationManager, + tokens: &StoredOAuthTokens, + exposure: CredentialExposure, +) -> Result<()> { + let store = InMemoryCredentialStore::new(); + store + .save(stored_credentials_from_tokens(tokens, exposure)) + .await + .context("failed to stage OAuth tokens for authorization manager")?; + + authorization_manager.set_credential_store(store); + authorization_manager + .initialize_from_store() + .await + .context("failed to adopt refreshed OAuth tokens")?; + Ok(()) +} + +#[derive(Clone, Copy)] +enum CredentialExposure { + Request, + Refresh, +} + +fn stored_credentials_from_tokens( + tokens: &StoredOAuthTokens, + exposure: CredentialExposure, +) -> StoredCredentials { + let token_response = match exposure { + CredentialExposure::Request => request_oauth_token_response(tokens), + CredentialExposure::Refresh => tokens.token_response.0.clone(), + }; + let granted_scopes = token_response + .scopes() + .map(|scopes| scopes.iter().map(|scope| scope.to_string()).collect()) + .unwrap_or_default(); + let token_received_at = match exposure { + CredentialExposure::Request => None, + CredentialExposure::Refresh => SystemTime::now() + .duration_since(UNIX_EPOCH) + .ok() + .map(|duration| duration.as_secs()), + }; + + StoredCredentials::new( + tokens.client_id.clone(), + Some(token_response), + granted_scopes, + token_received_at, + ) +} + +pub(crate) fn request_oauth_token_response(tokens: &StoredOAuthTokens) -> OAuthTokenResponse { + let mut token_response = tokens.token_response.0.clone(); + token_response.set_refresh_token(None); + token_response.set_expires_in(None); + token_response +} - self.persist_if_needed().await +fn refreshed_tokens( + mut token_response: OAuthTokenResponse, + previous: &StoredOAuthTokens, + inner: &OAuthPersistorInner, +) -> StoredOAuthTokens { + if token_response.refresh_token().is_none() { + token_response.set_refresh_token(previous.token_response.0.refresh_token().cloned()); + } + if token_response.scopes().is_none() { + token_response.set_scopes(previous.token_response.0.scopes().cloned()); + } + let expires_at = compute_expires_at_millis(&token_response); + StoredOAuthTokens { + server_name: inner.server_name.clone(), + url: inner.url.clone(), + client_id: previous.client_id.clone(), + token_response: WrappedOAuthTokenResponse(token_response), + expires_at, } } @@ -435,7 +869,8 @@ struct FallbackTokenEntry { } fn load_oauth_tokens_from_file(server_name: &str, url: &str) -> Result> { - let Some(store) = read_fallback_file()? else { + let _store_lock = OAuthStoreLock::acquire(OAuthStore::File)?; + let Some(store) = read_fallback_file_unlocked()? else { return Ok(None); }; @@ -478,8 +913,13 @@ fn load_oauth_tokens_from_file(server_name: &str, url: &str) -> Result Result<()> { + let _store_lock = OAuthStoreLock::acquire(OAuthStore::File)?; + save_oauth_tokens_to_file_unlocked(tokens) +} + +fn save_oauth_tokens_to_file_unlocked(tokens: &StoredOAuthTokens) -> Result<()> { let key = compute_store_key(&tokens.server_name, &tokens.url)?; - let mut store = read_fallback_file()?.unwrap_or_default(); + let mut store = read_fallback_file_unlocked()?.unwrap_or_default(); let token_response = &tokens.token_response.0; let expires_at = tokens @@ -507,7 +947,8 @@ fn save_oauth_tokens_to_file(tokens: &StoredOAuthTokens) -> Result<()> { } fn delete_oauth_tokens_from_file(key: &str) -> Result { - let mut store = match read_fallback_file()? { + let _store_lock = OAuthStoreLock::acquire(OAuthStore::File)?; + let mut store = match read_fallback_file_unlocked()? { Some(store) => store, None => return Ok(false), }; @@ -579,6 +1020,11 @@ fn fallback_file_path() -> Result { } fn read_fallback_file() -> Result> { + let _store_lock = OAuthStoreLock::acquire(OAuthStore::File)?; + read_fallback_file_unlocked() +} + +fn read_fallback_file_unlocked() -> Result> { let path = fallback_file_path()?; let contents = match fs::read_to_string(&path) { Ok(contents) => contents, @@ -604,6 +1050,23 @@ fn read_fallback_file() -> Result> { } } +fn refresh_lock_path(store_key: &str) -> Result { + let mut hasher = Sha256::new(); + hasher.update(store_key.as_bytes()); + let digest = hasher.finalize(); + Ok(find_codex_home()? + .join(REFRESH_LOCK_DIR) + .join(format!("{digest:x}.lock")) + .to_path_buf()) +} + +fn oauth_store_lock_path(store: OAuthStore) -> Result { + Ok(find_codex_home()? + .join(REFRESH_LOCK_DIR) + .join(store.lock_filename()) + .to_path_buf()) +} + fn write_fallback_file(store: &FallbackFile) -> Result<()> { let path = fallback_file_path()?; @@ -735,6 +1198,7 @@ mod tests { use std::sync::OnceLock; use std::sync::PoisonError; use tempfile::tempdir; + use tokio::time; use codex_keyring_store::tests::MockKeyringStore; @@ -1086,6 +1550,56 @@ mod tests { assert!(!super::oauth_tokens_are_usable(&tokens)); } + #[test] + fn request_oauth_token_response_strips_refresh_material() { + let tokens = sample_tokens(); + let request_tokens = super::request_oauth_token_response(&tokens); + + assert_eq!( + request_tokens.access_token().secret(), + tokens.token_response.0.access_token().secret() + ); + assert!(request_tokens.refresh_token().is_none()); + assert!(request_tokens.expires_in().is_none()); + assert_eq!(request_tokens.scopes(), tokens.token_response.0.scopes()); + } + + #[tokio::test(flavor = "current_thread")] + async fn locked_save_waits_for_refresh_transaction_lock() -> Result<()> { + let _env = TempCodexHome::new(); + let tokens = sample_tokens(); + let key = super::compute_store_key(&tokens.server_name, &tokens.url)?; + let lock = super::RefreshCredentialLock::acquire_for_server_with_timeout( + &tokens.server_name, + &tokens.url, + Duration::from_secs(/*secs*/ 1), + ) + .await?; + + let tokens_for_save = tokens.clone(); + let save = tokio::spawn(async move { + super::save_oauth_tokens_locked( + &tokens_for_save.server_name, + &tokens_for_save, + OAuthCredentialsStoreMode::File, + ) + .await + }); + + time::sleep(Duration::from_millis(/*millis*/ 50)).await; + assert!( + !super::fallback_file_path()?.exists(), + "login save should wait for the refresh transaction lock" + ); + + drop(lock); + save.await??; + + let saved = super::read_fallback_file()?.expect("fallback file should load"); + assert!(saved.contains_key(&key)); + Ok(()) + } + fn assert_tokens_match_without_expiry( actual: &StoredOAuthTokens, expected: &StoredOAuthTokens, diff --git a/codex-rs/rmcp-client/src/perform_oauth_device_login.rs b/codex-rs/rmcp-client/src/perform_oauth_device_login.rs index 306701cc935a..f6e2e2cb99e0 100644 --- a/codex-rs/rmcp-client/src/perform_oauth_device_login.rs +++ b/codex-rs/rmcp-client/src/perform_oauth_device_login.rs @@ -19,7 +19,7 @@ use crate::StoredOAuthTokens; use crate::WrappedOAuthTokenResponse; use crate::oauth::compute_expires_at_millis; use crate::perform_oauth_login::OAuthProviderError; -use crate::save_oauth_tokens; +use crate::save_oauth_tokens_locked; use crate::utils::build_default_headers; use crate::utils::build_reqwest_client; use codex_config::types::OAuthCredentialsStoreMode; @@ -109,7 +109,7 @@ pub async fn perform_oauth_device_login( token_response: WrappedOAuthTokenResponse(token_response), expires_at, }; - save_oauth_tokens(server_name, &stored, store_mode) + save_oauth_tokens_locked(server_name, &stored, store_mode).await } async fn resolve_device_oauth_client_id( diff --git a/codex-rs/rmcp-client/src/perform_oauth_login.rs b/codex-rs/rmcp-client/src/perform_oauth_login.rs index 9ec5394293a2..1de1925e5bb9 100644 --- a/codex-rs/rmcp-client/src/perform_oauth_login.rs +++ b/codex-rs/rmcp-client/src/perform_oauth_login.rs @@ -26,7 +26,7 @@ use urlencoding::decode; use crate::StoredOAuthTokens; use crate::WrappedOAuthTokenResponse; use crate::oauth::compute_expires_at_millis; -use crate::save_oauth_tokens; +use crate::save_oauth_tokens_locked; use crate::utils::build_default_headers; use crate::utils::build_reqwest_client; use codex_config::types::OAuthCredentialsStoreMode; @@ -571,7 +571,7 @@ impl OauthLoginFlow { token_response: WrappedOAuthTokenResponse(credentials), expires_at, }; - save_oauth_tokens(&self.server_name, &stored, self.store_mode)?; + save_oauth_tokens_locked(&self.server_name, &stored, self.store_mode).await?; Ok(()) } diff --git a/codex-rs/rmcp-client/src/rmcp_client.rs b/codex-rs/rmcp-client/src/rmcp_client.rs index ee60e793fa3d..ba5ee05197e8 100644 --- a/codex-rs/rmcp-client/src/rmcp_client.rs +++ b/codex-rs/rmcp-client/src/rmcp_client.rs @@ -63,8 +63,10 @@ use crate::elicitation_client_service::ElicitationClientService; use crate::http_client_adapter::StreamableHttpClientAdapter; use crate::http_client_adapter::StreamableHttpClientAdapterError; use crate::in_process_transport::InProcessTransportFactory; -use crate::load_oauth_tokens; +use crate::load_oauth_tokens_with_source; +use crate::oauth::LoadedOAuthTokens; use crate::oauth::OAuthPersistor; +use crate::oauth::ResolvedOAuthCredentialStore; use crate::oauth::StoredOAuthTokens; use crate::stdio_server_launcher::StdioServerCommand; use crate::stdio_server_launcher::StdioServerLauncher; @@ -79,6 +81,8 @@ mod streamable_http_retry; use self::streamable_http_retry::HandshakeError; use self::streamable_http_retry::STREAMABLE_HTTP_RETRY_DELAYS_MS; +use self::streamable_http_retry::initialize_timeout_error; +use self::streamable_http_retry::remaining_initialize_timeout; use self::streamable_http_retry::sleep_with_retry_deadline; enum PendingTransport { @@ -430,7 +434,7 @@ impl RmcpClient { }; let (service, oauth_persistor) = self - .connect_pending_transport_with_initialize_retries( + .connect_pending_transport_with_oauth_recovery( pending_transport, client_service.clone(), timeout, @@ -458,16 +462,10 @@ impl RmcpClient { } *guard = ClientState::Ready { service, - oauth: oauth_persistor.clone(), + oauth: oauth_persistor, }; } - if let Some(runtime) = oauth_persistor - && let Err(error) = runtime.persist_if_needed().await - { - warn!("failed to persist OAuth tokens after initialize: {error}"); - } - Ok(initialize_result) } @@ -476,14 +474,13 @@ impl RmcpClient { params: Option, timeout: Option, ) -> Result { - self.refresh_oauth_if_needed().await; + self.refresh_oauth_if_needed().await?; let result = self .run_service_operation("tools/list", timeout, move |service| { let params = params.clone(); async move { service.list_tools(params).await }.boxed() }) .await?; - self.persist_oauth_tokens().await; Ok(result) } @@ -492,7 +489,7 @@ impl RmcpClient { params: Option, timeout: Option, ) -> Result { - self.refresh_oauth_if_needed().await; + self.refresh_oauth_if_needed().await?; let result = self .run_service_operation("tools/list", timeout, move |service| { let params = params.clone(); @@ -517,7 +514,6 @@ impl RmcpClient { }) }) .collect::>>()?; - self.persist_oauth_tokens().await; Ok(ListToolsWithConnectorIdResult { next_cursor: result.next_cursor, tools, @@ -537,14 +533,13 @@ impl RmcpClient { params: Option, timeout: Option, ) -> Result { - self.refresh_oauth_if_needed().await; + self.refresh_oauth_if_needed().await?; let result = self .run_service_operation("resources/list", timeout, move |service| { let params = params.clone(); async move { service.list_resources(params).await }.boxed() }) .await?; - self.persist_oauth_tokens().await; Ok(result) } @@ -553,14 +548,13 @@ impl RmcpClient { params: Option, timeout: Option, ) -> Result { - self.refresh_oauth_if_needed().await; + self.refresh_oauth_if_needed().await?; let result = self .run_service_operation("resources/templates/list", timeout, move |service| { let params = params.clone(); async move { service.list_resource_templates(params).await }.boxed() }) .await?; - self.persist_oauth_tokens().await; Ok(result) } @@ -569,14 +563,13 @@ impl RmcpClient { params: ReadResourceRequestParams, timeout: Option, ) -> Result { - self.refresh_oauth_if_needed().await; + self.refresh_oauth_if_needed().await?; let result = self .run_service_operation("resources/read", timeout, move |service| { let params = params.clone(); async move { service.read_resource(params).await }.boxed() }) .await?; - self.persist_oauth_tokens().await; Ok(result) } @@ -587,7 +580,7 @@ impl RmcpClient { meta: Option, timeout: Option, ) -> Result { - self.refresh_oauth_if_needed().await; + self.refresh_oauth_if_needed().await?; let arguments = match arguments { Some(Value::Object(map)) => Some(map), Some(other) => { @@ -636,7 +629,6 @@ impl RmcpClient { .boxed() }) .await?; - self.persist_oauth_tokens().await; Ok(result) } @@ -645,7 +637,7 @@ impl RmcpClient { method: &str, params: Option, ) -> Result<()> { - self.refresh_oauth_if_needed().await; + self.refresh_oauth_if_needed().await?; self.run_service_operation( "notifications/custom", /*timeout*/ None, @@ -666,7 +658,6 @@ impl RmcpClient { }, ) .await?; - self.persist_oauth_tokens().await; Ok(()) } @@ -675,7 +666,7 @@ impl RmcpClient { method: &str, params: Option, ) -> Result { - self.refresh_oauth_if_needed().await; + self.refresh_oauth_if_needed().await?; let response = self .run_service_operation("requests/custom", /*timeout*/ None, move |service| { let params = params.clone(); @@ -689,7 +680,6 @@ impl RmcpClient { .boxed() }) .await?; - self.persist_oauth_tokens().await; Ok(response) } @@ -729,22 +719,11 @@ impl RmcpClient { drop(previous_state); } - /// This should be called after every tool call so that if a given tool call triggered - /// a refresh of the OAuth tokens, they are persisted. - async fn persist_oauth_tokens(&self) { - if let Some(runtime) = self.oauth_persistor().await - && let Err(error) = runtime.persist_if_needed().await - { - warn!("failed to persist OAuth tokens: {error}"); - } - } - - async fn refresh_oauth_if_needed(&self) { - if let Some(runtime) = self.oauth_persistor().await - && let Err(error) = runtime.refresh_if_needed().await - { - warn!("failed to refresh OAuth tokens: {error}"); + async fn refresh_oauth_if_needed(&self) -> Result<()> { + if let Some(runtime) = self.oauth_persistor().await { + runtime.refresh_if_needed().await?; } + Ok(()) } async fn create_pending_transport( @@ -776,7 +755,7 @@ impl RmcpClient { && auth_provider.is_none() && !default_headers.contains_key(AUTHORIZATION) { - match load_oauth_tokens(server_name, url, *store_mode) { + match load_oauth_tokens_with_source(server_name, url, *store_mode) { Ok(tokens) => tokens, Err(err) => { warn!("failed to read tokens for server `{server_name}`: {err}"); @@ -787,12 +766,16 @@ impl RmcpClient { None }; - if let Some(initial_tokens) = initial_oauth_tokens.clone() { + if let Some(LoadedOAuthTokens { + tokens: initial_tokens, + store: credential_store, + }) = initial_oauth_tokens + { match create_oauth_transport_and_runtime( server_name, url, initial_tokens.clone(), - *store_mode, + credential_store, default_headers.clone(), Arc::clone(http_client), ) @@ -862,6 +845,7 @@ impl RmcpClient { Arc>, Option, )> { + let deadline = timeout.map(|duration| Instant::now() + duration); let (transport, oauth_persistor) = match pending_transport { PendingTransport::InProcess { transport } => ( service::serve_client(client_service, transport).boxed(), @@ -878,13 +862,30 @@ impl RmcpClient { PendingTransport::StreamableHttpWithOAuth { transport, oauth_persistor, - } => ( - service::serve_client(client_service, transport).boxed(), - Some(oauth_persistor), - ), + } => { + match remaining_initialize_timeout(timeout, deadline)? { + Some(remaining) => { + if let Err(error) = oauth_persistor + .refresh_if_needed_with_timeout(remaining) + .await + { + if remaining_initialize_timeout(timeout, deadline).is_err() { + return Err(initialize_timeout_error(timeout, remaining)); + } + return Err(error); + } + } + None => oauth_persistor.refresh_if_needed().await?, + } + ( + service::serve_client(client_service, transport).boxed(), + Some(oauth_persistor), + ) + } }; - let service_result = match timeout { + let handshake_timeout = remaining_initialize_timeout(timeout, deadline)?; + let service_result = match handshake_timeout { Some(duration) => match time::timeout(duration, transport).await { Ok(result) => { result.map_err(|source| anyhow::Error::from(HandshakeError { source })) @@ -897,19 +898,7 @@ impl RmcpClient { .await .map_err(|source| anyhow::Error::from(HandshakeError { source })), }; - let service = match service_result { - Ok(service) => service, - Err(error) => { - if let Some(runtime) = oauth_persistor.as_ref() - && let Err(persist_error) = runtime.persist_if_needed().await - { - warn!( - "failed to persist OAuth tokens after failed initialize: {persist_error}" - ); - } - return Err(error); - } - }; + let service = service_result?; Ok((Arc::new(service), oauth_persistor)) } @@ -925,32 +914,46 @@ impl RmcpClient { Fut: std::future::Future>, { let service = self.service().await?; - match Self::run_service_operation_with_transient_retries( + let mut result = Self::run_service_operation_with_transient_retries( Arc::clone(&service), label, timeout, self.elicitation_pause_state.clone(), &operation, ) - .await + .await; + + if result + .as_ref() + .is_err_and(Self::is_unauthorized_operation_error) + && let Some(oauth_persistor) = self.oauth_persistor().await { - Ok(result) => Ok(result), - Err(error) if Self::is_session_expired_404(&error) => { - self.reinitialize_after_recoverable_transport_error(&service) - .await?; - let recovered_service = self.service().await?; - Self::run_service_operation_with_transient_retries( - recovered_service, - label, - timeout, - self.elicitation_pause_state.clone(), - &operation, - ) - .await - .map_err(Into::into) - } - Err(error) => Err(error.into()), + oauth_persistor.refresh_after_unauthorized().await?; + result = Self::run_service_operation_with_transient_retries( + Arc::clone(&service), + label, + timeout, + self.elicitation_pause_state.clone(), + &operation, + ) + .await; } + + if result.as_ref().is_err_and(Self::is_session_expired_404) { + self.reinitialize_after_recoverable_transport_error(&service) + .await?; + let recovered_service = self.service().await?; + result = Self::run_service_operation_with_transient_retries( + recovered_service, + label, + timeout, + self.elicitation_pause_state.clone(), + &operation, + ) + .await; + } + + result.map_err(Into::into) } async fn run_service_operation_with_transient_retries( @@ -1065,12 +1068,24 @@ impl RmcpClient { error, StreamableHttpError::Client( StreamableHttpClientAdapterError::SessionExpired404 - ) | StreamableHttpError::Auth(AuthError::AuthorizationRequired) - | StreamableHttpError::AuthRequired(_) + ) ) }) } + fn is_unauthorized_operation_error(error: &ClientOperationError) -> bool { + let ClientOperationError::Service(rmcp::service::ServiceError::TransportSend(error)) = + error + else { + return false; + }; + + error + .error + .downcast_ref::>() + .is_some_and(Self::is_unauthorized_streamable_http_error) + } + async fn reinitialize_after_recoverable_transport_error( &self, failed_service: &Arc>, @@ -1105,7 +1120,7 @@ impl RmcpClient { .ok_or_else(|| anyhow!("MCP client cannot recover before initialize succeeds"))?; let pending_transport = Self::create_pending_transport(&self.transport_recipe).await?; let (service, oauth_persistor) = self - .connect_pending_transport_with_initialize_retries( + .connect_pending_transport_with_oauth_recovery( pending_transport, initialize_context.client_service, initialize_context.timeout, @@ -1119,16 +1134,10 @@ impl RmcpClient { } *guard = ClientState::Ready { service, - oauth: oauth_persistor.clone(), + oauth: oauth_persistor, }; } - if let Some(runtime) = oauth_persistor - && let Err(error) = runtime.persist_if_needed().await - { - warn!("failed to persist OAuth tokens after session recovery: {error}"); - } - Ok(()) } } @@ -1137,7 +1146,7 @@ async fn create_oauth_transport_and_runtime( server_name: &str, url: &str, initial_tokens: StoredOAuthTokens, - credentials_store: OAuthCredentialsStoreMode, + credential_store: ResolvedOAuthCredentialStore, default_headers: HeaderMap, http_client: Arc, ) -> Result<( @@ -1181,12 +1190,12 @@ async fn create_oauth_transport_and_runtime( server_name.to_string(), url.to_string(), auth_manager, - credentials_store, - Some(initial_tokens), + credential_store, + Some(initial_tokens.clone()), ); - if let Err(error) = runtime.refresh_if_needed().await { - warn!("failed to refresh OAuth tokens during transport creation: {error}"); - } + runtime + .adopt_credentials(initial_tokens) + .await?; Ok((transport, runtime)) } @@ -1231,20 +1240,22 @@ mod tests { } #[test] - fn oauth_authorization_required_reinitializes_mcp_client() { + fn oauth_authorization_required_refreshes_oauth() { let error = transport_send_error(StreamableHttpError::Auth(AuthError::AuthorizationRequired)); - assert!(RmcpClient::is_session_expired_404(&error)); + assert!(RmcpClient::is_unauthorized_operation_error(&error)); + assert!(!RmcpClient::is_session_expired_404(&error)); } #[test] - fn streamable_http_auth_required_reinitializes_mcp_client() { + fn streamable_http_auth_required_refreshes_oauth() { let error = transport_send_error(StreamableHttpError::AuthRequired( AuthRequiredError::new("Bearer resource_metadata=\"https://example.test\"".to_string()), )); - assert!(RmcpClient::is_session_expired_404(&error)); + assert!(RmcpClient::is_unauthorized_operation_error(&error)); + assert!(!RmcpClient::is_session_expired_404(&error)); } #[test] diff --git a/codex-rs/rmcp-client/src/streamable_http_retry.rs b/codex-rs/rmcp-client/src/streamable_http_retry.rs index 73da95de58ec..4ec219ce9d89 100644 --- a/codex-rs/rmcp-client/src/streamable_http_retry.rs +++ b/codex-rs/rmcp-client/src/streamable_http_retry.rs @@ -8,6 +8,7 @@ use codex_exec_server::ExecServerError; use reqwest::StatusCode; use rmcp::service::RoleClient; use rmcp::service::RunningService; +use rmcp::transport::auth::AuthError; use rmcp::transport::streamable_http_client::StreamableHttpError; use tokio::time; use tracing::warn; @@ -22,8 +23,13 @@ use super::RmcpClient; const JSON_RPC_INTERNAL_ERROR_CODE: i64 = -32603; pub(super) const STREAMABLE_HTTP_RETRY_DELAYS_MS: [u64; 2] = [250, 1_000]; +#[derive(Default)] +struct InitializeAttemptContext { + oauth_persistor: Option, +} + impl RmcpClient { - pub(super) async fn connect_pending_transport_with_initialize_retries( + pub(super) async fn connect_pending_transport_with_oauth_recovery( &self, initial_transport: PendingTransport, client_service: ElicitationClientService, @@ -31,6 +37,68 @@ impl RmcpClient { ) -> Result<( Arc>, Option, + )> { + let mut attempt_context = InitializeAttemptContext::default(); + let deadline = timeout.map(|duration| Instant::now() + duration); + match self + .connect_pending_transport_with_initialize_retries( + initial_transport, + client_service.clone(), + timeout, + &mut attempt_context, + ) + .await + { + Ok(result) => Ok(result), + Err(error) if Self::is_unauthorized_initialize_error(&error) => { + let Some(oauth_persistor) = attempt_context.oauth_persistor else { + return Err(error); + }; + let refresh_result = match remaining_initialize_timeout(timeout, deadline)? { + Some(remaining) => { + oauth_persistor + .refresh_after_unauthorized_with_timeout(remaining) + .await + } + None => oauth_persistor.refresh_after_unauthorized().await, + }; + if let Err(error) = refresh_result { + remaining_initialize_timeout(timeout, deadline)?; + return Err(error); + } + let remaining = remaining_initialize_timeout(timeout, deadline)?; + let transport = match remaining { + Some(remaining) => time::timeout( + remaining, + Self::create_pending_transport(&self.transport_recipe), + ) + .await + .map_err(|_| initialize_timeout_error(timeout, remaining))??, + None => Self::create_pending_transport(&self.transport_recipe).await?, + }; + let remaining = remaining_initialize_timeout(timeout, deadline)?; + let mut retry_context = InitializeAttemptContext::default(); + self.connect_pending_transport_with_initialize_retries( + transport, + client_service, + remaining, + &mut retry_context, + ) + .await + } + Err(error) => Err(error), + } + } + + async fn connect_pending_transport_with_initialize_retries( + &self, + initial_transport: PendingTransport, + client_service: ElicitationClientService, + timeout: Option, + attempt_context: &mut InitializeAttemptContext, + ) -> Result<( + Arc>, + Option, )> { let should_retry = match &initial_transport { PendingTransport::InProcess { .. } | PendingTransport::Stdio { .. } => false, @@ -62,6 +130,14 @@ impl RmcpClient { } } }; + attempt_context.oauth_persistor = match &transport { + PendingTransport::StreamableHttpWithOAuth { + oauth_persistor, .. + } => Some(oauth_persistor.clone()), + PendingTransport::InProcess { .. } + | PendingTransport::Stdio { .. } + | PendingTransport::StreamableHttp { .. } => None, + }; let attempt_timeout = remaining_initialize_timeout(timeout, retry_deadline)?; match Self::connect_pending_transport( @@ -109,6 +185,29 @@ impl RmcpClient { }) } + fn is_unauthorized_initialize_error(error: &anyhow::Error) -> bool { + error.chain().any(|source| { + source + .downcast_ref::() + .is_some_and(|error| Self::is_unauthorized_client_initialize_error(&error.source)) + || source + .downcast_ref::() + .is_some_and(Self::is_unauthorized_client_initialize_error) + }) + } + + fn is_unauthorized_client_initialize_error( + error: &rmcp::service::ClientInitializeError, + ) -> bool { + match error { + rmcp::service::ClientInitializeError::TransportError { error, .. } => error + .error + .downcast_ref::>() + .is_some_and(Self::is_unauthorized_streamable_http_error), + _ => false, + } + } + fn is_retryable_client_initialize_error(error: &rmcp::service::ClientInitializeError) -> bool { match error { rmcp::service::ClientInitializeError::TransportError { error, context } @@ -163,6 +262,19 @@ impl RmcpClient { _ => false, } } + + pub(super) fn is_unauthorized_streamable_http_error( + error: &StreamableHttpError, + ) -> bool { + match error { + StreamableHttpError::AuthRequired(_) + | StreamableHttpError::Auth(AuthError::AuthorizationRequired) => true, + StreamableHttpError::UnexpectedServerResponse(message) => { + message.starts_with("HTTP 401") + } + _ => false, + } + } } fn is_retryable_unexpected_server_response(message: &str) -> bool { @@ -194,7 +306,7 @@ fn is_retryable_http_status(status: StatusCode) -> bool { ) } -fn remaining_initialize_timeout( +pub(super) fn remaining_initialize_timeout( timeout: Option, deadline: Option, ) -> Result> { @@ -209,7 +321,10 @@ fn remaining_initialize_timeout( } } -fn initialize_timeout_error(timeout: Option, fallback: Duration) -> anyhow::Error { +pub(super) fn initialize_timeout_error( + timeout: Option, + fallback: Duration, +) -> anyhow::Error { let duration = timeout.unwrap_or(fallback); anyhow!("timed out handshaking with MCP server after {duration:?}") } diff --git a/codex-rs/rmcp-client/tests/streamable_http_oauth_startup.rs b/codex-rs/rmcp-client/tests/streamable_http_oauth_startup.rs index d7d8cf42323b..841e9a903664 100644 --- a/codex-rs/rmcp-client/tests/streamable_http_oauth_startup.rs +++ b/codex-rs/rmcp-client/tests/streamable_http_oauth_startup.rs @@ -176,7 +176,7 @@ async fn persisted_credentials_auth_status_child() -> anyhow::Result<()> { url: UNEXPIRED_SERVER_URL.to_string(), client_id: "test-client-id".to_string(), token_response: WrappedOAuthTokenResponse(response), - expires_at: Some(now.saturating_add(/*rhs*/ 60_000)), + expires_at: Some(now.saturating_add(/*rhs*/ 120_000)), }; save_oauth_tokens(SERVER_NAME, &tokens, OAuthCredentialsStoreMode::File)?; From 662b07225d8600fa697d229a7377e0f0afb3b065 Mon Sep 17 00:00:00 2001 From: GraciousGazelles <55840159+GraciousGazelles@users.noreply.github.com> Date: Sat, 20 Jun 2026 16:36:43 +1000 Subject: [PATCH 2/4] Fix OAuth refresh lock compile issues --- codex-rs/rmcp-client/src/oauth.rs | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/codex-rs/rmcp-client/src/oauth.rs b/codex-rs/rmcp-client/src/oauth.rs index d007cec83fe3..f1b544c2e23d 100644 --- a/codex-rs/rmcp-client/src/oauth.rs +++ b/codex-rs/rmcp-client/src/oauth.rs @@ -151,14 +151,14 @@ fn load_oauth_tokens_with_keyring_store( store: ResolvedOAuthCredentialStore::File, }, )), - OAuthCredentialsStoreMode::Keyring => { - Ok(load_oauth_tokens_from_keyring(keyring_store, server_name, url) + OAuthCredentialsStoreMode::Keyring => Ok(Some( + load_oauth_tokens_from_keyring(keyring_store, server_name, url) .with_context(|| "failed to read OAuth tokens from keyring".to_string()) .map(|tokens| LoadedOAuthTokens { tokens, store: ResolvedOAuthCredentialStore::Keyring, - })?) - } + })?, + )), } } @@ -1056,7 +1056,7 @@ fn refresh_lock_path(store_key: &str) -> Result { let digest = hasher.finalize(); Ok(find_codex_home()? .join(REFRESH_LOCK_DIR) - .join(format!("{digest:x}.lock")) + .join(format!("{}.lock", digest_hex(digest))) .to_path_buf()) } From 9dc2b7f54ed0d9f133dd4b7627a6f007452a752c Mon Sep 17 00:00:00 2001 From: GraciousGazelles <55840159+GraciousGazelles@users.noreply.github.com> Date: Sat, 20 Jun 2026 16:42:27 +1000 Subject: [PATCH 3/4] Handle missing keyring OAuth tokens --- codex-rs/rmcp-client/src/oauth.rs | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/codex-rs/rmcp-client/src/oauth.rs b/codex-rs/rmcp-client/src/oauth.rs index f1b544c2e23d..ec3d0ef0e0a4 100644 --- a/codex-rs/rmcp-client/src/oauth.rs +++ b/codex-rs/rmcp-client/src/oauth.rs @@ -151,14 +151,14 @@ fn load_oauth_tokens_with_keyring_store( store: ResolvedOAuthCredentialStore::File, }, )), - OAuthCredentialsStoreMode::Keyring => Ok(Some( - load_oauth_tokens_from_keyring(keyring_store, server_name, url) - .with_context(|| "failed to read OAuth tokens from keyring".to_string()) - .map(|tokens| LoadedOAuthTokens { - tokens, - store: ResolvedOAuthCredentialStore::Keyring, - })?, - )), + OAuthCredentialsStoreMode::Keyring => { + let tokens = load_oauth_tokens_from_keyring(keyring_store, server_name, url) + .with_context(|| "failed to read OAuth tokens from keyring".to_string())?; + Ok(tokens.map(|tokens| LoadedOAuthTokens { + tokens, + store: ResolvedOAuthCredentialStore::Keyring, + })) + } } } From 211b4d38197bdc8618718fbb9514e0d96bb7adcb Mon Sep 17 00:00:00 2001 From: GraciousGazelles <55840159+GraciousGazelles@users.noreply.github.com> Date: Sat, 20 Jun 2026 18:46:10 +1000 Subject: [PATCH 4/4] Fix MCP OAuth formatting --- codex-rs/rmcp-client/src/oauth.rs | 23 ++++++++++++----------- codex-rs/rmcp-client/src/rmcp_client.rs | 4 +--- 2 files changed, 13 insertions(+), 14 deletions(-) diff --git a/codex-rs/rmcp-client/src/oauth.rs b/codex-rs/rmcp-client/src/oauth.rs index ec3d0ef0e0a4..516e082cc5d6 100644 --- a/codex-rs/rmcp-client/src/oauth.rs +++ b/codex-rs/rmcp-client/src/oauth.rs @@ -218,12 +218,14 @@ fn load_oauth_tokens_from_keyring_with_fallback_to_file( server_name: &str, url: &str, ) -> Result> { - Ok(load_oauth_tokens_from_keyring_with_fallback_to_file_with_source( - keyring_store, - server_name, - url, - )? - .map(|loaded| loaded.tokens)) + Ok( + load_oauth_tokens_from_keyring_with_fallback_to_file_with_source( + keyring_store, + server_name, + url, + )? + .map(|loaded| loaded.tokens), + ) } fn load_oauth_tokens_from_keyring_with_fallback_to_file_with_source( @@ -727,9 +729,7 @@ impl OAuthStoreLock { .create(true) .truncate(false) .open(&path) - .with_context(|| { - format!("failed to open MCP OAuth store lock {}", path.display()) - })?; + .with_context(|| format!("failed to open MCP OAuth store lock {}", path.display()))?; let started = Instant::now(); loop { @@ -747,8 +747,9 @@ impl OAuthStoreLock { std::thread::sleep(STORE_LOCK_RETRY_SLEEP); } Err(error) => { - return Err(std::io::Error::from(error)) - .with_context(|| format!("failed to lock MCP OAuth store lock {}", path.display())); + return Err(std::io::Error::from(error)).with_context(|| { + format!("failed to lock MCP OAuth store lock {}", path.display()) + }); } } } diff --git a/codex-rs/rmcp-client/src/rmcp_client.rs b/codex-rs/rmcp-client/src/rmcp_client.rs index ba5ee05197e8..cde252083a16 100644 --- a/codex-rs/rmcp-client/src/rmcp_client.rs +++ b/codex-rs/rmcp-client/src/rmcp_client.rs @@ -1193,9 +1193,7 @@ async fn create_oauth_transport_and_runtime( credential_store, Some(initial_tokens.clone()), ); - runtime - .adopt_credentials(initial_tokens) - .await?; + runtime.adopt_credentials(initial_tokens).await?; Ok((transport, runtime)) }