diff --git a/codex-rs/rmcp-client/src/oauth.rs b/codex-rs/rmcp-client/src/oauth.rs index bfd3c411e7c0..410e0016c5d4 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,101 @@ 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. - persistor - .persist_if_needed_with_keyring_store(&store) - .await?; + // 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 .refresh_if_needed_with_keyring_store(&store) .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) + .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_token_response_match_without_expiry( + &second_tokens.token_response, + &stored.token_response, + ); Ok(()) } @@ -1461,19 +1428,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 +1451,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 +1562,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 +1625,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 +1659,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 +1667,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 +1697,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 +1742,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 +1805,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 +1814,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 +1852,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 +1884,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 +1893,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 +2215,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..ae558c729007 100644 --- a/codex-rs/rmcp-client/src/oauth/persistor.rs +++ b/codex-rs/rmcp-client/src/oauth/persistor.rs @@ -43,7 +43,6 @@ const REFRESH_REQUEST_TIMEOUT: Duration = Duration::from_secs(45); pub(crate) struct OAuthPersistor { inner: Arc, } - struct OAuthPersistorInner { server_name: String, url: String, @@ -111,8 +110,10 @@ impl OAuthPersistor { } } - /// Persists the latest stored credentials if they have changed. - /// Deletes the credentials if they are no longer present. + /// 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 @@ -243,9 +244,8 @@ impl OAuthPersistor { 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. + // 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. @@ -329,6 +329,11 @@ impl OAuthPersistor { .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 +361,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 +402,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 +416,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 +449,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 +500,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 +625,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 +663,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 +692,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..77d3834a6ab2 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, @@ -507,7 +507,6 @@ impl RmcpClient { async move { service.list_tools(params).await }.boxed() }) .await?; - self.persist_oauth_tokens().await; Ok(result) } @@ -542,7 +541,6 @@ impl RmcpClient { }) }) .collect::>>()?; - self.persist_oauth_tokens().await; Ok(ListToolsWithConnectorIdResult { next_cursor: result.next_cursor, tools, @@ -569,7 +567,6 @@ impl RmcpClient { async move { service.list_resources(params).await }.boxed() }) .await?; - self.persist_oauth_tokens().await; Ok(result) } @@ -585,7 +582,6 @@ impl RmcpClient { async move { service.list_resource_templates(params).await }.boxed() }) .await?; - self.persist_oauth_tokens().await; Ok(result) } @@ -601,7 +597,6 @@ impl RmcpClient { async move { service.read_resource(params).await }.boxed() }) .await?; - self.persist_oauth_tokens().await; Ok(result) } @@ -659,7 +654,6 @@ impl RmcpClient { .boxed() }) .await?; - self.persist_oauth_tokens().await; Ok(result) } @@ -689,7 +683,6 @@ impl RmcpClient { }, ) .await?; - self.persist_oauth_tokens().await; Ok(()) } @@ -712,7 +705,6 @@ impl RmcpClient { .boxed() }) .await?; - self.persist_oauth_tokens().await; Ok(response) } @@ -736,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 = { @@ -752,16 +755,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?; @@ -995,38 +988,73 @@ 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; + } + + 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( service: Arc>, label: &str, timeout: Option, + retry_deadline: Option, pause_state: ElicitationPauseState, operation: &F, ) -> std::result::Result @@ -1034,7 +1062,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 +1167,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 +1229,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, 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 } 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,