From 835f4a1e2b2c8d3927ed762854bc62a07c401482 Mon Sep 17 00:00:00 2001 From: GraciousGazelles <55840159+GraciousGazelles@users.noreply.github.com> Date: Sat, 20 Jun 2026 21:55:18 +1000 Subject: [PATCH 1/2] test(mcp): preserve OAuth manager on keyring reread errors --- codex-rs/rmcp-client/src/oauth.rs | 129 +++++++++++++++++++++++++----- 1 file changed, 107 insertions(+), 22 deletions(-) diff --git a/codex-rs/rmcp-client/src/oauth.rs b/codex-rs/rmcp-client/src/oauth.rs index 516e082cc5d6..9c8082b0a80a 100644 --- a/codex-rs/rmcp-client/src/oauth.rs +++ b/codex-rs/rmcp-client/src/oauth.rs @@ -162,6 +162,26 @@ fn load_oauth_tokens_with_keyring_store( } } +fn load_oauth_tokens_for_refresh_with_keyring_store( + keyring_store: &K, + server_name: &str, + url: &str, + resolved_store: ResolvedOAuthCredentialStore, +) -> Result> { + match resolved_store { + ResolvedOAuthCredentialStore::File => load_oauth_tokens_from_file(server_name, url) + .context("failed to reread OAuth tokens from resolved file storage"), + ResolvedOAuthCredentialStore::Keyring => load_oauth_tokens_from_keyring( + keyring_store, + server_name, + url, + ) + .context( + "failed to reread OAuth tokens from resolved keyring storage; refusing file fallback", + ), + } +} + pub(crate) fn oauth_token_status( server_name: &str, url: &str, @@ -472,26 +492,6 @@ impl OAuthPersistor { .await } - fn load_resolved_credentials(&self) -> Result> { - let keyring_store = DefaultKeyringStore; - match self.inner.credential_store { - ResolvedOAuthCredentialStore::File => { - load_oauth_tokens_from_file(&self.inner.server_name, &self.inner.url) - .context("failed to reread OAuth tokens from resolved file storage") - } - ResolvedOAuthCredentialStore::Keyring => { - load_oauth_tokens_from_keyring( - &keyring_store, - &self.inner.server_name, - &self.inner.url, - ) - .context( - "failed to reread OAuth tokens from resolved keyring storage; refusing file fallback", - ) - } - } - } - fn save_resolved_credentials(&self, tokens: &StoredOAuthTokens) -> Result<()> { let keyring_store = DefaultKeyringStore; match self.inner.credential_store { @@ -502,14 +502,25 @@ impl OAuthPersistor { } } + async fn refresh_transaction( + &self, + reason: RefreshReason, + overall_timeout: Duration, + ) -> Result<()> { + let keyring_store = DefaultKeyringStore; + self.refresh_transaction_with_keyring_store(reason, overall_timeout, &keyring_store) + .await + } + #[expect( clippy::await_holding_invalid_type, reason = "AuthorizationManager async access must be serialized through its mutex" )] - async fn refresh_transaction( + async fn refresh_transaction_with_keyring_store( &self, reason: RefreshReason, overall_timeout: Duration, + keyring_store: &K, ) -> Result<()> { let local_access_token = { let last_credentials = self.inner.last_credentials.lock().await; @@ -531,7 +542,13 @@ impl OAuthPersistor { // Reread after acquiring the cross-process lock. If another Codex // client refreshed the rotating token first, this process must adopt // that newer state instead of replaying its stale refresh token. - let Some(latest) = self.load_resolved_credentials()? else { + let Some(latest) = load_oauth_tokens_for_refresh_with_keyring_store( + keyring_store, + &self.inner.server_name, + &self.inner.url, + self.inner.credential_store, + )? + else { self.clear_manager_credentials().await; let mut last_credentials = self.inner.last_credentials.lock().await; *last_credentials = None; @@ -1565,6 +1582,50 @@ mod tests { assert_eq!(request_tokens.scopes(), tokens.token_response.0.scopes()); } + #[tokio::test(flavor = "current_thread")] + async fn refresh_transaction_keeps_manager_credentials_on_keyring_reread_error() -> Result<()> { + let _env = TempCodexHome::new(); + let store = MockKeyringStore::default(); + let mut tokens = sample_tokens(); + tokens.expires_at = Some(0); + tokens + .token_response + .0 + .set_expires_in(Some(&Duration::ZERO)); + let key = super::compute_store_key(&tokens.server_name, &tokens.url)?; + store.set_error(&key, KeyringError::Invalid("error".into(), "load".into())); + + let manager = authorization_manager_for(&tokens).await?; + let persistor = OAuthPersistor::new( + tokens.server_name.clone(), + tokens.url.clone(), + manager.clone(), + ResolvedOAuthCredentialStore::Keyring, + Some(tokens.clone()), + ); + + let error = persistor + .refresh_transaction_with_keyring_store( + RefreshReason::Expiry, + Duration::from_secs(/*secs*/ 1), + &store, + ) + .await + .expect_err("keyring reread failure should abort refresh"); + + assert!( + format!("{error:#}") + .contains("failed to reread OAuth tokens from resolved keyring storage"), + "unexpected error: {error:#}" + ); + let access_token = manager.lock().await.get_access_token().await?; + assert_eq!( + access_token, + tokens.token_response.0.access_token().secret() + ); + Ok(()) + } + #[tokio::test(flavor = "current_thread")] async fn locked_save_waits_for_refresh_transaction_lock() -> Result<()> { let _env = TempCodexHome::new(); @@ -1601,6 +1662,30 @@ mod tests { Ok(()) } + async fn authorization_manager_for( + tokens: &StoredOAuthTokens, + ) -> Result>> { + let manager = std::sync::Arc::new(tokio::sync::Mutex::new( + AuthorizationManager::new(&tokens.url).await?, + )); + { + let mut guard = manager.lock().await; + guard.set_metadata(rmcp::transport::auth::AuthorizationMetadata { + authorization_endpoint: "https://auth.example.test/authorize".to_string(), + token_endpoint: "https://auth.example.test/token".to_string(), + registration_endpoint: None, + issuer: None, + jwks_uri: None, + scopes_supported: None, + response_types_supported: None, + code_challenge_methods_supported: None, + additional_fields: Default::default(), + }); + } + super::install_tokens_in_manager(&manager, tokens).await?; + Ok(manager) + } + fn assert_tokens_match_without_expiry( actual: &StoredOAuthTokens, expected: &StoredOAuthTokens, From 588caf26dd65a225256ced24b2aa722cecff5716 Mon Sep 17 00:00:00 2001 From: GraciousGazelles <55840159+GraciousGazelles@users.noreply.github.com> Date: Sat, 20 Jun 2026 22:16:11 +1000 Subject: [PATCH 2/2] Fix MCP OAuth reread regression test --- codex-rs/rmcp-client/src/oauth.rs | 29 +++++++++++------------------ 1 file changed, 11 insertions(+), 18 deletions(-) diff --git a/codex-rs/rmcp-client/src/oauth.rs b/codex-rs/rmcp-client/src/oauth.rs index 9c8082b0a80a..a1bc2d091ed5 100644 --- a/codex-rs/rmcp-client/src/oauth.rs +++ b/codex-rs/rmcp-client/src/oauth.rs @@ -1603,13 +1603,10 @@ mod tests { ResolvedOAuthCredentialStore::Keyring, Some(tokens.clone()), ); + let refresh_timeout = Duration::from_secs(1); let error = persistor - .refresh_transaction_with_keyring_store( - RefreshReason::Expiry, - Duration::from_secs(/*secs*/ 1), - &store, - ) + .refresh_transaction_with_keyring_store(RefreshReason::Expiry, refresh_timeout, &store) .await .expect_err("keyring reread failure should abort refresh"); @@ -1621,7 +1618,7 @@ mod tests { let access_token = manager.lock().await.get_access_token().await?; assert_eq!( access_token, - tokens.token_response.0.access_token().secret() + tokens.token_response.0.access_token().secret().as_str() ); Ok(()) } @@ -1668,21 +1665,17 @@ mod tests { let manager = std::sync::Arc::new(tokio::sync::Mutex::new( AuthorizationManager::new(&tokens.url).await?, )); + let store = InMemoryCredentialStore::new(); + store + .save(super::stored_credentials_from_tokens( + tokens, + CredentialExposure::Request, + )) + .await?; { let mut guard = manager.lock().await; - guard.set_metadata(rmcp::transport::auth::AuthorizationMetadata { - authorization_endpoint: "https://auth.example.test/authorize".to_string(), - token_endpoint: "https://auth.example.test/token".to_string(), - registration_endpoint: None, - issuer: None, - jwks_uri: None, - scopes_supported: None, - response_types_supported: None, - code_challenge_methods_supported: None, - additional_fields: Default::default(), - }); + guard.set_credential_store(store); } - super::install_tokens_in_manager(&manager, tokens).await?; Ok(manager) }