diff --git a/codex-rs/cli/src/mcp_cmd.rs b/codex-rs/cli/src/mcp_cmd.rs index 1466f1039479..fdff6bf494ea 100644 --- a/codex-rs/cli/src/mcp_cmd.rs +++ b/codex-rs/cli/src/mcp_cmd.rs @@ -26,7 +26,7 @@ use codex_mcp::oauth_login_support; use codex_mcp::resolve_oauth_scopes; use codex_mcp::should_retry_without_scopes; use codex_protocol::protocol::McpAuthStatus; -use codex_rmcp_client::delete_oauth_tokens; +use codex_rmcp_client::delete_oauth_tokens_locked; use codex_rmcp_client::perform_oauth_login; use codex_utils_cli::CliConfigOverrides; use codex_utils_cli::format_env_display; @@ -519,12 +519,14 @@ async fn run_logout(config_overrides: &CliConfigOverrides, logout_args: LogoutAr _ => bail!("OAuth logout is only supported for streamable_http transports."), }; - match delete_oauth_tokens( + match delete_oauth_tokens_locked( &name, &url, config.mcp_oauth_credentials_store_mode, config.auth_keyring_backend_kind(), - ) { + ) + .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 e1ee18c75324..e58c344aca87 100644 --- a/codex-rs/rmcp-client/src/lib.rs +++ b/codex-rs/rmcp-client/src/lib.rs @@ -20,6 +20,7 @@ pub use in_process_transport::InProcessTransportFactory; pub use oauth::StoredOAuthTokens; pub use oauth::WrappedOAuthTokenResponse; pub use oauth::delete_oauth_tokens; +pub use oauth::delete_oauth_tokens_locked; pub(crate) use oauth::load_oauth_tokens; pub use oauth::save_oauth_tokens; 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 d0005d20a6d0..07ad99fe12a2 100644 --- a/codex-rs/rmcp-client/src/oauth.rs +++ b/codex-rs/rmcp-client/src/oauth.rs @@ -41,10 +41,13 @@ 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::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; @@ -52,13 +55,24 @@ 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 MCP_OAUTH_SECRET_PREFIX: &str = "MCP_OAUTH"; -const REFRESH_SKEW_MILLIS: u64 = 30_000; +// Refresh proactively so ordinary requests do not race token expiry. +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_REQUEST_TIMEOUT: Duration = Duration::from_secs(45); #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] pub struct StoredOAuthTokens { @@ -97,21 +111,72 @@ pub(crate) fn load_oauth_tokens( keyring_backend_kind: AuthKeyringBackendKind, ) -> Result> { let keyring_store = DefaultKeyringStore; + load_oauth_tokens_with_keyring_store( + &keyring_store, + server_name, + url, + store_mode, + keyring_backend_kind, + ) +} + +fn load_oauth_tokens_with_keyring_store( + keyring_store: &K, + server_name: &str, + url: &str, + store_mode: OAuthCredentialsStoreMode, + keyring_backend_kind: AuthKeyringBackendKind, +) -> Result> { match store_mode { OAuthCredentialsStoreMode::Auto => load_oauth_tokens_from_keyring_with_fallback_to_file( - &keyring_store, + 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) + load_oauth_tokens_from_keyring(keyring_store, keyring_backend_kind, server_name, url) .with_context(|| "failed to read OAuth tokens from keyring".to_string()) } } } +fn load_oauth_tokens_for_refresh_with_keyring_store( + keyring_store: &K, + server_name: &str, + url: &str, + store_mode: OAuthCredentialsStoreMode, + keyring_backend_kind: AuthKeyringBackendKind, +) -> Result> { + if store_mode != OAuthCredentialsStoreMode::Auto { + return load_oauth_tokens_with_keyring_store( + keyring_store, + server_name, + url, + store_mode, + keyring_backend_kind, + ); + } + + 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(keyring_error) => { + warn!("failed to reread OAuth tokens from keyring: {keyring_error}"); + match load_oauth_tokens_from_file(server_name, url) { + Ok(Some(tokens)) => Ok(Some(tokens)), + Ok(None) => Err(keyring_error).context( + "failed to reread OAuth tokens from keyring and no fallback credential was found", + ), + Err(fallback_error) => Err(fallback_error).with_context(|| { + format!("failed to reread OAuth tokens from keyring: {keyring_error}") + }), + } + } + } +} + pub(crate) fn oauth_token_status( server_name: &str, url: &str, @@ -220,6 +285,7 @@ fn load_oauth_tokens_from_secrets_keyring( server_name: &str, url: &str, ) -> Result> { + let _store_lock = OAuthStoreLock::acquire(OAuthStore::Secrets)?; let codex_home = find_codex_home()?; let manager = SecretsManager::new_with_keyring_store_and_namespace( codex_home.to_path_buf(), @@ -249,20 +315,67 @@ 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, + ) +} + +pub(crate) async fn save_oauth_tokens_locked( + server_name: &str, + tokens: &StoredOAuthTokens, + store_mode: OAuthCredentialsStoreMode, + keyring_backend_kind: AuthKeyringBackendKind, +) -> Result<()> { + let keyring_store = DefaultKeyringStore; + save_oauth_tokens_locked_with_keyring_store( + &keyring_store, + server_name, + tokens, + store_mode, + keyring_backend_kind, + ) + .await +} + +async fn save_oauth_tokens_locked_with_keyring_store( + keyring_store: &K, + server_name: &str, + tokens: &StoredOAuthTokens, + store_mode: OAuthCredentialsStoreMode, + keyring_backend_kind: AuthKeyringBackendKind, +) -> Result<()> { + let _lock = RefreshCredentialLock::acquire_for_server(server_name, &tokens.url).await?; + 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, - keyring_backend_kind, - server_name, - tokens, - ), + OAuthCredentialsStoreMode::Keyring => { + save_oauth_tokens_with_keyring(keyring_store, keyring_backend_kind, server_name, tokens) + } } } @@ -314,6 +427,29 @@ fn save_oauth_tokens_to_secrets_keyring( tokens: &StoredOAuthTokens, ) -> Result<()> { let serialized = serde_json::to_string(tokens).context("failed to serialize OAuth tokens")?; + { + let _store_lock = OAuthStoreLock::acquire(OAuthStore::Secrets)?; + save_oauth_tokens_to_secrets_keyring_unlocked( + keyring_store, + server_name, + tokens, + &serialized, + )?; + } + + 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_secrets_keyring_unlocked( + keyring_store: &K, + server_name: &str, + tokens: &StoredOAuthTokens, + serialized: &str, +) -> Result<()> { let codex_home = find_codex_home()?; let manager = SecretsManager::new_with_keyring_store_and_namespace( codex_home.to_path_buf(), @@ -323,14 +459,8 @@ 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(()) + .set(&SecretScope::Global, &secret_name, serialized) + .context("failed to write OAuth tokens to encrypted storage") } fn save_oauth_tokens_with_keyring_with_fallback_to_file( @@ -366,6 +496,40 @@ pub fn delete_oauth_tokens( ) } +pub async fn delete_oauth_tokens_locked( + server_name: &str, + url: &str, + store_mode: OAuthCredentialsStoreMode, + keyring_backend_kind: AuthKeyringBackendKind, +) -> Result { + let keyring_store = DefaultKeyringStore; + delete_oauth_tokens_locked_with_keyring_store( + &keyring_store, + store_mode, + keyring_backend_kind, + server_name, + url, + ) + .await +} + +async fn delete_oauth_tokens_locked_with_keyring_store( + keyring_store: &K, + store_mode: OAuthCredentialsStoreMode, + keyring_backend_kind: AuthKeyringBackendKind, + server_name: &str, + url: &str, +) -> Result { + let _lock = RefreshCredentialLock::acquire_for_server(server_name, url).await?; + delete_oauth_tokens_from_keyring_and_file( + keyring_store, + store_mode, + keyring_backend_kind, + server_name, + url, + ) +} + fn delete_oauth_tokens_from_keyring_and_file( keyring_store: &K, store_mode: OAuthCredentialsStoreMode, @@ -430,6 +594,7 @@ fn delete_oauth_tokens_from_secrets_keyring( server_name: &str, url: &str, ) -> Result { + let _store_lock = OAuthStoreLock::acquire(OAuthStore::Secrets)?; let codex_home = find_codex_home()?; let manager = SecretsManager::new_with_keyring_store_and_namespace( codex_home.to_path_buf(), @@ -479,96 +644,406 @@ impl OAuthPersistor { } } - /// 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 - ); - } - } + pub(crate) async fn refresh_if_needed(&self) -> Result<()> { + self.refresh_if_needed_with_keyring_store(&DefaultKeyringStore) + .await + } + + pub(crate) async fn refresh_after_unauthorized(&self) -> Result<()> { + self.refresh_after_unauthorized_with_keyring_store(&DefaultKeyringStore) + .await + } + + 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 + } + + async fn refresh_if_needed_with_keyring_store_and_timeout( + &self, + keyring_store: &K, + refresh_request_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(()); } - Ok(()) + self.refresh_transaction( + keyring_store, + RefreshReason::Expiry, + refresh_request_timeout, + ) + .await + } + + async fn refresh_after_unauthorized_with_keyring_store( + &self, + keyring_store: &K, + ) -> Result<()> { + self.refresh_transaction( + keyring_store, + RefreshReason::Unauthorized, + REFRESH_REQUEST_TIMEOUT, + ) + .await } #[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, + keyring_store: &K, + reason: RefreshReason, + refresh_request_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 _lock = + RefreshCredentialLock::acquire_for_server(&self.inner.server_name, &self.inner.url) + .await?; + let latest = load_oauth_tokens_for_refresh_with_keyring_store( + keyring_store, + &self.inner.server_name, + &self.inner.url, + self.inner.store_mode, + self.inner.keyring_backend_kind, + )?; + + let Some(latest) = latest 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_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 {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"); + } + drop(guard); + + let refreshed = refresh_result?; + save_oauth_tokens_with_keyring_store( + keyring_store, + &self.inner.server_name, + &refreshed, + self.inner.store_mode, + self.inner.keyring_backend_kind, + )?; + let mut last_credentials = self.inner.last_credentials.lock().await; + *last_credentials = Some(refreshed); + Ok(()) + } + + 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, +} + +struct RefreshCredentialLock { + _file: File, +} + +impl RefreshCredentialLock { + async fn acquire_for_server(server_name: &str, url: &str) -> Result { + let key = compute_store_key(server_name, url)?; + Self::acquire(&key) + .await + .with_context(|| format!("failed to acquire OAuth credential lock for {server_name}")) + } + + async fn acquire(store_key: &str) -> Result { + Self::acquire_with_timeout(store_key, LOCK_ACQUIRE_TIMEOUT).await + } + + 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, + Secrets, +} + +impl OAuthStore { + fn lock_filename(self) -> &'static str { + match self { + Self::File => "file-store.lock", + Self::Secrets => "secrets-store.lock", + } + } + + fn description(self) -> &'static str { + match self { + Self::File => "fallback file", + Self::Secrets => "encrypted secrets", + } + } +} + +struct OAuthStoreLock { + _file: File, +} + +impl OAuthStoreLock { + fn acquire(store: OAuthStore) -> Result { + Self::acquire_with_timeout(store, LOCK_ACQUIRE_TIMEOUT) + } + + fn acquire_with_timeout(store: OAuthStore, acquire_timeout: Duration) -> 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 {}", + store.description(), + 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() >= acquire_timeout => { + anyhow::bail!( + "timed out after {acquire_timeout:?} waiting for MCP OAuth {} store lock {}", + store.description(), + path.display() + ); + } + Err(std::fs::TryLockError::WouldBlock) => { + std::thread::sleep(STORE_LOCK_RETRY_SLEEP.min(acquire_timeout)); + } + Err(error) => { + return Err(std::io::Error::from(error)).with_context(|| { + format!( + "failed to lock MCP OAuth {} store lock {}", + store.description(), + 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, } } @@ -592,7 +1067,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); }; @@ -635,8 +1111,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 @@ -664,7 +1145,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), }; @@ -750,9 +1232,26 @@ fn fallback_file_path() -> Result { Ok(find_codex_home()?.join(FALLBACK_FILENAME).to_path_buf()) } -fn read_fallback_file() -> Result> { - let path = fallback_file_path()?; - let contents = match fs::read_to_string(&path) { +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 read_fallback_file_unlocked() -> Result> { + let path = fallback_file_path()?; + let contents = match fs::read_to_string(&path) { Ok(contents) => contents, Err(err) if err.kind() == ErrorKind::NotFound => return Ok(None), Err(err) => { @@ -817,12 +1316,22 @@ mod tests { use codex_secrets::compute_keyring_account; use keyring::Error as KeyringError; use pretty_assertions::assert_eq; + use rmcp::transport::auth::OAuthState; + use serde_json::json; use std::sync::Arc; use std::sync::Mutex; use std::sync::MutexGuard; use std::sync::OnceLock; use std::sync::PoisonError; + use std::sync::mpsc; use tempfile::tempdir; + use tokio::sync::Mutex as TokioMutex; + use wiremock::Mock; + use wiremock::MockServer; + use wiremock::ResponseTemplate; + use wiremock::matchers::body_string_contains; + use wiremock::matchers::method; + use wiremock::matchers::path; use codex_keyring_store::tests::MockKeyringStore; @@ -964,7 +1473,7 @@ mod tests { let fallback_path = super::fallback_file_path()?; assert!(fallback_path.exists(), "fallback file should be created"); - let saved = super::read_fallback_file()?.expect("fallback file should load"); + let saved = read_fallback_file()?.expect("fallback file should load"); let key = super::compute_store_key(&tokens.server_name, &tokens.url)?; let entry = saved.get(&key).expect("entry for key"); assert_eq!(entry.server_name, tokens.server_name); @@ -978,6 +1487,45 @@ mod tests { Ok(()) } + #[test] + fn file_store_lock_preserves_updates_for_different_servers() -> Result<()> { + let _env = TempCodexHome::new(); + let first = sample_tokens(); + let mut second = sample_tokens(); + second.server_name = "second-server".to_string(); + second.url = "https://second.example.test".to_string(); + + let held_lock = + OAuthStoreLock::acquire_with_timeout(OAuthStore::File, Duration::from_millis(100))?; + let (started_tx, started_rx) = mpsc::channel(); + let (result_tx, result_rx) = mpsc::channel(); + let second_for_writer = second.clone(); + let writer = std::thread::spawn(move || { + started_tx.send(()).expect("signal writer start"); + result_tx + .send(super::save_oauth_tokens_to_file(&second_for_writer)) + .expect("send writer result"); + }); + + started_rx.recv_timeout(Duration::from_secs(1))?; + assert!(matches!( + result_rx.recv_timeout(Duration::from_millis(100)), + Err(mpsc::RecvTimeoutError::Timeout) + )); + super::save_oauth_tokens_to_file_unlocked(&first)?; + drop(held_lock); + + result_rx.recv_timeout(Duration::from_secs(10))??; + writer.join().expect("file store writer should finish"); + let loaded_first = super::load_oauth_tokens_from_file(&first.server_name, &first.url)? + .expect("first server tokens should remain stored"); + let loaded_second = super::load_oauth_tokens_from_file(&second.server_name, &second.url)? + .expect("second server tokens should be stored"); + assert_tokens_match_without_expiry(&loaded_first, &first); + assert_tokens_match_without_expiry(&loaded_second, &second); + Ok(()) + } + #[test] fn save_oauth_tokens_with_secrets_backend_writes_encrypted_storage() -> Result<()> { let env = TempCodexHome::new(); @@ -1038,6 +1586,68 @@ mod tests { Ok(()) } + #[test] + fn secrets_store_lock_preserves_updates_for_different_servers() -> Result<()> { + let _env = TempCodexHome::new(); + let store = MockKeyringStore::default(); + let first = sample_tokens(); + let mut second = sample_tokens(); + second.server_name = "second-server".to_string(); + second.url = "https://second.example.test".to_string(); + + let held_lock = + OAuthStoreLock::acquire_with_timeout(OAuthStore::Secrets, Duration::from_millis(100))?; + let (started_tx, started_rx) = mpsc::channel(); + let (result_tx, result_rx) = mpsc::channel(); + let store_for_writer = store.clone(); + let second_for_writer = second.clone(); + let writer = std::thread::spawn(move || { + started_tx.send(()).expect("signal writer start"); + result_tx + .send(super::save_oauth_tokens_with_keyring( + &store_for_writer, + AuthKeyringBackendKind::Secrets, + &second_for_writer.server_name, + &second_for_writer, + )) + .expect("send writer result"); + }); + + started_rx.recv_timeout(Duration::from_secs(1))?; + assert!(matches!( + result_rx.recv_timeout(Duration::from_millis(100)), + Err(mpsc::RecvTimeoutError::Timeout) + )); + let first_serialized = serde_json::to_string(&first)?; + super::save_oauth_tokens_to_secrets_keyring_unlocked( + &store, + &first.server_name, + &first, + &first_serialized, + )?; + drop(held_lock); + + result_rx.recv_timeout(Duration::from_secs(10))??; + writer.join().expect("secrets store writer should finish"); + let loaded_first = super::load_oauth_tokens_from_keyring( + &store, + AuthKeyringBackendKind::Secrets, + &first.server_name, + &first.url, + )? + .expect("first server tokens should remain stored"); + let loaded_second = super::load_oauth_tokens_from_keyring( + &store, + AuthKeyringBackendKind::Secrets, + &second.server_name, + &second.url, + )? + .expect("second server tokens should be stored"); + assert_tokens_match_without_expiry(&loaded_first, &first); + assert_tokens_match_without_expiry(&loaded_second, &second); + Ok(()) + } + #[test] fn load_oauth_tokens_with_secrets_backend_ignores_direct_entry() -> Result<()> { let _env = TempCodexHome::new(); @@ -1058,6 +1668,739 @@ mod tests { Ok(()) } + #[tokio::test] + async fn refresh_transaction_preserves_credentials_when_auto_keyring_reread_fails() -> Result<()> + { + let _env = TempCodexHome::new(); + let server = MockServer::start().await; + mount_oauth_metadata(&server).await; + let store = MockKeyringStore::default(); + let initial_tokens = expired_sample_tokens(&format!("{}/mcp", server.uri())); + let key = super::compute_store_key(&initial_tokens.server_name, &initial_tokens.url)?; + store.set_error(&key, KeyringError::Invalid("error".into(), "load".into())); + + let manager = authorization_manager_for(&initial_tokens).await?; + let persistor = OAuthPersistor::new( + initial_tokens.server_name.clone(), + initial_tokens.url.clone(), + manager.clone(), + OAuthCredentialsStoreMode::Auto, + AuthKeyringBackendKind::Direct, + Some(initial_tokens.clone()), + ); + + let error = persistor + .refresh_if_needed_with_keyring_store(&store) + .await + .expect_err("keyring reread failure should abort refresh"); + + assert!( + error + .to_string() + .contains("failed to reread OAuth tokens from keyring"), + "unexpected error: {error:#}" + ); + let manager_tokens = tokens_from_manager(&manager).await?; + assert_eq!( + manager_tokens.token_response, + WrappedOAuthTokenResponse(request_oauth_token_response(&initial_tokens)) + ); + Ok(()) + } + + #[tokio::test] + async fn refresh_transaction_adopts_keyring_update_without_second_provider_refresh() + -> Result<()> { + let _env = TempCodexHome::new(); + let server = MockServer::start().await; + mount_oauth_metadata(&server).await; + mount_refresh_response( + &server, + "refresh-token", + "refreshed-access-token", + "rotated-refresh-token", + ) + .await; + + let store = MockKeyringStore::default(); + let initial_tokens = expired_sample_tokens(&format!("{}/mcp", server.uri())); + super::save_oauth_tokens_with_keyring_store( + &store, + &initial_tokens.server_name, + &initial_tokens, + OAuthCredentialsStoreMode::Keyring, + AuthKeyringBackendKind::Direct, + )?; + + let winner_manager = authorization_manager_for(&initial_tokens).await?; + let loser_manager = authorization_manager_for(&initial_tokens).await?; + let winner = OAuthPersistor::new( + initial_tokens.server_name.clone(), + initial_tokens.url.clone(), + winner_manager, + OAuthCredentialsStoreMode::Keyring, + AuthKeyringBackendKind::Direct, + Some(initial_tokens.clone()), + ); + let loser = OAuthPersistor::new( + initial_tokens.server_name.clone(), + initial_tokens.url.clone(), + loser_manager.clone(), + OAuthCredentialsStoreMode::Keyring, + AuthKeyringBackendKind::Direct, + Some(initial_tokens.clone()), + ); + + winner.refresh_if_needed_with_keyring_store(&store).await?; + loser.refresh_if_needed_with_keyring_store(&store).await?; + + server.verify().await; + let stored = super::load_oauth_tokens_with_keyring_store( + &store, + &initial_tokens.server_name, + &initial_tokens.url, + OAuthCredentialsStoreMode::Keyring, + AuthKeyringBackendKind::Direct, + )? + .expect("refreshed tokens should be persisted"); + assert_eq!(access_token(&stored), "refreshed-access-token"); + assert_eq!( + refresh_token(&stored), + Some("rotated-refresh-token".to_string()) + ); + + let loser_tokens = tokens_from_manager(&loser_manager).await?; + assert_eq!( + loser_tokens.token_response, + WrappedOAuthTokenResponse(request_oauth_token_response(&stored)) + ); + Ok(()) + } + + #[tokio::test] + async fn unauthorized_transaction_without_expiry_refreshes_once_across_clients() -> Result<()> { + let _env = TempCodexHome::new(); + let server = MockServer::start().await; + mount_oauth_metadata(&server).await; + mount_refresh_response( + &server, + "refresh-token", + "refreshed-access-token", + "rotated-refresh-token", + ) + .await; + + let store = MockKeyringStore::default(); + let mut initial_tokens = sample_tokens(); + initial_tokens.url = format!("{}/mcp", server.uri()); + initial_tokens.expires_at = None; + initial_tokens.token_response.0.set_expires_in(None); + super::save_oauth_tokens_with_keyring_store( + &store, + &initial_tokens.server_name, + &initial_tokens, + OAuthCredentialsStoreMode::Keyring, + AuthKeyringBackendKind::Direct, + )?; + + let first_manager = authorization_manager_for(&initial_tokens).await?; + let second_manager = authorization_manager_for(&initial_tokens).await?; + let first = OAuthPersistor::new( + initial_tokens.server_name.clone(), + initial_tokens.url.clone(), + first_manager, + OAuthCredentialsStoreMode::Keyring, + AuthKeyringBackendKind::Direct, + Some(initial_tokens.clone()), + ); + let second = OAuthPersistor::new( + initial_tokens.server_name.clone(), + initial_tokens.url.clone(), + second_manager.clone(), + OAuthCredentialsStoreMode::Keyring, + AuthKeyringBackendKind::Direct, + Some(initial_tokens.clone()), + ); + + first + .refresh_after_unauthorized_with_keyring_store(&store) + .await?; + second + .refresh_after_unauthorized_with_keyring_store(&store) + .await?; + + server.verify().await; + let stored = super::load_oauth_tokens_with_keyring_store( + &store, + &initial_tokens.server_name, + &initial_tokens.url, + OAuthCredentialsStoreMode::Keyring, + AuthKeyringBackendKind::Direct, + )? + .expect("refreshed tokens should be persisted"); + assert_eq!(access_token(&stored), "refreshed-access-token"); + assert_eq!( + refresh_token(&stored), + Some("rotated-refresh-token".to_string()) + ); + let second_tokens = tokens_from_manager(&second_manager).await?; + assert_eq!( + second_tokens.token_response, + WrappedOAuthTokenResponse(request_oauth_token_response(&stored)) + ); + Ok(()) + } + + #[tokio::test] + async fn refresh_transaction_adopts_valid_reread_without_provider_refresh() -> Result<()> { + let _env = TempCodexHome::new(); + let server = MockServer::start().await; + mount_oauth_metadata(&server).await; + let store = MockKeyringStore::default(); + let mut initial_tokens = expired_sample_tokens(&format!("{}/mcp", server.uri())); + initial_tokens + .token_response + .0 + .set_refresh_token(Some(RefreshToken::new("stale-refresh-token".to_string()))); + + let mut latest_tokens = sample_tokens(); + latest_tokens.url.clone_from(&initial_tokens.url); + latest_tokens + .token_response + .0 + .set_access_token(AccessToken::new( + "already-refreshed-access-token".to_string(), + )); + latest_tokens + .token_response + .0 + .set_refresh_token(Some(RefreshToken::new( + "already-rotated-refresh-token".to_string(), + ))); + + super::save_oauth_tokens_with_keyring_store( + &store, + &latest_tokens.server_name, + &latest_tokens, + OAuthCredentialsStoreMode::Keyring, + AuthKeyringBackendKind::Direct, + )?; + + let manager = authorization_manager_for(&initial_tokens).await?; + let persistor = OAuthPersistor::new( + initial_tokens.server_name.clone(), + initial_tokens.url.clone(), + manager.clone(), + OAuthCredentialsStoreMode::Keyring, + AuthKeyringBackendKind::Direct, + Some(initial_tokens), + ); + + persistor + .refresh_if_needed_with_keyring_store(&store) + .await?; + + let manager_tokens = tokens_from_manager(&manager).await?; + assert_eq!( + manager_tokens.token_response, + WrappedOAuthTokenResponse(request_oauth_token_response(&latest_tokens)) + ); + Ok(()) + } + + #[tokio::test] + async fn refresh_transaction_refreshes_in_guard_band_despite_expires_in_drift() -> Result<()> { + let _env = TempCodexHome::new(); + let server = MockServer::start().await; + mount_oauth_metadata(&server).await; + mount_refresh_response( + &server, + "refresh-token", + "refreshed-after-expiry-drift", + "rotated-after-expiry-drift", + ) + .await; + + let store = MockKeyringStore::default(); + let mut initial_tokens = sample_tokens(); + initial_tokens.url = format!("{}/mcp", server.uri()); + initial_tokens.expires_at = Some(now_millis().saturating_add(45_000)); + initial_tokens + .token_response + .0 + .set_expires_in(Some(&Duration::from_secs(3600))); + super::save_oauth_tokens_with_keyring_store( + &store, + &initial_tokens.server_name, + &initial_tokens, + OAuthCredentialsStoreMode::Keyring, + AuthKeyringBackendKind::Direct, + )?; + + let manager = authorization_manager_for(&initial_tokens).await?; + let persistor = OAuthPersistor::new( + initial_tokens.server_name.clone(), + initial_tokens.url.clone(), + manager, + OAuthCredentialsStoreMode::Keyring, + AuthKeyringBackendKind::Direct, + Some(initial_tokens.clone()), + ); + + persistor + .refresh_if_needed_with_keyring_store(&store) + .await?; + + server.verify().await; + let stored = super::load_oauth_tokens_with_keyring_store( + &store, + &initial_tokens.server_name, + &initial_tokens.url, + OAuthCredentialsStoreMode::Keyring, + AuthKeyringBackendKind::Direct, + )? + .expect("refreshed tokens should be persisted"); + assert_eq!(access_token(&stored), "refreshed-after-expiry-drift"); + assert_eq!( + refresh_token(&stored), + Some("rotated-after-expiry-drift".to_string()) + ); + Ok(()) + } + + #[tokio::test] + async fn refresh_transaction_uses_latest_refresh_token_when_reread_is_expired() -> Result<()> { + let _env = TempCodexHome::new(); + let server = MockServer::start().await; + mount_oauth_metadata(&server).await; + mount_refresh_response( + &server, + "latest-refresh-token", + "refreshed-from-latest-token", + "rotated-from-latest-token", + ) + .await; + + let store = MockKeyringStore::default(); + let mut initial_tokens = expired_sample_tokens(&format!("{}/mcp", server.uri())); + initial_tokens + .token_response + .0 + .set_refresh_token(Some(RefreshToken::new("stale-refresh-token".to_string()))); + + let mut latest_tokens = initial_tokens.clone(); + latest_tokens + .token_response + .0 + .set_access_token(AccessToken::new("latest-expired-access-token".to_string())); + latest_tokens + .token_response + .0 + .set_refresh_token(Some(RefreshToken::new("latest-refresh-token".to_string()))); + super::save_oauth_tokens_with_keyring_store( + &store, + &latest_tokens.server_name, + &latest_tokens, + OAuthCredentialsStoreMode::Keyring, + AuthKeyringBackendKind::Direct, + )?; + + let manager = authorization_manager_for(&initial_tokens).await?; + let persistor = OAuthPersistor::new( + initial_tokens.server_name.clone(), + initial_tokens.url.clone(), + manager, + OAuthCredentialsStoreMode::Keyring, + AuthKeyringBackendKind::Direct, + Some(initial_tokens.clone()), + ); + + persistor + .refresh_if_needed_with_keyring_store(&store) + .await?; + + server.verify().await; + let stored = super::load_oauth_tokens_with_keyring_store( + &store, + &initial_tokens.server_name, + &initial_tokens.url, + OAuthCredentialsStoreMode::Keyring, + AuthKeyringBackendKind::Direct, + )? + .expect("refreshed tokens should be persisted"); + assert_eq!(access_token(&stored), "refreshed-from-latest-token"); + assert_eq!( + refresh_token(&stored), + Some("rotated-from-latest-token".to_string()) + ); + Ok(()) + } + + #[tokio::test] + async fn refresh_lock_acquisition_times_out_without_stealing() -> Result<()> { + let _env = TempCodexHome::new(); + let store_key = "test-store-key"; + let held_lock = + RefreshCredentialLock::acquire_with_timeout(store_key, Duration::from_millis(100)) + .await?; + + let error = + match RefreshCredentialLock::acquire_with_timeout(store_key, Duration::from_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_with_timeout(store_key, Duration::from_millis(100)) + .await?; + Ok(()) + } + + #[tokio::test] + async fn provider_refresh_timeout_releases_lock_without_persisting() -> Result<()> { + let _env = TempCodexHome::new(); + let server = MockServer::start().await; + mount_oauth_metadata(&server).await; + let refresh_started = mount_refresh_response_with_signal( + &server, + "refresh-token", + "late-access-token", + "late-refresh-token", + Duration::from_secs(1), + ) + .await; + + let store = MockKeyringStore::default(); + let initial_tokens = expired_sample_tokens(&format!("{}/mcp", server.uri())); + super::save_oauth_tokens_with_keyring_store( + &store, + &initial_tokens.server_name, + &initial_tokens, + OAuthCredentialsStoreMode::Keyring, + AuthKeyringBackendKind::Direct, + )?; + + let manager = authorization_manager_for(&initial_tokens).await?; + let persistor = OAuthPersistor::new( + initial_tokens.server_name.clone(), + initial_tokens.url.clone(), + manager, + OAuthCredentialsStoreMode::Keyring, + AuthKeyringBackendKind::Direct, + Some(initial_tokens.clone()), + ); + let refresh_task = tokio::spawn({ + let persistor = persistor.clone(); + let store = store.clone(); + async move { + persistor + .refresh_if_needed_with_keyring_store_and_timeout( + &store, + Duration::from_millis(200), + ) + .await + } + }); + + wait_for_signal(refresh_started).await?; + let error = refresh_task + .await? + .expect_err("delayed provider response should time out"); + assert_eq!( + error.to_string(), + "timed out after 200ms refreshing OAuth tokens for server test-server" + ); + + let store_key = super::compute_store_key(&initial_tokens.server_name, &initial_tokens.url)?; + let _lock = + RefreshCredentialLock::acquire_with_timeout(&store_key, Duration::from_millis(100)) + .await?; + let stored = super::load_oauth_tokens_with_keyring_store( + &store, + &initial_tokens.server_name, + &initial_tokens.url, + OAuthCredentialsStoreMode::Keyring, + AuthKeyringBackendKind::Direct, + )? + .expect("original tokens should remain persisted"); + assert_eq!(access_token(&stored), "access-token"); + assert_eq!(refresh_token(&stored), Some("refresh-token".to_string())); + server.verify().await; + Ok(()) + } + + #[tokio::test] + async fn locked_login_save_before_refresh_prevents_overwrite() -> Result<()> { + let _env = TempCodexHome::new(); + let server = MockServer::start().await; + mount_oauth_metadata(&server).await; + + let store = MockKeyringStore::default(); + let initial_tokens = expired_sample_tokens(&format!("{}/mcp", server.uri())); + super::save_oauth_tokens_with_keyring_store( + &store, + &initial_tokens.server_name, + &initial_tokens, + OAuthCredentialsStoreMode::Keyring, + AuthKeyringBackendKind::Direct, + )?; + + let mut login_tokens = sample_tokens(); + login_tokens.url.clone_from(&initial_tokens.url); + login_tokens + .token_response + .0 + .set_access_token(AccessToken::new("login-access-token".to_string())); + login_tokens + .token_response + .0 + .set_refresh_token(Some(RefreshToken::new("login-refresh-token".to_string()))); + super::save_oauth_tokens_locked_with_keyring_store( + &store, + &login_tokens.server_name, + &login_tokens, + OAuthCredentialsStoreMode::Keyring, + AuthKeyringBackendKind::Direct, + ) + .await?; + + let manager = authorization_manager_for(&initial_tokens).await?; + let persistor = OAuthPersistor::new( + initial_tokens.server_name.clone(), + initial_tokens.url.clone(), + manager, + OAuthCredentialsStoreMode::Keyring, + AuthKeyringBackendKind::Direct, + Some(initial_tokens.clone()), + ); + + persistor + .refresh_if_needed_with_keyring_store(&store) + .await?; + + let stored = super::load_oauth_tokens_with_keyring_store( + &store, + &initial_tokens.server_name, + &initial_tokens.url, + OAuthCredentialsStoreMode::Keyring, + AuthKeyringBackendKind::Direct, + )? + .expect("login tokens should remain persisted"); + assert_eq!(access_token(&stored), "login-access-token"); + assert_eq!( + refresh_token(&stored), + Some("login-refresh-token".to_string()) + ); + Ok(()) + } + + #[tokio::test] + async fn locked_login_save_after_refresh_still_wins() -> Result<()> { + let _env = TempCodexHome::new(); + let server = MockServer::start().await; + mount_oauth_metadata(&server).await; + let refresh_started = mount_refresh_response_with_signal( + &server, + "refresh-token", + "refreshed-before-login", + "rotated-before-login", + Duration::from_millis(200), + ) + .await; + + let store = MockKeyringStore::default(); + let initial_tokens = expired_sample_tokens(&format!("{}/mcp", server.uri())); + super::save_oauth_tokens_with_keyring_store( + &store, + &initial_tokens.server_name, + &initial_tokens, + OAuthCredentialsStoreMode::Keyring, + AuthKeyringBackendKind::Direct, + )?; + + let manager = authorization_manager_for(&initial_tokens).await?; + let persistor = OAuthPersistor::new( + initial_tokens.server_name.clone(), + initial_tokens.url.clone(), + manager, + OAuthCredentialsStoreMode::Keyring, + AuthKeyringBackendKind::Direct, + Some(initial_tokens.clone()), + ); + let refresh_task = tokio::spawn({ + let persistor = persistor.clone(); + let store = store.clone(); + async move { persistor.refresh_if_needed_with_keyring_store(&store).await } + }); + + wait_for_signal(refresh_started).await?; + + let mut login_tokens = sample_tokens(); + login_tokens.url.clone_from(&initial_tokens.url); + login_tokens + .token_response + .0 + .set_access_token(AccessToken::new("login-after-refresh-access".to_string())); + login_tokens + .token_response + .0 + .set_refresh_token(Some(RefreshToken::new( + "login-after-refresh-token".to_string(), + ))); + super::save_oauth_tokens_locked_with_keyring_store( + &store, + &login_tokens.server_name, + &login_tokens, + OAuthCredentialsStoreMode::Keyring, + AuthKeyringBackendKind::Direct, + ) + .await?; + refresh_task.await??; + + server.verify().await; + let stored = super::load_oauth_tokens_with_keyring_store( + &store, + &initial_tokens.server_name, + &initial_tokens.url, + OAuthCredentialsStoreMode::Keyring, + AuthKeyringBackendKind::Direct, + )? + .expect("login tokens should remain persisted"); + assert_eq!(access_token(&stored), "login-after-refresh-access"); + assert_eq!( + refresh_token(&stored), + Some("login-after-refresh-token".to_string()) + ); + Ok(()) + } + + #[tokio::test] + async fn locked_logout_before_refresh_prevents_resurrection() -> Result<()> { + let _env = TempCodexHome::new(); + let server = MockServer::start().await; + mount_oauth_metadata(&server).await; + + let store = MockKeyringStore::default(); + let initial_tokens = expired_sample_tokens(&format!("{}/mcp", server.uri())); + super::save_oauth_tokens_with_keyring_store( + &store, + &initial_tokens.server_name, + &initial_tokens, + OAuthCredentialsStoreMode::Keyring, + AuthKeyringBackendKind::Direct, + )?; + let removed = super::delete_oauth_tokens_locked_with_keyring_store( + &store, + OAuthCredentialsStoreMode::Keyring, + AuthKeyringBackendKind::Direct, + &initial_tokens.server_name, + &initial_tokens.url, + ) + .await?; + assert!(removed); + + let manager = authorization_manager_for(&initial_tokens).await?; + let persistor = OAuthPersistor::new( + initial_tokens.server_name.clone(), + initial_tokens.url.clone(), + manager, + OAuthCredentialsStoreMode::Keyring, + AuthKeyringBackendKind::Direct, + Some(initial_tokens.clone()), + ); + + let error = persistor + .refresh_if_needed_with_keyring_store(&store) + .await + .expect_err("missing latest tokens should abort refresh"); + assert!( + error.to_string().contains("authorization required"), + "unexpected error: {error:#}" + ); + let stored = super::load_oauth_tokens_with_keyring_store( + &store, + &initial_tokens.server_name, + &initial_tokens.url, + OAuthCredentialsStoreMode::Keyring, + AuthKeyringBackendKind::Direct, + )?; + assert!(stored.is_none()); + Ok(()) + } + + #[tokio::test] + async fn locked_logout_after_refresh_still_deletes() -> Result<()> { + let _env = TempCodexHome::new(); + let server = MockServer::start().await; + mount_oauth_metadata(&server).await; + let refresh_started = mount_refresh_response_with_signal( + &server, + "refresh-token", + "refreshed-before-logout", + "rotated-before-logout", + Duration::from_millis(200), + ) + .await; + + let store = MockKeyringStore::default(); + let initial_tokens = expired_sample_tokens(&format!("{}/mcp", server.uri())); + super::save_oauth_tokens_with_keyring_store( + &store, + &initial_tokens.server_name, + &initial_tokens, + OAuthCredentialsStoreMode::Keyring, + AuthKeyringBackendKind::Direct, + )?; + + let manager = authorization_manager_for(&initial_tokens).await?; + let persistor = OAuthPersistor::new( + initial_tokens.server_name.clone(), + initial_tokens.url.clone(), + manager, + OAuthCredentialsStoreMode::Keyring, + AuthKeyringBackendKind::Direct, + Some(initial_tokens.clone()), + ); + let refresh_task = tokio::spawn({ + let persistor = persistor.clone(); + let store = store.clone(); + async move { persistor.refresh_if_needed_with_keyring_store(&store).await } + }); + + wait_for_signal(refresh_started).await?; + + let removed = super::delete_oauth_tokens_locked_with_keyring_store( + &store, + OAuthCredentialsStoreMode::Keyring, + AuthKeyringBackendKind::Direct, + &initial_tokens.server_name, + &initial_tokens.url, + ) + .await?; + refresh_task.await??; + + server.verify().await; + assert!(removed); + let stored = super::load_oauth_tokens_with_keyring_store( + &store, + &initial_tokens.server_name, + &initial_tokens.url, + OAuthCredentialsStoreMode::Keyring, + AuthKeyringBackendKind::Direct, + )?; + assert!(stored.is_none()); + Ok(()) + } + #[test] fn save_oauth_tokens_with_secrets_backend_falls_back_to_file_when_keyring_fails() -> Result<()> { @@ -1076,7 +2419,7 @@ mod tests { &tokens, )?; - let saved = super::read_fallback_file()?.expect("fallback file should load"); + let saved = read_fallback_file()?.expect("fallback file should load"); let key = super::compute_store_key(&tokens.server_name, &tokens.url)?; assert!(saved.contains_key(&key)); Ok(()) @@ -1317,6 +2660,11 @@ mod tests { ); } + fn read_fallback_file() -> Result> { + let _store_lock = OAuthStoreLock::acquire(OAuthStore::File)?; + super::read_fallback_file_unlocked() + } + fn assert_token_response_match_without_expiry( actual: &WrappedOAuthTokenResponse, expected: &WrappedOAuthTokenResponse, @@ -1344,6 +2692,152 @@ mod tests { ); } + async fn authorization_manager_for( + tokens: &StoredOAuthTokens, + ) -> Result>> { + let mut state = OAuthState::new(tokens.url.clone(), Some(reqwest::Client::new())).await?; + state + .set_credentials(&tokens.client_id, request_oauth_token_response(tokens)) + .await?; + let manager = match state { + OAuthState::Authorized(manager) | OAuthState::Unauthorized(manager) => manager, + OAuthState::Session(_) | OAuthState::AuthorizedHttpClient(_) => { + anyhow::bail!("unexpected OAuth state") + } + _ => anyhow::bail!("unexpected OAuth state"), + }; + Ok(Arc::new(TokioMutex::new(manager))) + } + + async fn mount_oauth_metadata(server: &MockServer) { + Mock::given(method("GET")) + .and(path("/.well-known/oauth-authorization-server/mcp")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "authorization_endpoint": format!("{}/oauth/authorize", server.uri()), + "token_endpoint": format!("{}/oauth/token", server.uri()), + "scopes_supported": ["scope-a", "scope-b"], + }))) + .mount(server) + .await; + } + + async fn mount_refresh_response( + server: &MockServer, + request_refresh_token: &str, + response_access_token: &str, + response_refresh_token: &str, + ) { + Mock::given(method("POST")) + .and(path("/oauth/token")) + .and(body_string_contains("grant_type=refresh_token")) + .and(body_string_contains(format!( + "refresh_token={request_refresh_token}" + ))) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "access_token": response_access_token, + "token_type": "Bearer", + "expires_in": 3600, + "refresh_token": response_refresh_token, + "scope": "scope-a scope-b", + }))) + .expect(1) + .mount(server) + .await; + } + + async fn mount_refresh_response_with_signal( + server: &MockServer, + request_refresh_token: &str, + response_access_token: &str, + response_refresh_token: &str, + response_delay: Duration, + ) -> mpsc::Receiver<()> { + let (tx, rx) = mpsc::channel(); + let response_access_token = response_access_token.to_string(); + let response_refresh_token = response_refresh_token.to_string(); + Mock::given(method("POST")) + .and(path("/oauth/token")) + .and(body_string_contains("grant_type=refresh_token")) + .and(body_string_contains(format!( + "refresh_token={request_refresh_token}" + ))) + .respond_with(move |_request: &wiremock::Request| { + let _ = tx.send(()); + let access_token = response_access_token.clone(); + let refresh_token = response_refresh_token.clone(); + ResponseTemplate::new(200) + .set_delay(response_delay) + .set_body_json(json!({ + "access_token": access_token, + "token_type": "Bearer", + "expires_in": 3600, + "refresh_token": refresh_token, + "scope": "scope-a scope-b", + })) + }) + .expect(1) + .mount(server) + .await; + rx + } + + async fn wait_for_signal(rx: mpsc::Receiver<()>) -> Result<()> { + tokio::task::spawn_blocking(move || { + rx.recv_timeout(Duration::from_secs(5)) + .context("timed out waiting for refresh request") + }) + .await? + } + + #[expect( + clippy::await_holding_invalid_type, + reason = "AuthorizationManager async access must be serialized through its mutex" + )] + async fn tokens_from_manager( + manager: &Arc>, + ) -> Result { + let guard = manager.lock().await; + let (client_id, token_response) = guard.get_credentials().await?; + let token_response = token_response.expect("manager should have token response"); + Ok(StoredOAuthTokens { + server_name: "test-server".to_string(), + url: "https://example.test".to_string(), + client_id, + token_response: WrappedOAuthTokenResponse(token_response), + expires_at: None, + }) + } + + fn access_token(tokens: &StoredOAuthTokens) -> &str { + tokens.token_response.0.access_token().secret() + } + + fn refresh_token(tokens: &StoredOAuthTokens) -> Option { + tokens + .token_response + .0 + .refresh_token() + .map(|token| token.secret().to_string()) + } + + fn expired_sample_tokens(url: &str) -> StoredOAuthTokens { + let mut tokens = sample_tokens(); + tokens.url = url.to_string(); + tokens.expires_at = Some(0); + tokens + .token_response + .0 + .set_expires_in(Some(&Duration::ZERO)); + tokens + } + + fn now_millis() -> u64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_else(|_| Duration::from_secs(0)) + .as_millis() as u64 + } + fn sample_tokens() -> StoredOAuthTokens { let mut response = OAuthTokenResponse::new( AccessToken::new("access-token".to_string()), diff --git a/codex-rs/rmcp-client/src/perform_oauth_login.rs b/codex-rs/rmcp-client/src/perform_oauth_login.rs index 8c72f320bc2c..ab74caf9e79f 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::oauth::save_oauth_tokens_locked; use crate::utils::apply_default_headers; use crate::utils::build_default_headers; use codex_config::types::AuthKeyringBackendKind; @@ -583,12 +583,13 @@ impl OauthLoginFlow { token_response: WrappedOAuthTokenResponse(credentials), expires_at, }; - save_oauth_tokens( + save_oauth_tokens_locked( &self.server_name, &stored, self.store_mode, self.keyring_backend_kind, - )?; + ) + .await?; Ok(()) } diff --git a/codex-rs/rmcp-client/src/rmcp_client.rs b/codex-rs/rmcp-client/src/rmcp_client.rs index 4a9b89de3c32..3bd83e63f741 100644 --- a/codex-rs/rmcp-client/src/rmcp_client.rs +++ b/codex-rs/rmcp-client/src/rmcp_client.rs @@ -68,6 +68,7 @@ use crate::in_process_transport::InProcessTransportFactory; use crate::load_oauth_tokens; use crate::oauth::OAuthPersistor; use crate::oauth::StoredOAuthTokens; +use crate::oauth::request_oauth_token_response; use crate::stdio_server_launcher::StdioServerCommand; use crate::stdio_server_launcher::StdioServerLauncher; use crate::stdio_server_launcher::StdioServerProcessHandle; @@ -445,7 +446,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, @@ -477,12 +478,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 +486,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) } @@ -507,7 +501,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(); @@ -532,7 +526,6 @@ impl RmcpClient { }) }) .collect::>>()?; - self.persist_oauth_tokens().await; Ok(ListToolsWithConnectorIdResult { next_cursor: result.next_cursor, tools, @@ -552,14 +545,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) } @@ -568,14 +560,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) } @@ -584,14 +575,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) } @@ -602,7 +592,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) => { @@ -649,7 +639,6 @@ impl RmcpClient { .boxed() }) .await?; - self.persist_oauth_tokens().await; Ok(result) } @@ -658,7 +647,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, @@ -679,7 +668,6 @@ impl RmcpClient { }, ) .await?; - self.persist_oauth_tokens().await; Ok(()) } @@ -688,7 +676,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(); @@ -702,7 +690,6 @@ impl RmcpClient { .boxed() }) .await?; - self.persist_oauth_tokens().await; Ok(response) } @@ -742,22 +729,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( @@ -893,10 +869,13 @@ impl RmcpClient { PendingTransport::StreamableHttpWithOAuth { transport, oauth_persistor, - } => ( - service::serve_client(client_service, transport).boxed(), - Some(oauth_persistor), - ), + } => { + oauth_persistor.refresh_if_needed().await?; + ( + service::serve_client(client_service, transport).boxed(), + Some(oauth_persistor), + ) + } }; let service_result = match timeout { @@ -912,19 +891,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)) } @@ -950,6 +917,21 @@ impl RmcpClient { .await { Ok(result) => Ok(result), + Err(error) if Self::is_unauthorized_operation_error(&error) => { + let Some(oauth_persistor) = self.oauth_persistor().await else { + return Err(error.into()); + }; + oauth_persistor.refresh_after_unauthorized().await?; + Self::run_service_operation_with_transient_retries( + service, + label, + timeout, + self.elicitation_pause_state.clone(), + &operation, + ) + .await + .map_err(Into::into) + } Err(error) if Self::is_session_expired_404(&error) => { self.reinitialize_after_session_expiry(&service).await?; let recovered_service = self.service().await?; @@ -1084,6 +1066,31 @@ impl RmcpClient { }) } + 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) + } + + pub(super) fn is_unauthorized_streamable_http_error( + error: &StreamableHttpError, + ) -> bool { + match error { + StreamableHttpError::AuthRequired(_) => true, + StreamableHttpError::UnexpectedServerResponse(message) => { + message.starts_with("HTTP 401") + } + _ => false, + } + } + async fn reinitialize_after_session_expiry( &self, failed_service: &Arc>, @@ -1118,7 +1125,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, @@ -1136,12 +1143,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(()) } } @@ -1172,7 +1173,7 @@ async fn create_oauth_transport_and_runtime( oauth_state .set_credentials( &initial_tokens.client_id, - initial_tokens.token_response.0.clone(), + request_oauth_token_response(&initial_tokens), ) .await?; diff --git a/codex-rs/rmcp-client/src/streamable_http_retry.rs b/codex-rs/rmcp-client/src/streamable_http_retry.rs index 73da95de58ec..24b73afab392 100644 --- a/codex-rs/rmcp-client/src/streamable_http_retry.rs +++ b/codex-rs/rmcp-client/src/streamable_http_retry.rs @@ -23,6 +23,60 @@ const JSON_RPC_INTERNAL_ERROR_CODE: i64 = -32603; pub(super) const STREAMABLE_HTTP_RETRY_DELAYS_MS: [u64; 2] = [250, 1_000]; impl RmcpClient { + pub(super) async fn connect_pending_transport_with_oauth_recovery( + &self, + initial_transport: PendingTransport, + client_service: ElicitationClientService, + timeout: Option, + ) -> Result<( + Arc>, + Option, + )> { + let oauth_persistor = match &initial_transport { + PendingTransport::StreamableHttpWithOAuth { + oauth_persistor, .. + } => Some(oauth_persistor.clone()), + _ => None, + }; + let deadline = timeout.map(|duration| Instant::now() + duration); + match self + .connect_pending_transport_with_initialize_retries( + initial_transport, + client_service.clone(), + timeout, + ) + .await + { + Ok(result) => Ok(result), + Err(error) + if oauth_persistor.is_some() && Self::is_unauthorized_initialize_error(&error) => + { + let Some(oauth_persistor) = oauth_persistor else { + return Err(error); + }; + oauth_persistor.refresh_after_unauthorized().await?; + 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)?; + self.connect_pending_transport_with_initialize_retries( + transport, + client_service, + remaining, + ) + .await + } + Err(error) => Err(error), + } + } + pub(super) async fn connect_pending_transport_with_initialize_retries( &self, initial_transport: PendingTransport, @@ -109,6 +163,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 } 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 866daa67a4c0..4ec89ddc8952 100644 --- a/codex-rs/rmcp-client/tests/streamable_http_oauth_startup.rs +++ b/codex-rs/rmcp-client/tests/streamable_http_oauth_startup.rs @@ -38,6 +38,10 @@ const SERVER_NAME: &str = "test-streamable-http-oauth-startup"; const EXPIRED_ACCESS_TOKEN: &str = "expired-access-token"; const REFRESH_TOKEN: &str = "valid-refresh-token"; const REFRESHED_ACCESS_TOKEN: &str = "refreshed-access-token"; +const ROTATED_REFRESH_TOKEN: &str = "rotated-refresh-token"; +const FINAL_ACCESS_TOKEN: &str = "final-access-token"; +const FINAL_REFRESH_TOKEN: &str = "final-refresh-token"; +const PREFLIGHT_REFRESH_ERROR: &str = "preflight refresh failed distinctly"; const CHILD_SERVER_URL_ENV: &str = "MCP_TEST_OAUTH_STARTUP_SERVER_URL"; const UNREFRESHABLE_SERVER_URL: &str = "https://unrefreshable.example/mcp"; const UNEXPIRED_SERVER_URL: &str = "https://unexpired.example/mcp"; @@ -121,6 +125,316 @@ async fn refreshes_expired_persisted_token_before_initialize() -> anyhow::Result Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn recovers_initialization_and_operation_401_without_rmcp_refresh() -> anyhow::Result<()> { + let server = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/.well-known/oauth-authorization-server/mcp")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "authorization_endpoint": format!("{}/oauth/authorize", server.uri()), + "token_endpoint": format!("{}/oauth/token", server.uri()), + "scopes_supported": [""], + }))) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/oauth/token")) + .and(body_string_contains("grant_type=refresh_token")) + .and(body_string_contains(format!( + "refresh_token={REFRESH_TOKEN}" + ))) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "access_token": REFRESHED_ACCESS_TOKEN, + "token_type": "Bearer", + "expires_in": 7200, + "refresh_token": ROTATED_REFRESH_TOKEN, + }))) + .expect(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/oauth/token")) + .and(body_string_contains("grant_type=refresh_token")) + .and(body_string_contains(format!( + "refresh_token={ROTATED_REFRESH_TOKEN}" + ))) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "access_token": FINAL_ACCESS_TOKEN, + "token_type": "Bearer", + "expires_in": 7200, + "refresh_token": FINAL_REFRESH_TOKEN, + }))) + .expect(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/mcp")) + .and(header( + "authorization", + format!("Bearer {EXPIRED_ACCESS_TOKEN}"), + )) + .respond_with( + ResponseTemplate::new(401).insert_header("www-authenticate", "Bearer realm=test"), + ) + .expect(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/mcp")) + .and(header( + "authorization", + format!("Bearer {REFRESHED_ACCESS_TOKEN}"), + )) + .respond_with(|request: &Request| { + let body: Value = request.body_json().expect("valid JSON-RPC request"); + match body.get("method").and_then(Value::as_str) { + Some("initialize") => initialize_response(&body, "oauth-401-recovery-test"), + Some("notifications/initialized") => ResponseTemplate::new(202), + Some("tools/list") => ResponseTemplate::new(401) + .insert_header("www-authenticate", "Bearer realm=test"), + method => ResponseTemplate::new(400) + .set_body_string(format!("unexpected JSON-RPC method: {method:?}")), + } + }) + .expect(3) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/mcp")) + .and(header( + "authorization", + format!("Bearer {FINAL_ACCESS_TOKEN}"), + )) + .respond_with(|request: &Request| { + let body: Value = request.body_json().expect("valid JSON-RPC request"); + match body.get("method").and_then(Value::as_str) { + Some("initialize") => initialize_response(&body, "oauth-401-persisted-test"), + Some("notifications/initialized") => ResponseTemplate::new(202), + Some("tools/list") => ResponseTemplate::new(200).set_body_json(json!({ + "jsonrpc": "2.0", + "id": body.get("id").cloned().unwrap_or(Value::Null), + "result": { "tools": [] }, + })), + method => ResponseTemplate::new(400) + .set_body_string(format!("unexpected JSON-RPC method: {method:?}")), + } + }) + .expect(3) + .mount(&server) + .await; + + let codex_home = TempDir::new()?; + let server_url = format!("{}/mcp", server.uri()); + let status = Command::new(std::env::current_exe()?) + .args([ + "oauth_401_recovery_child", + "--exact", + "--ignored", + "--nocapture", + ]) + .env("CODEX_HOME", codex_home.path()) + .env(CHILD_SERVER_URL_ENV, server_url) + .status() + .await?; + assert!( + status.success(), + "OAuth 401 recovery child failed: {status}" + ); + server.verify().await; + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn concurrent_file_mode_startup_refreshes_once() -> anyhow::Result<()> { + let server = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/.well-known/oauth-authorization-server/mcp")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "authorization_endpoint": format!("{}/oauth/authorize", server.uri()), + "token_endpoint": format!("{}/oauth/token", server.uri()), + "scopes_supported": [""], + }))) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/oauth/token")) + .and(body_string_contains("grant_type=refresh_token")) + .and(body_string_contains(format!( + "refresh_token={REFRESH_TOKEN}" + ))) + .respond_with( + ResponseTemplate::new(200) + .set_delay(Duration::from_millis(250)) + .set_body_json(json!({ + "access_token": REFRESHED_ACCESS_TOKEN, + "token_type": "Bearer", + "expires_in": 7200, + "refresh_token": ROTATED_REFRESH_TOKEN, + })), + ) + .expect(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/mcp")) + .and(header( + "authorization", + format!("Bearer {REFRESHED_ACCESS_TOKEN}"), + )) + .respond_with(|request: &Request| { + let body: Value = request.body_json().expect("valid JSON-RPC request"); + match body.get("method").and_then(Value::as_str) { + Some("initialize") => ResponseTemplate::new(200).set_body_json(json!({ + "jsonrpc": "2.0", + "id": body.get("id").cloned().unwrap_or(Value::Null), + "result": { + "protocolVersion": body + .pointer("/params/protocolVersion") + .cloned() + .unwrap_or_else(|| json!("2025-06-18")), + "capabilities": {}, + "serverInfo": { + "name": "oauth-startup-test", + "version": "0.0.0-test", + }, + }, + })), + Some("notifications/initialized") => ResponseTemplate::new(202), + method => ResponseTemplate::new(400) + .set_body_string(format!("unexpected JSON-RPC method: {method:?}")), + } + }) + .expect(4) + .mount(&server) + .await; + + let codex_home = TempDir::new()?; + let server_url = format!("{}/mcp", server.uri()); + let seed_status = Command::new(std::env::current_exe()?) + .args(["oauth_concurrency_seed_child", "--exact", "--ignored"]) + .env("CODEX_HOME", codex_home.path()) + .env(CHILD_SERVER_URL_ENV, &server_url) + .status() + .await?; + assert!( + seed_status.success(), + "OAuth concurrency seed child failed: {seed_status}" + ); + + let first_status = Command::new(std::env::current_exe()?) + .args([ + "oauth_concurrency_client_child", + "--exact", + "--ignored", + "--nocapture", + ]) + .env("CODEX_HOME", codex_home.path()) + .env(CHILD_SERVER_URL_ENV, &server_url) + .status(); + let second_status = Command::new(std::env::current_exe()?) + .args([ + "oauth_concurrency_client_child", + "--exact", + "--ignored", + "--nocapture", + ]) + .env("CODEX_HOME", codex_home.path()) + .env(CHILD_SERVER_URL_ENV, &server_url) + .status(); + let (first_status, second_status) = tokio::try_join!(first_status, second_status)?; + assert!( + first_status.success(), + "first OAuth concurrency child failed: {first_status}" + ); + assert!( + second_status.success(), + "second OAuth concurrency child failed: {second_status}" + ); + + server.verify().await; + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn operation_preflight_refresh_failure_blocks_rmcp_request() -> anyhow::Result<()> { + let server = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/.well-known/oauth-authorization-server/mcp")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "authorization_endpoint": format!("{}/oauth/authorize", server.uri()), + "token_endpoint": format!("{}/oauth/token", server.uri()), + "scopes_supported": [""], + }))) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/oauth/token")) + .and(body_string_contains("grant_type=refresh_token")) + .and(body_string_contains(format!( + "refresh_token={REFRESH_TOKEN}" + ))) + .respond_with(ResponseTemplate::new(400).set_body_json(json!({ + "error": "invalid_grant", + "error_description": PREFLIGHT_REFRESH_ERROR, + }))) + .expect(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/mcp")) + .and(header( + "authorization", + format!("Bearer {EXPIRED_ACCESS_TOKEN}"), + )) + .respond_with(|request: &Request| { + let body: Value = request.body_json().expect("valid JSON-RPC request"); + match body.get("method").and_then(Value::as_str) { + Some("initialize") => ResponseTemplate::new(200).set_body_json(json!({ + "jsonrpc": "2.0", + "id": body.get("id").cloned().unwrap_or(Value::Null), + "result": { + "protocolVersion": body + .pointer("/params/protocolVersion") + .cloned() + .unwrap_or_else(|| json!("2025-06-18")), + "capabilities": {}, + "serverInfo": { + "name": "oauth-preflight-test", + "version": "0.0.0-test", + }, + }, + })), + Some("notifications/initialized") => ResponseTemplate::new(202), + method => ResponseTemplate::new(400) + .set_body_string(format!("unexpected JSON-RPC method: {method:?}")), + } + }) + .expect(2) + .mount(&server) + .await; + + let codex_home = TempDir::new()?; + let server_url = format!("{}/mcp", server.uri()); + let status = Command::new(std::env::current_exe()?) + .args([ + "operation_preflight_refresh_failure_child", + "--exact", + "--ignored", + "--nocapture", + ]) + .env("CODEX_HOME", codex_home.path()) + .env(CHILD_SERVER_URL_ENV, &server_url) + .status() + .await?; + assert!( + status.success(), + "OAuth preflight failure child failed: {status}" + ); + + server.verify().await; + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 1)] async fn reports_auth_status_for_persisted_credentials() -> anyhow::Result<()> { let codex_home = TempDir::new()?; @@ -182,7 +496,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, @@ -279,3 +593,197 @@ async fn oauth_startup_child() -> anyhow::Result<()> { initialize_client(&client).await?; Ok(()) } + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +#[ignore = "spawned by recovers_initialization_and_operation_401_without_rmcp_refresh"] +async fn oauth_401_recovery_child() -> anyhow::Result<()> { + let server_url = std::env::var(CHILD_SERVER_URL_ENV)?; + save_no_expiry_file_mode_tokens(&server_url)?; + + let client = new_file_mode_oauth_client(&server_url).await?; + initialize_client(&client).await?; + let tools = client + .list_tools(/*params*/ None, Some(Duration::from_secs(5))) + .await?; + assert!(tools.tools.is_empty()); + client.shutdown().await; + + let reloaded_client = new_file_mode_oauth_client(&server_url).await?; + initialize_client(&reloaded_client).await?; + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +#[ignore = "spawned by operation_preflight_refresh_failure_blocks_rmcp_request"] +async fn operation_preflight_refresh_failure_child() -> anyhow::Result<()> { + let server_url = std::env::var(CHILD_SERVER_URL_ENV)?; + save_near_expiry_file_mode_tokens(&server_url)?; + + let client = RmcpClient::new_streamable_http_client( + SERVER_NAME, + &server_url, + /*bearer_token*/ None, + /*http_headers*/ None, + /*env_http_headers*/ None, + OAuthCredentialsStoreMode::File, + AuthKeyringBackendKind::default(), + Environment::default_for_tests().get_http_client(), + /*auth_provider*/ None, + ) + .await?; + + initialize_client(&client).await?; + + tokio::time::sleep(Duration::from_millis(2_200)).await; + let error = client + .list_tools(/*params*/ None, Some(Duration::from_secs(5))) + .await + .expect_err("preflight refresh failure should abort the operation"); + let message = format!("{error:#}"); + assert!( + message.contains("failed to refresh OAuth tokens for server"), + "unexpected error: {message}" + ); + assert!( + message.contains(PREFLIGHT_REFRESH_ERROR), + "unexpected error: {message}" + ); + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +#[ignore = "spawned by concurrent_file_mode_startup_refreshes_once"] +async fn oauth_concurrency_seed_child() -> anyhow::Result<()> { + let server_url = std::env::var(CHILD_SERVER_URL_ENV)?; + save_expired_file_mode_tokens(&server_url)?; + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +#[ignore = "spawned by concurrent_file_mode_startup_refreshes_once"] +async fn oauth_concurrency_client_child() -> anyhow::Result<()> { + let server_url = std::env::var(CHILD_SERVER_URL_ENV)?; + let client = RmcpClient::new_streamable_http_client( + SERVER_NAME, + &server_url, + /*bearer_token*/ None, + /*http_headers*/ None, + /*env_http_headers*/ None, + OAuthCredentialsStoreMode::File, + AuthKeyringBackendKind::default(), + Environment::default_for_tests().get_http_client(), + /*auth_provider*/ None, + ) + .await?; + + initialize_client(&client).await?; + Ok(()) +} + +fn save_expired_file_mode_tokens(server_url: &str) -> anyhow::Result<()> { + let mut response = OAuthTokenResponse::new( + AccessToken::new(EXPIRED_ACCESS_TOKEN.to_string()), + BasicTokenType::Bearer, + VendorExtraTokenFields::default(), + ); + response.set_refresh_token(Some(RefreshToken::new(REFRESH_TOKEN.to_string()))); + response.set_expires_in(Some(&Duration::from_secs(7200))); + let tokens = StoredOAuthTokens { + server_name: SERVER_NAME.to_string(), + url: server_url.to_string(), + client_id: "test-client-id".to_string(), + token_response: WrappedOAuthTokenResponse(response), + expires_at: Some(0), + }; + save_oauth_tokens( + SERVER_NAME, + &tokens, + OAuthCredentialsStoreMode::File, + AuthKeyringBackendKind::default(), + )?; + Ok(()) +} + +fn save_near_expiry_file_mode_tokens(server_url: &str) -> anyhow::Result<()> { + let mut response = OAuthTokenResponse::new( + AccessToken::new(EXPIRED_ACCESS_TOKEN.to_string()), + BasicTokenType::Bearer, + VendorExtraTokenFields::default(), + ); + response.set_refresh_token(Some(RefreshToken::new(REFRESH_TOKEN.to_string()))); + response.set_expires_in(Some(&Duration::from_secs(62))); + let now = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_else(|_| Duration::from_secs(0)) + .as_millis() as u64; + let tokens = StoredOAuthTokens { + server_name: SERVER_NAME.to_string(), + url: server_url.to_string(), + client_id: "test-client-id".to_string(), + token_response: WrappedOAuthTokenResponse(response), + expires_at: Some(now.saturating_add(/*rhs*/ 62_000)), + }; + save_oauth_tokens( + SERVER_NAME, + &tokens, + OAuthCredentialsStoreMode::File, + AuthKeyringBackendKind::default(), + )?; + Ok(()) +} + +fn save_no_expiry_file_mode_tokens(server_url: &str) -> anyhow::Result<()> { + let mut response = OAuthTokenResponse::new( + AccessToken::new(EXPIRED_ACCESS_TOKEN.to_string()), + BasicTokenType::Bearer, + VendorExtraTokenFields::default(), + ); + response.set_refresh_token(Some(RefreshToken::new(REFRESH_TOKEN.to_string()))); + let tokens = StoredOAuthTokens { + server_name: SERVER_NAME.to_string(), + url: server_url.to_string(), + client_id: "test-client-id".to_string(), + token_response: WrappedOAuthTokenResponse(response), + expires_at: None, + }; + save_oauth_tokens( + SERVER_NAME, + &tokens, + OAuthCredentialsStoreMode::File, + AuthKeyringBackendKind::default(), + )?; + Ok(()) +} + +async fn new_file_mode_oauth_client(server_url: &str) -> anyhow::Result { + RmcpClient::new_streamable_http_client( + SERVER_NAME, + server_url, + /*bearer_token*/ None, + /*http_headers*/ None, + /*env_http_headers*/ None, + OAuthCredentialsStoreMode::File, + AuthKeyringBackendKind::default(), + Environment::default_for_tests().get_http_client(), + /*auth_provider*/ None, + ) + .await +} + +fn initialize_response(body: &Value, name: &str) -> ResponseTemplate { + ResponseTemplate::new(200).set_body_json(json!({ + "jsonrpc": "2.0", + "id": body.get("id").cloned().unwrap_or(Value::Null), + "result": { + "protocolVersion": body + .pointer("/params/protocolVersion") + .cloned() + .unwrap_or_else(|| json!("2025-06-18")), + "capabilities": {}, + "serverInfo": { + "name": name, + "version": "0.0.0-test", + }, + }, + })) +}