diff --git a/codex-rs/cli/src/mcp_cmd.rs b/codex-rs/cli/src/mcp_cmd.rs index 586f9c184b..07e304be03 100644 --- a/codex-rs/cli/src/mcp_cmd.rs +++ b/codex-rs/cli/src/mcp_cmd.rs @@ -29,7 +29,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; @@ -523,12 +523,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 5aae4493ce..df78794508 100644 --- a/codex-rs/rmcp-client/src/lib.rs +++ b/codex-rs/rmcp-client/src/lib.rs @@ -23,6 +23,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 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 7ca5b28fd2..65908f9932 100644 --- a/codex-rs/rmcp-client/src/oauth.rs +++ b/codex-rs/rmcp-client/src/oauth.rs @@ -58,6 +58,9 @@ use codex_keyring_store::KeyringStore; use codex_utils_home_dir::find_codex_home; pub(crate) use self::persistor::OAuthPersistor; +#[cfg(test)] +use self::persistor::install_tokens_in_manager; +use self::refresh_lock::RefreshCredentialLock; pub(crate) use self::resolved_store::LoadedOAuthTokens; pub(crate) use self::resolved_store::ResolvedOAuthCredentialStore; #[cfg(test)] @@ -244,6 +247,43 @@ pub fn save_oauth_tokens( ) } +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<()> { + // Login persistence shares the refresh transaction lock so a completed login always becomes + // authoritative: it either lands before refresh's reread or waits and overwrites the refresh + // result afterward. + 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, @@ -380,6 +420,42 @@ 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 { + // Logout shares the refresh transaction lock so refresh cannot resurrect credentials after a + // completed delete: it either observes the deletion or finishes before logout removes it. + 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, @@ -1264,6 +1340,232 @@ mod tests { Ok(()) } + #[tokio::test] + async fn completed_login_prevents_delayed_post_operation_overwrite() -> Result<()> { + let _env = TempCodexHome::new(); + let server = MockServer::start().await; + let store = MockKeyringStore::default(); + let (initial_tokens, manager, persistor) = + persistor_after_rmcp_refresh(&store, &server).await?; + + let login_tokens = tokens_with_credentials( + initial_tokens.clone(), + "login-access-token", + "login-refresh-token", + ); + super::save_oauth_tokens_locked_with_keyring_store( + &store, + &login_tokens.server_name, + &login_tokens, + OAuthCredentialsStoreMode::Keyring, + AuthKeyringBackendKind::Direct, + ) + .await?; + + persistor + .persist_if_needed_locked_with_keyring_store(&store) + .await?; + + let stored = super::load_oauth_tokens_with_source_and_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.tokens), "login-access-token"); + assert_eq!( + refresh_token(&stored.tokens), + Some("login-refresh-token".to_string()) + ); + let manager_tokens = tokens_from_manager(&manager).await?; + assert_eq!(access_token(&manager_tokens), "login-access-token"); + 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, + ResolvedOAuthCredentialStore::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_source_and_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.tokens), "login-after-refresh-access"); + assert_eq!( + refresh_token(&stored.tokens), + Some("login-after-refresh-token".to_string()) + ); + Ok(()) + } + + #[tokio::test] + #[expect( + clippy::await_holding_invalid_type, + reason = "AuthorizationManager async access must be serialized through its mutex" + )] + async fn completed_logout_prevents_delayed_post_operation_resurrection() -> Result<()> { + let _env = TempCodexHome::new(); + let server = MockServer::start().await; + let store = MockKeyringStore::default(); + let (initial_tokens, manager, persistor) = + persistor_after_rmcp_refresh(&store, &server).await?; + + 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); + + persistor + .persist_if_needed_locked_with_keyring_store(&store) + .await?; + + let stored = super::load_oauth_tokens_with_source_and_keyring_store( + &store, + &initial_tokens.server_name, + &initial_tokens.url, + OAuthCredentialsStoreMode::Keyring, + AuthKeyringBackendKind::Direct, + )?; + assert!(stored.is_none()); + let guard = manager.lock().await; + let (_, manager_tokens) = guard.get_credentials().await?; + assert!(manager_tokens.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, + ResolvedOAuthCredentialStore::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_source_and_keyring_store( + &store, + &initial_tokens.server_name, + &initial_tokens.url, + OAuthCredentialsStoreMode::Keyring, + AuthKeyringBackendKind::Direct, + )?; + assert!(stored.is_none()); + Ok(()) + } + #[tokio::test] async fn provider_refresh_timeout_permits_a_later_serialized_retry() -> Result<()> { let _env = TempCodexHome::new(); @@ -1696,6 +1998,57 @@ mod tests { ); } + async fn persistor_after_rmcp_refresh( + store: &MockKeyringStore, + server: &MockServer, + ) -> Result<( + StoredOAuthTokens, + Arc>, + OAuthPersistor, + )> { + mount_oauth_metadata(server).await; + let mut initial_tokens = sample_tokens(); + initial_tokens.url = 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.clone(), + ResolvedOAuthCredentialStore::Keyring(AuthKeyringBackendKind::Direct), + Some(initial_tokens.clone()), + ); + let rmcp_tokens = tokens_with_credentials( + initial_tokens.clone(), + "rmcp-access-token", + "rmcp-refresh-token", + ); + install_tokens_in_manager(&manager, &rmcp_tokens).await?; + Ok((initial_tokens, manager, persistor)) + } + + fn tokens_with_credentials( + mut tokens: StoredOAuthTokens, + access_token: &str, + refresh_token: &str, + ) -> StoredOAuthTokens { + tokens + .token_response + .0 + .set_access_token(AccessToken::new(access_token.to_string())); + tokens + .token_response + .0 + .set_refresh_token(Some(RefreshToken::new(refresh_token.to_string()))); + tokens + } + async fn authorization_manager_for( tokens: &StoredOAuthTokens, ) -> Result>> { diff --git a/codex-rs/rmcp-client/src/oauth/persistor.rs b/codex-rs/rmcp-client/src/oauth/persistor.rs index 9a740885bb..ff7f6265d8 100644 --- a/codex-rs/rmcp-client/src/oauth/persistor.rs +++ b/codex-rs/rmcp-client/src/oauth/persistor.rs @@ -15,6 +15,7 @@ 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; @@ -73,10 +74,105 @@ impl OAuthPersistor { /// Persists the latest stored credentials if they have changed. /// Deletes the credentials if they are no longer present. pub(crate) async fn persist_if_needed(&self) -> Result<()> { - self.persist_if_needed_with_keyring_store(&DefaultKeyringStore) + self.persist_if_needed_locked_with_keyring_store(&DefaultKeyringStore) .await } + pub(super) async fn persist_if_needed_locked_with_keyring_store< + K: KeyringStore + Clone + 'static, + >( + &self, + keyring_store: &K, + ) -> Result<()> { + let snapshot = { + let last_credentials = self.inner.last_credentials.lock().await; + last_credentials.clone() + }; + let (client_id, current_credentials) = self.manager_credentials().await?; + let manager_changed = match (&snapshot, current_credentials.as_ref()) { + (Some(previous), Some(current)) => { + previous.client_id != client_id + || previous.token_response != WrappedOAuthTokenResponse(current.clone()) + } + (None, None) => false, + (Some(_), None) | (None, Some(_)) => true, + }; + if !manager_changed { + return Ok(()); + } + + let _lock = + RefreshCredentialLock::acquire_for_server(&self.inner.server_name, &self.inner.url) + .await?; + let latest = self.load_resolved_credentials(keyring_store)?; + + let latest_matches_snapshot = match (&latest, &snapshot) { + (Some(latest), Some(snapshot)) => { + // `expires_in` is reconstructed from the durable `expires_at` timestamp on every + // load, so elapsed time alone must not look like a concurrent credential change. + let mut comparable_latest = latest.clone(); + comparable_latest + .token_response + .0 + .set_expires_in(snapshot.token_response.0.expires_in().as_ref()); + comparable_latest == *snapshot + } + (None, None) => true, + (Some(_), None) | (None, Some(_)) => false, + }; + + if !latest_matches_snapshot { + // A completed login or logout is authoritative over tokens refreshed inside an RMCP + // operation. Adopt that state instead of allowing delayed post-operation persistence + // to overwrite the login or resurrect the logout. + match latest { + Some(latest) => self.adopt_credentials(latest).await?, + None => { + self.clear_manager_credentials().await; + let mut last_credentials = self.inner.last_credentials.lock().await; + *last_credentials = None; + } + } + return Ok(()); + } + + self.persist_if_needed_with_keyring_store(keyring_store) + .await + } + + #[expect( + clippy::await_holding_invalid_type, + reason = "AuthorizationManager async access must be serialized through its mutex" + )] + async fn manager_credentials(&self) -> Result<(String, Option)> { + let manager = self.inner.authorization_manager.clone(); + let guard = manager.lock().await; + guard.get_credentials().await.map_err(Into::into) + } + + fn load_resolved_credentials( + &self, + keyring_store: &K, + ) -> Result> { + match self.inner.credential_store { + ResolvedOAuthCredentialStore::File => { + load_oauth_tokens_from_file(&self.inner.server_name, &self.inner.url) + .context("failed to reread OAuth tokens from resolved file storage") + } + ResolvedOAuthCredentialStore::Keyring(keyring_backend_kind) => { + load_oauth_tokens_from_keyring( + keyring_store, + keyring_backend_kind, + &self.inner.server_name, + &self.inner.url, + ) + .context( + "failed to reread OAuth tokens from resolved keyring storage; refusing file fallback", + ) + } + } + } + #[expect( clippy::await_holding_invalid_type, reason = "AuthorizationManager async access must be serialized through its mutex" @@ -241,8 +337,9 @@ impl OAuthPersistor { let transaction_started_at = Instant::now(); let lock_started_at = Instant::now(); debug!("waiting for the MCP OAuth credential transaction lock"); - let key = compute_store_key(&self.inner.server_name, &self.inner.url)?; - let _lock = RefreshCredentialLock::acquire(&key).await?; + 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" @@ -251,23 +348,7 @@ impl OAuthPersistor { // 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. - let latest = match self.inner.credential_store { - ResolvedOAuthCredentialStore::File => { - load_oauth_tokens_from_file(&self.inner.server_name, &self.inner.url) - .context("failed to reread OAuth tokens from resolved file storage")? - } - ResolvedOAuthCredentialStore::Keyring(keyring_backend_kind) => { - load_oauth_tokens_from_keyring( - keyring_store, - keyring_backend_kind, - &self.inner.server_name, - &self.inner.url, - ) - .context( - "failed to reread OAuth tokens from resolved keyring storage; refusing file fallback", - )? - } - }; + let latest = self.load_resolved_credentials(keyring_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 @@ -363,7 +444,7 @@ impl OAuthPersistor { clippy::await_holding_invalid_type, reason = "AuthorizationManager async access must be serialized through its mutex" )] -async fn install_tokens_in_manager( +pub(super) async fn install_tokens_in_manager( authorization_manager: &Arc>, tokens: &StoredOAuthTokens, ) -> Result<()> { diff --git a/codex-rs/rmcp-client/src/oauth/refresh_lock.rs b/codex-rs/rmcp-client/src/oauth/refresh_lock.rs index b540066cd1..17bdb6acb1 100644 --- a/codex-rs/rmcp-client/src/oauth/refresh_lock.rs +++ b/codex-rs/rmcp-client/src/oauth/refresh_lock.rs @@ -30,6 +30,13 @@ pub(super) struct RefreshCredentialLock { } impl RefreshCredentialLock { + pub(super) async fn acquire_for_server(server_name: &str, url: &str) -> Result { + let key = super::compute_store_key(server_name, url)?; + Self::acquire(&key) + .await + .with_context(|| format!("failed to acquire OAuth credential lock for {server_name}")) + } + pub(super) async fn acquire(store_key: &str) -> Result { let codex_home = find_codex_home()?; Self::acquire_in(&codex_home, store_key, REFRESH_LOCK_ACQUIRE_TIMEOUT).await diff --git a/codex-rs/rmcp-client/src/perform_oauth_login.rs b/codex-rs/rmcp-client/src/perform_oauth_login.rs index 7a9d242bc4..cd2b961b81 100644 --- a/codex-rs/rmcp-client/src/perform_oauth_login.rs +++ b/codex-rs/rmcp-client/src/perform_oauth_login.rs @@ -28,8 +28,8 @@ use urlencoding::decode; use crate::StoredOAuthTokens; use crate::WrappedOAuthTokenResponse; use crate::oauth::compute_expires_at_millis; +use crate::oauth::save_oauth_tokens_locked; use crate::oauth_http_client::OAuthHttpClientAdapter; -use crate::save_oauth_tokens; use crate::utils::build_default_headers; use codex_config::types::AuthKeyringBackendKind; use codex_config::types::OAuthCredentialsStoreMode; @@ -623,12 +623,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(()) }