From 448dcbd8c1bf11d07155c292c5e02f3533a082b5 Mon Sep 17 00:00:00 2001 From: Steven Lee Date: Thu, 25 Jun 2026 17:07:00 +0000 Subject: [PATCH 1/6] Add Codex-owned MCP OAuth recovery --- codex-rs/rmcp-client/src/oauth.rs | 623 ++++++++---------- codex-rs/rmcp-client/src/oauth/persistor.rs | 450 ++++++------- .../rmcp-client/src/oauth/refresh_lock.rs | 9 +- codex-rs/rmcp-client/src/rmcp_client.rs | 137 ++-- .../rmcp-client/src/streamable_http_retry.rs | 104 ++- 5 files changed, 642 insertions(+), 681 deletions(-) diff --git a/codex-rs/rmcp-client/src/oauth.rs b/codex-rs/rmcp-client/src/oauth.rs index bfd3c411e7c0..191ac1d54230 100644 --- a/codex-rs/rmcp-client/src/oauth.rs +++ b/codex-rs/rmcp-client/src/oauth.rs @@ -61,8 +61,6 @@ 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; @@ -77,7 +75,8 @@ use rmcp::transport::auth::AuthorizationManager; 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; #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] pub struct StoredOAuthTokens { @@ -1200,15 +1199,7 @@ mod tests { -> Result<()> { let _env = TempCodexHome::new(); 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": ["scope-a", "scope-b"], - }))) - .mount(&server) - .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)?; @@ -1244,29 +1235,43 @@ mod tests { 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", + "updated-access-token", + "updated-refresh-token", + Duration::from_millis(200), + ) + .await; let store = MockKeyringStore::default(); - let mut initial_tokens = sample_tokens(); - initial_tokens.url = format!("{}/mcp", server.uri()); - let mut updated_tokens = initial_tokens.clone(); - updated_tokens - .token_response - .0 - .set_access_token(AccessToken::new("updated-access-token".to_string())); + 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(&updated_tokens).await?; + 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), + Some(initial_tokens.clone()), ); - let key = super::compute_store_key(&updated_tokens.server_name, &updated_tokens.url)?; - store.set_error(&key, KeyringError::Invalid("error".into(), "save".into())); + let refresh_task = tokio::spawn({ + let persistor = persistor.clone(); + let store = store.clone(); + async move { persistor.refresh_if_needed_with_keyring_store(&store).await } + }); - let error = persistor - .persist_if_needed_with_keyring_store(&store) - .await + wait_for_signal(refresh_started).await?; + let key = super::compute_store_key(&initial_tokens.server_name, &initial_tokens.url)?; + store.set_error(&key, KeyringError::Invalid("error".into(), "save".into())); + let error = refresh_task + .await? .expect_err("resolved keyring write should fail instead of falling back"); assert!( @@ -1276,139 +1281,107 @@ mod tests { "unexpected error: {error:#}" ); assert!(!super::fallback_file_path()?.exists()); + let manager_tokens = tokens_from_manager(&manager).await?; + assert_eq!(access_token(&manager_tokens), "updated-access-token"); + assert_eq!( + refresh_token(&manager_tokens), + Some("updated-refresh-token".to_string()) + ); - // B remains authoritative in memory without retrying the failed durable write. + // B is now the in-process authority even though its save failed. A later operation may use + // B while it remains healthy, but we deliberately do not retry persistence here. If B + // later needs refresh, the persistor refuses to reread and replay durable A. persistor - .persist_if_needed_with_keyring_store(&store) + .refresh_if_needed_with_keyring_store(&store) .await?; persistor - .refresh_if_needed_with_keyring_store(&store) + .refresh_after_unauthorized_with_keyring_store( + &store, + initial_tokens.token_response.0.access_token().clone(), + ) .await?; - let manager_tokens = tokens_from_manager(&manager).await?; - assert_eq!(access_token(&manager_tokens), "updated-access-token"); + let error = persistor + .refresh_after_unauthorized_with_keyring_store( + &store, + AccessToken::new("updated-access-token".to_string()), + ) + .await + .expect_err("B cannot be refreshed by rereading durable A"); + assert!( + error + .to_string() + .contains("previous refresh succeeded but its credentials were not persisted"), + "unexpected error: {error:#}" + ); + server.verify().await; Ok(()) } #[tokio::test] - async fn unpersisted_refresh_does_not_reinstall_unchanged_durable_credentials() -> Result<()> { + 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; - Mock::given(method("POST")) - .and(path("/oauth/token")) - .respond_with(ResponseTemplate::new(500)) - .expect(0) - .mount(&server) - .await; + mount_refresh_response( + &server, + "refresh-token", + "refreshed-access-token", + "rotated-refresh-token", + ) + .await; - let failing_store = MockKeyringStore::default(); + 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( - &failing_store, + &store, &initial_tokens.server_name, &initial_tokens, OAuthCredentialsStoreMode::Keyring, AuthKeyringBackendKind::Direct, )?; - let mut unpersisted_tokens = initial_tokens.clone(); - unpersisted_tokens - .token_response - .0 - .set_access_token(AccessToken::new("unpersisted-access-token".to_string())); - let manager = authorization_manager_for(&unpersisted_tokens).await?; - let persistor = OAuthPersistor::new( + 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(), - manager.clone(), + first_manager, ResolvedOAuthCredentialStore::Keyring(AuthKeyringBackendKind::Direct), Some(initial_tokens.clone()), ); - let key = super::compute_store_key(&initial_tokens.server_name, &initial_tokens.url)?; - failing_store.set_error(&key, KeyringError::Invalid("error".into(), "save".into())); - persistor - .persist_if_needed_with_keyring_store(&failing_store) - .await - .expect_err("the refreshed credential should remain unpersisted"); - - let unchanged_store = MockKeyringStore::default(); - super::save_oauth_tokens_with_keyring_store( - &unchanged_store, - &initial_tokens.server_name, - &initial_tokens, - OAuthCredentialsStoreMode::Keyring, - AuthKeyringBackendKind::Direct, - )?; - let mut subsequent_tokens = unpersisted_tokens.clone(); - subsequent_tokens - .token_response - .0 - .set_access_token(AccessToken::new("subsequent-access-token".to_string())); - install_tokens_in_manager(&manager, &subsequent_tokens).await?; - - persistor - .persist_if_needed_locked_with_keyring_store(&unchanged_store) - .await?; - - let manager_tokens = tokens_from_manager(&manager).await?; - assert_eq!( - manager_tokens.token_response, - subsequent_tokens.token_response - ); - let stored = super::load_oauth_tokens_from_keyring( - &unchanged_store, - AuthKeyringBackendKind::Direct, - &subsequent_tokens.server_name, - &subsequent_tokens.url, - )? - .expect("the subsequent credentials should replace the unchanged stale snapshot"); - assert_eq!(stored.token_response, subsequent_tokens.token_response); - server.verify().await; - Ok(()) - } - - #[tokio::test] - async fn resolved_secrets_write_does_not_cleanup_fallback_file() -> Result<()> { - let _env = TempCodexHome::new(); - let server = MockServer::start().await; - mount_oauth_metadata(&server).await; - let store = MockKeyringStore::default(); - let mut initial_tokens = sample_tokens(); - initial_tokens.url = format!("{}/mcp", server.uri()); - super::save_oauth_tokens_to_file(&initial_tokens)?; - let fallback_path = super::fallback_file_path()?; - let fallback_before = fs::read(&fallback_path)?; - - let mut updated_tokens = initial_tokens.clone(); - updated_tokens - .token_response - .0 - .set_access_token(AccessToken::new("updated-access-token".to_string())); - let manager = authorization_manager_for(&updated_tokens).await?; - let persistor = OAuthPersistor::new( + let second = OAuthPersistor::new( initial_tokens.server_name.clone(), initial_tokens.url.clone(), - manager, - ResolvedOAuthCredentialStore::Keyring(AuthKeyringBackendKind::Secrets), + second_manager.clone(), + ResolvedOAuthCredentialStore::Keyring(AuthKeyringBackendKind::Direct), Some(initial_tokens.clone()), ); - persistor - .persist_if_needed_with_keyring_store(&store) + first + .refresh_after_unauthorized_with_keyring_store(&store) + .await?; + second + .refresh_after_unauthorized_with_keyring_store(&store) .await?; - assert_eq!(fs::read(&fallback_path)?, fallback_before); - let secrets_tokens = super::load_oauth_tokens_from_keyring( + server.verify().await; + let stored = super::load_oauth_tokens_from_keyring( &store, - AuthKeyringBackendKind::Secrets, - &updated_tokens.server_name, - &updated_tokens.url, + AuthKeyringBackendKind::Direct, + &initial_tokens.server_name, + &initial_tokens.url, )? - .expect("updated tokens should be persisted to resolved Secrets storage"); - let mut expected_tokens = updated_tokens; - expected_tokens.expires_at = secrets_tokens.expires_at; - assert_tokens_match_without_expiry(&secrets_tokens, &expected_tokens); + .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, stored.token_response); Ok(()) } @@ -1461,19 +1434,15 @@ mod tests { .await?; let manager_tokens = tokens_from_manager(&manager).await?; - assert_eq!( - access_token(&manager_tokens), - "already-refreshed-access-token" - ); - assert_eq!( - refresh_token(&manager_tokens), - Some("already-rotated-refresh-token".to_string()) + assert_token_response_match_without_expiry( + &manager_tokens.token_response, + &latest_tokens.token_response, ); Ok(()) } #[tokio::test] - async fn refresh_transaction_refreshes_when_only_derived_expires_in_drifted() -> Result<()> { + 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; @@ -1488,7 +1457,7 @@ mod tests { 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(5_000)); + initial_tokens.expires_at = Some(now_millis().saturating_add(45_000)); initial_tokens .token_response .0 @@ -1599,29 +1568,60 @@ mod tests { } #[tokio::test] - async fn completed_login_prevents_delayed_post_operation_overwrite() -> Result<()> { + async fn provider_refresh_timeout_permits_a_later_serialized_retry() -> 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?; + mount_oauth_metadata(&server).await; + let request_count = Arc::new(AtomicUsize::new(/*v*/ 0)); + let request_count_for_response = Arc::clone(&request_count); + Mock::given(method("POST")) + .and(path("/oauth/token")) + .and(body_string_contains("grant_type=refresh_token")) + .and(body_string_contains("refresh_token=refresh-token")) + .respond_with(move |_request: &wiremock::Request| { + let response = ResponseTemplate::new(200).set_body_json(json!({ + "access_token": "retried-access-token", + "token_type": "Bearer", + "expires_in": 3600, + "refresh_token": "retried-refresh-token", + "scope": "scope-a scope-b", + })); + if request_count_for_response.fetch_add(1, Ordering::SeqCst) == 0 { + response.set_delay(Duration::from_millis(500)) + } else { + response + } + }) + .expect(2) + .mount(&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( + let store = MockKeyringStore::default(); + let initial_tokens = expired_sample_tokens(&format!("{}/mcp", server.uri())); + super::save_oauth_tokens_with_keyring_store( &store, - &login_tokens.server_name, - &login_tokens, + &initial_tokens.server_name, + &initial_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, + ResolvedOAuthCredentialStore::Keyring(AuthKeyringBackendKind::Direct), + Some(initial_tokens.clone()), + ); + + let first_error = persistor + .refresh_if_needed_with_keyring_store_and_timeout(&store, Duration::from_millis(100)) + .await + .expect_err("the first provider request should reach its explicit timeout"); + assert!(first_error.to_string().contains("timed out after 100ms")); persistor - .persist_if_needed_locked_with_keyring_store(&store) + .refresh_if_needed_with_keyring_store_and_timeout(&store, Duration::from_secs(1)) .await?; let stored = super::load_oauth_tokens_with_source_and_keyring_store( @@ -1631,28 +1631,28 @@ mod tests { OAuthCredentialsStoreMode::Keyring, AuthKeyringBackendKind::Direct, )? - .expect("login tokens should remain persisted"); - assert_eq!(access_token(&stored.tokens), "login-access-token"); + .expect("the later retry should persist the rotated credentials"); + assert_eq!(access_token(&stored.tokens), "retried-access-token"); assert_eq!( refresh_token(&stored.tokens), - Some("login-refresh-token".to_string()) + Some("retried-refresh-token".to_string()) ); - let manager_tokens = tokens_from_manager(&manager).await?; - assert_eq!(access_token(&manager_tokens), "login-access-token"); + assert_eq!(request_count.load(Ordering::SeqCst), 2); + server.verify().await; Ok(()) } #[tokio::test] - async fn locked_login_save_after_refresh_still_wins() -> Result<()> { + async fn caller_cancellation_does_not_cancel_refresh_and_persistence() -> 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), + "cancel-safe-access-token", + "cancel-safe-refresh-token", + Duration::from_millis(300), ) .await; @@ -1665,7 +1665,6 @@ mod tests { OAuthCredentialsStoreMode::Keyring, AuthKeyringBackendKind::Direct, )?; - let manager = authorization_manager_for(&initial_tokens).await?; let persistor = OAuthPersistor::new( initial_tokens.server_name.clone(), @@ -1674,37 +1673,29 @@ mod tests { ResolvedOAuthCredentialStore::Keyring(AuthKeyringBackendKind::Direct), Some(initial_tokens.clone()), ); - let refresh_task = tokio::spawn({ + let caller = 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?; + caller.abort(); + let caller_error = caller + .await + .expect_err("the caller task should observe cancellation"); + assert!(caller_error.is_cancelled()); - 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, + // The provider handler fires only after the owned transaction has acquired this lock. + // Reacquiring it therefore waits deterministically for refresh and persistence to finish, + // without relying on a scheduler-sensitive sleep after aborting the caller. + let store_key = super::compute_store_key(&initial_tokens.server_name, &initial_tokens.url)?; + let _lock = tokio::time::timeout( + Duration::from_secs(/*secs*/ 2), + RefreshCredentialLock::acquire(&store_key), ) - .await?; - refresh_task.await??; - - server.verify().await; + .await + .expect("the independently owned refresh should release its credential lock")?; let stored = super::load_oauth_tokens_with_source_and_keyring_store( &store, &initial_tokens.server_name, @@ -1712,66 +1703,27 @@ mod tests { OAuthCredentialsStoreMode::Keyring, AuthKeyringBackendKind::Direct, )? - .expect("login tokens should remain persisted"); - assert_eq!(access_token(&stored.tokens), "login-after-refresh-access"); + .expect("the independently owned refresh should still persist credentials"); + assert_eq!(access_token(&stored.tokens), "cancel-safe-access-token"); assert_eq!( refresh_token(&stored.tokens), - Some("login-after-refresh-token".to_string()) + Some("cancel-safe-refresh-token".to_string()) ); + server.verify().await; 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<()> { + 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", - "refreshed-before-logout", - "rotated-before-logout", - Duration::from_millis(200), + "late-access-token", + "late-refresh-token", + Duration::from_secs(1), ) .await; @@ -1796,62 +1748,59 @@ mod tests { let refresh_task = tokio::spawn({ let persistor = persistor.clone(); let store = store.clone(); - async move { persistor.refresh_if_needed_with_keyring_store(&store).await } + 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 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 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_source_and_keyring_store( &store, &initial_tokens.server_name, &initial_tokens.url, OAuthCredentialsStoreMode::Keyring, AuthKeyringBackendKind::Direct, - )?; - assert!(stored.is_none()); + )? + .expect("original tokens should remain persisted"); + assert_eq!(access_token(&stored.tokens), "access-token"); + assert_eq!( + refresh_token(&stored.tokens), + Some("refresh-token".to_string()) + ); + server.verify().await; Ok(()) } #[tokio::test] - async fn provider_refresh_timeout_permits_a_later_serialized_retry() -> Result<()> { + 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 request_count = Arc::new(AtomicUsize::new(/*v*/ 0)); - let request_count_for_response = Arc::clone(&request_count); - Mock::given(method("POST")) - .and(path("/oauth/token")) - .and(body_string_contains("grant_type=refresh_token")) - .and(body_string_contains("refresh_token=refresh-token")) - .respond_with(move |_request: &wiremock::Request| { - let response = ResponseTemplate::new(200).set_body_json(json!({ - "access_token": "retried-access-token", - "token_type": "Bearer", - "expires_in": 3600, - "refresh_token": "retried-refresh-token", - "scope": "scope-a scope-b", - })); - if request_count_for_response.fetch_add(1, Ordering::SeqCst) == 0 { - response.set_delay(Duration::from_millis(500)) - } else { - response - } - }) - .expect(2) - .mount(&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())); @@ -1862,6 +1811,7 @@ mod tests { OAuthCredentialsStoreMode::Keyring, AuthKeyringBackendKind::Direct, )?; + let manager = authorization_manager_for(&initial_tokens).await?; let persistor = OAuthPersistor::new( initial_tokens.server_name.clone(), @@ -1870,17 +1820,37 @@ mod tests { 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 } + }); - let first_error = persistor - .refresh_if_needed_with_keyring_store_and_timeout(&store, Duration::from_millis(100)) - .await - .expect_err("the first provider request should reach its explicit timeout"); - assert!(first_error.to_string().contains("timed out after 100ms")); + wait_for_signal(refresh_started).await?; - persistor - .refresh_if_needed_with_keyring_store_and_timeout(&store, Duration::from_secs(1)) - .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, @@ -1888,28 +1858,26 @@ mod tests { OAuthCredentialsStoreMode::Keyring, AuthKeyringBackendKind::Direct, )? - .expect("the later retry should persist the rotated credentials"); - assert_eq!(access_token(&stored.tokens), "retried-access-token"); + .expect("login tokens should remain persisted"); + assert_eq!(access_token(&stored.tokens), "login-after-refresh-access"); assert_eq!( refresh_token(&stored.tokens), - Some("retried-refresh-token".to_string()) + Some("login-after-refresh-token".to_string()) ); - assert_eq!(request_count.load(Ordering::SeqCst), 2); - server.verify().await; Ok(()) } #[tokio::test] - async fn caller_cancellation_does_not_cancel_refresh_and_persistence() -> Result<()> { + 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", - "cancel-safe-access-token", - "cancel-safe-refresh-token", - Duration::from_millis(300), + "refreshed-before-logout", + "rotated-before-logout", + Duration::from_millis(200), ) .await; @@ -1922,6 +1890,7 @@ mod tests { OAuthCredentialsStoreMode::Keyring, AuthKeyringBackendKind::Direct, )?; + let manager = authorization_manager_for(&initial_tokens).await?; let persistor = OAuthPersistor::new( initial_tokens.server_name.clone(), @@ -1930,43 +1899,34 @@ mod tests { ResolvedOAuthCredentialStore::Keyring(AuthKeyringBackendKind::Direct), Some(initial_tokens.clone()), ); - let caller = tokio::spawn({ + 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?; - caller.abort(); - let caller_error = caller - .await - .expect_err("the caller task should observe cancellation"); - assert!(caller_error.is_cancelled()); - // The provider handler fires only after the owned transaction has acquired this lock. - // Reacquiring it therefore waits deterministically for refresh and persistence to finish, - // without relying on a scheduler-sensitive sleep after aborting the caller. - let store_key = super::compute_store_key(&initial_tokens.server_name, &initial_tokens.url)?; - let _lock = tokio::time::timeout( - Duration::from_secs(/*secs*/ 2), - RefreshCredentialLock::acquire(&store_key), + let removed = super::delete_oauth_tokens_locked_with_keyring_store( + &store, + OAuthCredentialsStoreMode::Keyring, + AuthKeyringBackendKind::Direct, + &initial_tokens.server_name, + &initial_tokens.url, ) - .await - .expect("the independently owned refresh should release its credential lock")?; + .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, - )? - .expect("the independently owned refresh should still persist credentials"); - assert_eq!(access_token(&stored.tokens), "cancel-safe-access-token"); - assert_eq!( - refresh_token(&stored.tokens), - Some("cancel-safe-refresh-token".to_string()) - ); - server.verify().await; + )?; + assert!(stored.is_none()); Ok(()) } @@ -2261,57 +2221,6 @@ 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 3748da5dc218..ff3273a68abc 100644 --- a/codex-rs/rmcp-client/src/oauth/persistor.rs +++ b/codex-rs/rmcp-client/src/oauth/persistor.rs @@ -8,7 +8,6 @@ use std::time::UNIX_EPOCH; use anyhow::Context; use anyhow::Result; -use codex_config::types::AuthKeyringBackendKind; use codex_keyring_store::DefaultKeyringStore; use codex_keyring_store::KeyringStore; use oauth2::TokenResponse; @@ -26,10 +25,6 @@ use super::ResolvedOAuthCredentialStore; use super::StoredOAuthTokens; use super::WrappedOAuthTokenResponse; use super::compute_expires_at_millis; -use super::compute_store_key; -use super::delete_oauth_tokens_from_direct_keyring; -use super::delete_oauth_tokens_from_file; -use super::delete_oauth_tokens_from_secrets_keyring; use super::load_oauth_tokens_from_file; use super::load_oauth_tokens_from_keyring; use super::refresh_lock::RefreshCredentialLock; @@ -43,7 +38,6 @@ const REFRESH_REQUEST_TIMEOUT: Duration = Duration::from_secs(45); pub(crate) struct OAuthPersistor { inner: Arc, } - struct OAuthPersistorInner { server_name: String, url: String, @@ -111,79 +105,6 @@ 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_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, unpersisted_refresh) = { - let state = self.inner.credential_state.lock().await; - (state.current.clone(), state.unpersisted_refresh.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 = durable_credentials_match_snapshot(&latest, &snapshot) - || unpersisted_refresh.as_ref().is_some_and(|pending| { - // The failed write never changed durable authority. If storage still contains the - // last known durable snapshot, persist the manager's newer result instead of - // adopting and replaying that stale snapshot. - durable_credentials_match_snapshot(&latest, &pending.previously_persisted) - }); - - 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 state = self.inner.credential_state.lock().await; - state.current = None; - state.unpersisted_refresh = 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, @@ -207,128 +128,16 @@ impl OAuthPersistor { } } - #[expect( - clippy::await_holding_invalid_type, - reason = "AuthorizationManager async access must be serialized through its mutex" - )] - pub(super) async fn persist_if_needed_with_keyring_store( - &self, - keyring_store: &K, - ) -> 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 state = self.inner.credential_state.lock().await; - let new_token_response = WrappedOAuthTokenResponse(credentials.clone()); - let same_token = state - .current - .as_ref() - .map(|prev| prev.token_response == new_token_response) - .unwrap_or(false); - let expires_at = if same_token { - state.current.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 state.current.as_ref() != Some(&stored) { - // The provider may already have consumed the old rotating refresh token. Make - // B authoritative in this process before the fallible save so a later public - // operation cannot reinstall A from the last snapshot. - // Preserve the last snapshot known to be durable across repeated in-memory - // changes. Using the immediately preceding in-memory token here could make an - // unchanged stale store look like an external login or refresh. - let previously_persisted = state - .unpersisted_refresh - .take() - .map(|pending| pending.previously_persisted) - .unwrap_or_else(|| state.current.clone()); - state.current = Some(stored.clone()); - state.unpersisted_refresh = Some(UnpersistedRefresh { - previously_persisted, - }); - debug!("persisting refreshed MCP OAuth credentials to the resolved store"); - let persistence_started_at = Instant::now(); - let persistence_result = match self.inner.credential_store { - ResolvedOAuthCredentialStore::File => save_oauth_tokens_to_file(&stored), - ResolvedOAuthCredentialStore::Keyring(keyring_backend_kind) => { - save_oauth_tokens_with_keyring( - keyring_store, - keyring_backend_kind, - &self.inner.server_name, - &stored, - ) - } - }; - if let Err(error) = persistence_result { - warn!( - persistence_elapsed_ms = persistence_started_at.elapsed().as_millis(), - error = %error, - "failed to persist refreshed MCP OAuth credentials; retaining them as the in-process authority without retrying persistence" - ); - return Err(error); - } - state.unpersisted_refresh = None; - debug!( - persistence_elapsed_ms = persistence_started_at.elapsed().as_millis(), - "persisted refreshed MCP OAuth credentials" - ); - } - } - None => { - let mut state = self.inner.credential_state.lock().await; - if state.current.take().is_some() - && let Err(error) = match self.inner.credential_store { - ResolvedOAuthCredentialStore::File => { - let key = compute_store_key(&self.inner.server_name, &self.inner.url)?; - delete_oauth_tokens_from_file(&key).map(|_| ()) - } - ResolvedOAuthCredentialStore::Keyring(AuthKeyringBackendKind::Direct) => { - delete_oauth_tokens_from_direct_keyring( - keyring_store, - &self.inner.server_name, - &self.inner.url, - ) - .map(|_| ()) - } - ResolvedOAuthCredentialStore::Keyring(AuthKeyringBackendKind::Secrets) => { - delete_oauth_tokens_from_secrets_keyring( - keyring_store, - &self.inner.server_name, - &self.inner.url, - ) - .map(|_| ()) - } - } - { - warn!( - "failed to remove OAuth tokens for server {}: {error}", - self.inner.server_name - ); - } - state.unpersisted_refresh = None; - } - } - - Ok(()) - } - 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 + } + pub(super) async fn refresh_if_needed_with_keyring_store( &self, keyring_store: &K, @@ -356,13 +165,32 @@ impl OAuthPersistor { return Ok(()); } - self.run_owned_refresh_transaction(keyring_store.clone(), refresh_request_timeout) - .await + self.run_owned_refresh_transaction( + keyring_store.clone(), + RefreshReason::Expiry, + refresh_request_timeout, + ) + .await + } + + pub(super) async fn refresh_after_unauthorized_with_keyring_store< + K: KeyringStore + Clone + 'static, + >( + &self, + keyring_store: &K, + ) -> Result<()> { + self.run_owned_refresh_transaction( + keyring_store.clone(), + RefreshReason::Unauthorized, + REFRESH_REQUEST_TIMEOUT, + ) + .await } async fn run_owned_refresh_transaction( &self, keyring_store: K, + reason: RefreshReason, refresh_request_timeout: Duration, ) -> Result<()> { let persistor = self.clone(); @@ -378,7 +206,7 @@ impl OAuthPersistor { // token-family-revocation risk rather than holding the lock indefinitely. tokio::spawn(async move { persistor - .refresh_transaction(&keyring_store, refresh_request_timeout) + .refresh_transaction(&keyring_store, reason, refresh_request_timeout) .await }) .await @@ -392,16 +220,27 @@ impl OAuthPersistor { #[tracing::instrument( level = "debug", skip_all, - fields(server_name = %self.inner.server_name), + fields( + server_name = %self.inner.server_name, + refresh_reason = reason.as_str(), + ), err )] async fn refresh_transaction( &self, keyring_store: &K, + reason: RefreshReason, refresh_request_timeout: Duration, ) -> Result<()> { let transaction_started_at = Instant::now(); let lock_started_at = Instant::now(); + let local_access_token = { + let state = self.inner.credential_state.lock().await; + state + .current + .as_ref() + .map(|tokens| tokens.token_response.0.access_token().secret().to_string()) + }; debug!("waiting for the MCP OAuth credential transaction lock"); let _lock = RefreshCredentialLock::acquire_for_server(&self.inner.server_name, &self.inner.url) @@ -414,10 +253,11 @@ impl OAuthPersistor { let unpersisted_refresh = { let state = self.inner.credential_state.lock().await; if let Some(unpersisted_refresh) = state.unpersisted_refresh.as_ref() { - if state - .current - .as_ref() - .is_some_and(|tokens| !token_needs_refresh(tokens.expires_at)) + if matches!(reason, RefreshReason::Expiry) + && state + .current + .as_ref() + .is_some_and(|tokens| !token_needs_refresh(tokens.expires_at)) { debug!( "using the memory-authoritative MCP OAuth credentials from a refresh whose persistence failed" @@ -464,67 +304,114 @@ impl OAuthPersistor { ); }; - if !token_needs_refresh(latest.expires_at) { + let latest_access_token = latest.token_response.0.access_token().secret(); + // Expiry refresh can adopt any reread that is now healthy. A 401 is different: an + // unexpired token is still rejected, so adopt only when another actor has already changed + // the access token; otherwise force one serialized provider refresh. + 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(()); } - self.adopt_credentials(latest).await?; - - { - let manager = self.inner.authorization_manager.clone(); - let guard = manager.lock().await; - let provider_started_at = Instant::now(); - debug!( - timeout_ms = refresh_request_timeout.as_millis(), - "requesting refreshed MCP OAuth credentials from the provider" - ); - match timeout(refresh_request_timeout, guard.refresh_token()).await { - Ok(Ok(_token_response)) => { - debug!( - provider_elapsed_ms = provider_started_at.elapsed().as_millis(), - "received refreshed MCP OAuth credentials from the provider" - ); - } - Ok(Err(error)) => { - warn!( - provider_elapsed_ms = provider_started_at.elapsed().as_millis(), - error = %error, - "MCP OAuth provider refresh failed" - ); - return Err(error).with_context(|| { - format!( - "failed to refresh OAuth tokens for server {}", - self.inner.server_name - ) - }); - } - Err(_) => { - warn!( - provider_elapsed_ms = provider_started_at.elapsed().as_millis(), - timeout_ms = refresh_request_timeout.as_millis(), - "MCP OAuth provider refresh timed out; the outcome is unknown and a later serialized retry is permitted" - ); - anyhow::bail!( - "timed out after {refresh_request_timeout:?} refreshing OAuth tokens for server {}", + let manager = self.inner.authorization_manager.clone(); + let mut guard = manager.lock().await; + install_tokens_in_manager_guard(&mut guard, &latest) + .await + .context("failed to stage OAuth credentials for refresh")?; + let provider_started_at = Instant::now(); + debug!( + timeout_ms = refresh_request_timeout.as_millis(), + "requesting refreshed MCP OAuth credentials from the provider" + ); + let refreshed = match timeout(refresh_request_timeout, guard.refresh_token()).await { + Ok(Ok(token_response)) => { + debug!( + provider_elapsed_ms = provider_started_at.elapsed().as_millis(), + "received refreshed MCP OAuth credentials from the provider" + ); + refreshed_tokens(token_response, &latest, &self.inner) + } + Ok(Err(error)) => { + warn!( + provider_elapsed_ms = provider_started_at.elapsed().as_millis(), + error = %error, + "MCP OAuth provider refresh failed" + ); + return Err(error).with_context(|| { + format!( + "failed to refresh OAuth tokens for server {}", self.inner.server_name - ); - } + ) + }); } - } + Err(_) => { + warn!( + provider_elapsed_ms = provider_started_at.elapsed().as_millis(), + timeout_ms = refresh_request_timeout.as_millis(), + "MCP OAuth provider refresh timed out; the outcome is unknown and a later serialized retry is permitted" + ); + anyhow::bail!( + "timed out after {refresh_request_timeout:?} refreshing OAuth tokens for server {}", + self.inner.server_name + ); + } + }; + install_tokens_in_manager_guard(&mut guard, &refreshed) + .await + .context("failed to install refreshed OAuth credentials")?; + drop(guard); - // Once the provider returns a rotated token, persistence must finish before the credential - // lock is released. In particular, caller startup deadlines must not cancel this step. - let result = self - .persist_if_needed_with_keyring_store(keyring_store) - .await; - if result.is_ok() { - debug!( + { + let mut state = self.inner.credential_state.lock().await; + state.current = Some(refreshed.clone()); + state.unpersisted_refresh = Some(UnpersistedRefresh { + previously_persisted: Some(latest.clone()), + }); + } + debug!( + "installed refreshed MCP OAuth credentials in memory and marked persistence pending" + ); + // Once the provider rotates a refresh token, persistence must complete even if the caller's + // deadline expires in the meantime. Returning early here would lose the only usable token. + // Refresh persistence stays on the source resolved at client startup. In particular, a + // keyring failure must surface instead of writing the rotated token to fallback File. + debug!("persisting refreshed MCP OAuth credentials to the resolved store"); + let persistence_started_at = Instant::now(); + let persistence_result = match self.inner.credential_store { + ResolvedOAuthCredentialStore::File => save_oauth_tokens_to_file(&refreshed), + ResolvedOAuthCredentialStore::Keyring(keyring_backend_kind) => { + save_oauth_tokens_with_keyring( + keyring_store, + keyring_backend_kind, + &self.inner.server_name, + &refreshed, + ) + } + }; + if let Err(error) = persistence_result { + warn!( + persistence_elapsed_ms = persistence_started_at.elapsed().as_millis(), transaction_elapsed_ms = transaction_started_at.elapsed().as_millis(), - "completed the MCP OAuth refresh transaction" + error = %error, + "failed to persist refreshed MCP OAuth credentials; retaining them as the in-process authority without retrying persistence" ); + return Err(error); } - result + let mut state = self.inner.credential_state.lock().await; + state.unpersisted_refresh = None; + debug!( + persistence_elapsed_ms = persistence_started_at.elapsed().as_millis(), + transaction_elapsed_ms = transaction_started_at.elapsed().as_millis(), + "persisted refreshed MCP OAuth credentials and completed the transaction" + ); + Ok(()) } async fn adopt_credentials(&self, tokens: StoredOAuthTokens) -> Result<()> { @@ -542,13 +429,37 @@ impl OAuthPersistor { } } +#[derive(Clone, Copy)] +enum RefreshReason { + Expiry, + Unauthorized, +} + +impl RefreshReason { + fn as_str(&self) -> &'static str { + match self { + Self::Expiry => "expiry", + Self::Unauthorized => "unauthorized", + } + } +} + #[expect( clippy::await_holding_invalid_type, reason = "AuthorizationManager async access must be serialized through its mutex" )] -pub(super) async fn install_tokens_in_manager( +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).await +} + +async fn install_tokens_in_manager_guard( + authorization_manager: &mut AuthorizationManager, + tokens: &StoredOAuthTokens, ) -> Result<()> { let store = InMemoryCredentialStore::new(); store @@ -556,13 +467,11 @@ pub(super) async fn install_tokens_in_manager( .await .context("failed to stage OAuth tokens for authorization manager")?; - let manager = authorization_manager.clone(); - let mut guard = manager.lock().await; - guard.set_credential_store(store); + authorization_manager.set_credential_store(store); // TODO(stevenlee): RMCP's `initialize_from_store` updates the credential store and client ID // but not its private `current_scopes`. Credential adoption can therefore leave scope-upgrade // state stale until RMCP exposes an adoption API that synchronizes both. - guard + authorization_manager .initialize_from_store() .await .context("failed to adopt refreshed OAuth tokens")?; @@ -587,3 +496,24 @@ fn stored_credentials_from_tokens(tokens: &StoredOAuthTokens) -> StoredCredentia token_received_at, ) } + +fn refreshed_tokens( + mut token_response: OAuthTokenResponse, + previous: &StoredOAuthTokens, + inner: &OAuthPersistorInner, +) -> StoredOAuthTokens { + if token_response.refresh_token().is_none() { + token_response.set_refresh_token(previous.token_response.0.refresh_token().cloned()); + } + if token_response.scopes().is_none() { + token_response.set_scopes(previous.token_response.0.scopes().cloned()); + } + let expires_at = compute_expires_at_millis(&token_response); + StoredOAuthTokens { + server_name: inner.server_name.clone(), + url: inner.url.clone(), + client_id: previous.client_id.clone(), + token_response: WrappedOAuthTokenResponse(token_response), + expires_at, + } +} diff --git a/codex-rs/rmcp-client/src/oauth/refresh_lock.rs b/codex-rs/rmcp-client/src/oauth/refresh_lock.rs index 17bdb6acb143..51577606d986 100644 --- a/codex-rs/rmcp-client/src/oauth/refresh_lock.rs +++ b/codex-rs/rmcp-client/src/oauth/refresh_lock.rs @@ -38,8 +38,15 @@ impl RefreshCredentialLock { } pub(super) async fn acquire(store_key: &str) -> Result { + Self::acquire_with_timeout(store_key, REFRESH_LOCK_ACQUIRE_TIMEOUT).await + } + + pub(super) async fn acquire_with_timeout( + store_key: &str, + acquire_timeout: Duration, + ) -> Result { let codex_home = find_codex_home()?; - Self::acquire_in(&codex_home, store_key, REFRESH_LOCK_ACQUIRE_TIMEOUT).await + Self::acquire_in(&codex_home, store_key, acquire_timeout).await } async fn acquire_in( diff --git a/codex-rs/rmcp-client/src/rmcp_client.rs b/codex-rs/rmcp-client/src/rmcp_client.rs index eb3e0deb53f6..9b74721f6d05 100644 --- a/codex-rs/rmcp-client/src/rmcp_client.rs +++ b/codex-rs/rmcp-client/src/rmcp_client.rs @@ -453,7 +453,7 @@ impl RmcpClient { let mut initialize_deadline = timeout.map(|duration| Instant::now() + duration); let (service, oauth_persistor) = self - .connect_pending_transport_with_initialize_retries( + .connect_pending_transport_with_oauth_recovery( pending_transport, client_service.clone(), timeout, @@ -486,12 +486,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) } @@ -507,7 +501,6 @@ impl RmcpClient { async move { service.list_tools(params).await }.boxed() }) .await?; - self.persist_oauth_tokens().await; Ok(result) } @@ -542,7 +535,6 @@ impl RmcpClient { }) }) .collect::>>()?; - self.persist_oauth_tokens().await; Ok(ListToolsWithConnectorIdResult { next_cursor: result.next_cursor, tools, @@ -569,7 +561,6 @@ impl RmcpClient { async move { service.list_resources(params).await }.boxed() }) .await?; - self.persist_oauth_tokens().await; Ok(result) } @@ -585,7 +576,6 @@ impl RmcpClient { async move { service.list_resource_templates(params).await }.boxed() }) .await?; - self.persist_oauth_tokens().await; Ok(result) } @@ -601,7 +591,6 @@ impl RmcpClient { async move { service.read_resource(params).await }.boxed() }) .await?; - self.persist_oauth_tokens().await; Ok(result) } @@ -659,7 +648,6 @@ impl RmcpClient { .boxed() }) .await?; - self.persist_oauth_tokens().await; Ok(result) } @@ -689,7 +677,6 @@ impl RmcpClient { }, ) .await?; - self.persist_oauth_tokens().await; Ok(()) } @@ -712,7 +699,6 @@ impl RmcpClient { .boxed() }) .await?; - self.persist_oauth_tokens().await; Ok(response) } @@ -752,16 +738,6 @@ 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) -> Result<()> { if let Some(runtime) = self.oauth_persistor().await { runtime.refresh_if_needed().await?; @@ -968,19 +944,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)) } @@ -995,38 +959,69 @@ impl RmcpClient { F: Fn(Arc>) -> Fut, Fut: std::future::Future>, { + let deadline = timeout.map(|duration| Instant::now() + duration); let service = self.service().await?; - match Self::run_service_operation_with_transient_retries( + let mut result = Self::run_service_operation_with_transient_retries( Arc::clone(&service), label, timeout, + deadline, self.elicitation_pause_state.clone(), &operation, ) - .await + .await; + + if result + .as_ref() + .is_err_and(Self::is_unauthorized_operation_error) + && let Some(oauth_persistor) = self.oauth_persistor().await { - Ok(result) => Ok(result), - Err(error) if Self::is_session_expired_404(&error) => { - self.reinitialize_after_session_expiry(&service).await?; - let recovered_service = self.service().await?; - Self::run_service_operation_with_transient_retries( - recovered_service, - label, - timeout, - self.elicitation_pause_state.clone(), - &operation, - ) - .await - .map_err(Into::into) + // OAuth recovery is intentionally one-shot. The caller deadline gates whether recovery + // starts and whether the retry runs, but it must not cancel an in-flight provider + // refresh before a rotated token is persisted. + remaining_operation_timeout(label, timeout, deadline)?; + let refresh_result = oauth_persistor.refresh_after_unauthorized().await; + if let Err(error) = refresh_result { + if let Err(timeout_error) = remaining_operation_timeout(label, timeout, deadline) { + return Err(timeout_error.into()); + } + return Err(error); } - Err(error) => Err(error.into()), + result = Self::run_service_operation_with_transient_retries( + Arc::clone(&service), + label, + timeout, + deadline, + self.elicitation_pause_state.clone(), + &operation, + ) + .await; } + + if result.as_ref().is_err_and(Self::is_session_expired_404) { + // Session recovery remains one-shot and runs after the optional OAuth retry, so a 401 + // followed by the old session's 404 still reconstructs the transport before retrying. + self.reinitialize_after_session_expiry(&service).await?; + let recovered_service = self.service().await?; + result = Self::run_service_operation_with_transient_retries( + recovered_service, + label, + timeout, + deadline, + self.elicitation_pause_state.clone(), + &operation, + ) + .await; + } + + result.map_err(Into::into) } async fn run_service_operation_with_transient_retries( service: Arc>, label: &str, timeout: Option, + retry_deadline: Option, pause_state: ElicitationPauseState, operation: &F, ) -> std::result::Result @@ -1034,7 +1029,6 @@ impl RmcpClient { F: Fn(Arc>) -> Fut, Fut: std::future::Future>, { - let retry_deadline = timeout.map(|duration| Instant::now() + duration); for (attempt, retry_delay_ms) in STREAMABLE_HTTP_RETRY_DELAYS_MS .iter() .copied() @@ -1140,6 +1134,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>, @@ -1177,7 +1196,7 @@ impl RmcpClient { .timeout .map(|duration| Instant::now() + duration); 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, @@ -1196,12 +1215,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(()) } } diff --git a/codex-rs/rmcp-client/src/streamable_http_retry.rs b/codex-rs/rmcp-client/src/streamable_http_retry.rs index 2794489b2a76..bd7a4cb7f800 100644 --- a/codex-rs/rmcp-client/src/streamable_http_retry.rs +++ b/codex-rs/rmcp-client/src/streamable_http_retry.rs @@ -22,8 +22,13 @@ use super::RmcpClient; const JSON_RPC_INTERNAL_ERROR_CODE: i64 = -32603; pub(super) const STREAMABLE_HTTP_RETRY_DELAYS_MS: [u64; 2] = [250, 1_000]; +#[derive(Default)] +struct InitializeAttemptContext { + oauth_persistor: Option, +} + impl RmcpClient { - pub(super) async fn connect_pending_transport_with_initialize_retries( + pub(super) async fn connect_pending_transport_with_oauth_recovery( &self, initial_transport: PendingTransport, client_service: ElicitationClientService, @@ -32,6 +37,69 @@ impl RmcpClient { ) -> Result<( Arc>, Option, + )> { + let mut attempt_context = InitializeAttemptContext::default(); + match self + .connect_pending_transport_with_initialize_retries( + initial_transport, + client_service.clone(), + timeout, + initialize_deadline, + &mut attempt_context, + ) + .await + { + Ok(result) => Ok(result), + Err(error) if Self::is_unauthorized_initialize_error(&error) => { + let Some(oauth_persistor) = attempt_context.oauth_persistor else { + return Err(error); + }; + // Initialization gets one OAuth refresh and one reconstructed transport. Reusing + // this wrapper for the retry would turn persistent 401s into a refresh loop. The + // startup deadline gates whether recovery starts and bounds transport setup plus + // the retry handshake, but the refresh transaction has its own bounds and is + // deliberately excluded from the startup budget. + remaining_initialize_timeout(timeout, *initialize_deadline)?; + let refresh_started_at = Instant::now(); + let refresh_result = oauth_persistor.refresh_after_unauthorized().await; + if let Some(deadline) = initialize_deadline.as_mut() { + *deadline += refresh_started_at.elapsed(); + } + refresh_result?; + let remaining = remaining_initialize_timeout(timeout, *initialize_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 mut retry_context = InitializeAttemptContext::default(); + self.connect_pending_transport_with_initialize_retries( + transport, + client_service, + timeout, + initialize_deadline, + &mut retry_context, + ) + .await + } + Err(error) => Err(error), + } + } + + async fn connect_pending_transport_with_initialize_retries( + &self, + initial_transport: PendingTransport, + client_service: ElicitationClientService, + timeout: Option, + initialize_deadline: &mut Option, + attempt_context: &mut InitializeAttemptContext, + ) -> Result<( + Arc>, + Option, )> { let should_retry = match &initial_transport { PendingTransport::InProcess { .. } | PendingTransport::Stdio { .. } => false, @@ -62,6 +130,17 @@ impl RmcpClient { } } }; + // Keep the persistor paired with the transport attempt that returned 401. Rebuilt + // transports reuse the recipe's lifecycle-pinned credential source, and this pairing + // also keeps the authorization manager and snapshot aligned with the failed attempt. + attempt_context.oauth_persistor = match &transport { + PendingTransport::StreamableHttpWithOAuth { + oauth_persistor, .. + } => Some(oauth_persistor.clone()), + PendingTransport::InProcess { .. } + | PendingTransport::Stdio { .. } + | PendingTransport::StreamableHttp { .. } => None, + }; match Self::connect_pending_transport( transport, client_service.clone(), @@ -108,6 +187,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 } From 93c547454a442a78cb891394b057a1ca8ede58da Mon Sep 17 00:00:00 2001 From: Steven Lee Date: Thu, 25 Jun 2026 17:10:22 +0000 Subject: [PATCH 2/6] Test memory authority after persistence failure --- codex-rs/rmcp-client/src/oauth.rs | 11 +---------- 1 file changed, 1 insertion(+), 10 deletions(-) diff --git a/codex-rs/rmcp-client/src/oauth.rs b/codex-rs/rmcp-client/src/oauth.rs index 191ac1d54230..df5506647ec0 100644 --- a/codex-rs/rmcp-client/src/oauth.rs +++ b/codex-rs/rmcp-client/src/oauth.rs @@ -1294,17 +1294,8 @@ mod tests { persistor .refresh_if_needed_with_keyring_store(&store) .await?; - persistor - .refresh_after_unauthorized_with_keyring_store( - &store, - initial_tokens.token_response.0.access_token().clone(), - ) - .await?; let error = persistor - .refresh_after_unauthorized_with_keyring_store( - &store, - AccessToken::new("updated-access-token".to_string()), - ) + .refresh_after_unauthorized_with_keyring_store(&store) .await .expect_err("B cannot be refreshed by rereading durable A"); assert!( From 5f5c41d6b9fa59732d178eb0accfdaeaeceba2ba Mon Sep 17 00:00:00 2001 From: Steven Lee Date: Thu, 25 Jun 2026 17:12:40 +0000 Subject: [PATCH 3/6] Stabilize OAuth recovery assertions --- codex-rs/rmcp-client/src/oauth.rs | 5 ++++- codex-rs/rmcp-client/tests/streamable_http_oauth_startup.rs | 4 +++- 2 files changed, 7 insertions(+), 2 deletions(-) diff --git a/codex-rs/rmcp-client/src/oauth.rs b/codex-rs/rmcp-client/src/oauth.rs index df5506647ec0..410e0016c5d4 100644 --- a/codex-rs/rmcp-client/src/oauth.rs +++ b/codex-rs/rmcp-client/src/oauth.rs @@ -1372,7 +1372,10 @@ mod tests { Some("rotated-refresh-token".to_string()) ); let second_tokens = tokens_from_manager(&second_manager).await?; - assert_eq!(second_tokens.token_response, stored.token_response); + assert_token_response_match_without_expiry( + &second_tokens.token_response, + &stored.token_response, + ); Ok(()) } 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 d578f40a3ebd..c0c691980e2f 100644 --- a/codex-rs/rmcp-client/tests/streamable_http_oauth_startup.rs +++ b/codex-rs/rmcp-client/tests/streamable_http_oauth_startup.rs @@ -587,7 +587,9 @@ 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)), + // Keep this outside the 60-second proactive refresh guard band. The test is checking a + // healthy persisted access token, not the boundary where a refresh becomes necessary. + expires_at: Some(now.saturating_add(/*rhs*/ 120_000)), }; save_oauth_tokens( SERVER_NAME, From c73188bb1003f90417c766b305e0ad29f2f564e0 Mon Sep 17 00:00:00 2001 From: Steven Lee Date: Thu, 25 Jun 2026 18:11:34 +0000 Subject: [PATCH 4/6] Preserve RMCP OAuth persistence until ownership switches --- codex-rs/rmcp-client/src/oauth/persistor.rs | 193 ++++++++++++++++++++ codex-rs/rmcp-client/src/rmcp_client.rs | 43 ++++- 2 files changed, 234 insertions(+), 2 deletions(-) diff --git a/codex-rs/rmcp-client/src/oauth/persistor.rs b/codex-rs/rmcp-client/src/oauth/persistor.rs index ff3273a68abc..4e6c02c8b82e 100644 --- a/codex-rs/rmcp-client/src/oauth/persistor.rs +++ b/codex-rs/rmcp-client/src/oauth/persistor.rs @@ -8,6 +8,7 @@ use std::time::UNIX_EPOCH; use anyhow::Context; use anyhow::Result; +use codex_config::types::AuthKeyringBackendKind; use codex_keyring_store::DefaultKeyringStore; use codex_keyring_store::KeyringStore; use oauth2::TokenResponse; @@ -25,6 +26,10 @@ use super::ResolvedOAuthCredentialStore; use super::StoredOAuthTokens; use super::WrappedOAuthTokenResponse; use super::compute_expires_at_millis; +use super::compute_store_key; +use super::delete_oauth_tokens_from_direct_keyring; +use super::delete_oauth_tokens_from_file; +use super::delete_oauth_tokens_from_secrets_keyring; use super::load_oauth_tokens_from_file; use super::load_oauth_tokens_from_keyring; use super::refresh_lock::RefreshCredentialLock; @@ -105,6 +110,88 @@ impl OAuthPersistor { } } + /// Persists credentials that RMCP may still refresh before the ownership switch lands. + /// + /// The next stack layer removes this compatibility path atomically with request-only RMCP + /// credentials. Keeping it here makes this intermediate tip preserve the existing contract. + pub(crate) async fn persist_if_needed(&self) -> Result<()> { + 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 state = self.inner.credential_state.lock().await; + state.current.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 state = self.inner.credential_state.lock().await; + state.current = None; + state.has_unpersisted_refresh = false; + } + } + 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, @@ -128,6 +215,112 @@ impl OAuthPersistor { } } + #[expect( + clippy::await_holding_invalid_type, + reason = "AuthorizationManager async access must be serialized through its mutex" + )] + pub(super) async fn persist_if_needed_with_keyring_store( + &self, + keyring_store: &K, + ) -> 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 state = self.inner.credential_state.lock().await; + let new_token_response = WrappedOAuthTokenResponse(credentials.clone()); + let same_token = state + .current + .as_ref() + .map(|prev| prev.token_response == new_token_response) + .unwrap_or(false); + let expires_at = if same_token { + state.current.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 state.current.as_ref() != Some(&stored) { + // A provider may already have consumed the old rotating refresh token. Make B + // authoritative in this process before the fallible save. + state.current = Some(stored.clone()); + state.has_unpersisted_refresh = true; + debug!("persisting refreshed MCP OAuth credentials to the resolved store"); + let persistence_started_at = Instant::now(); + let persistence_result = match self.inner.credential_store { + ResolvedOAuthCredentialStore::File => save_oauth_tokens_to_file(&stored), + ResolvedOAuthCredentialStore::Keyring(keyring_backend_kind) => { + save_oauth_tokens_with_keyring( + keyring_store, + keyring_backend_kind, + &self.inner.server_name, + &stored, + ) + } + }; + if let Err(error) = persistence_result { + warn!( + persistence_elapsed_ms = persistence_started_at.elapsed().as_millis(), + error = %error, + "failed to persist refreshed MCP OAuth credentials; retaining them as the in-process authority without retrying persistence" + ); + return Err(error); + } + state.has_unpersisted_refresh = false; + debug!( + persistence_elapsed_ms = persistence_started_at.elapsed().as_millis(), + "persisted refreshed MCP OAuth credentials" + ); + } + } + None => { + let mut state = self.inner.credential_state.lock().await; + if state.current.take().is_some() + && let Err(error) = match self.inner.credential_store { + ResolvedOAuthCredentialStore::File => { + let key = compute_store_key(&self.inner.server_name, &self.inner.url)?; + delete_oauth_tokens_from_file(&key).map(|_| ()) + } + ResolvedOAuthCredentialStore::Keyring(AuthKeyringBackendKind::Direct) => { + delete_oauth_tokens_from_direct_keyring( + keyring_store, + &self.inner.server_name, + &self.inner.url, + ) + .map(|_| ()) + } + ResolvedOAuthCredentialStore::Keyring(AuthKeyringBackendKind::Secrets) => { + delete_oauth_tokens_from_secrets_keyring( + keyring_store, + &self.inner.server_name, + &self.inner.url, + ) + .map(|_| ()) + } + } + { + warn!( + "failed to remove OAuth tokens for server {}: {error}", + self.inner.server_name + ); + } + state.has_unpersisted_refresh = false; + } + } + + Ok(()) + } + pub(crate) async fn refresh_if_needed(&self) -> Result<()> { self.refresh_if_needed_with_keyring_store(&DefaultKeyringStore) .await diff --git a/codex-rs/rmcp-client/src/rmcp_client.rs b/codex-rs/rmcp-client/src/rmcp_client.rs index 9b74721f6d05..77d3834a6ab2 100644 --- a/codex-rs/rmcp-client/src/rmcp_client.rs +++ b/codex-rs/rmcp-client/src/rmcp_client.rs @@ -486,6 +486,12 @@ 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) } @@ -722,6 +728,17 @@ impl RmcpClient { } } + // RMCP still receives refresh-capable credentials in this intermediate layer. Preserve its + // legacy best-effort persistence until the next layer atomically installs request-only + // credentials and the Codex-owned all-traffic wrapper. + 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}"); + } + } + /// Stop the MCP transport and any stdio server process owned by this client. pub async fn shutdown(&self) { let previous_state = { @@ -944,7 +961,19 @@ impl RmcpClient { .await .map_err(|source| anyhow::Error::from(HandshakeError { source })), }; - let service = service_result?; + 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); + } + }; Ok((Arc::new(service), oauth_persistor)) } @@ -1014,7 +1043,11 @@ impl RmcpClient { .await; } - result.map_err(Into::into) + let result = result.map_err(Into::into); + if result.is_ok() { + self.persist_oauth_tokens().await; + } + result } async fn run_service_operation_with_transient_retries( @@ -1215,6 +1248,12 @@ 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(()) } } From 0e139f87c9481c733427eb64eae140e7228b3749 Mon Sep 17 00:00:00 2001 From: Steven Lee Date: Thu, 25 Jun 2026 18:36:07 +0000 Subject: [PATCH 5/6] Preserve pending OAuth authority in compatibility path --- codex-rs/rmcp-client/src/oauth/persistor.rs | 43 +++++++++++---------- 1 file changed, 23 insertions(+), 20 deletions(-) diff --git a/codex-rs/rmcp-client/src/oauth/persistor.rs b/codex-rs/rmcp-client/src/oauth/persistor.rs index 4e6c02c8b82e..53144518c2d2 100644 --- a/codex-rs/rmcp-client/src/oauth/persistor.rs +++ b/codex-rs/rmcp-client/src/oauth/persistor.rs @@ -125,9 +125,9 @@ impl OAuthPersistor { &self, keyring_store: &K, ) -> Result<()> { - let snapshot = { + let (snapshot, unpersisted_refresh) = { let state = self.inner.credential_state.lock().await; - state.current.clone() + (state.current.clone(), state.unpersisted_refresh.clone()) }; let (client_id, current_credentials) = self.manager_credentials().await?; let manager_changed = match (&snapshot, current_credentials.as_ref()) { @@ -147,20 +147,13 @@ impl OAuthPersistor { .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, - }; + let latest_matches_snapshot = durable_credentials_match_snapshot(&latest, &snapshot) + || unpersisted_refresh.as_ref().is_some_and(|pending| { + // The failed write never changed durable authority. If storage still contains the + // last known durable snapshot, persist the manager's newer result instead of + // adopting and replaying that stale snapshot. + durable_credentials_match_snapshot(&latest, &pending.previously_persisted) + }); if !latest_matches_snapshot { // A completed login or logout is authoritative over tokens refreshed inside an RMCP @@ -172,7 +165,7 @@ impl OAuthPersistor { self.clear_manager_credentials().await; let mut state = self.inner.credential_state.lock().await; state.current = None; - state.has_unpersisted_refresh = false; + state.unpersisted_refresh = None; } } return Ok(()); @@ -253,8 +246,18 @@ impl OAuthPersistor { if state.current.as_ref() != Some(&stored) { // A provider may already have consumed the old rotating refresh token. Make B // authoritative in this process before the fallible save. + // Preserve the last snapshot known to be durable across repeated in-memory + // changes. Using the immediately preceding in-memory token here could make an + // unchanged stale store look like an external login or refresh. + let previously_persisted = state + .unpersisted_refresh + .take() + .map(|pending| pending.previously_persisted) + .unwrap_or_else(|| state.current.clone()); state.current = Some(stored.clone()); - state.has_unpersisted_refresh = true; + state.unpersisted_refresh = Some(UnpersistedRefresh { + previously_persisted, + }); debug!("persisting refreshed MCP OAuth credentials to the resolved store"); let persistence_started_at = Instant::now(); let persistence_result = match self.inner.credential_store { @@ -276,7 +279,7 @@ impl OAuthPersistor { ); return Err(error); } - state.has_unpersisted_refresh = false; + state.unpersisted_refresh = None; debug!( persistence_elapsed_ms = persistence_started_at.elapsed().as_millis(), "persisted refreshed MCP OAuth credentials" @@ -314,7 +317,7 @@ impl OAuthPersistor { self.inner.server_name ); } - state.has_unpersisted_refresh = false; + state.unpersisted_refresh = None; } } From 78a6a3e3604a2aa6b0a16f6ddc49390ecd302717 Mon Sep 17 00:00:00 2001 From: Steven Lee Date: Thu, 25 Jun 2026 19:07:36 +0000 Subject: [PATCH 6/6] Fix OAuth recovery Clippy lint --- codex-rs/rmcp-client/src/oauth/persistor.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/codex-rs/rmcp-client/src/oauth/persistor.rs b/codex-rs/rmcp-client/src/oauth/persistor.rs index 53144518c2d2..ae558c729007 100644 --- a/codex-rs/rmcp-client/src/oauth/persistor.rs +++ b/codex-rs/rmcp-client/src/oauth/persistor.rs @@ -632,7 +632,7 @@ enum RefreshReason { } impl RefreshReason { - fn as_str(&self) -> &'static str { + fn as_str(self) -> &'static str { match self { Self::Expiry => "expiry", Self::Unauthorized => "unauthorized",