diff --git a/codex-rs/rmcp-client/src/oauth.rs b/codex-rs/rmcp-client/src/oauth.rs index a582e39aea23..c4c2c966944b 100644 --- a/codex-rs/rmcp-client/src/oauth.rs +++ b/codex-rs/rmcp-client/src/oauth.rs @@ -59,7 +59,6 @@ use self::resolution_state::StoreResolutionReason; use self::resolution_state::record_store_resolution; use self::store_lock::OAuthStore; use self::store_lock::OAuthStoreLock; - use codex_keyring_store::DefaultKeyringStore; use codex_keyring_store::KeyringStore; use codex_utils_home_dir::find_codex_home; @@ -71,6 +70,8 @@ pub(crate) use self::resolved_store::ResolvedOAuthCredentialStore; pub(crate) use self::resolved_store::ResolvedOAuthTokens; pub(crate) use self::resolved_store::load_oauth_tokens_from_store; pub(crate) use self::resolved_store::resolve_oauth_tokens; +#[cfg(test)] +use rmcp::transport::auth::AuthorizationManager; const KEYRING_SERVICE: &str = "Codex MCP Credentials"; const MCP_OAUTH_SECRET_PREFIX: &str = "MCP_OAUTH"; @@ -822,20 +823,90 @@ fn sha_256_prefix(value: &Value) -> Result { #[cfg(test)] mod tests { + use super::refresh_lock::RefreshCredentialLock; use super::*; use anyhow::Result; use codex_secrets::compute_keyring_account; use keyring::Error as KeyringError; use pretty_assertions::assert_eq; + use rmcp::transport::auth::OAuthState; + use serde_json::json; use std::sync::Arc; use std::sync::Mutex; use std::sync::MutexGuard; use std::sync::OnceLock; use std::sync::PoisonError; + use std::sync::atomic::AtomicUsize; + use std::sync::atomic::Ordering; + use std::sync::mpsc; use tempfile::tempdir; + use tokio::sync::Mutex as TokioMutex; + use tracing::Event; + use tracing::Id; + use tracing::Metadata; + use tracing::Subscriber; + use tracing::span::Attributes; + use tracing::span::Record; + use tracing::subscriber::Interest; + use wiremock::Mock; + use wiremock::MockServer; + use wiremock::ResponseTemplate; + use wiremock::matchers::body_string_contains; + use wiremock::matchers::method; + use wiremock::matchers::path; use codex_keyring_store::tests::MockKeyringStore; + const STORE_LOCK_CONTENTION_EVENT_TARGET: &str = + "codex_rmcp_client::oauth::store_lock::contention"; + const REFRESH_LOCK_CONTENTION_EVENT_TARGET: &str = + "codex_rmcp_client::oauth::refresh_lock::contention"; + + struct LockContentionSubscriber { + target: &'static str, + contended_tx: mpsc::Sender<()>, + } + + impl Subscriber for LockContentionSubscriber { + fn enabled(&self, metadata: &Metadata<'_>) -> bool { + metadata.target() == self.target + } + + fn register_callsite(&self, metadata: &'static Metadata<'static>) -> Interest { + if self.enabled(metadata) { + Interest::always() + } else { + Interest::never() + } + } + + fn max_level_hint(&self) -> Option { + Some(tracing::level_filters::LevelFilter::DEBUG) + } + + fn new_span(&self, _span: &Attributes<'_>) -> Id { + // This subscriber enables only the contention event callsite, so it never observes + // spans. + Id::from_u64(/*u*/ 1) + } + + fn record(&self, _span: &Id, _values: &Record<'_>) {} + + fn record_follows_from(&self, _span: &Id, _follows: &Id) {} + + fn event(&self, event: &Event<'_>) { + if event.metadata().target() == self.target { + self.contended_tx + .send(()) + .expect("signal actual OAuth store lock contention"); + } + } + + fn enter(&self, _span: &Id) {} + + fn exit(&self, _span: &Id) {} + } + struct TempCodexHome { _guard: MutexGuard<'static, ()>, _dir: tempfile::TempDir, @@ -871,6 +942,41 @@ mod tests { } } + fn complete_after_store_lock_contention( + codex_home: &std::path::Path, + store: OAuthStore, + operation: impl FnOnce() -> Result + Send + 'static, + ) -> Result + where + T: Send + 'static, + { + let held_lock = + OAuthStoreLock::acquire_in(codex_home, store, Duration::from_millis(/*millis*/ 100))?; + let (contended_tx, contended_rx) = mpsc::channel(); + let (result_tx, result_rx) = mpsc::channel(); + let worker = std::thread::spawn(move || { + tracing::subscriber::with_default( + LockContentionSubscriber { + target: STORE_LOCK_CONTENTION_EVENT_TARGET, + contended_tx, + }, + || { + result_tx + .send(operation()) + .expect("send contending OAuth store operation result"); + }, + ); + }); + + contended_rx.recv_timeout(Duration::from_secs(/*secs*/ 1))?; + drop(held_lock); + let result = result_rx.recv_timeout(Duration::from_secs(/*secs*/ 10))??; + worker + .join() + .expect("contending OAuth store worker should finish"); + Ok(result) + } + #[test] fn load_oauth_tokens_reads_from_keyring_when_available() -> Result<()> { let _env = TempCodexHome::new(); @@ -909,6 +1015,33 @@ mod tests { AuthKeyringBackendKind::Direct, )? .expect("tokens should load from fallback"); + assert_eq!(resolved.store, ResolvedOAuthCredentialStore::File); + assert_tokens_match_without_expiry(&resolved.tokens, &expected); + Ok(()) + } + + #[test] + fn resolution_diagnostic_failure_does_not_change_auto_selection() -> Result<()> { + let env = TempCodexHome::new(); + let store = MockKeyringStore::default(); + let tokens = sample_tokens(); + let expected = tokens.clone(); + super::save_oauth_tokens_to_file(&tokens)?; + fs::write( + env.path().join(".mcp-oauth-store-resolutions.json"), + "not valid JSON", + )?; + + let resolved = super::resolve_oauth_tokens( + &store, + &tokens.server_name, + &tokens.url, + OAuthCredentialsStoreMode::Auto, + AuthKeyringBackendKind::Direct, + )? + .expect("diagnostic failure must not hide fallback credentials"); + + assert_eq!(resolved.store, ResolvedOAuthCredentialStore::File); assert_tokens_match_without_expiry(&resolved.tokens, &expected); Ok(()) } @@ -932,10 +1065,46 @@ mod tests { AuthKeyringBackendKind::Direct, )? .expect("tokens should load from fallback"); + assert_eq!(resolved.store, ResolvedOAuthCredentialStore::File); assert_tokens_match_without_expiry(&resolved.tokens, &expected); Ok(()) } + #[test] + fn auto_resolution_prioritizes_keyring_and_tracks_its_source() -> Result<()> { + let _env = TempCodexHome::new(); + let store = MockKeyringStore::default(); + let keyring_tokens = sample_tokens(); + let mut file_tokens = sample_tokens(); + file_tokens + .token_response + .0 + .set_access_token(AccessToken::new("file-access-token".to_string())); + super::save_oauth_tokens_to_file(&file_tokens)?; + super::save_oauth_tokens_with_keyring( + &store, + AuthKeyringBackendKind::Direct, + &keyring_tokens.server_name, + &keyring_tokens, + )?; + + let resolved = super::resolve_oauth_tokens( + &store, + &keyring_tokens.server_name, + &keyring_tokens.url, + OAuthCredentialsStoreMode::Auto, + AuthKeyringBackendKind::Direct, + )? + .expect("Auto should load keyring credentials"); + + assert_eq!( + resolved.store, + ResolvedOAuthCredentialStore::Keyring(AuthKeyringBackendKind::Direct) + ); + assert_tokens_match_without_expiry(&resolved.tokens, &keyring_tokens); + Ok(()) + } + #[test] fn save_oauth_tokens_prefers_keyring_when_available() -> Result<()> { let _env = TempCodexHome::new(); @@ -976,7 +1145,7 @@ mod tests { let fallback_path = super::fallback_file_path()?; assert!(fallback_path.exists(), "fallback file should be created"); - let saved = super::read_fallback_file_unlocked()?.expect("fallback file should load"); + let saved = read_fallback_file()?.expect("fallback file should load"); let key = super::compute_store_key(&tokens.server_name, &tokens.url)?; let entry = saved.get(&key).expect("entry for key"); assert_eq!(entry.server_name, tokens.server_name); @@ -990,6 +1159,78 @@ mod tests { Ok(()) } + #[test] + fn file_store_lock_preserves_updates_for_different_servers() -> Result<()> { + let env = TempCodexHome::new(); + let first = sample_tokens(); + let mut second = sample_tokens(); + second.server_name = "second-server".to_string(); + second.url = "https://second.example.test".to_string(); + + let held_lock = OAuthStoreLock::acquire_in( + env.path(), + OAuthStore::File, + Duration::from_millis(/*millis*/ 100), + )?; + let (contended_tx, contended_rx) = mpsc::channel(); + let (result_tx, result_rx) = mpsc::channel(); + let second_for_writer = second.clone(); + let writer = std::thread::spawn(move || { + tracing::subscriber::with_default( + LockContentionSubscriber { + target: STORE_LOCK_CONTENTION_EVENT_TARGET, + contended_tx, + }, + || { + result_tx + .send(super::save_oauth_tokens_to_file(&second_for_writer)) + .expect("send writer result"); + }, + ); + }); + + // The writer reports only from OAuthStoreLock's real try_lock WouldBlock branch. This + // fails deterministically if save stops acquiring the aggregate File lock. + contended_rx.recv_timeout(Duration::from_secs(/*secs*/ 1))?; + super::save_oauth_tokens_to_file_unlocked(&first)?; + drop(held_lock); + + result_rx.recv_timeout(Duration::from_secs(10))??; + writer.join().expect("file store writer should finish"); + let loaded_first = super::load_oauth_tokens_from_file(&first.server_name, &first.url)? + .expect("first server tokens should remain stored"); + let loaded_second = super::load_oauth_tokens_from_file(&second.server_name, &second.url)? + .expect("second server tokens should be stored"); + assert_tokens_match_without_expiry(&loaded_first, &first); + assert_tokens_match_without_expiry(&loaded_second, &second); + Ok(()) + } + + #[test] + fn file_store_load_and_delete_observe_aggregate_lock() -> Result<()> { + let env = TempCodexHome::new(); + let tokens = sample_tokens(); + super::save_oauth_tokens_to_file(&tokens)?; + + let server_name = tokens.server_name.clone(); + let url = tokens.url.clone(); + let loaded = + complete_after_store_lock_contention(env.path(), OAuthStore::File, move || { + super::load_oauth_tokens_from_file(&server_name, &url) + })? + .expect("file credentials should remain readable after contention"); + assert_tokens_match_without_expiry(&loaded, &tokens); + + let key = super::compute_store_key(&tokens.server_name, &tokens.url)?; + let removed = + complete_after_store_lock_contention(env.path(), OAuthStore::File, move || { + super::delete_oauth_tokens_from_file(&key) + })?; + assert!(removed); + assert!(super::load_oauth_tokens_from_file(&tokens.server_name, &tokens.url)?.is_none()); + Ok(()) + } + #[test] fn save_oauth_tokens_with_secrets_backend_writes_encrypted_storage() -> Result<()> { let env = TempCodexHome::new(); @@ -1050,6 +1291,127 @@ mod tests { Ok(()) } + #[test] + fn secrets_store_lock_preserves_updates_for_different_servers() -> Result<()> { + let env = TempCodexHome::new(); + let store = MockKeyringStore::default(); + let first = sample_tokens(); + let mut second = sample_tokens(); + second.server_name = "second-server".to_string(); + second.url = "https://second.example.test".to_string(); + + let held_lock = OAuthStoreLock::acquire_in( + env.path(), + OAuthStore::Secrets, + Duration::from_millis(/*millis*/ 100), + )?; + let (contended_tx, contended_rx) = mpsc::channel(); + let (result_tx, result_rx) = mpsc::channel(); + let store_for_writer = store.clone(); + let second_for_writer = second.clone(); + let writer = std::thread::spawn(move || { + tracing::subscriber::with_default( + LockContentionSubscriber { + target: STORE_LOCK_CONTENTION_EVENT_TARGET, + contended_tx, + }, + || { + result_tx + .send(super::save_oauth_tokens_with_keyring( + &store_for_writer, + AuthKeyringBackendKind::Secrets, + &second_for_writer.server_name, + &second_for_writer, + )) + .expect("send writer result"); + }, + ); + }); + + // As above, wait for the actual WouldBlock branch rather than inferring contention from a + // scheduler-sensitive period in which the writer has not returned. + contended_rx.recv_timeout(Duration::from_secs(/*secs*/ 1))?; + let first_serialized = serde_json::to_string(&first)?; + super::save_oauth_tokens_to_secrets_keyring_unlocked( + &store, + &first.server_name, + &first, + &first_serialized, + )?; + drop(held_lock); + + result_rx.recv_timeout(Duration::from_secs(10))??; + writer.join().expect("secrets store writer should finish"); + let loaded_first = super::load_oauth_tokens_from_keyring( + &store, + AuthKeyringBackendKind::Secrets, + &first.server_name, + &first.url, + )? + .expect("first server tokens should remain stored"); + let loaded_second = super::load_oauth_tokens_from_keyring( + &store, + AuthKeyringBackendKind::Secrets, + &second.server_name, + &second.url, + )? + .expect("second server tokens should be stored"); + assert_tokens_match_without_expiry(&loaded_first, &first); + assert_tokens_match_without_expiry(&loaded_second, &second); + Ok(()) + } + + #[test] + fn secrets_store_load_and_delete_observe_aggregate_lock() -> Result<()> { + let env = TempCodexHome::new(); + let store = MockKeyringStore::default(); + let tokens = sample_tokens(); + super::save_oauth_tokens_with_keyring( + &store, + AuthKeyringBackendKind::Secrets, + &tokens.server_name, + &tokens, + )?; + + let store_for_load = store.clone(); + let server_name = tokens.server_name.clone(); + let url = tokens.url.clone(); + let loaded = + complete_after_store_lock_contention(env.path(), OAuthStore::Secrets, move || { + super::load_oauth_tokens_from_keyring( + &store_for_load, + AuthKeyringBackendKind::Secrets, + &server_name, + &url, + ) + })? + .expect("encrypted credentials should remain readable after contention"); + assert_tokens_match_without_expiry(&loaded, &tokens); + + let store_for_delete = store.clone(); + let server_name = tokens.server_name.clone(); + let url = tokens.url.clone(); + let removed = + complete_after_store_lock_contention(env.path(), OAuthStore::Secrets, move || { + super::delete_oauth_tokens_from_secrets_keyring( + &store_for_delete, + &server_name, + &url, + ) + })?; + assert!(removed); + assert!( + super::load_oauth_tokens_from_keyring( + &store, + AuthKeyringBackendKind::Secrets, + &tokens.server_name, + &tokens.url, + )? + .is_none() + ); + Ok(()) + } + #[test] fn load_oauth_tokens_with_secrets_backend_ignores_direct_entry() -> Result<()> { let _env = TempCodexHome::new(); @@ -1070,6 +1432,714 @@ mod tests { Ok(()) } + #[tokio::test] + async fn refresh_transaction_preserves_credentials_when_resolved_keyring_reread_fails() + -> Result<()> { + let _env = TempCodexHome::new(); + let server = MockServer::start().await; + mount_oauth_metadata(&server).await; + let store = MockKeyringStore::default(); + let initial_tokens = expired_sample_tokens(&format!("{}/mcp", server.uri())); + let key = super::compute_store_key(&initial_tokens.server_name, &initial_tokens.url)?; + store.set_error(&key, KeyringError::Invalid("error".into(), "load".into())); + + let manager = authorization_manager_for(&initial_tokens).await?; + let persistor = OAuthPersistor::new( + initial_tokens.server_name.clone(), + initial_tokens.url.clone(), + manager.clone(), + ResolvedOAuthCredentialStore::Keyring(AuthKeyringBackendKind::Direct), + Some(initial_tokens.clone()), + ); + + let error = persistor + .refresh_if_needed_with_keyring_store(&store) + .await + .expect_err("keyring reread failure should abort refresh"); + + assert!( + error + .to_string() + .contains("failed to reread OAuth tokens from resolved keyring storage"), + "unexpected error: {error:#}" + ); + let manager_tokens = tokens_from_manager(&manager).await?; + assert_eq!( + manager_tokens.token_response, + WrappedOAuthTokenResponse(request_oauth_token_response(&initial_tokens)) + ); + Ok(()) + } + + #[tokio::test] + async fn resolved_keyring_write_failure_errors_and_restores_previous_credentials() -> 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", + "updated-access-token", + "updated-refresh-token", + Duration::from_millis(200), + ) + .await; + let store = MockKeyringStore::default(); + let initial_tokens = expired_sample_tokens(&format!("{}/mcp", server.uri())); + super::save_oauth_tokens_with_keyring_store( + &store, + &initial_tokens.server_name, + &initial_tokens, + OAuthCredentialsStoreMode::Keyring, + AuthKeyringBackendKind::Direct, + )?; + + let manager = authorization_manager_for(&initial_tokens).await?; + let persistor = OAuthPersistor::new( + initial_tokens.server_name.clone(), + initial_tokens.url.clone(), + manager.clone(), + ResolvedOAuthCredentialStore::Keyring(AuthKeyringBackendKind::Direct), + Some(initial_tokens.clone()), + ); + let refresh_task = tokio::spawn({ + let persistor = persistor.clone(); + let store = store.clone(); + async move { persistor.refresh_if_needed_with_keyring_store(&store).await } + }); + + wait_for_signal(refresh_started).await?; + let 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!( + error + .to_string() + .contains("failed to write OAuth tokens to keyring"), + "unexpected error: {error:#}" + ); + assert!(!super::fallback_file_path()?.exists()); + let manager_tokens = tokens_from_manager(&manager).await?; + assert_eq!( + manager_tokens.token_response, + WrappedOAuthTokenResponse(request_oauth_token_response(&initial_tokens)) + ); + server.verify().await; + Ok(()) + } + + #[tokio::test] + async fn unauthorized_transaction_refreshes_once_for_delayed_same_and_cross_client_401s() + -> Result<()> { + let _env = TempCodexHome::new(); + let server = MockServer::start().await; + mount_oauth_metadata(&server).await; + mount_refresh_response( + &server, + "refresh-token", + "refreshed-access-token", + "rotated-refresh-token", + ) + .await; + + let store = MockKeyringStore::default(); + let mut initial_tokens = sample_tokens(); + initial_tokens.url = format!("{}/mcp", server.uri()); + initial_tokens.expires_at = None; + initial_tokens.token_response.0.set_expires_in(None); + super::save_oauth_tokens_with_keyring_store( + &store, + &initial_tokens.server_name, + &initial_tokens, + OAuthCredentialsStoreMode::Keyring, + AuthKeyringBackendKind::Direct, + )?; + + let first_manager = authorization_manager_for(&initial_tokens).await?; + let second_manager = authorization_manager_for(&initial_tokens).await?; + let first = OAuthPersistor::new( + initial_tokens.server_name.clone(), + initial_tokens.url.clone(), + first_manager, + ResolvedOAuthCredentialStore::Keyring(AuthKeyringBackendKind::Direct), + Some(initial_tokens.clone()), + ); + let second = OAuthPersistor::new( + initial_tokens.server_name.clone(), + initial_tokens.url.clone(), + second_manager.clone(), + ResolvedOAuthCredentialStore::Keyring(AuthKeyringBackendKind::Direct), + Some(initial_tokens.clone()), + ); + let rejected_access_token = initial_tokens.token_response.0.access_token().clone(); + + first + .refresh_after_unauthorized_with_keyring_store(&store, rejected_access_token.clone()) + .await?; + // A second request from this same client may have sent the original token before the first + // request completed recovery. Its delayed 401 must adopt the rotated credentials rather + // than trigger another provider refresh. + first + .refresh_after_unauthorized_with_keyring_store(&store, rejected_access_token.clone()) + .await?; + second + .refresh_after_unauthorized_with_keyring_store(&store, rejected_access_token) + .await?; + + server.verify().await; + let stored = super::load_oauth_tokens_from_keyring( + &store, + AuthKeyringBackendKind::Direct, + &initial_tokens.server_name, + &initial_tokens.url, + )? + .expect("refreshed tokens should be persisted"); + assert_eq!(access_token(&stored), "refreshed-access-token"); + assert_eq!( + refresh_token(&stored), + Some("rotated-refresh-token".to_string()) + ); + let second_tokens = tokens_from_manager(&second_manager).await?; + assert_eq!( + second_tokens.token_response, + WrappedOAuthTokenResponse(request_oauth_token_response(&stored)) + ); + Ok(()) + } + + #[tokio::test] + async fn refresh_transaction_adopts_valid_reread_without_provider_refresh() -> Result<()> { + let _env = TempCodexHome::new(); + let server = MockServer::start().await; + mount_oauth_metadata(&server).await; + let store = MockKeyringStore::default(); + let mut initial_tokens = expired_sample_tokens(&format!("{}/mcp", server.uri())); + initial_tokens + .token_response + .0 + .set_refresh_token(Some(RefreshToken::new("stale-refresh-token".to_string()))); + + let mut latest_tokens = sample_tokens(); + latest_tokens.url.clone_from(&initial_tokens.url); + latest_tokens + .token_response + .0 + .set_access_token(AccessToken::new( + "already-refreshed-access-token".to_string(), + )); + latest_tokens + .token_response + .0 + .set_refresh_token(Some(RefreshToken::new( + "already-rotated-refresh-token".to_string(), + ))); + + super::save_oauth_tokens_with_keyring_store( + &store, + &latest_tokens.server_name, + &latest_tokens, + OAuthCredentialsStoreMode::Keyring, + AuthKeyringBackendKind::Direct, + )?; + + let manager = authorization_manager_for(&initial_tokens).await?; + let persistor = OAuthPersistor::new( + initial_tokens.server_name.clone(), + initial_tokens.url.clone(), + manager.clone(), + ResolvedOAuthCredentialStore::Keyring(AuthKeyringBackendKind::Direct), + Some(initial_tokens), + ); + + persistor + .refresh_if_needed_with_keyring_store(&store) + .await?; + + let manager_tokens = tokens_from_manager(&manager).await?; + assert_eq!( + manager_tokens.token_response, + WrappedOAuthTokenResponse(request_oauth_token_response(&latest_tokens)) + ); + Ok(()) + } + + #[tokio::test] + async fn refresh_transaction_refreshes_in_guard_band_despite_expires_in_drift() -> Result<()> { + let _env = TempCodexHome::new(); + let server = MockServer::start().await; + mount_oauth_metadata(&server).await; + mount_refresh_response( + &server, + "refresh-token", + "refreshed-after-expiry-drift", + "rotated-after-expiry-drift", + ) + .await; + + let store = MockKeyringStore::default(); + let mut initial_tokens = sample_tokens(); + initial_tokens.url = format!("{}/mcp", server.uri()); + initial_tokens.expires_at = Some(now_millis().saturating_add(45_000)); + initial_tokens + .token_response + .0 + .set_expires_in(Some(&Duration::from_secs(3600))); + super::save_oauth_tokens_with_keyring_store( + &store, + &initial_tokens.server_name, + &initial_tokens, + OAuthCredentialsStoreMode::Keyring, + AuthKeyringBackendKind::Direct, + )?; + + let manager = authorization_manager_for(&initial_tokens).await?; + let persistor = OAuthPersistor::new( + initial_tokens.server_name.clone(), + initial_tokens.url.clone(), + manager, + ResolvedOAuthCredentialStore::Keyring(AuthKeyringBackendKind::Direct), + Some(initial_tokens.clone()), + ); + + persistor + .refresh_if_needed_with_keyring_store(&store) + .await?; + + server.verify().await; + let stored = super::resolve_oauth_tokens( + &store, + &initial_tokens.server_name, + &initial_tokens.url, + OAuthCredentialsStoreMode::Keyring, + AuthKeyringBackendKind::Direct, + )? + .expect("refreshed tokens should be persisted"); + assert_eq!(access_token(&stored.tokens), "refreshed-after-expiry-drift"); + assert_eq!( + refresh_token(&stored.tokens), + Some("rotated-after-expiry-drift".to_string()) + ); + Ok(()) + } + + #[tokio::test] + async fn refresh_transaction_uses_latest_refresh_token_when_reread_is_expired() -> Result<()> { + let _env = TempCodexHome::new(); + let server = MockServer::start().await; + mount_oauth_metadata(&server).await; + mount_refresh_response( + &server, + "latest-refresh-token", + "refreshed-from-latest-token", + "rotated-from-latest-token", + ) + .await; + + let store = MockKeyringStore::default(); + let mut initial_tokens = expired_sample_tokens(&format!("{}/mcp", server.uri())); + initial_tokens + .token_response + .0 + .set_refresh_token(Some(RefreshToken::new("stale-refresh-token".to_string()))); + + let mut latest_tokens = initial_tokens.clone(); + latest_tokens + .token_response + .0 + .set_access_token(AccessToken::new("latest-expired-access-token".to_string())); + latest_tokens + .token_response + .0 + .set_refresh_token(Some(RefreshToken::new("latest-refresh-token".to_string()))); + super::save_oauth_tokens_with_keyring_store( + &store, + &latest_tokens.server_name, + &latest_tokens, + OAuthCredentialsStoreMode::Keyring, + AuthKeyringBackendKind::Direct, + )?; + + let manager = authorization_manager_for(&initial_tokens).await?; + let persistor = OAuthPersistor::new( + initial_tokens.server_name.clone(), + initial_tokens.url.clone(), + manager, + ResolvedOAuthCredentialStore::Keyring(AuthKeyringBackendKind::Direct), + Some(initial_tokens.clone()), + ); + + persistor + .refresh_if_needed_with_keyring_store(&store) + .await?; + + server.verify().await; + let stored = super::resolve_oauth_tokens( + &store, + &initial_tokens.server_name, + &initial_tokens.url, + OAuthCredentialsStoreMode::Keyring, + AuthKeyringBackendKind::Direct, + )? + .expect("refreshed tokens should be persisted"); + assert_eq!(access_token(&stored.tokens), "refreshed-from-latest-token"); + assert_eq!( + refresh_token(&stored.tokens), + Some("rotated-from-latest-token".to_string()) + ); + Ok(()) + } + + #[tokio::test] + async fn provider_refresh_timeout_releases_lock_without_persisting_and_permits_retry() + -> 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 store = MockKeyringStore::default(); + let initial_tokens = expired_sample_tokens(&format!("{}/mcp", server.uri())); + super::save_oauth_tokens_with_keyring_store( + &store, + &initial_tokens.server_name, + &initial_tokens, + OAuthCredentialsStoreMode::Keyring, + AuthKeyringBackendKind::Direct, + )?; + let manager = authorization_manager_for(&initial_tokens).await?; + let persistor = OAuthPersistor::new( + initial_tokens.server_name.clone(), + initial_tokens.url.clone(), + manager, + ResolvedOAuthCredentialStore::Keyring(AuthKeyringBackendKind::Direct), + Some(initial_tokens.clone()), + ); + + let 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")); + + let lock = tokio::time::timeout( + Duration::from_millis(100), + RefreshCredentialLock::acquire_for_server( + &initial_tokens.server_name, + &initial_tokens.url, + ), + ) + .await + .expect("the timed-out provider request should release its credential lock")?; + let stored = super::resolve_oauth_tokens( + &store, + &initial_tokens.server_name, + &initial_tokens.url, + OAuthCredentialsStoreMode::Keyring, + AuthKeyringBackendKind::Direct, + )? + .expect("a timed-out provider request must not update durable credentials"); + assert_eq!(access_token(&stored.tokens), "access-token"); + assert_eq!( + refresh_token(&stored.tokens), + Some("refresh-token".to_string()) + ); + drop(lock); + + persistor + .refresh_if_needed_with_keyring_store_and_timeout(&store, Duration::from_secs(1)) + .await?; + + let stored = super::resolve_oauth_tokens( + &store, + &initial_tokens.server_name, + &initial_tokens.url, + OAuthCredentialsStoreMode::Keyring, + AuthKeyringBackendKind::Direct, + )? + .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("retried-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<()> { + 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), + ) + .await; + + let store = MockKeyringStore::default(); + let initial_tokens = expired_sample_tokens(&format!("{}/mcp", server.uri())); + super::save_oauth_tokens_with_keyring_store( + &store, + &initial_tokens.server_name, + &initial_tokens, + OAuthCredentialsStoreMode::Keyring, + AuthKeyringBackendKind::Direct, + )?; + let manager = authorization_manager_for(&initial_tokens).await?; + let persistor = OAuthPersistor::new( + initial_tokens.server_name.clone(), + initial_tokens.url.clone(), + manager, + ResolvedOAuthCredentialStore::Keyring(AuthKeyringBackendKind::Direct), + Some(initial_tokens.clone()), + ); + let 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()); + + // 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 _lock = tokio::time::timeout( + Duration::from_secs(/*secs*/ 2), + RefreshCredentialLock::acquire_for_server( + &initial_tokens.server_name, + &initial_tokens.url, + ), + ) + .await + .expect("the independently owned refresh should release its credential lock")?; + let stored = super::resolve_oauth_tokens( + &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; + Ok(()) + } + + #[tokio::test] + async fn locked_login_save_after_refresh_still_wins() -> Result<()> { + let _env = TempCodexHome::new(); + let server = MockServer::start().await; + mount_oauth_metadata(&server).await; + let refresh_started = mount_refresh_response_with_signal( + &server, + "refresh-token", + "refreshed-before-login", + "rotated-before-login", + Duration::from_millis(200), + ) + .await; + + let store = MockKeyringStore::default(); + let initial_tokens = expired_sample_tokens(&format!("{}/mcp", server.uri())); + super::save_oauth_tokens_with_keyring_store( + &store, + &initial_tokens.server_name, + &initial_tokens, + OAuthCredentialsStoreMode::Keyring, + AuthKeyringBackendKind::Direct, + )?; + + let manager = authorization_manager_for(&initial_tokens).await?; + let persistor = OAuthPersistor::new( + initial_tokens.server_name.clone(), + initial_tokens.url.clone(), + manager, + ResolvedOAuthCredentialStore::Keyring(AuthKeyringBackendKind::Direct), + Some(initial_tokens.clone()), + ); + let refresh_task = tokio::spawn({ + let persistor = persistor.clone(); + let store = store.clone(); + async move { persistor.refresh_if_needed_with_keyring_store(&store).await } + }); + + wait_for_signal(refresh_started).await?; + + let mut login_tokens = sample_tokens(); + login_tokens.url.clone_from(&initial_tokens.url); + login_tokens + .token_response + .0 + .set_access_token(AccessToken::new("login-after-refresh-access".to_string())); + login_tokens + .token_response + .0 + .set_refresh_token(Some(RefreshToken::new( + "login-after-refresh-token".to_string(), + ))); + let (contended_tx, contended_rx) = mpsc::channel(); + let _subscriber_guard = tracing::subscriber::set_default(LockContentionSubscriber { + target: REFRESH_LOCK_CONTENTION_EVENT_TARGET, + contended_tx, + }); + let login_task = tokio::spawn({ + let store = store.clone(); + async move { + super::save_oauth_tokens_locked_with_keyring_store( + &store, + &login_tokens.server_name, + &login_tokens, + OAuthCredentialsStoreMode::Keyring, + AuthKeyringBackendKind::Direct, + ) + .await + } + }); + + // Prove the login save actually reaches RefreshCredentialLock's WouldBlock branch while + // refresh owns the transaction. Without login serialization this marker never arrives. + wait_for_signal(contended_rx).await?; + refresh_task.await??; + login_task.await??; + + server.verify().await; + let stored = super::resolve_oauth_tokens( + &store, + &initial_tokens.server_name, + &initial_tokens.url, + OAuthCredentialsStoreMode::Keyring, + AuthKeyringBackendKind::Direct, + )? + .expect("login tokens should remain persisted"); + assert_eq!(access_token(&stored.tokens), "login-after-refresh-access"); + assert_eq!( + refresh_token(&stored.tokens), + Some("login-after-refresh-token".to_string()) + ); + Ok(()) + } + + #[tokio::test] + async fn locked_logout_after_refresh_still_deletes() -> Result<()> { + let _env = TempCodexHome::new(); + let server = MockServer::start().await; + mount_oauth_metadata(&server).await; + let refresh_started = mount_refresh_response_with_signal( + &server, + "refresh-token", + "refreshed-before-logout", + "rotated-before-logout", + Duration::from_millis(200), + ) + .await; + + let store = MockKeyringStore::default(); + let initial_tokens = expired_sample_tokens(&format!("{}/mcp", server.uri())); + super::save_oauth_tokens_with_keyring_store( + &store, + &initial_tokens.server_name, + &initial_tokens, + OAuthCredentialsStoreMode::Keyring, + AuthKeyringBackendKind::Direct, + )?; + + let manager = authorization_manager_for(&initial_tokens).await?; + let persistor = OAuthPersistor::new( + initial_tokens.server_name.clone(), + initial_tokens.url.clone(), + manager, + ResolvedOAuthCredentialStore::Keyring(AuthKeyringBackendKind::Direct), + Some(initial_tokens.clone()), + ); + let refresh_task = tokio::spawn({ + let persistor = persistor.clone(); + let store = store.clone(); + async move { persistor.refresh_if_needed_with_keyring_store(&store).await } + }); + + wait_for_signal(refresh_started).await?; + + let (contended_tx, contended_rx) = mpsc::channel(); + let _subscriber_guard = tracing::subscriber::set_default(LockContentionSubscriber { + target: REFRESH_LOCK_CONTENTION_EVENT_TARGET, + contended_tx, + }); + let logout_task = tokio::spawn({ + let store = store.clone(); + let server_name = initial_tokens.server_name.clone(); + let url = initial_tokens.url.clone(); + async move { + super::delete_oauth_tokens_locked_with_keyring_store( + &store, + OAuthCredentialsStoreMode::Keyring, + AuthKeyringBackendKind::Direct, + &server_name, + &url, + ) + .await + } + }); + + // As with login, synchronize on the lock implementation's actual contention branch so a + // scheduler delay cannot accidentally turn this into a non-racing logout. + wait_for_signal(contended_rx).await?; + refresh_task.await??; + let removed = logout_task.await??; + + server.verify().await; + assert!(removed); + let stored = super::resolve_oauth_tokens( + &store, + &initial_tokens.server_name, + &initial_tokens.url, + OAuthCredentialsStoreMode::Keyring, + AuthKeyringBackendKind::Direct, + )?; + assert!(stored.is_none()); + Ok(()) + } + #[test] fn save_oauth_tokens_with_secrets_backend_falls_back_to_file_when_keyring_fails() -> Result<()> { @@ -1088,7 +2158,7 @@ mod tests { &tokens, )?; - let saved = super::read_fallback_file_unlocked()?.expect("fallback file should load"); + let saved = read_fallback_file()?.expect("fallback file should load"); let key = super::compute_store_key(&tokens.server_name, &tokens.url)?; assert!(saved.contains_key(&key)); Ok(()) @@ -1329,6 +2399,11 @@ mod tests { ); } + fn read_fallback_file() -> Result> { + let _store_lock = OAuthStoreLock::acquire(OAuthStore::File)?; + super::read_fallback_file_unlocked() + } + fn assert_token_response_match_without_expiry( actual: &WrappedOAuthTokenResponse, expected: &WrappedOAuthTokenResponse, @@ -1356,6 +2431,152 @@ mod tests { ); } + async fn authorization_manager_for( + tokens: &StoredOAuthTokens, + ) -> Result>> { + let mut state = OAuthState::new(tokens.url.clone(), Some(reqwest::Client::new())).await?; + state + .set_credentials(&tokens.client_id, request_oauth_token_response(tokens)) + .await?; + let manager = match state { + OAuthState::Authorized(manager) | OAuthState::Unauthorized(manager) => manager, + OAuthState::Session(_) | OAuthState::AuthorizedHttpClient(_) => { + anyhow::bail!("unexpected OAuth state") + } + _ => anyhow::bail!("unexpected OAuth state"), + }; + Ok(Arc::new(TokioMutex::new(manager))) + } + + async fn mount_oauth_metadata(server: &MockServer) { + Mock::given(method("GET")) + .and(path("/.well-known/oauth-authorization-server/mcp")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "authorization_endpoint": format!("{}/oauth/authorize", server.uri()), + "token_endpoint": format!("{}/oauth/token", server.uri()), + "scopes_supported": ["scope-a", "scope-b"], + }))) + .mount(server) + .await; + } + + async fn mount_refresh_response( + server: &MockServer, + request_refresh_token: &str, + response_access_token: &str, + response_refresh_token: &str, + ) { + Mock::given(method("POST")) + .and(path("/oauth/token")) + .and(body_string_contains("grant_type=refresh_token")) + .and(body_string_contains(format!( + "refresh_token={request_refresh_token}" + ))) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "access_token": response_access_token, + "token_type": "Bearer", + "expires_in": 3600, + "refresh_token": response_refresh_token, + "scope": "scope-a scope-b", + }))) + .expect(1) + .mount(server) + .await; + } + + async fn mount_refresh_response_with_signal( + server: &MockServer, + request_refresh_token: &str, + response_access_token: &str, + response_refresh_token: &str, + response_delay: Duration, + ) -> mpsc::Receiver<()> { + let (tx, rx) = mpsc::channel(); + let response_access_token = response_access_token.to_string(); + let response_refresh_token = response_refresh_token.to_string(); + Mock::given(method("POST")) + .and(path("/oauth/token")) + .and(body_string_contains("grant_type=refresh_token")) + .and(body_string_contains(format!( + "refresh_token={request_refresh_token}" + ))) + .respond_with(move |_request: &wiremock::Request| { + let _ = tx.send(()); + let access_token = response_access_token.clone(); + let refresh_token = response_refresh_token.clone(); + ResponseTemplate::new(200) + .set_delay(response_delay) + .set_body_json(json!({ + "access_token": access_token, + "token_type": "Bearer", + "expires_in": 3600, + "refresh_token": refresh_token, + "scope": "scope-a scope-b", + })) + }) + .expect(1) + .mount(server) + .await; + rx + } + + async fn wait_for_signal(rx: mpsc::Receiver<()>) -> Result<()> { + tokio::task::spawn_blocking(move || { + rx.recv_timeout(Duration::from_secs(5)) + .context("timed out waiting for refresh request") + }) + .await? + } + + #[expect( + clippy::await_holding_invalid_type, + reason = "AuthorizationManager async access must be serialized through its mutex" + )] + async fn tokens_from_manager( + manager: &Arc>, + ) -> Result { + let guard = manager.lock().await; + let (client_id, token_response) = guard.get_credentials().await?; + let token_response = token_response.expect("manager should have token response"); + Ok(StoredOAuthTokens { + server_name: "test-server".to_string(), + url: "https://example.test".to_string(), + client_id, + token_response: WrappedOAuthTokenResponse(token_response), + expires_at: None, + }) + } + + fn access_token(tokens: &StoredOAuthTokens) -> &str { + tokens.token_response.0.access_token().secret() + } + + fn refresh_token(tokens: &StoredOAuthTokens) -> Option { + tokens + .token_response + .0 + .refresh_token() + .map(|token| token.secret().to_string()) + } + + fn expired_sample_tokens(url: &str) -> StoredOAuthTokens { + let mut tokens = sample_tokens(); + tokens.url = url.to_string(); + tokens.expires_at = Some(0); + tokens + .token_response + .0 + .set_expires_in(Some(&Duration::ZERO)); + tokens + } + + fn now_millis() -> u64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_else(|_| Duration::from_secs(0)) + .as_millis() as u64 + } + fn sample_tokens() -> StoredOAuthTokens { let mut response = OAuthTokenResponse::new( AccessToken::new("access-token".to_string()), diff --git a/codex-rs/rmcp-client/src/oauth/resolution_state.rs b/codex-rs/rmcp-client/src/oauth/resolution_state.rs index d188c1d60e00..33ec910b5a15 100644 --- a/codex-rs/rmcp-client/src/oauth/resolution_state.rs +++ b/codex-rs/rmcp-client/src/oauth/resolution_state.rs @@ -249,3 +249,7 @@ fn read_resolution_state(path: &Path) -> Result { ) }) } + +#[cfg(test)] +#[path = "resolution_state_tests.rs"] +mod tests; diff --git a/codex-rs/rmcp-client/src/oauth/resolution_state_tests.rs b/codex-rs/rmcp-client/src/oauth/resolution_state_tests.rs new file mode 100644 index 000000000000..876747777e99 --- /dev/null +++ b/codex-rs/rmcp-client/src/oauth/resolution_state_tests.rs @@ -0,0 +1,131 @@ +use codex_config::types::AuthKeyringBackendKind; +use codex_config::types::OAuthCredentialsStoreMode; +use pretty_assertions::assert_eq; + +use super::ObservedStore; +use super::ResolutionState; +use super::StoreResolution; +use super::StoreResolutionChange; +use super::record_store_resolution_in; +use crate::oauth::ResolvedOAuthCredentialStore; +use crate::oauth::compute_store_key; + +#[test] +fn auto_resolution_change_is_reported_and_persisted() -> anyhow::Result<()> { + let codex_home = tempfile::tempdir()?; + let server_name = "test-server"; + let url = "https://example.test/mcp"; + + assert_eq!( + record_store_resolution_in( + codex_home.path(), + server_name, + url, + OAuthCredentialsStoreMode::Auto, + AuthKeyringBackendKind::Direct, + ResolvedOAuthCredentialStore::Keyring(AuthKeyringBackendKind::Direct), + )?, + None + ); + assert_eq!( + record_store_resolution_in( + codex_home.path(), + server_name, + url, + OAuthCredentialsStoreMode::Auto, + AuthKeyringBackendKind::Direct, + ResolvedOAuthCredentialStore::File, + )?, + Some(StoreResolutionChange { + previous: ObservedStore::Keyring, + current: ObservedStore::File, + }) + ); + + let state: ResolutionState = serde_json::from_str(&std::fs::read_to_string( + codex_home.path().join(super::RESOLUTION_STATE_FILENAME), + )?)?; + assert_eq!( + state, + ResolutionState { + version: super::RESOLUTION_STATE_VERSION, + resolutions: [( + compute_store_key(server_name, url)?, + StoreResolution { + store_mode: OAuthCredentialsStoreMode::Auto, + keyring_backend: AuthKeyringBackendKind::Direct, + resolved_store: ObservedStore::File, + }, + )] + .into(), + } + ); + Ok(()) +} + +#[test] +fn intentional_configuration_changes_reset_the_auto_comparison() -> anyhow::Result<()> { + let codex_home = tempfile::tempdir()?; + let server_name = "test-server"; + let url = "https://example.test/mcp"; + + assert_eq!( + record_store_resolution_in( + codex_home.path(), + server_name, + url, + OAuthCredentialsStoreMode::Auto, + AuthKeyringBackendKind::Direct, + ResolvedOAuthCredentialStore::Keyring(AuthKeyringBackendKind::Direct), + )?, + None + ); + assert_eq!( + record_store_resolution_in( + codex_home.path(), + server_name, + url, + OAuthCredentialsStoreMode::File, + AuthKeyringBackendKind::Direct, + ResolvedOAuthCredentialStore::File, + )?, + None + ); + assert_eq!( + record_store_resolution_in( + codex_home.path(), + server_name, + url, + OAuthCredentialsStoreMode::Auto, + AuthKeyringBackendKind::Direct, + ResolvedOAuthCredentialStore::File, + )?, + None + ); + assert_eq!( + record_store_resolution_in( + codex_home.path(), + server_name, + url, + OAuthCredentialsStoreMode::Auto, + AuthKeyringBackendKind::Secrets, + ResolvedOAuthCredentialStore::Keyring(AuthKeyringBackendKind::Secrets), + )?, + None + ); + assert_eq!( + record_store_resolution_in( + codex_home.path(), + server_name, + url, + OAuthCredentialsStoreMode::Auto, + AuthKeyringBackendKind::Secrets, + ResolvedOAuthCredentialStore::File, + )?, + Some(StoreResolutionChange { + previous: ObservedStore::Keyring, + current: ObservedStore::File, + }) + ); + Ok(()) +} diff --git a/codex-rs/rmcp-client/src/oauth_transport.rs b/codex-rs/rmcp-client/src/oauth_transport.rs index fbdda2641647..b6e9556de888 100644 --- a/codex-rs/rmcp-client/src/oauth_transport.rs +++ b/codex-rs/rmcp-client/src/oauth_transport.rs @@ -215,3 +215,7 @@ fn oauth_transport_error( ) -> StreamableHttpError { StreamableHttpError::Client(StreamableHttpClientAdapterError::OAuth(error)) } + +#[cfg(test)] +#[path = "oauth_transport_tests.rs"] +mod tests; diff --git a/codex-rs/rmcp-client/src/oauth_transport_tests.rs b/codex-rs/rmcp-client/src/oauth_transport_tests.rs new file mode 100644 index 000000000000..508c3dd913b4 --- /dev/null +++ b/codex-rs/rmcp-client/src/oauth_transport_tests.rs @@ -0,0 +1,190 @@ +use std::collections::HashMap; +use std::sync::Arc; + +use codex_config::types::AuthKeyringBackendKind; +use codex_config::types::OAuthCredentialsStoreMode; +use codex_exec_server::Environment; +use oauth2::AccessToken; +use oauth2::RefreshToken; +use oauth2::basic::BasicTokenType; +use reqwest::header::HeaderMap; +use rmcp::model::ClientJsonRpcMessage; +use rmcp::transport::auth::AuthClient; +use rmcp::transport::auth::OAuthState; +use rmcp::transport::auth::OAuthTokenResponse; +use rmcp::transport::auth::VendorExtraTokenFields; +use rmcp::transport::streamable_http_client::StreamableHttpClient; +use rmcp::transport::streamable_http_client::StreamableHttpPostResponse; +use serde_json::json; +use tempfile::TempDir; +use tokio::process::Command; +use wiremock::Mock; +use wiremock::MockServer; +use wiremock::ResponseTemplate; +use wiremock::matchers::body_string_contains; +use wiremock::matchers::header; +use wiremock::matchers::method; +use wiremock::matchers::path; + +use super::OAuthTransportClient; +use crate::http_client_adapter::StreamableHttpClientAdapter; +use crate::oauth::OAuthPersistor; +use crate::oauth::ResolvedOAuthCredentialStore; +use crate::oauth::StoredOAuthTokens; +use crate::oauth::WrappedOAuthTokenResponse; +use crate::oauth::request_oauth_token_response; +use crate::oauth::save_oauth_tokens; +use crate::oauth_http_client::OAuthHttpClientAdapter; + +const SERVER_NAME: &str = "oauth-transport-response-test"; +const SERVER_URL_ENV: &str = "MCP_TEST_OAUTH_RESPONSE_SERVER_URL"; +const ACCESS_TOKEN_A: &str = "response-access-a"; +const REFRESH_TOKEN_A: &str = "response-refresh-a"; +const ACCESS_TOKEN_B: &str = "response-access-b"; +const REFRESH_TOKEN_B: &str = "response-refresh-b"; + +#[tokio::test] +async fn server_response_post_receives_one_shot_oauth_recovery() -> anyhow::Result<()> { + let server = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/.well-known/oauth-authorization-server/mcp")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "authorization_endpoint": format!("{}/oauth/authorize", server.uri()), + "token_endpoint": format!("{}/oauth/token", server.uri()), + "scopes_supported": ["scope-a"], + }))) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/oauth/token")) + .and(body_string_contains("grant_type=refresh_token")) + .and(body_string_contains(format!( + "refresh_token={REFRESH_TOKEN_A}" + ))) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "access_token": ACCESS_TOKEN_B, + "token_type": "Bearer", + "expires_in": 3600, + "refresh_token": REFRESH_TOKEN_B, + "scope": "scope-a", + }))) + .expect(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/mcp")) + .and(header("authorization", format!("Bearer {ACCESS_TOKEN_A}"))) + .respond_with(ResponseTemplate::new(401)) + .expect(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/mcp")) + .and(header("authorization", format!("Bearer {ACCESS_TOKEN_B}"))) + .respond_with(ResponseTemplate::new(202)) + .expect(1) + .mount(&server) + .await; + + let codex_home = TempDir::new()?; + let server_url = format!("{}/mcp", server.uri()); + let status = Command::new(std::env::current_exe()?) + .args([ + "oauth_transport::tests::server_response_post_child", + "--exact", + "--ignored", + "--nocapture", + ]) + .env("CODEX_HOME", codex_home.path()) + .env(SERVER_URL_ENV, server_url) + .status() + .await?; + anyhow::ensure!(status.success(), "OAuth response child failed: {status}"); + server.verify().await; + Ok(()) +} + +#[tokio::test] +#[ignore = "spawned by server_response_post_receives_one_shot_oauth_recovery"] +async fn server_response_post_child() -> anyhow::Result<()> { + let server_url = std::env::var(SERVER_URL_ENV)?; + let initial_tokens = initial_tokens(&server_url); + save_oauth_tokens( + SERVER_NAME, + &initial_tokens, + OAuthCredentialsStoreMode::File, + AuthKeyringBackendKind::default(), + )?; + + let http_client = Environment::default_for_tests().get_http_client(); + let oauth_http_client = Arc::new(OAuthHttpClientAdapter::new( + Arc::clone(&http_client), + HeaderMap::new(), + )); + let mut oauth_state = + OAuthState::new_with_oauth_http_client(server_url.clone(), oauth_http_client).await?; + oauth_state + .set_credentials( + &initial_tokens.client_id, + request_oauth_token_response(&initial_tokens), + ) + .await?; + let manager = match oauth_state { + OAuthState::Authorized(manager) | OAuthState::Unauthorized(manager) => manager, + _ => anyhow::bail!("unexpected OAuth state during response test setup"), + }; + let auth_client = AuthClient::new( + StreamableHttpClientAdapter::new( + Arc::clone(&http_client), + HeaderMap::new(), + /*auth_provider*/ None, + ), + manager, + ); + let persistor = OAuthPersistor::new( + SERVER_NAME.to_string(), + server_url.clone(), + Arc::clone(&auth_client.auth_manager), + ResolvedOAuthCredentialStore::File, + Some(initial_tokens), + ); + let client = OAuthTransportClient::new(auth_client, persistor); + let response_message: ClientJsonRpcMessage = serde_json::from_value(json!({ + "jsonrpc": "2.0", + "id": "server-request-1", + "result": { + "action": "accept", + "content": { "confirmed": true } + } + }))?; + + let response = client + .post_message( + Arc::from(server_url), + response_message, + Some(Arc::from("response-session")), + /*auth_token*/ None, + HashMap::new(), + ) + .await?; + + assert!(matches!(response, StreamableHttpPostResponse::Accepted)); + Ok(()) +} + +fn initial_tokens(server_url: &str) -> StoredOAuthTokens { + let mut response = OAuthTokenResponse::new( + AccessToken::new(ACCESS_TOKEN_A.to_string()), + BasicTokenType::Bearer, + VendorExtraTokenFields::default(), + ); + response.set_refresh_token(Some(RefreshToken::new(REFRESH_TOKEN_A.to_string()))); + response.set_expires_in(None); + StoredOAuthTokens { + server_name: SERVER_NAME.to_string(), + url: server_url.to_string(), + client_id: "test-client-id".to_string(), + token_response: WrappedOAuthTokenResponse(response), + expires_at: None, + } +} diff --git a/codex-rs/rmcp-client/tests/streamable_http_oauth_internal.rs b/codex-rs/rmcp-client/tests/streamable_http_oauth_internal.rs new file mode 100644 index 000000000000..ceee3625cd54 --- /dev/null +++ b/codex-rs/rmcp-client/tests/streamable_http_oauth_internal.rs @@ -0,0 +1,219 @@ +mod streamable_http_test_support; + +use std::time::Duration; + +use codex_config::types::AuthKeyringBackendKind; +use codex_config::types::OAuthCredentialsStoreMode; +use codex_exec_server::Environment; +use codex_rmcp_client::RmcpClient; +use codex_rmcp_client::StoredOAuthTokens; +use codex_rmcp_client::WrappedOAuthTokenResponse; +use codex_rmcp_client::save_oauth_tokens; +use oauth2::AccessToken; +use oauth2::RefreshToken; +use oauth2::basic::BasicTokenType; +use rmcp::transport::auth::OAuthTokenResponse; +use rmcp::transport::auth::VendorExtraTokenFields; +use serde_json::Value; +use serde_json::json; +use tempfile::TempDir; +use tokio::process::Command; +use wiremock::Mock; +use wiremock::MockServer; +use wiremock::Request; +use wiremock::ResponseTemplate; +use wiremock::matchers::body_string_contains; +use wiremock::matchers::header; +use wiremock::matchers::method; +use wiremock::matchers::path; + +use streamable_http_test_support::initialize_client; + +const SERVER_NAME: &str = "test-streamable-http-oauth-internal"; +const SERVER_URL_ENV: &str = "MCP_TEST_OAUTH_INTERNAL_SERVER_URL"; +const ACCESS_TOKEN_A: &str = "internal-access-a"; +const REFRESH_TOKEN_A: &str = "internal-refresh-a"; +const ACCESS_TOKEN_B: &str = "internal-access-b"; +const REFRESH_TOKEN_B: &str = "internal-refresh-b"; +const ACCESS_TOKEN_C: &str = "internal-access-c"; +const REFRESH_TOKEN_C: &str = "internal-refresh-c"; + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn rmcp_owned_get_and_delete_receive_oauth_recovery() -> anyhow::Result<()> { + let server = MockServer::start().await; + mount_oauth_metadata(&server).await; + mount_refresh(&server, REFRESH_TOKEN_A, ACCESS_TOKEN_B, REFRESH_TOKEN_B).await; + mount_refresh(&server, REFRESH_TOKEN_B, ACCESS_TOKEN_C, REFRESH_TOKEN_C).await; + + Mock::given(method("POST")) + .and(path("/mcp")) + .and(header("authorization", format!("Bearer {ACCESS_TOKEN_A}"))) + .respond_with(|request: &Request| { + let body: Value = request.body_json().expect("valid JSON-RPC request"); + match body.get("method").and_then(Value::as_str) { + Some("initialize") => initialize_response(&body), + Some("notifications/initialized") => ResponseTemplate::new(202), + method => ResponseTemplate::new(400) + .set_body_string(format!("unexpected JSON-RPC method: {method:?}")), + } + }) + .expect(2) + .mount(&server) + .await; + Mock::given(method("GET")) + .and(path("/mcp")) + .and(header("authorization", format!("Bearer {ACCESS_TOKEN_A}"))) + .respond_with(ResponseTemplate::new(401)) + .expect(1) + .mount(&server) + .await; + Mock::given(method("GET")) + .and(path("/mcp")) + .and(header("authorization", format!("Bearer {ACCESS_TOKEN_B}"))) + // A 405 tells RMCP that the optional common SSE stream is unsupported. Reaching this + // response proves that the wrapper retried the RMCP-owned GET with B after refreshing A. + .respond_with(ResponseTemplate::new(405)) + .expect(1) + .mount(&server) + .await; + Mock::given(method("DELETE")) + .and(path("/mcp")) + .and(header("authorization", format!("Bearer {ACCESS_TOKEN_B}"))) + .respond_with(ResponseTemplate::new(401)) + .expect(1) + .mount(&server) + .await; + Mock::given(method("DELETE")) + .and(path("/mcp")) + .and(header("authorization", format!("Bearer {ACCESS_TOKEN_C}"))) + .respond_with(ResponseTemplate::new(204)) + .expect(1) + .mount(&server) + .await; + + run_child( + "oauth_internal_get_delete_child", + &format!("{}/mcp", server.uri()), + ) + .await?; + server.verify().await; + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +#[ignore = "spawned by rmcp_owned_get_and_delete_receive_oauth_recovery"] +async fn oauth_internal_get_delete_child() -> anyhow::Result<()> { + let client = create_oauth_client().await?; + initialize_client(&client).await?; + tokio::time::sleep(Duration::from_millis(500)).await; + client.shutdown().await; + tokio::time::sleep(Duration::from_millis(750)).await; + Ok(()) +} + +async fn create_oauth_client() -> anyhow::Result { + let server_url = std::env::var(SERVER_URL_ENV)?; + save_initial_tokens(&server_url)?; + RmcpClient::new_streamable_http_client( + SERVER_NAME, + &server_url, + /*bearer_token*/ None, + /*http_headers*/ None, + /*env_http_headers*/ None, + OAuthCredentialsStoreMode::File, + AuthKeyringBackendKind::default(), + Environment::default_for_tests().get_http_client(), + /*auth_provider*/ None, + ) + .await +} + +fn save_initial_tokens(server_url: &str) -> anyhow::Result<()> { + let mut response = OAuthTokenResponse::new( + AccessToken::new(ACCESS_TOKEN_A.to_string()), + BasicTokenType::Bearer, + VendorExtraTokenFields::default(), + ); + response.set_refresh_token(Some(RefreshToken::new(REFRESH_TOKEN_A.to_string()))); + response.set_expires_in(None); + save_oauth_tokens( + SERVER_NAME, + &StoredOAuthTokens { + server_name: SERVER_NAME.to_string(), + url: server_url.to_string(), + client_id: "test-client-id".to_string(), + token_response: WrappedOAuthTokenResponse(response), + expires_at: None, + }, + OAuthCredentialsStoreMode::File, + AuthKeyringBackendKind::default(), + ) +} + +async fn mount_oauth_metadata(server: &MockServer) { + Mock::given(method("GET")) + .and(path("/.well-known/oauth-authorization-server/mcp")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "authorization_endpoint": format!("{}/oauth/authorize", server.uri()), + "token_endpoint": format!("{}/oauth/token", server.uri()), + "scopes_supported": ["scope-a"], + }))) + .mount(server) + .await; +} + +async fn mount_refresh( + server: &MockServer, + request_refresh_token: &str, + response_access_token: &str, + response_refresh_token: &str, +) { + Mock::given(method("POST")) + .and(path("/oauth/token")) + .and(body_string_contains("grant_type=refresh_token")) + .and(body_string_contains(format!( + "refresh_token={request_refresh_token}" + ))) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "access_token": response_access_token, + "token_type": "Bearer", + "expires_in": 3600, + "refresh_token": response_refresh_token, + "scope": "scope-a", + }))) + .expect(1) + .mount(server) + .await; +} + +fn initialize_response(body: &Value) -> ResponseTemplate { + ResponseTemplate::new(200) + .insert_header("mcp-session-id", "oauth-internal-session") + .set_body_json(json!({ + "jsonrpc": "2.0", + "id": body.get("id").cloned().unwrap_or(Value::Null), + "result": { + "protocolVersion": body + .pointer("/params/protocolVersion") + .cloned() + .unwrap_or_else(|| json!("2025-06-18")), + "capabilities": {}, + "serverInfo": { + "name": "oauth-internal-test", + "version": "0.0.0-test" + } + } + })) +} + +async fn run_child(test_name: &str, server_url: &str) -> anyhow::Result<()> { + let codex_home = TempDir::new()?; + let status = Command::new(std::env::current_exe()?) + .args([test_name, "--exact", "--ignored", "--nocapture"]) + .env("CODEX_HOME", codex_home.path()) + .env(SERVER_URL_ENV, server_url) + .status() + .await?; + anyhow::ensure!(status.success(), "OAuth internal child failed: {status}"); + 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 ab193e7358d2..d4d2f12ecb4b 100644 --- a/codex-rs/rmcp-client/tests/streamable_http_oauth_startup.rs +++ b/codex-rs/rmcp-client/tests/streamable_http_oauth_startup.rs @@ -1,9 +1,20 @@ mod streamable_http_test_support; +use std::fs; +use std::path::Path; +use std::path::PathBuf; +use std::sync::Arc; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; use std::time::Duration; +use std::time::Instant; use std::time::SystemTime; use std::time::UNIX_EPOCH; +use axum::Json; +use axum::Router; +use axum::extract::State; +use axum::routing::post; use codex_config::types::AuthKeyringBackendKind; use codex_config::types::OAuthCredentialsStoreMode; use codex_exec_server::Environment; @@ -22,7 +33,18 @@ use rmcp::transport::auth::VendorExtraTokenFields; use serde_json::Value; use serde_json::json; use tempfile::TempDir; +use tokio::net::TcpListener; use tokio::process::Command; +use tokio::sync::Notify; +use tokio::sync::Semaphore; +use tokio::task::JoinHandle; +use tracing::Event; +use tracing::Id; +use tracing::Metadata; +use tracing::Subscriber; +use tracing::span::Attributes; +use tracing::span::Record; +use tracing::subscriber::Interest; use wiremock::Mock; use wiremock::MockServer; use wiremock::Request; @@ -33,18 +55,184 @@ use wiremock::matchers::method; use wiremock::matchers::path; use streamable_http_test_support::initialize_client; +use streamable_http_test_support::initialize_client_with_timeout; const SERVER_NAME: &str = "test-streamable-http-oauth-startup"; const EXPIRED_ACCESS_TOKEN: &str = "expired-access-token"; const REFRESH_TOKEN: &str = "valid-refresh-token"; const REFRESHED_ACCESS_TOKEN: &str = "refreshed-access-token"; +const ROTATED_REFRESH_TOKEN: &str = "rotated-refresh-token"; +const CHILD_CONTENTION_FILE_ENV: &str = "MCP_TEST_OAUTH_STARTUP_CONTENTION_FILE"; +const CHILD_READY_FILE_ENV: &str = "MCP_TEST_OAUTH_STARTUP_READY_FILE"; +const CHILD_RELEASE_FILE_ENV: &str = "MCP_TEST_OAUTH_STARTUP_RELEASE_FILE"; +const FINAL_ACCESS_TOKEN: &str = "final-access-token"; +const FINAL_REFRESH_TOKEN: &str = "final-refresh-token"; +const POST_DEADLINE_ACCESS_TOKEN: &str = "post-deadline-access-token"; +const POST_DEADLINE_REFRESH_TOKEN: &str = "post-deadline-refresh-token"; +const PREFLIGHT_REFRESH_ERROR: &str = "preflight refresh failed distinctly"; const CHILD_SERVER_URL_ENV: &str = "MCP_TEST_OAUTH_STARTUP_SERVER_URL"; +// This mirrors the private event target in oauth::refresh_lock without exposing test-only crate API. +const LOCK_CONTENTION_EVENT_TARGET: &str = "codex_rmcp_client::oauth::refresh_lock::contention"; const UNREFRESHABLE_SERVER_URL: &str = "https://unrefreshable.example/mcp"; const UNEXPIRED_SERVER_URL: &str = "https://unexpired.example/mcp"; const REFRESHABLE_SERVER_URL: &str = "https://refreshable.example/mcp"; +#[derive(Clone)] +struct GatedRefreshState { + request_count: Arc, + request_started: Arc, + response_release: Arc, +} + +struct GatedRefreshServer { + token_endpoint: String, + state: GatedRefreshState, + task: JoinHandle<()>, +} + +impl GatedRefreshServer { + async fn start() -> anyhow::Result { + let state = GatedRefreshState { + request_count: Arc::new(AtomicUsize::new(/*v*/ 0)), + request_started: Arc::new(Notify::new()), + response_release: Arc::new(Semaphore::new(/*permits*/ 0)), + }; + let router = Router::new() + .route("/oauth/token", post(gated_refresh_response)) + .with_state(state.clone()); + let listener = TcpListener::bind("127.0.0.1:0").await?; + let address = listener.local_addr()?; + let task = tokio::spawn(async move { + if let Err(error) = axum::serve(listener, router).await { + panic!("gated refresh server failed: {error}"); + } + }); + Ok(Self { + token_endpoint: format!("http://{address}/oauth/token"), + state, + task, + }) + } + + async fn wait_until_request_started(&self) -> anyhow::Result<()> { + let notified = self.state.request_started.notified(); + if self.request_count() == 0 { + tokio::time::timeout(Duration::from_secs(/*secs*/ 10), notified) + .await + .map_err(|_| anyhow::anyhow!("provider refresh request did not start"))?; + } + Ok(()) + } + + fn release_responses(&self) { + // Two permits also let the no-lock negative control exit cleanly after issuing two refresh + // requests. The passing path consumes one permit because only the lock owner calls out. + self.state.response_release.add_permits(/*n*/ 2); + } + + fn request_count(&self) -> usize { + self.state.request_count.load(Ordering::SeqCst) + } +} + +impl Drop for GatedRefreshServer { + fn drop(&mut self) { + self.task.abort(); + } +} + +async fn gated_refresh_response( + State(state): State, + body: String, +) -> Json { + assert!(body.contains("grant_type=refresh_token")); + assert!(body.contains(&format!("refresh_token={REFRESH_TOKEN}"))); + state.request_count.fetch_add(/*val*/ 1, Ordering::SeqCst); + state.request_started.notify_one(); + let Ok(permit) = state.response_release.acquire().await else { + panic!("gated refresh server closed its response semaphore"); + }; + permit.forget(); + Json(json!({ + "access_token": REFRESHED_ACCESS_TOKEN, + "token_type": "Bearer", + "expires_in": 7200, + "refresh_token": ROTATED_REFRESH_TOKEN, + })) +} + +struct LockContentionMarkerSubscriber { + marker_file: PathBuf, +} + +impl Subscriber for LockContentionMarkerSubscriber { + fn enabled(&self, metadata: &Metadata<'_>) -> bool { + metadata.target() == LOCK_CONTENTION_EVENT_TARGET + } + + fn register_callsite(&self, metadata: &'static Metadata<'static>) -> Interest { + if self.enabled(metadata) { + Interest::always() + } else { + Interest::never() + } + } + + fn max_level_hint(&self) -> Option { + Some(tracing::level_filters::LevelFilter::DEBUG) + } + + fn new_span(&self, _span: &Attributes<'_>) -> Id { + // This subscriber enables only the contention event callsite, so it never observes spans. + Id::from_u64(/*u*/ 1) + } + + fn record(&self, _span: &Id, _values: &Record<'_>) {} + + fn record_follows_from(&self, _span: &Id, _follows: &Id) {} + + fn event(&self, event: &Event<'_>) { + if event.metadata().target() == LOCK_CONTENTION_EVENT_TARGET + && let Err(error) = fs::write(&self.marker_file, b"contended") + { + panic!("failed to write refresh-lock contention marker: {error}"); + } + } + + fn enter(&self, _span: &Id) {} + + fn exit(&self, _span: &Id) {} +} + +async fn wait_for_marker(path: &Path, timeout_message: &str) -> anyhow::Result<()> { + tokio::time::timeout(Duration::from_secs(/*secs*/ 10), async { + while !path.exists() { + tokio::time::sleep(Duration::from_millis(/*millis*/ 10)).await; + } + }) + .await + .map_err(|_| anyhow::anyhow!("{timeout_message}")) +} + +fn oauth_concurrency_child_command( + codex_home: &Path, + server_url: &str, + ready_file: &Path, + release_file: &Path, +) -> anyhow::Result { + let mut command = Command::new(std::env::current_exe()?); + command + .args(["oauth_concurrency_client_child", "--exact", "--ignored"]) + .env("CODEX_HOME", codex_home) + .env(CHILD_SERVER_URL_ENV, server_url) + .env(CHILD_READY_FILE_ENV, ready_file) + .env(CHILD_RELEASE_FILE_ENV, release_file) + .kill_on_drop(true); + Ok(command) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 1)] -async fn refreshes_expired_persisted_token_before_initialize() -> anyhow::Result<()> { +async fn startup_refresh_does_not_consume_handshake_timeout() -> anyhow::Result<()> { let server = MockServer::start().await; Mock::given(method("GET")) .and(path("/.well-known/oauth-authorization-server/mcp")) @@ -62,12 +250,18 @@ async fn refreshes_expired_persisted_token_before_initialize() -> anyhow::Result .and(body_string_contains(format!( "refresh_token={REFRESH_TOKEN}" ))) - .respond_with(ResponseTemplate::new(200).set_body_json(json!({ - "access_token": REFRESHED_ACCESS_TOKEN, - "token_type": "Bearer", - "expires_in": 7200, - "refresh_token": REFRESH_TOKEN, - }))) + // The provider takes longer than the configured MCP handshake timeout. Refresh has its own + // bound, so this delay must not leave the subsequent handshake with an expired budget. + .respond_with( + ResponseTemplate::new(200) + .set_delay(Duration::from_millis(1_500)) + .set_body_json(json!({ + "access_token": REFRESHED_ACCESS_TOKEN, + "token_type": "Bearer", + "expires_in": 7200, + "refresh_token": REFRESH_TOKEN, + })), + ) .expect(1) .mount(&server) .await; @@ -121,6 +315,403 @@ async fn refreshes_expired_persisted_token_before_initialize() -> anyhow::Result Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn recovers_initialization_and_operation_401_without_rmcp_refresh() -> anyhow::Result<()> { + let server = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/.well-known/oauth-authorization-server/mcp")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "authorization_endpoint": format!("{}/oauth/authorize", server.uri()), + "token_endpoint": format!("{}/oauth/token", server.uri()), + "scopes_supported": [""], + }))) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/oauth/token")) + .and(body_string_contains("grant_type=refresh_token")) + .and(body_string_contains(format!( + "refresh_token={REFRESH_TOKEN}" + ))) + // Initialization receives a 401 before this refresh. Its response arrives after the MCP + // handshake timeout, proving that refresh time is excluded before the retry handshake. + .respond_with( + ResponseTemplate::new(200) + .set_delay(Duration::from_millis(1_500)) + .set_body_json(json!({ + "access_token": REFRESHED_ACCESS_TOKEN, + "token_type": "Bearer", + "expires_in": 7200, + "refresh_token": ROTATED_REFRESH_TOKEN, + })), + ) + .expect(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/oauth/token")) + .and(body_string_contains("grant_type=refresh_token")) + .and(body_string_contains(format!( + "refresh_token={FINAL_REFRESH_TOKEN}" + ))) + .respond_with( + ResponseTemplate::new(200) + .set_delay(Duration::from_secs(2)) + .set_body_json(json!({ + "access_token": POST_DEADLINE_ACCESS_TOKEN, + "token_type": "Bearer", + "expires_in": 7200, + "refresh_token": POST_DEADLINE_REFRESH_TOKEN, + })), + ) + .expect(1) + .mount(&server) + .await; + let final_tools_requests = Arc::new(AtomicUsize::new(0)); + let final_tools_requests_for_response = Arc::clone(&final_tools_requests); + Mock::given(method("POST")) + .and(path("/oauth/token")) + .and(body_string_contains("grant_type=refresh_token")) + .and(body_string_contains(format!( + "refresh_token={ROTATED_REFRESH_TOKEN}" + ))) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "access_token": FINAL_ACCESS_TOKEN, + "token_type": "Bearer", + "expires_in": 7200, + "refresh_token": FINAL_REFRESH_TOKEN, + }))) + .expect(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/mcp")) + .and(header( + "authorization", + format!("Bearer {EXPIRED_ACCESS_TOKEN}"), + )) + .respond_with( + ResponseTemplate::new(401).insert_header("www-authenticate", "Bearer realm=test"), + ) + .expect(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/mcp")) + .and(header( + "authorization", + format!("Bearer {REFRESHED_ACCESS_TOKEN}"), + )) + .respond_with(|request: &Request| { + let body: Value = request.body_json().expect("valid JSON-RPC request"); + match body.get("method").and_then(Value::as_str) { + Some("initialize") => initialize_response(&body, "oauth-401-recovery-test"), + Some("notifications/initialized") => ResponseTemplate::new(202), + Some("tools/list") => ResponseTemplate::new(401) + .insert_header("www-authenticate", "Bearer realm=test"), + method => ResponseTemplate::new(400) + .set_body_string(format!("unexpected JSON-RPC method: {method:?}")), + } + }) + .expect(3) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/mcp")) + .and(header( + "authorization", + format!("Bearer {FINAL_ACCESS_TOKEN}"), + )) + .respond_with(move |request: &Request| { + let body: Value = request.body_json().expect("valid JSON-RPC request"); + match body.get("method").and_then(Value::as_str) { + Some("initialize") => initialize_response(&body, "oauth-401-persisted-test"), + Some("notifications/initialized") => ResponseTemplate::new(202), + Some("tools/list") => { + match final_tools_requests_for_response.fetch_add(1, Ordering::SeqCst) { + 0 => ResponseTemplate::new(404), + 1 => ResponseTemplate::new(200).set_body_json(json!({ + "jsonrpc": "2.0", + "id": body.get("id").cloned().unwrap_or(Value::Null), + "result": { "tools": [] }, + })), + 2 => ResponseTemplate::new(401) + .insert_header("www-authenticate", "Bearer realm=test"), + _ => ResponseTemplate::new(200).set_body_json(json!({ + "jsonrpc": "2.0", + "id": body.get("id").cloned().unwrap_or(Value::Null), + "result": { "tools": [] }, + })), + } + } + method => ResponseTemplate::new(400) + .set_body_string(format!("unexpected JSON-RPC method: {method:?}")), + } + }) + .expect(5) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/mcp")) + .and(header( + "authorization", + format!("Bearer {POST_DEADLINE_ACCESS_TOKEN}"), + )) + .respond_with(|request: &Request| { + let body: Value = request.body_json().expect("valid JSON-RPC request"); + match body.get("method").and_then(Value::as_str) { + Some("initialize") => initialize_response(&body, "oauth-post-deadline-test"), + Some("notifications/initialized") => ResponseTemplate::new(202), + Some("tools/list") => ResponseTemplate::new(200).set_body_json(json!({ + "jsonrpc": "2.0", + "id": body.get("id").cloned().unwrap_or(Value::Null), + "result": { "tools": [] }, + })), + method => ResponseTemplate::new(400) + .set_body_string(format!("unexpected JSON-RPC method: {method:?}")), + } + }) + .expect(3) + .mount(&server) + .await; + + let codex_home = TempDir::new()?; + let server_url = format!("{}/mcp", server.uri()); + let status = Command::new(std::env::current_exe()?) + .args([ + "oauth_401_recovery_child", + "--exact", + "--ignored", + "--nocapture", + ]) + .env("CODEX_HOME", codex_home.path()) + .env(CHILD_SERVER_URL_ENV, server_url) + .status() + .await?; + assert!( + status.success(), + "OAuth 401 recovery child failed: {status}" + ); + server.verify().await; + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn concurrent_file_mode_startup_refreshes_once() -> anyhow::Result<()> { + let server = MockServer::start().await; + let refresh_server = GatedRefreshServer::start().await?; + let codex_home = TempDir::new()?; + let server_url = format!("{}/mcp", server.uri()); + 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": refresh_server.token_endpoint.clone(), + "scopes_supported": [""], + }))) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/mcp")) + .and(header( + "authorization", + format!("Bearer {REFRESHED_ACCESS_TOKEN}"), + )) + .respond_with(|request: &Request| { + let body: Value = request.body_json().expect("valid JSON-RPC request"); + match body.get("method").and_then(Value::as_str) { + Some("initialize") => ResponseTemplate::new(200).set_body_json(json!({ + "jsonrpc": "2.0", + "id": body.get("id").cloned().unwrap_or(Value::Null), + "result": { + "protocolVersion": body + .pointer("/params/protocolVersion") + .cloned() + .unwrap_or_else(|| json!("2025-06-18")), + "capabilities": {}, + "serverInfo": { + "name": "oauth-startup-test", + "version": "0.0.0-test", + }, + }, + })), + Some("notifications/initialized") => ResponseTemplate::new(202), + method => ResponseTemplate::new(400) + .set_body_string(format!("unexpected JSON-RPC method: {method:?}")), + } + }) + .expect(4) + .mount(&server) + .await; + + let seed_status = Command::new(std::env::current_exe()?) + .args(["oauth_concurrency_seed_child", "--exact", "--ignored"]) + .env("CODEX_HOME", codex_home.path()) + .env(CHILD_SERVER_URL_ENV, &server_url) + .status() + .await?; + assert!( + seed_status.success(), + "OAuth concurrency seed child failed: {seed_status}" + ); + + let first_ready_file = codex_home.path().join("oauth-client-first.ready"); + let first_release_file = codex_home.path().join("oauth-client-first.release"); + let second_ready_file = codex_home.path().join("oauth-client-second.ready"); + let second_release_file = codex_home.path().join("oauth-client-second.release"); + let contention_file = codex_home.path().join("oauth-client-second.contended"); + let mut first_child = oauth_concurrency_child_command( + codex_home.path(), + &server_url, + &first_ready_file, + &first_release_file, + )? + .spawn()?; + let mut second_command = oauth_concurrency_child_command( + codex_home.path(), + &server_url, + &second_ready_file, + &second_release_file, + )?; + second_command.env(CHILD_CONTENTION_FILE_ENV, &contention_file); + let mut second_child = second_command.spawn()?; + + wait_for_marker( + &first_ready_file, + "first OAuth concurrency child did not become ready", + ) + .await?; + wait_for_marker( + &second_ready_file, + "second OAuth concurrency child did not become ready", + ) + .await?; + + fs::write(&first_release_file, b"release")?; + refresh_server.wait_until_request_started().await?; + + // The first child is now inside the provider request while retaining the credential lock. + // Releasing the second child must make its first try_lock call observe WouldBlock. Keep the + // provider response gated until that exact branch emits the contention marker. + fs::write(&second_release_file, b"release")?; + let contention_result = wait_for_marker( + &contention_file, + "second OAuth concurrency child did not observe refresh-lock contention", + ) + .await; + refresh_server.release_responses(); + + let (first_status, second_status) = tokio::try_join!(first_child.wait(), second_child.wait())?; + assert!( + first_status.success(), + "first OAuth concurrency child failed: {first_status}" + ); + assert!( + second_status.success(), + "second OAuth concurrency child failed: {second_status}" + ); + + server.verify().await; + contention_result?; + assert_eq!(refresh_server.request_count(), 1); + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn operation_preflight_refresh_failure_blocks_rmcp_request() -> anyhow::Result<()> { + let server = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/.well-known/oauth-authorization-server/mcp")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "authorization_endpoint": format!("{}/oauth/authorize", server.uri()), + "token_endpoint": format!("{}/oauth/token", server.uri()), + "scopes_supported": [""], + }))) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/oauth/token")) + .and(body_string_contains("grant_type=refresh_token")) + .and(body_string_contains(format!( + "refresh_token={REFRESH_TOKEN}" + ))) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "access_token": REFRESHED_ACCESS_TOKEN, + "token_type": "Bearer", + "expires_in": 31, + "refresh_token": ROTATED_REFRESH_TOKEN, + }))) + .expect(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/oauth/token")) + .and(body_string_contains("grant_type=refresh_token")) + .and(body_string_contains(format!( + "refresh_token={ROTATED_REFRESH_TOKEN}" + ))) + .respond_with(ResponseTemplate::new(400).set_body_json(json!({ + "error": "invalid_grant", + "error_description": PREFLIGHT_REFRESH_ERROR, + }))) + .expect(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/mcp")) + .and(header( + "authorization", + format!("Bearer {REFRESHED_ACCESS_TOKEN}"), + )) + .respond_with(|request: &Request| { + let body: Value = request.body_json().expect("valid JSON-RPC request"); + match body.get("method").and_then(Value::as_str) { + Some("initialize") => ResponseTemplate::new(200).set_body_json(json!({ + "jsonrpc": "2.0", + "id": body.get("id").cloned().unwrap_or(Value::Null), + "result": { + "protocolVersion": body + .pointer("/params/protocolVersion") + .cloned() + .unwrap_or_else(|| json!("2025-06-18")), + "capabilities": {}, + "serverInfo": { + "name": "oauth-preflight-test", + "version": "0.0.0-test", + }, + }, + })), + Some("notifications/initialized") => ResponseTemplate::new(202), + method => ResponseTemplate::new(400) + .set_body_string(format!("unexpected JSON-RPC method: {method:?}")), + } + }) + .expect(2) + .mount(&server) + .await; + + let codex_home = TempDir::new()?; + let server_url = format!("{}/mcp", server.uri()); + let status = Command::new(std::env::current_exe()?) + .args([ + "operation_preflight_refresh_failure_child", + "--exact", + "--ignored", + "--nocapture", + ]) + .env("CODEX_HOME", codex_home.path()) + .env(CHILD_SERVER_URL_ENV, &server_url) + .status() + .await?; + assert!( + status.success(), + "OAuth preflight failure child failed: {status}" + ); + + server.verify().await; + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 1)] async fn reports_auth_status_for_persisted_credentials() -> anyhow::Result<()> { let codex_home = TempDir::new()?; @@ -235,7 +826,7 @@ async fn auth_status(server_url: &str) -> anyhow::Result { } #[tokio::test(flavor = "multi_thread", worker_threads = 1)] -#[ignore = "spawned by refreshes_expired_persisted_token_before_initialize"] +#[ignore = "spawned by startup_refresh_does_not_consume_handshake_timeout"] async fn oauth_startup_child() -> anyhow::Result<()> { let server_url = std::env::var(CHILD_SERVER_URL_ENV)?; @@ -278,6 +869,205 @@ async fn oauth_startup_child() -> anyhow::Result<()> { ) .await?; + initialize_client_with_timeout(&client, Duration::from_secs(1)).await?; + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +#[ignore = "spawned by recovers_initialization_and_operation_401_without_rmcp_refresh"] +async fn oauth_401_recovery_child() -> anyhow::Result<()> { + let server_url = std::env::var(CHILD_SERVER_URL_ENV)?; + save_no_expiry_file_mode_tokens(&server_url)?; + + let client = new_file_mode_oauth_client(&server_url).await?; + initialize_client_with_timeout(&client, Duration::from_secs(1)).await?; + let tools = client + .list_tools(/*params*/ None, Some(Duration::from_secs(5))) + .await?; + assert!(tools.tools.is_empty()); + + // The provider response arrives after the caller's operation deadline. Recovery must finish + // receiving and persisting that rotated token, then return the caller timeout without retrying. + let started = Instant::now(); + let error = client + .list_tools(/*params*/ None, Some(Duration::from_millis(500))) + .await + .expect_err("the operation deadline should prevent a retry after refresh"); + assert!( + error.to_string().contains("timed out"), + "unexpected timeout error: {error:#}" + ); + assert!(started.elapsed() >= Duration::from_millis(1_500)); + + let tools = client + .list_tools(/*params*/ None, Some(Duration::from_secs(5))) + .await?; + assert!(tools.tools.is_empty()); + client.shutdown().await; + + let reloaded_client = new_file_mode_oauth_client(&server_url).await?; + initialize_client(&reloaded_client).await?; + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +#[ignore = "spawned by operation_preflight_refresh_failure_blocks_rmcp_request"] +async fn operation_preflight_refresh_failure_child() -> anyhow::Result<()> { + let server_url = std::env::var(CHILD_SERVER_URL_ENV)?; + save_expired_file_mode_tokens(&server_url)?; + + let client = RmcpClient::new_streamable_http_client( + SERVER_NAME, + &server_url, + /*bearer_token*/ None, + /*http_headers*/ None, + /*env_http_headers*/ None, + OAuthCredentialsStoreMode::File, + AuthKeyringBackendKind::default(), + Environment::default_for_tests().get_http_client(), + /*auth_provider*/ None, + ) + .await?; + + initialize_client(&client).await?; + + tokio::time::sleep(Duration::from_millis(2_200)).await; + let error = client + .list_tools(/*params*/ None, Some(Duration::from_secs(5))) + .await + .expect_err("preflight refresh failure should abort the operation"); + let message = format!("{error:#}"); + assert!( + message.contains("failed to refresh OAuth tokens for server"), + "unexpected error: {message}" + ); + assert!( + message.contains(PREFLIGHT_REFRESH_ERROR), + "unexpected error: {message}" + ); + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +#[ignore = "spawned by concurrent_file_mode_startup_refreshes_once"] +async fn oauth_concurrency_seed_child() -> anyhow::Result<()> { + let server_url = std::env::var(CHILD_SERVER_URL_ENV)?; + save_expired_file_mode_tokens(&server_url)?; + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +#[ignore = "spawned by concurrent_file_mode_startup_refreshes_once"] +async fn oauth_concurrency_client_child() -> anyhow::Result<()> { + let server_url = std::env::var(CHILD_SERVER_URL_ENV)?; + let ready_file = PathBuf::from(std::env::var(CHILD_READY_FILE_ENV)?); + let release_file = PathBuf::from(std::env::var(CHILD_RELEASE_FILE_ENV)?); + if let Ok(marker_file) = std::env::var(CHILD_CONTENTION_FILE_ENV) { + tracing::subscriber::set_global_default(LockContentionMarkerSubscriber { + marker_file: PathBuf::from(marker_file), + }) + .map_err(|error| anyhow::anyhow!("failed to install contention subscriber: {error}"))?; + } + let client = RmcpClient::new_streamable_http_client( + SERVER_NAME, + &server_url, + /*bearer_token*/ None, + /*http_headers*/ None, + /*env_http_headers*/ None, + OAuthCredentialsStoreMode::File, + AuthKeyringBackendKind::default(), + Environment::default_for_tests().get_http_client(), + /*auth_provider*/ None, + ) + .await?; + + // Both processes must construct their OAuth client from the same expired snapshot before the + // parent releases either one. The parent then gates them separately to force lock contention. + fs::write(ready_file, b"ready")?; + while !release_file.exists() { + tokio::time::sleep(Duration::from_millis(/*millis*/ 10)).await; + } initialize_client(&client).await?; Ok(()) } + +fn save_expired_file_mode_tokens(server_url: &str) -> anyhow::Result<()> { + let mut response = OAuthTokenResponse::new( + AccessToken::new(EXPIRED_ACCESS_TOKEN.to_string()), + BasicTokenType::Bearer, + VendorExtraTokenFields::default(), + ); + response.set_refresh_token(Some(RefreshToken::new(REFRESH_TOKEN.to_string()))); + response.set_expires_in(Some(&Duration::from_secs(7200))); + let tokens = StoredOAuthTokens { + server_name: SERVER_NAME.to_string(), + url: server_url.to_string(), + client_id: "test-client-id".to_string(), + token_response: WrappedOAuthTokenResponse(response), + expires_at: Some(0), + }; + save_oauth_tokens( + SERVER_NAME, + &tokens, + OAuthCredentialsStoreMode::File, + AuthKeyringBackendKind::default(), + )?; + Ok(()) +} +fn save_no_expiry_file_mode_tokens(server_url: &str) -> anyhow::Result<()> { + let mut response = OAuthTokenResponse::new( + AccessToken::new(EXPIRED_ACCESS_TOKEN.to_string()), + BasicTokenType::Bearer, + VendorExtraTokenFields::default(), + ); + response.set_refresh_token(Some(RefreshToken::new(REFRESH_TOKEN.to_string()))); + let tokens = StoredOAuthTokens { + server_name: SERVER_NAME.to_string(), + url: server_url.to_string(), + client_id: "test-client-id".to_string(), + token_response: WrappedOAuthTokenResponse(response), + expires_at: None, + }; + save_oauth_tokens( + SERVER_NAME, + &tokens, + OAuthCredentialsStoreMode::File, + AuthKeyringBackendKind::default(), + )?; + Ok(()) +} + +async fn new_file_mode_oauth_client(server_url: &str) -> anyhow::Result { + RmcpClient::new_streamable_http_client( + SERVER_NAME, + server_url, + /*bearer_token*/ None, + /*http_headers*/ None, + /*env_http_headers*/ None, + OAuthCredentialsStoreMode::File, + AuthKeyringBackendKind::default(), + Environment::default_for_tests().get_http_client(), + /*auth_provider*/ None, + ) + .await +} + +fn initialize_response(body: &Value, name: &str) -> ResponseTemplate { + ResponseTemplate::new(200) + .insert_header("mcp-session-id", format!("{name}-session")) + .set_body_json(json!({ + "jsonrpc": "2.0", + "id": body.get("id").cloned().unwrap_or(Value::Null), + "result": { + "protocolVersion": body + .pointer("/params/protocolVersion") + .cloned() + .unwrap_or_else(|| json!("2025-06-18")), + "capabilities": {}, + "serverInfo": { + "name": name, + "version": "0.0.0-test", + }, + }, + })) +} diff --git a/codex-rs/rmcp-client/tests/streamable_http_test_support.rs b/codex-rs/rmcp-client/tests/streamable_http_test_support.rs index 87a52f55dff6..532861de7ddf 100644 --- a/codex-rs/rmcp-client/tests/streamable_http_test_support.rs +++ b/codex-rs/rmcp-client/tests/streamable_http_test_support.rs @@ -106,10 +106,17 @@ pub(crate) async fn create_client_with_http_client( } pub(crate) async fn initialize_client(client: &RmcpClient) -> anyhow::Result<()> { + initialize_client_with_timeout(client, Duration::from_secs(5)).await +} + +pub(crate) async fn initialize_client_with_timeout( + client: &RmcpClient, + timeout: Duration, +) -> anyhow::Result<()> { client .initialize( init_params(), - Some(Duration::from_secs(5)), + Some(timeout), Box::new(|_, _| { async { Ok(ElicitationResponse {