diff --git a/codex-rs/config/src/types.rs b/codex-rs/config/src/types.rs index 5303fccd43b6..ccb647a7bf16 100644 --- a/codex-rs/config/src/types.rs +++ b/codex-rs/config/src/types.rs @@ -103,7 +103,10 @@ pub enum AuthCredentialsStoreMode { #[derive(Debug, Default, Copy, Clone, PartialEq, Eq, Serialize, Deserialize, JsonSchema)] #[serde(rename_all = "lowercase")] pub enum OAuthCredentialsStoreMode { - /// `Keyring` when available; otherwise, `File`. + /// Prefer `Keyring` and use `File` when keyring storage is unavailable. + /// Once an MCP client loads credentials from one store, that client keeps the resolved store + /// for its lifetime: refresh reads and writes surface store failures instead of switching to + /// the other store and risking replay of an older rotating refresh token. /// Credentials stored in the keyring will only be readable by Codex unless the user explicitly grants access via OS-level keyring access. #[default] Auto, diff --git a/codex-rs/core/config.schema.json b/codex-rs/core/config.schema.json index 512c050a20c1..b3c69d4e40c0 100644 --- a/codex-rs/core/config.schema.json +++ b/codex-rs/core/config.schema.json @@ -2120,7 +2120,7 @@ "description": "Determine where Codex should store and read MCP credentials.", "oneOf": [ { - "description": "`Keyring` when available; otherwise, `File`. Credentials stored in the keyring will only be readable by Codex unless the user explicitly grants access via OS-level keyring access.", + "description": "Prefer `Keyring` and use `File` when keyring storage is unavailable. Once an MCP client loads credentials from one store, that client keeps the resolved store for its lifetime: refresh reads and writes surface store failures instead of switching to the other store and risking replay of an older rotating refresh token. Credentials stored in the keyring will only be readable by Codex unless the user explicitly grants access via OS-level keyring access.", "enum": [ "auto" ], @@ -5519,4 +5519,4 @@ }, "title": "ConfigToml", "type": "object" -} \ No newline at end of file +} diff --git a/codex-rs/rmcp-client/src/lib.rs b/codex-rs/rmcp-client/src/lib.rs index e1041fdafbbf..5aae4493ce34 100644 --- a/codex-rs/rmcp-client/src/lib.rs +++ b/codex-rs/rmcp-client/src/lib.rs @@ -23,7 +23,6 @@ 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::save_oauth_tokens; pub use perform_oauth_login::OAuthProviderError; pub use perform_oauth_login::OauthLoginHandle; diff --git a/codex-rs/rmcp-client/src/oauth.rs b/codex-rs/rmcp-client/src/oauth.rs index d0005d20a6d0..9759981f59d4 100644 --- a/codex-rs/rmcp-client/src/oauth.rs +++ b/codex-rs/rmcp-client/src/oauth.rs @@ -16,6 +16,10 @@ //! //! If the keyring is not available or fails, we fall back to CODEX_HOME/.credentials.json which is consistent with other coding CLI agents. +mod persistor; +mod refresh_lock; +mod resolved_store; + use anyhow::Context; use anyhow::Error; use anyhow::Result; @@ -51,11 +55,14 @@ use tracing::warn; use codex_keyring_store::DefaultKeyringStore; use codex_keyring_store::KeyringStore; -use rmcp::transport::auth::AuthorizationManager; -use tokio::sync::Mutex; - use codex_utils_home_dir::find_codex_home; +pub(crate) use self::persistor::OAuthPersistor; +pub(crate) use self::resolved_store::ResolvedOAuthCredentialStore; +pub(crate) use self::resolved_store::ResolvedOAuthTokens; +pub(crate) use self::resolved_store::load_oauth_tokens_from_store; +pub(crate) use self::resolved_store::resolve_oauth_tokens; + const KEYRING_SERVICE: &str = "Codex MCP Credentials"; const MCP_OAUTH_SECRET_PREFIX: &str = "MCP_OAUTH"; const REFRESH_SKEW_MILLIS: u64 = 30_000; @@ -90,41 +97,24 @@ pub(crate) enum StoredOAuthTokenStatus { AuthorizationRequired, } -pub(crate) fn load_oauth_tokens( - server_name: &str, - url: &str, - store_mode: OAuthCredentialsStoreMode, - keyring_backend_kind: AuthKeyringBackendKind, -) -> Result> { - let keyring_store = DefaultKeyringStore; - match store_mode { - OAuthCredentialsStoreMode::Auto => load_oauth_tokens_from_keyring_with_fallback_to_file( - &keyring_store, - keyring_backend_kind, - server_name, - url, - ), - OAuthCredentialsStoreMode::File => load_oauth_tokens_from_file(server_name, url), - OAuthCredentialsStoreMode::Keyring => { - load_oauth_tokens_from_keyring(&keyring_store, keyring_backend_kind, server_name, url) - .with_context(|| "failed to read OAuth tokens from keyring".to_string()) - } - } -} - pub(crate) fn oauth_token_status( server_name: &str, url: &str, store_mode: OAuthCredentialsStoreMode, keyring_backend_kind: AuthKeyringBackendKind, ) -> Result { - Ok( - match load_oauth_tokens(server_name, url, store_mode, keyring_backend_kind)?.as_ref() { - None => StoredOAuthTokenStatus::Missing, - Some(tokens) if oauth_tokens_are_usable(tokens) => StoredOAuthTokenStatus::Usable, - Some(_) => StoredOAuthTokenStatus::AuthorizationRequired, - }, - ) + let resolved = resolve_oauth_tokens( + &DefaultKeyringStore, + server_name, + url, + store_mode, + keyring_backend_kind, + )?; + Ok(match resolved.as_ref().map(|resolved| &resolved.tokens) { + None => StoredOAuthTokenStatus::Missing, + Some(tokens) if oauth_tokens_are_usable(tokens) => StoredOAuthTokenStatus::Usable, + Some(_) => StoredOAuthTokenStatus::AuthorizationRequired, + }) } fn oauth_tokens_are_usable(tokens: &StoredOAuthTokens) -> bool { @@ -164,23 +154,6 @@ fn refresh_expires_in_from_timestamp(tokens: &mut StoredOAuthTokens) { } } -fn load_oauth_tokens_from_keyring_with_fallback_to_file( - keyring_store: &K, - keyring_backend_kind: AuthKeyringBackendKind, - server_name: &str, - url: &str, -) -> Result> { - match load_oauth_tokens_from_keyring(keyring_store, keyring_backend_kind, server_name, url) { - Ok(Some(tokens)) => Ok(Some(tokens)), - Ok(None) => load_oauth_tokens_from_file(server_name, url), - Err(error) => { - warn!("failed to read OAuth tokens from keyring: {error}"); - load_oauth_tokens_from_file(server_name, url) - .with_context(|| format!("failed to read OAuth tokens from keyring: {error}")) - } - } -} - fn load_oauth_tokens_from_keyring( keyring_store: &K, keyring_backend_kind: AuthKeyringBackendKind, @@ -249,16 +222,32 @@ pub fn save_oauth_tokens( keyring_backend_kind: AuthKeyringBackendKind, ) -> Result<()> { let keyring_store = DefaultKeyringStore; + save_oauth_tokens_with_keyring_store( + &keyring_store, + server_name, + tokens, + store_mode, + keyring_backend_kind, + ) +} + +fn save_oauth_tokens_with_keyring_store( + keyring_store: &K, + server_name: &str, + tokens: &StoredOAuthTokens, + store_mode: OAuthCredentialsStoreMode, + keyring_backend_kind: AuthKeyringBackendKind, +) -> Result<()> { match store_mode { OAuthCredentialsStoreMode::Auto => save_oauth_tokens_with_keyring_with_fallback_to_file( - &keyring_store, + keyring_store, keyring_backend_kind, server_name, tokens, ), OAuthCredentialsStoreMode::File => save_oauth_tokens_to_file(tokens), - OAuthCredentialsStoreMode::Keyring => save_oauth_tokens_with_keyring( - &keyring_store, + OAuthCredentialsStoreMode::Keyring => save_oauth_tokens_with_keyring_and_cleanup_file( + keyring_store, keyring_backend_kind, server_name, tokens, @@ -282,6 +271,24 @@ fn save_oauth_tokens_with_keyring( } } +/// Saves to the selected keyring backend, then best-effort removes the fallback file entry. +/// +/// A cleanup failure does not change the current client's selected authority, but it can leave +/// legacy residue that a different `Auto` process may discover if keyring availability changes. +fn save_oauth_tokens_with_keyring_and_cleanup_file( + keyring_store: &K, + keyring_backend_kind: AuthKeyringBackendKind, + server_name: &str, + tokens: &StoredOAuthTokens, +) -> Result<()> { + save_oauth_tokens_with_keyring(keyring_store, keyring_backend_kind, server_name, tokens)?; + let key = compute_store_key(server_name, &tokens.url)?; + if let Err(error) = delete_oauth_tokens_from_file(&key) { + warn!("failed to remove OAuth tokens from fallback storage: {error:?}"); + } + Ok(()) +} + fn save_oauth_tokens_to_direct_keyring( keyring_store: &K, server_name: &str, @@ -291,12 +298,7 @@ fn save_oauth_tokens_to_direct_keyring( let key = compute_store_key(server_name, &tokens.url)?; match keyring_store.save(KEYRING_SERVICE, &key, &serialized) { - Ok(()) => { - if let Err(error) = delete_oauth_tokens_from_file(&key) { - warn!("failed to remove OAuth tokens from fallback storage: {error:?}"); - } - Ok(()) - } + Ok(()) => Ok(()), Err(error) => { let message = format!( "failed to write OAuth tokens to keyring: {}", @@ -324,13 +326,7 @@ fn save_oauth_tokens_to_secrets_keyring( let secret_name = compute_secret_name(server_name, &tokens.url)?; manager .set(&SecretScope::Global, &secret_name, &serialized) - .context("failed to write OAuth tokens to encrypted storage")?; - - let key = compute_store_key(server_name, &tokens.url)?; - if let Err(error) = delete_oauth_tokens_from_file(&key) { - warn!("failed to remove OAuth tokens from fallback storage: {error:?}"); - } - Ok(()) + .context("failed to write OAuth tokens to encrypted storage") } fn save_oauth_tokens_with_keyring_with_fallback_to_file( @@ -339,7 +335,12 @@ fn save_oauth_tokens_with_keyring_with_fallback_to_file Result<()> { - match save_oauth_tokens_with_keyring(keyring_store, keyring_backend_kind, server_name, tokens) { + match save_oauth_tokens_with_keyring_and_cleanup_file( + keyring_store, + keyring_backend_kind, + server_name, + tokens, + ) { Ok(()) => Ok(()), Err(error) => { let message = error.to_string(); @@ -444,134 +445,6 @@ fn delete_oauth_tokens_from_secrets_keyring( Ok(secrets_removed) } -#[derive(Clone)] -pub(crate) struct OAuthPersistor { - inner: Arc, -} - -struct OAuthPersistorInner { - server_name: String, - url: String, - authorization_manager: Arc>, - store_mode: OAuthCredentialsStoreMode, - keyring_backend_kind: AuthKeyringBackendKind, - last_credentials: Mutex>, -} - -impl OAuthPersistor { - pub(crate) fn new( - server_name: String, - url: String, - authorization_manager: Arc>, - store_mode: OAuthCredentialsStoreMode, - keyring_backend_kind: AuthKeyringBackendKind, - initial_credentials: Option, - ) -> Self { - Self { - inner: Arc::new(OAuthPersistorInner { - server_name, - url, - authorization_manager, - store_mode, - keyring_backend_kind, - 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, - self.inner.keyring_backend_kind, - )?; - *last_credentials = Some(stored); - } - } - 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, - self.inner.keyring_backend_kind, - ) - { - warn!( - "failed to remove OAuth tokens for server {}: {error}", - self.inner.server_name - ); - } - } - } - - Ok(()) - } - - #[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) - }; - - if !token_needs_refresh(expires_at) { - return Ok(()); - } - - { - let manager = self.inner.authorization_manager.clone(); - let guard = manager.lock().await; - guard.refresh_token().await.with_context(|| { - format!( - "failed to refresh OAuth tokens for server {}", - self.inner.server_name - ) - })?; - } - - self.persist_if_needed().await - } -} - const FALLBACK_FILENAME: &str = ".credentials.json"; const MCP_SERVER_TYPE: &str = "http"; @@ -891,14 +764,15 @@ mod tests { super::save_oauth_tokens_to_file(&tokens)?; - let loaded = super::load_oauth_tokens_from_keyring_with_fallback_to_file( + let resolved = super::resolve_oauth_tokens( &store, - AuthKeyringBackendKind::Direct, &tokens.server_name, &tokens.url, + OAuthCredentialsStoreMode::Auto, + AuthKeyringBackendKind::Direct, )? .expect("tokens should load from fallback"); - assert_tokens_match_without_expiry(&loaded, &expected); + assert_tokens_match_without_expiry(&resolved.tokens, &expected); Ok(()) } @@ -913,14 +787,15 @@ mod tests { super::save_oauth_tokens_to_file(&tokens)?; - let loaded = super::load_oauth_tokens_from_keyring_with_fallback_to_file( + let resolved = super::resolve_oauth_tokens( &store, - AuthKeyringBackendKind::Direct, &tokens.server_name, &tokens.url, + OAuthCredentialsStoreMode::Auto, + AuthKeyringBackendKind::Direct, )? .expect("tokens should load from fallback"); - assert_tokens_match_without_expiry(&loaded, &expected); + assert_tokens_match_without_expiry(&resolved.tokens, &expected); Ok(()) } diff --git a/codex-rs/rmcp-client/src/oauth/persistor.rs b/codex-rs/rmcp-client/src/oauth/persistor.rs new file mode 100644 index 000000000000..7fc54198e794 --- /dev/null +++ b/codex-rs/rmcp-client/src/oauth/persistor.rs @@ -0,0 +1,485 @@ +//! Lifecycle-local persistence and serialized refresh transactions for MCP OAuth credentials. + +use std::sync::Arc; +use std::time::Duration; +use std::time::Instant; +use std::time::SystemTime; +use std::time::UNIX_EPOCH; + +use anyhow::Context; +use anyhow::Result; +use codex_keyring_store::DefaultKeyringStore; +use codex_keyring_store::KeyringStore; +use oauth2::AccessToken; +use oauth2::TokenResponse; +use rmcp::transport::auth::AuthorizationManager; +use rmcp::transport::auth::CredentialStore as _; +use rmcp::transport::auth::InMemoryCredentialStore; +use rmcp::transport::auth::OAuthTokenResponse; +use rmcp::transport::auth::StoredCredentials; +use tokio::sync::Mutex; +use tokio::time::timeout; +use tracing::debug; +use tracing::warn; + +use super::ResolvedOAuthCredentialStore; +use super::StoredOAuthTokens; +use super::WrappedOAuthTokenResponse; +use super::compute_expires_at_millis; +use super::load_oauth_tokens_from_store; +use super::refresh_lock::RefreshCredentialLock; +use super::save_oauth_tokens_to_file; +use super::save_oauth_tokens_with_keyring; +use super::token_needs_refresh; + +const REFRESH_REQUEST_TIMEOUT: Duration = Duration::from_secs(45); + +#[derive(Clone)] +pub(crate) struct OAuthPersistor { + inner: Arc, +} + +struct OAuthPersistorInner { + server_name: String, + url: String, + authorization_manager: Arc>, + credential_store: ResolvedOAuthCredentialStore, + current_credentials: Mutex>, +} + +impl OAuthPersistor { + pub(crate) fn new( + server_name: String, + url: String, + authorization_manager: Arc>, + credential_store: ResolvedOAuthCredentialStore, + initial_credentials: Option, + ) -> Self { + Self { + inner: Arc::new(OAuthPersistorInner { + server_name, + url, + authorization_manager, + credential_store, + current_credentials: Mutex::new(initial_credentials), + }), + } + } + + pub(crate) async fn refresh_if_needed(&self) -> Result<()> { + self.refresh_if_needed_with_keyring_store(&DefaultKeyringStore) + .await + } + + #[expect(dead_code, reason = "wired by part 4 of this stack")] + pub(crate) async fn refresh_after_unauthorized( + &self, + rejected_access_token: AccessToken, + ) -> Result<()> { + self.refresh_after_unauthorized_with_keyring_store( + &DefaultKeyringStore, + rejected_access_token, + ) + .await + } + + pub(super) async fn refresh_if_needed_with_keyring_store( + &self, + keyring_store: &K, + ) -> Result<()> { + self.refresh_if_needed_with_keyring_store_and_timeout( + keyring_store, + REFRESH_REQUEST_TIMEOUT, + ) + .await + } + + pub(super) async fn refresh_if_needed_with_keyring_store_and_timeout< + K: KeyringStore + Clone + 'static, + >( + &self, + keyring_store: &K, + refresh_request_timeout: Duration, + ) -> Result<()> { + let expires_at = { + let guard = self.inner.current_credentials.lock().await; + guard.as_ref().and_then(|tokens| tokens.expires_at) + }; + + if !token_needs_refresh(expires_at) { + return Ok(()); + } + + self.run_owned_refresh_transaction( + keyring_store.clone(), + RefreshReason::Expiry, + refresh_request_timeout, + ) + .await + } + + pub(super) async fn refresh_after_unauthorized_with_keyring_store< + K: KeyringStore + Clone + 'static, + >( + &self, + keyring_store: &K, + rejected_access_token: AccessToken, + ) -> Result<()> { + self.run_owned_refresh_transaction( + keyring_store.clone(), + RefreshReason::Unauthorized { + rejected_access_token, + }, + REFRESH_REQUEST_TIMEOUT, + ) + .await + } + + async fn run_owned_refresh_transaction( + &self, + keyring_store: K, + reason: RefreshReason, + refresh_request_timeout: Duration, + ) -> Result<()> { + let persistor = self.clone(); + let server_name = self.inner.server_name.clone(); + // Once a provider request can consume a rotating refresh token, dropping the caller's + // future must not also drop the refresh-and-persist transaction. Dropping this JoinHandle + // detaches the task, so it continues while the provider request remains bounded by + // `refresh_request_timeout` and the credential lock remains bounded by its own timeout. + // + // A provider timeout deliberately leaves the outcome unknown, releases the lock, and + // permits a later serialized retry. Some providers accept the previous token during a + // grace period; otherwise that retry surfaces reauthorization. We accept that residual + // token-family-revocation risk rather than holding the lock indefinitely. + let refresh_reason = reason.as_str(); + tokio::spawn(async move { + let result = persistor + .refresh_transaction(&keyring_store, reason, refresh_request_timeout) + .await; + + // Keep this summary inside the owned task so caller cancellation cannot suppress it. + if let Err(error) = &result { + warn!( + server_name = %persistor.inner.server_name, + refresh_reason, + error = %error, + "MCP OAuth refresh transaction failed" + ); + } + + result + }) + .await + .with_context(|| format!("OAuth refresh task failed for server {server_name}"))? + } + + #[expect( + clippy::await_holding_invalid_type, + reason = "AuthorizationManager async access must be serialized through its mutex" + )] + #[tracing::instrument( + level = "debug", + skip_all, + fields( + server_name = %self.inner.server_name, + refresh_reason = reason.as_str(), + ), + err + )] + async fn refresh_transaction( + &self, + keyring_store: &K, + reason: RefreshReason, + refresh_request_timeout: Duration, + ) -> Result<()> { + let transaction_started_at = Instant::now(); + let lock_started_at = Instant::now(); + debug!("waiting for the MCP OAuth credential transaction lock"); + let _lock = + RefreshCredentialLock::acquire_for_server(&self.inner.server_name, &self.inner.url) + .await?; + debug!( + lock_wait_ms = lock_started_at.elapsed().as_millis(), + "acquired the MCP OAuth credential transaction lock" + ); + + // The refresh transaction must stay on the store that supplied its snapshot. Falling back + // here could replay an older rotating refresh token from the other store. We assume store + // availability is stable for this client lifecycle and surface violations of that + // assumption instead of switching stores. + debug!("rereading authoritative MCP OAuth credentials"); + let latest = load_oauth_tokens_from_store( + keyring_store, + &self.inner.server_name, + &self.inner.url, + self.inner.credential_store, + )?; + + // The pre-lock snapshot only decides whether a refresh transaction might be needed. Once + // the lock is held, this reread is authoritative: adopt it before deciding whether to + // refresh so this process never sends a refresh token superseded by another process. + let Some(latest) = latest else { + self.clear_manager_credentials().await; + *self.inner.current_credentials.lock().await = 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(); + // Expiry refresh can adopt any reread that is now healthy. A 401 belongs to the access + // token sent with that specific request, not to this client's mutable current snapshot. + // If a delayed request rejected A after another request refreshed A to B, adopt B and let + // the caller retry instead of rotating B again. + let should_adopt = !token_needs_refresh(latest.expires_at) + && match &reason { + RefreshReason::Expiry => true, + RefreshReason::Unauthorized { + rejected_access_token, + } => rejected_access_token.secret() != latest_access_token, + }; + if should_adopt { + debug!("adopting newer MCP OAuth credentials without contacting the provider"); + 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 + { + 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"); + } + // The provider request has its own bound. The independently owned task prevents caller + // startup and operation deadlines from canceling this future after the provider may have + // rotated the refresh token. + let provider_started_at = Instant::now(); + debug!( + timeout_ms = refresh_request_timeout.as_millis(), + "requesting refreshed MCP OAuth credentials from the provider" + ); + let refresh_result = match timeout(refresh_request_timeout, guard.refresh_token()).await { + Ok(Ok(token_response)) => { + debug!( + provider_elapsed_ms = provider_started_at.elapsed().as_millis(), + "received refreshed MCP OAuth credentials from the provider" + ); + Ok(refreshed_tokens(token_response, &latest, &self.inner)) + } + Ok(Err(error)) => { + warn!( + provider_elapsed_ms = provider_started_at.elapsed().as_millis(), + error = %error, + "MCP OAuth provider refresh failed" + ); + Err(error).with_context(|| { + format!( + "failed to refresh OAuth tokens for server {}", + self.inner.server_name + ) + }) + } + Err(_) => { + warn!( + provider_elapsed_ms = provider_started_at.elapsed().as_millis(), + timeout_ms = refresh_request_timeout.as_millis(), + "MCP OAuth provider refresh timed out; the outcome is unknown and a later serialized retry is permitted" + ); + Err(anyhow::anyhow!( + "timed out after {refresh_request_timeout:?} refreshing OAuth tokens for server {}", + self.inner.server_name + )) + } + }; + let request_tokens = refresh_result.as_ref().unwrap_or(&latest); + if let Err(error) = + install_tokens_in_manager_guard(&mut guard, request_tokens, CredentialExposure::Request) + .await + { + return Err(error).context("failed to restore request-only OAuth credentials"); + } + let refreshed = refresh_result?; + // Once the provider rotates a refresh token, the owned task must attempt persistence even + // if the caller's deadline expires in the meantime. Refresh persistence stays on the + // source resolved at client startup. In particular, a keyring failure must surface instead + // of writing the rotated token to fallback File. + // + // A refreshed token becomes authoritative only after this write succeeds. If it fails, we + // deliberately restore the previous in-process credentials and return the error rather + // than serving an unpersisted token whose eventual loss would be difficult to correlate + // with this transaction. A later refresh may require reauthorization if the provider + // already consumed the previous token; that is the accepted fail-closed behavior. + // TODO: If persistence failures are common in practice, add an explicit bounded retry or + // reconciliation policy here. Do not silently continue with unpersisted credentials. + debug!("persisting refreshed MCP OAuth credentials to the resolved store"); + let persistence_started_at = Instant::now(); + let persistence_result = match self.inner.credential_store { + ResolvedOAuthCredentialStore::File => save_oauth_tokens_to_file(&refreshed), + ResolvedOAuthCredentialStore::Keyring(keyring_backend_kind) => { + save_oauth_tokens_with_keyring( + keyring_store, + keyring_backend_kind, + &self.inner.server_name, + &refreshed, + ) + } + }; + if let Err(error) = persistence_result { + warn!( + persistence_elapsed_ms = persistence_started_at.elapsed().as_millis(), + transaction_elapsed_ms = transaction_started_at.elapsed().as_millis(), + error = %error, + "failed to persist refreshed MCP OAuth credentials; returning the error and restoring the previous in-process credentials" + ); + install_tokens_in_manager_guard(&mut guard, &latest, CredentialExposure::Request) + .await + .context( + "failed to restore request-only OAuth credentials after refresh persistence failed", + )?; + return Err(error); + } + *self.inner.current_credentials.lock().await = Some(refreshed); + drop(guard); + debug!( + persistence_elapsed_ms = persistence_started_at.elapsed().as_millis(), + transaction_elapsed_ms = transaction_started_at.elapsed().as_millis(), + "persisted refreshed MCP OAuth credentials and completed the transaction" + ); + Ok(()) + } + + async fn adopt_credentials(&self, tokens: StoredOAuthTokens) -> Result<()> { + install_tokens_in_manager(&self.inner.authorization_manager, &tokens).await?; + *self.inner.current_credentials.lock().await = 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()); + } +} + +enum RefreshReason { + Expiry, + Unauthorized { rejected_access_token: AccessToken }, +} + +impl RefreshReason { + fn as_str(&self) -> &'static str { + match self { + Self::Expiry => "expiry", + Self::Unauthorized { .. } => "unauthorized", + } + } +} + +#[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); + // TODO(stevenlee): RMCP's `initialize_from_store` updates the credential store and client ID + // but not its private `current_scopes`. Credential adoption can therefore leave scope-upgrade + // state stale until RMCP exposes an adoption API that synchronizes both. + authorization_manager + .initialize_from_store() + .await + .context("failed to adopt refreshed OAuth tokens")?; + Ok(()) +} + +/// Controls which credentials are exposed to RMCP's authorization manager. +/// +/// Normal requests receive neither the refresh token nor expiry metadata, so RMCP cannot refresh +/// outside Codex's cross-process transaction. Full credentials are exposed only while that lock is +/// held, and request-only credentials are restored before the transaction returns unless that +/// restoration itself fails. +#[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 +} + +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, + } +} diff --git a/codex-rs/rmcp-client/src/oauth/refresh_lock.rs b/codex-rs/rmcp-client/src/oauth/refresh_lock.rs new file mode 100644 index 000000000000..5189f06ebdff --- /dev/null +++ b/codex-rs/rmcp-client/src/oauth/refresh_lock.rs @@ -0,0 +1,116 @@ +//! Cross-process serialization for one MCP OAuth credential's refresh transaction. +//! +//! The guard is intentionally acquired before the authoritative credential reread and retained +//! through provider refresh and persistence. This prevents two processes from replaying the same +//! rotating refresh token or observing a partially persisted transaction. + +use anyhow::Context; +use anyhow::Result; +use codex_utils_home_dir::find_codex_home; +use sha2::Digest; +use sha2::Sha256; +use std::fs; +use std::fs::File; +use std::fs::OpenOptions; +use std::path::Path; +use std::path::PathBuf; +use std::time::Duration; +use tokio::time::sleep; +use tokio::time::timeout; + +const REFRESH_LOCK_DIR: &str = "mcp-oauth-refresh-locks"; +const REFRESH_LOCK_ACQUIRE_TIMEOUT: Duration = Duration::from_secs(/*secs*/ 60); +const REFRESH_LOCK_RETRY_SLEEP: Duration = Duration::from_millis(/*millis*/ 50); +// Keep this internal target stable so diagnostics and cross-process tests can distinguish actual +// WouldBlock contention from a contender that merely started late and observed persisted tokens. +const LOCK_CONTENTION_EVENT_TARGET: &str = "codex_rmcp_client::oauth::refresh_lock::contention"; + +pub(super) struct RefreshCredentialLock { + _file: File, +} + +impl RefreshCredentialLock { + pub(super) async fn acquire_for_server(server_name: &str, url: &str) -> Result { + let store_key = super::compute_store_key(server_name, url)?; + let codex_home = find_codex_home()?; + Self::acquire_in(&codex_home, &store_key, REFRESH_LOCK_ACQUIRE_TIMEOUT) + .await + .with_context(|| format!("failed to acquire OAuth credential lock for {server_name}")) + } + + async fn acquire_in( + codex_home: &Path, + store_key: &str, + acquire_timeout: Duration, + ) -> Result { + let path = refresh_lock_path(codex_home, 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()))?; + + // Bound every contender, but keep the acquired lock for the full provider request and + // persistence transaction. Releasing it while awaiting the provider would allow concurrent + // use of a rotating refresh token. + let mut reported_contention = false; + match timeout(acquire_timeout, async { + loop { + match file.try_lock() { + Ok(()) => return Ok(()), + Err(std::fs::TryLockError::WouldBlock) => { + if !reported_contention { + tracing::debug!( + target: LOCK_CONTENTION_EVENT_TARGET, + lock_path = %path.display(), + "waiting for another process to finish refreshing MCP OAuth credentials" + ); + reported_contention = true; + } + sleep(REFRESH_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 }) + } +} + +fn refresh_lock_path(codex_home: &Path, store_key: &str) -> PathBuf { + // Credential coordination is deliberately scoped to the active CODEX_HOME, alongside File + // and Secrets state. Coordinating the process-global Direct keyring across distinct homes + // would require a separately defined global lock namespace and is outside this transaction. + // TODO(stevenlee): define a safe per-user, cross-platform rendezvous before extending Direct + // keyring coordination across distinct CODEX_HOME values. + let mut hasher = Sha256::new(); + hasher.update(store_key.as_bytes()); + let digest = hasher.finalize(); + codex_home + .join(REFRESH_LOCK_DIR) + .join(format!("{digest:x}.lock")) +} + +#[cfg(test)] +#[path = "refresh_lock_tests.rs"] +mod tests; diff --git a/codex-rs/rmcp-client/src/oauth/refresh_lock_tests.rs b/codex-rs/rmcp-client/src/oauth/refresh_lock_tests.rs new file mode 100644 index 000000000000..2b90e73996cc --- /dev/null +++ b/codex-rs/rmcp-client/src/oauth/refresh_lock_tests.rs @@ -0,0 +1,42 @@ +use super::RefreshCredentialLock; +use anyhow::Result; +use std::time::Duration; +use tempfile::tempdir; + +#[tokio::test] +async fn acquisition_times_out_without_stealing() -> Result<()> { + let codex_home = tempdir()?; + let store_key = "test-store-key"; + let held_lock = RefreshCredentialLock::acquire_in( + codex_home.path(), + store_key, + Duration::from_millis(/*millis*/ 100), + ) + .await?; + + let error = match RefreshCredentialLock::acquire_in( + codex_home.path(), + store_key, + Duration::from_millis(/*millis*/ 50), + ) + .await + { + Ok(_) => panic!("contending lock acquisition should time out"), + Err(error) => error, + }; + assert!( + error + .to_string() + .contains("timed out after 50ms waiting for OAuth refresh lock"), + "unexpected error: {error:#}" + ); + + drop(held_lock); + let _reacquired = RefreshCredentialLock::acquire_in( + codex_home.path(), + store_key, + Duration::from_millis(/*millis*/ 100), + ) + .await?; + Ok(()) +} diff --git a/codex-rs/rmcp-client/src/oauth/resolved_store.rs b/codex-rs/rmcp-client/src/oauth/resolved_store.rs new file mode 100644 index 000000000000..b384cc536047 --- /dev/null +++ b/codex-rs/rmcp-client/src/oauth/resolved_store.rs @@ -0,0 +1,120 @@ +//! Resolves the configured MCP OAuth store and pins that concrete source for one client lifecycle. + +use anyhow::Context; +use anyhow::Result; +use codex_config::types::AuthKeyringBackendKind; +use codex_config::types::OAuthCredentialsStoreMode; +use codex_keyring_store::KeyringStore; +use tracing::warn; + +use super::StoredOAuthTokens; +use super::load_oauth_tokens_from_file; +use super::load_oauth_tokens_from_keyring; + +/// Concrete credential store resolved for one MCP OAuth client lifecycle. +/// +/// This is intentionally not durable. `Auto` may resolve differently in a later process, but a +/// client that loaded credentials from one store must reread, refresh, persist, and remove only +/// through that store. A mid-lifecycle backend failure is unexpected and must return an error +/// rather than falling back to another possibly stale refresh token. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum ResolvedOAuthCredentialStore { + File, + Keyring(AuthKeyringBackendKind), +} + +#[derive(Debug)] +pub(crate) struct ResolvedOAuthTokens { + pub(crate) tokens: StoredOAuthTokens, + pub(crate) store: ResolvedOAuthCredentialStore, +} + +pub(crate) fn resolve_oauth_tokens( + keyring_store: &K, + server_name: &str, + url: &str, + store_mode: OAuthCredentialsStoreMode, + keyring_backend_kind: AuthKeyringBackendKind, +) -> Result> { + match store_mode { + OAuthCredentialsStoreMode::Auto => { + // Auto remains keyring-first at lifecycle startup. The returned source is then pinned + // by the client transport recipe and OAuth persistor so retries, recovery, and + // refresh work cannot hot-switch stores. + // TODO(stevenlee): Different processes can still resolve Auto to different stores + // when keyring availability differs. Solving that safely requires durable backend + // selection or reconciliation of legacy entries and is intentionally outside this + // stack. + match load_oauth_tokens_from_keyring( + keyring_store, + keyring_backend_kind, + server_name, + url, + ) { + Ok(Some(tokens)) => Ok(Some(ResolvedOAuthTokens { + tokens, + store: ResolvedOAuthCredentialStore::Keyring(keyring_backend_kind), + })), + Ok(None) => Ok( + load_oauth_tokens_from_file(server_name, url)?.map(|tokens| { + ResolvedOAuthTokens { + tokens, + store: ResolvedOAuthCredentialStore::File, + } + }), + ), + Err(error) => { + warn!("failed to read OAuth tokens from keyring: {error}"); + Ok(load_oauth_tokens_from_file(server_name, url) + .with_context(|| { + format!("failed to read OAuth tokens from keyring: {error}") + })? + .map(|tokens| ResolvedOAuthTokens { + tokens, + store: ResolvedOAuthCredentialStore::File, + })) + } + } + } + OAuthCredentialsStoreMode::File => Ok(load_oauth_tokens_from_file(server_name, url)?.map( + |tokens| ResolvedOAuthTokens { + tokens, + store: ResolvedOAuthCredentialStore::File, + }, + )), + OAuthCredentialsStoreMode::Keyring => Ok(load_oauth_tokens_from_keyring( + keyring_store, + keyring_backend_kind, + server_name, + url, + ) + .with_context(|| "failed to read OAuth tokens from keyring".to_string())? + .map(|tokens| ResolvedOAuthTokens { + tokens, + store: ResolvedOAuthCredentialStore::Keyring(keyring_backend_kind), + })), + } +} + +pub(crate) fn load_oauth_tokens_from_store( + keyring_store: &K, + server_name: &str, + url: &str, + store: ResolvedOAuthCredentialStore, +) -> Result> { + match store { + ResolvedOAuthCredentialStore::File => load_oauth_tokens_from_file(server_name, url) + .context("failed to reread OAuth tokens from resolved file storage"), + ResolvedOAuthCredentialStore::Keyring(keyring_backend_kind) => { + load_oauth_tokens_from_keyring( + keyring_store, + keyring_backend_kind, + server_name, + url, + ) + .context( + "failed to reread OAuth tokens from resolved keyring storage; refusing file fallback", + ) + } + } +} diff --git a/codex-rs/rmcp-client/src/rmcp_client.rs b/codex-rs/rmcp-client/src/rmcp_client.rs index c6527990fac0..55feb4a8ffd6 100644 --- a/codex-rs/rmcp-client/src/rmcp_client.rs +++ b/codex-rs/rmcp-client/src/rmcp_client.rs @@ -3,6 +3,7 @@ use std::ffi::OsString; use std::future::Future; use std::io; use std::sync::Arc; +use std::sync::OnceLock; use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; use std::time::Duration; @@ -14,6 +15,7 @@ use codex_api::SharedAuthProvider; use codex_config::types::AuthKeyringBackendKind; use codex_config::types::McpServerEnvVar; use codex_exec_server::HttpClient; +use codex_keyring_store::DefaultKeyringStore; use futures::FutureExt; use futures::future::BoxFuture; use oauth2::TokenResponse; @@ -64,9 +66,12 @@ 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::oauth::OAuthPersistor; +use crate::oauth::ResolvedOAuthCredentialStore; +use crate::oauth::ResolvedOAuthTokens; use crate::oauth::StoredOAuthTokens; +use crate::oauth::load_oauth_tokens_from_store; +use crate::oauth::resolve_oauth_tokens; use crate::oauth_http_client::OAuthHttpClientAdapter; use crate::stdio_server_launcher::StdioServerCommand; use crate::stdio_server_launcher::StdioServerLauncher; @@ -80,6 +85,7 @@ 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::remaining_initialize_timeout; use self::streamable_http_retry::sleep_with_retry_deadline; enum PendingTransport { @@ -126,6 +132,7 @@ enum TransportRecipe { env_http_headers: Option>, store_mode: OAuthCredentialsStoreMode, keyring_backend_kind: AuthKeyringBackendKind, + resolved_store: Arc>, http_client: Arc, auth_provider: Option, }, @@ -402,6 +409,7 @@ impl RmcpClient { env_http_headers, store_mode, keyring_backend_kind, + resolved_store: Arc::new(OnceLock::new()), http_client, auth_provider, }; @@ -444,11 +452,13 @@ impl RmcpClient { } }; + let mut initialize_deadline = timeout.map(|duration| Instant::now() + duration); let (service, oauth_persistor) = self .connect_pending_transport_with_initialize_retries( pending_transport, client_service.clone(), timeout, + &mut initialize_deadline, ) .await?; @@ -477,12 +487,6 @@ impl RmcpClient { }; } - 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) } @@ -491,14 +495,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) } @@ -508,7 +511,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(); @@ -533,7 +536,6 @@ impl RmcpClient { }) }) .collect::>>()?; - self.persist_oauth_tokens().await; Ok(ListToolsWithConnectorIdResult { next_cursor: result.next_cursor, tools, @@ -553,14 +555,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) } @@ -569,14 +570,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) } @@ -585,14 +585,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) } @@ -603,7 +602,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) => { @@ -650,7 +649,6 @@ impl RmcpClient { .boxed() }) .await?; - self.persist_oauth_tokens().await; Ok(result) } @@ -659,7 +657,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, @@ -680,7 +678,6 @@ impl RmcpClient { }, ) .await?; - self.persist_oauth_tokens().await; Ok(()) } @@ -689,7 +686,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(); @@ -703,7 +700,6 @@ impl RmcpClient { .boxed() }) .await?; - self.persist_oauth_tokens().await; Ok(response) } @@ -743,22 +739,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( @@ -781,6 +766,7 @@ impl RmcpClient { env_http_headers, store_mode, keyring_backend_kind, + resolved_store, http_client, auth_provider, } => { @@ -793,28 +779,55 @@ impl RmcpClient { auth_provider.clone() }; - let initial_oauth_tokens = if bearer_token.is_none() + let resolved_oauth_tokens = if bearer_token.is_none() && auth_provider.is_none() && !default_headers.contains_key(AUTHORIZATION) { - match load_oauth_tokens(server_name, url, *store_mode, *keyring_backend_kind) { - Ok(tokens) => tokens, - Err(err) => { - warn!("failed to read tokens for server `{server_name}`: {err}"); - None + if let Some(store) = resolved_store.get().copied() { + load_oauth_tokens_from_store(&DefaultKeyringStore, server_name, url, store)? + .map(|tokens| ResolvedOAuthTokens { tokens, store }) + } else { + match resolve_oauth_tokens( + &DefaultKeyringStore, + server_name, + url, + *store_mode, + *keyring_backend_kind, + ) { + Ok(tokens) => { + if let Some(resolved) = tokens.as_ref() { + // Transport retries and session recovery are part of the same + // client lifecycle. Pin the first concrete source in memory so + // rebuilding a transport never re-evaluates Auto and adopts a + // possibly stale credential from another store. + resolved_store.set(resolved.store).map_err(|_| { + anyhow!( + "OAuth credential store resolved concurrently for MCP server `{server_name}`" + ) + })?; + } + tokens + } + Err(err) => { + warn!("failed to read tokens for server `{server_name}`: {err}"); + None + } } } } else { None }; - if let Some(initial_tokens) = initial_oauth_tokens.clone() { + if let Some(ResolvedOAuthTokens { + tokens: initial_tokens, + store: credential_store, + }) = resolved_oauth_tokens + { match create_oauth_transport_and_runtime( server_name, url, initial_tokens.clone(), - *store_mode, - *keyring_backend_kind, + credential_store, default_headers.clone(), Arc::clone(http_client), ) @@ -880,6 +893,7 @@ impl RmcpClient { pending_transport: PendingTransport, client_service: ElicitationClientService, timeout: Option, + initialize_deadline: &mut Option, ) -> Result<( Arc>, Option, @@ -900,13 +914,26 @@ impl RmcpClient { PendingTransport::StreamableHttpWithOAuth { transport, oauth_persistor, - } => ( - service::serve_client(client_service, transport).boxed(), - Some(oauth_persistor), - ), + } => { + // `startup_timeout_sec` bounds MCP transport setup, retry delays, and the + // initialization handshake. OAuth refresh has independent lock and provider + // request bounds, and persistence must finish after a successful response, so the + // complete refresh transaction is deliberately excluded from that deadline. + let refresh_started_at = Instant::now(); + let refresh_result = oauth_persistor.refresh_if_needed().await; + if let Some(deadline) = initialize_deadline.as_mut() { + *deadline += refresh_started_at.elapsed(); + } + refresh_result?; + ( + service::serve_client(client_service, transport).boxed(), + Some(oauth_persistor), + ) + } }; - let service_result = match timeout { + let handshake_timeout = remaining_initialize_timeout(timeout, *initialize_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 })) @@ -919,19 +946,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)) } @@ -1124,11 +1139,15 @@ impl RmcpClient { .clone() .ok_or_else(|| anyhow!("MCP client cannot recover before initialize succeeds"))?; let pending_transport = Self::create_pending_transport(&self.transport_recipe).await?; + let mut initialize_deadline = initialize_context + .timeout + .map(|duration| Instant::now() + duration); let (service, oauth_persistor) = self .connect_pending_transport_with_initialize_retries( pending_transport, initialize_context.client_service, initialize_context.timeout, + &mut initialize_deadline, ) .await?; @@ -1143,12 +1162,6 @@ impl RmcpClient { }; } - 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(()) } } @@ -1157,8 +1170,7 @@ async fn create_oauth_transport_and_runtime( server_name: &str, url: &str, initial_tokens: StoredOAuthTokens, - credentials_store: OAuthCredentialsStoreMode, - keyring_backend_kind: AuthKeyringBackendKind, + credential_store: ResolvedOAuthCredentialStore, default_headers: HeaderMap, http_client: Arc, ) -> Result<( @@ -1202,8 +1214,7 @@ async fn create_oauth_transport_and_runtime( server_name.to_string(), url.to_string(), auth_manager, - credentials_store, - keyring_backend_kind, + credential_store, Some(initial_tokens), ); diff --git a/codex-rs/rmcp-client/src/streamable_http_retry.rs b/codex-rs/rmcp-client/src/streamable_http_retry.rs index 73da95de58ec..2794489b2a76 100644 --- a/codex-rs/rmcp-client/src/streamable_http_retry.rs +++ b/codex-rs/rmcp-client/src/streamable_http_retry.rs @@ -28,6 +28,7 @@ impl RmcpClient { initial_transport: PendingTransport, client_service: ElicitationClientService, timeout: Option, + initialize_deadline: &mut Option, ) -> Result<( Arc>, Option, @@ -37,7 +38,6 @@ impl RmcpClient { PendingTransport::StreamableHttp { .. } | PendingTransport::StreamableHttpWithOAuth { .. } => true, }; - let retry_deadline = timeout.map(|duration| Instant::now() + duration); let mut pending_transport = Some(initial_transport); for (attempt, retry_delay_ms) in STREAMABLE_HTTP_RETRY_DELAYS_MS @@ -50,7 +50,7 @@ impl RmcpClient { let transport = match pending_transport.take() { Some(transport) => transport, None => { - let remaining = remaining_initialize_timeout(timeout, retry_deadline)?; + let remaining = remaining_initialize_timeout(timeout, *initialize_deadline)?; match remaining { Some(remaining) => time::timeout( remaining, @@ -62,12 +62,11 @@ impl RmcpClient { } } }; - let attempt_timeout = remaining_initialize_timeout(timeout, retry_deadline)?; - match Self::connect_pending_transport( transport, client_service.clone(), - attempt_timeout, + timeout, + initialize_deadline, ) .await { @@ -84,7 +83,7 @@ impl RmcpClient { error = %error, "streamable HTTP MCP initialize failed with a retryable error; retrying" ); - if !sleep_with_retry_deadline(delay, retry_deadline).await { + if !sleep_with_retry_deadline(delay, *initialize_deadline).await { let duration = timeout.unwrap_or(delay); return Err(anyhow!( "timed out handshaking with MCP server after {duration:?}" @@ -194,7 +193,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 +208,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:?}") }