diff --git a/codex-rs/rmcp-client/src/oauth.rs b/codex-rs/rmcp-client/src/oauth.rs index 75bebddea7..6c2929e544 100644 --- a/codex-rs/rmcp-client/src/oauth.rs +++ b/codex-rs/rmcp-client/src/oauth.rs @@ -41,6 +41,7 @@ use std::io::ErrorKind; use std::path::Path; use std::path::PathBuf; use std::sync::Arc; +use std::sync::OnceLock; use std::time::Duration; use std::time::SystemTime; use std::time::UNIX_EPOCH; @@ -399,12 +400,22 @@ struct OAuthPersistorInner { credential_store: InMemoryCredentialStore, store_mode: OAuthCredentialsStoreMode, last_credentials: Mutex>, + refresh_attempt: Arc>, refresh_persistence_failure: Mutex>, } struct RefreshPersistenceFailure { message: String, - _refresh_locks: Vec, +} + +static FAILED_REFRESH_LOCKS: OnceLock>> = OnceLock::new(); + +fn retain_failed_refresh_locks(refresh_locks: Vec) { + FAILED_REFRESH_LOCKS + .get_or_init(std::sync::Mutex::default) + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .extend(refresh_locks); } impl OAuthPersistor { @@ -424,6 +435,7 @@ impl OAuthPersistor { credential_store, store_mode, last_credentials: Mutex::new(initial_credentials), + refresh_attempt: Arc::new(Mutex::new(())), refresh_persistence_failure: Mutex::new(None), }), } @@ -505,6 +517,7 @@ impl OAuthPersistor { reason = "AuthorizationManager async access must be serialized through its mutex" )] pub(crate) async fn refresh_if_needed(&self) -> Result<()> { + let refresh_attempt_guard = Arc::clone(&self.inner.refresh_attempt).lock_owned().await; let previous_failure = self .inner .refresh_persistence_failure @@ -578,6 +591,7 @@ impl OAuthPersistor { let persistor = self.clone(); tokio::spawn(async move { + let _refresh_attempt_guard = refresh_attempt_guard; let manager = persistor.inner.authorization_manager.clone(); let guard = manager.lock().await; match guard.refresh_token().await { @@ -604,10 +618,10 @@ impl OAuthPersistor { persistor.inner.server_name ); warn!("{message}"); + retain_failed_refresh_locks(refresh_locks); *persistor.inner.refresh_persistence_failure.lock().await = Some(RefreshPersistenceFailure { message: message.clone(), - _refresh_locks: refresh_locks, }); return Err(anyhow!(message)); } diff --git a/codex-rs/rmcp-client/src/rmcp_client.rs b/codex-rs/rmcp-client/src/rmcp_client.rs index f846406b15..62b7a28f31 100644 --- a/codex-rs/rmcp-client/src/rmcp_client.rs +++ b/codex-rs/rmcp-client/src/rmcp_client.rs @@ -966,18 +966,8 @@ impl RmcpClient { } fn persist_oauth_tokens_in_background(runtime: OAuthPersistor, error_context: &'static str) { - Self::spawn_oauth_persistence( - async move { runtime.persist_if_needed().await }, - error_context, - ); - } - - fn spawn_oauth_persistence( - persistence: impl Future> + Send + 'static, - error_context: &'static str, - ) { tokio::spawn(async move { - if let Err(error) = persistence.await { + if let Err(error) = runtime.persist_if_needed().await { warn!("{error_context}: {error}"); } }); @@ -1129,13 +1119,83 @@ async fn create_oauth_transport_and_runtime( #[cfg(test)] mod tests { + use std::ffi::OsString; + use std::fs::OpenOptions; + use std::path::Path; + use std::path::PathBuf; + use std::sync::Arc; use std::time::Duration; + use std::time::SystemTime; + use std::time::UNIX_EPOCH; + use anyhow::Context; + use codex_exec_server::Environment; + use oauth2::AccessToken; + use oauth2::RefreshToken; + use oauth2::basic::BasicTokenType; use pretty_assertions::assert_eq; - use tokio::sync::oneshot; + use rmcp::model::ClientCapabilities; + use rmcp::model::Implementation; + use rmcp::model::ProtocolVersion; + use rmcp::transport::auth::CredentialStore; + use rmcp::transport::auth::OAuthTokenResponse; + use rmcp::transport::auth::StoredCredentials; + use rmcp::transport::auth::VendorExtraTokenFields; + use serde_json::Value; + use serde_json::json; + use sha2::Digest; + use sha2::Sha256; + use tempfile::TempDir; use tokio::time; + use wiremock::Mock; + use wiremock::MockServer; + use wiremock::Request; + use wiremock::ResponseTemplate; + use wiremock::matchers::method; + use wiremock::matchers::path; use super::*; + use crate::oauth::WrappedOAuthTokenResponse; + use crate::oauth::load_oauth_tokens; + use crate::oauth::save_oauth_tokens; + + const OPERATION_TEST_SERVER_NAME: &str = "test-completed-operation-persistence"; + const OPERATION_TEST_CLIENT_ID: &str = "test-operation-client-id"; + const OPERATION_TEST_OLD_ACCESS_TOKEN: &str = "old-operation-access-token"; + const OPERATION_TEST_OLD_REFRESH_TOKEN: &str = "old-operation-refresh-token"; + const OPERATION_TEST_NEW_ACCESS_TOKEN: &str = "new-operation-access-token"; + const OPERATION_TEST_NEW_REFRESH_TOKEN: &str = "new-operation-refresh-token"; + + struct CodexHomeGuard { + original: Option, + } + + impl CodexHomeGuard { + fn set(path: &Path) -> Self { + let original = std::env::var_os("CODEX_HOME"); + // SAFETY: this test is serialized and restores the process environment on drop. + unsafe { + std::env::set_var("CODEX_HOME", path); + } + Self { original } + } + } + + impl Drop for CodexHomeGuard { + fn drop(&mut self) { + if let Some(original) = &self.original { + // SAFETY: this test is serialized and restores the prior value. + unsafe { + std::env::set_var("CODEX_HOME", original); + } + } else { + // SAFETY: this test is serialized and removes only the value it set. + unsafe { + std::env::remove_var("CODEX_HOME"); + } + } + } + } #[tokio::test] async fn active_time_timeout_pauses_while_elicitation_is_pending() { @@ -1156,36 +1216,273 @@ mod tests { assert_eq!(Ok("done"), result); } - #[tokio::test] - async fn completed_operation_is_not_delayed_by_oauth_persistence() { - let (started_tx, started_rx) = oneshot::channel(); - let (release_tx, release_rx) = oneshot::channel(); - let (finished_tx, mut finished_rx) = oneshot::channel(); + #[cfg(unix)] + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + #[serial_test::serial] + async fn completed_operation_returns_while_oauth_persistence_is_contended() -> 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": [""], + }))) + .expect(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/mcp")) + .respond_with(|request: &Request| { + let Ok(body) = request.body_json::() else { + return ResponseTemplate::new(400).set_body_string("invalid 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": "2025-06-18", + "capabilities": {}, + "serverInfo": { + "name": "completed-operation-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; - RmcpClient::spawn_oauth_persistence( - async move { - started_tx.send(()).expect("signal persistence start"); - release_rx.await.expect("release persistence"); - finished_tx.send(()).expect("signal persistence finish"); - Ok(()) - }, - "test OAuth persistence failure", + let codex_home = TempDir::new()?; + let _codex_home_guard = CodexHomeGuard::set(codex_home.path()); + let server_url = format!("{}/mcp", server.uri()); + let old_response = operation_test_token_response( + OPERATION_TEST_OLD_ACCESS_TOKEN, + OPERATION_TEST_OLD_REFRESH_TOKEN, ); - let completed_result = "completed tool response"; + let initial_tokens = StoredOAuthTokens { + server_name: OPERATION_TEST_SERVER_NAME.to_string(), + url: server_url.clone(), + client_id: OPERATION_TEST_CLIENT_ID.to_string(), + token_response: WrappedOAuthTokenResponse(old_response.clone()), + expires_at: Some(unix_millis().saturating_add(3_600_000)), + }; + save_oauth_tokens( + OPERATION_TEST_SERVER_NAME, + &initial_tokens, + OAuthCredentialsStoreMode::File, + )?; - time::timeout(Duration::from_secs(1), started_rx) - .await - .expect("persistence task did not start") - .expect("persistence start sender dropped"); - assert_eq!(completed_result, "completed tool response"); - assert!(matches!( - finished_rx.try_recv(), - Err(oneshot::error::TryRecvError::Empty) - )); - release_tx.send(()).expect("release persistence task"); - time::timeout(Duration::from_secs(1), finished_rx) - .await - .expect("persistence task did not finish") - .expect("persistence finish sender dropped"); + let client = RmcpClient::new_streamable_http_client( + OPERATION_TEST_SERVER_NAME, + &server_url, + Some("test-bearer".to_string()), + /*http_headers*/ None, + /*env_http_headers*/ None, + OAuthCredentialsStoreMode::File, + Environment::default_for_tests().get_http_client(), + /*auth_provider*/ None, + ) + .await?; + client + .initialize( + InitializeRequestParams::new( + ClientCapabilities::default(), + Implementation::new("codex-test", "0.0.0-test"), + ) + .with_protocol_version(ProtocolVersion::V_2025_06_18), + Some(Duration::from_secs(2)), + Box::new(|_, _| { + async { + Ok(ElicitationResponse { + action: ElicitationAction::Accept, + content: Some(json!({})), + meta: None, + }) + } + .boxed() + }), + ) + .await?; + let service = client.service().await?; + + let credential_store = InMemoryCredentialStore::new(); + let mut oauth_state = + OAuthState::new(server_url.clone(), Some(reqwest::Client::new())).await?; + match &mut oauth_state { + OAuthState::Unauthorized(manager) => { + manager.set_credential_store(credential_store.clone()); + } + _ => { + return Err(anyhow!( + "unexpected OAuth state during operation test setup" + )); + } + } + oauth_state + .set_credentials(OPERATION_TEST_CLIENT_ID, old_response) + .await?; + let authorization_manager = match oauth_state { + OAuthState::Authorized(manager) => manager, + _ => return Err(anyhow!("OAuth operation test setup did not authorize")), + }; + let oauth_persistor = OAuthPersistor::new( + OPERATION_TEST_SERVER_NAME.to_string(), + server_url.clone(), + Arc::new(Mutex::new(authorization_manager)), + credential_store.clone(), + OAuthCredentialsStoreMode::File, + Some(initial_tokens.clone()), + ); + + let fallback_lock = OpenOptions::new() + .read(true) + .write(true) + .create(true) + .truncate(false) + .open(operation_test_fallback_lock_path(codex_home.path()))?; + fallback_lock.lock_shared()?; + + let new_response = operation_test_token_response( + OPERATION_TEST_NEW_ACCESS_TOKEN, + OPERATION_TEST_NEW_REFRESH_TOKEN, + ); + let operation_store = credential_store.clone(); + let operation_response = new_response.clone(); + let operation = move |_service| { + let operation_store = operation_store.clone(); + let operation_response = operation_response.clone(); + async move { + if let Err(error) = operation_store + .save(StoredCredentials::new( + OPERATION_TEST_CLIENT_ID.to_string(), + Some(operation_response), + Vec::new(), + Some(unix_seconds()), + )) + .await + { + panic!("failed to simulate completed OAuth side effect: {error}"); + } + Ok::<_, rmcp::service::ServiceError>("side effect completed".to_string()) + } + .boxed() + }; + + let result = RmcpClient::run_service_operation_once( + service, + Some(oauth_persistor), + "tools/call", + Some(Duration::from_millis(100)), + ElicitationPauseState::new(), + &operation, + ) + .await?; + assert_eq!(result, "side effect completed"); + time::sleep(Duration::from_millis(200)).await; + let persisted_while_locked = load_oauth_tokens( + OPERATION_TEST_SERVER_NAME, + &server_url, + OAuthCredentialsStoreMode::File, + )? + .context("missing persisted OAuth credentials while fallback lock is held")?; + assert_eq!( + ( + persisted_while_locked + .token_response + .0 + .access_token() + .secret() + .as_str(), + persisted_while_locked + .token_response + .0 + .refresh_token() + .map(|token| token.secret().as_str()), + ), + ( + OPERATION_TEST_OLD_ACCESS_TOKEN, + Some(OPERATION_TEST_OLD_REFRESH_TOKEN), + ) + ); + + drop(fallback_lock); + let persisted = time::timeout(Duration::from_secs(2), async { + loop { + let loaded = load_oauth_tokens( + OPERATION_TEST_SERVER_NAME, + &server_url, + OAuthCredentialsStoreMode::File, + )?; + if loaded.as_ref().is_some_and(|tokens| { + tokens.token_response.0.access_token().secret() + == OPERATION_TEST_NEW_ACCESS_TOKEN + }) { + return Ok::<_, anyhow::Error>(loaded); + } + time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .context("timed out waiting for post-operation OAuth persistence")??; + assert_eq!( + persisted.as_ref().map(|tokens| tokens + .token_response + .0 + .access_token() + .secret() + .as_str()), + Some(OPERATION_TEST_NEW_ACCESS_TOKEN) + ); + + server.verify().await; + Ok(()) + } + + fn operation_test_token_response( + access_token: &str, + refresh_token: &str, + ) -> OAuthTokenResponse { + let mut response = OAuthTokenResponse::new( + AccessToken::new(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))); + response + } + + #[cfg(unix)] + fn operation_test_fallback_lock_path(codex_home: &Path) -> PathBuf { + let mut hasher = Sha256::new(); + hasher.update(codex_home.as_os_str().to_string_lossy().as_bytes()); + let digest = format!("{:x}", hasher.finalize()); + // SAFETY: getuid has no preconditions and returns the real UID. + let user_namespace = format!("uid-{}", unsafe { libc::getuid() }); + PathBuf::from("/tmp").join(format!( + "codex-mcp-oauth-fallback-{user_namespace}-{}.lock", + &digest[..16] + )) + } + + fn unix_millis() -> u64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + .as_millis() as u64 + } + + fn unix_seconds() -> u64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + .as_secs() } } diff --git a/codex-rs/rmcp-client/tests/streamable_http_oauth_refresh_race.rs b/codex-rs/rmcp-client/tests/streamable_http_oauth_refresh_race.rs index 69a7ce5619..b5725908c5 100644 --- a/codex-rs/rmcp-client/tests/streamable_http_oauth_refresh_race.rs +++ b/codex-rs/rmcp-client/tests/streamable_http_oauth_refresh_race.rs @@ -63,6 +63,7 @@ const SERVER_URL_ENV: &str = "MCP_TEST_OAUTH_RACE_SERVER_URL"; const READY_PATH_ENV: &str = "MCP_TEST_OAUTH_RACE_READY_PATH"; const GO_PATH_ENV: &str = "MCP_TEST_OAUTH_RACE_GO_PATH"; const RESULT_PATH_ENV: &str = "MCP_TEST_OAUTH_RACE_RESULT_PATH"; +const ATTEMPT_PATH_ENV: &str = "MCP_TEST_OAUTH_RACE_ATTEMPT_PATH"; const ACCESS_TOKEN_ENV: &str = "MCP_TEST_OAUTH_RACE_ACCESS_TOKEN"; const REFRESH_TOKEN_ENV: &str = "MCP_TEST_OAUTH_RACE_REFRESH_TOKEN"; const EXPIRES_AT_ENV: &str = "MCP_TEST_OAUTH_RACE_EXPIRES_AT"; @@ -510,6 +511,101 @@ async fn provider_refresh_persists_after_startup_timeout_before_next_refresh() - Ok(()) } +#[cfg(unix)] +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn canceled_startup_persistence_failure_retains_lock_until_process_exit() -> anyhow::Result<()> +{ + use std::os::unix::fs::PermissionsExt; + + let server = MockServer::start().await; + mount_oauth_metadata(&server, /*expected_requests*/ 2).await; + let refresh_attempts = Arc::new(AtomicUsize::new(0)); + let responder_attempts = Arc::clone(&refresh_attempts); + Mock::given(method("POST")) + .and(path("/oauth/token")) + .respond_with(move |_request: &Request| { + match responder_attempts.fetch_add(1, Ordering::SeqCst) { + 0 => ResponseTemplate::new(200) + .set_delay(Duration::from_secs(6)) + .set_body_json(json!({ + "access_token": NEW_ACCESS_TOKEN, + "token_type": "Bearer", + "expires_in": 7200, + "refresh_token": NEW_REFRESH_TOKEN, + })), + _ => ResponseTemplate::new(400).set_body_json(json!({ + "error": "invalid_grant", + "error_description": "refresh token already used", + })), + } + }) + .expect(2) + .mount(&server) + .await; + mount_mcp_server(&server, /*expected_requests*/ 0).await; + + let codex_home = TempDir::new()?; + let control_dir = TempDir::new()?; + let server_url = format!("{}/mcp", server.uri()); + seed_tokens(&codex_home, &server_url, /*expires_at*/ 0).await?; + let credentials_path = codex_home.path().join(".credentials.json"); + fs::set_permissions(&credentials_path, fs::Permissions::from_mode(0o400))?; + + let first_ready_path = control_dir.path().join("client-a.ready"); + let first_go_path = control_dir.path().join("client-a.go"); + let first_result_path = control_dir.path().join("client-a.result"); + let first_release_path = control_dir.path().join("client-a.release"); + let first_child = spawn_client_child_held_open( + &codex_home, + &server_url, + &first_ready_path, + &first_go_path, + &first_result_path, + &first_release_path, + )?; + wait_for_paths(std::slice::from_ref(&first_ready_path)).await?; + fs::write(&first_go_path, "go")?; + wait_for_paths(std::slice::from_ref(&first_result_path)).await?; + assert_eq!( + fs::read_to_string(&first_result_path)?, + "error:timed out handshaking with MCP server after 5s" + ); + sleep(Duration::from_secs(2)).await; + assert_eq!(refresh_attempts.load(Ordering::SeqCst), 1); + + let second_ready_path = control_dir.path().join("client-b.ready"); + let second_go_path = control_dir.path().join("client-b.go"); + let second_attempt_path = control_dir.path().join("client-b.attempt"); + let second_result_path = control_dir.path().join("client-b.result"); + let second_child = spawn_client_child_with_attempt( + &codex_home, + &server_url, + &second_ready_path, + &second_go_path, + &second_attempt_path, + &second_result_path, + )?; + wait_for_paths(std::slice::from_ref(&second_ready_path)).await?; + fs::write(&second_go_path, "go")?; + wait_for_paths(std::slice::from_ref(&second_attempt_path)).await?; + sleep(Duration::from_millis(250)).await; + assert_eq!(refresh_attempts.load(Ordering::SeqCst), 1); + assert!(!second_result_path.exists()); + + fs::set_permissions(&credentials_path, fs::Permissions::from_mode(0o600))?; + fs::write(first_release_path, "release")?; + wait_for_children(vec![first_child]).await?; + wait_for_children(vec![second_child]).await?; + assert_eq!( + fs::read_to_string(&second_result_path)?, + "error:Auth required: OAuth refresh token was rejected; reauthentication required" + ); + assert_eq!(refresh_attempts.load(Ordering::SeqCst), 2); + + server.verify().await; + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 1)] async fn fallback_read_lock_wait_uses_operation_timeout_budget() -> anyhow::Result<()> { let server = MockServer::start().await; @@ -831,7 +927,7 @@ async fn rotated_token_storage_failure_blocks_another_process_refresh() -> anyho })), } }) - .expect(2) + .expect(3) .mount(&server) .await; mount_mcp_server(&server, /*expected_requests*/ 2).await; @@ -871,27 +967,107 @@ async fn rotated_token_storage_failure_blocks_another_process_refresh() -> anyho let second_ready_path = control_dir.path().join("client-b.ready"); let second_go_path = control_dir.path().join("client-b.go"); + let second_attempt_path = control_dir.path().join("client-b.attempt"); let second_result_path = control_dir.path().join("client-b.result"); - let mut second_child = spawn_client_child( + let second_child = spawn_client_child_with_attempt( &codex_home, &server_url, &second_ready_path, &second_go_path, + &second_attempt_path, &second_result_path, )?; wait_for_paths(std::slice::from_ref(&second_ready_path)).await?; fs::write(&second_go_path, "go")?; + wait_for_paths(std::slice::from_ref(&second_attempt_path)).await?; sleep(Duration::from_millis(250)).await; assert_eq!(refresh_attempts.load(Ordering::SeqCst), 2); assert!(!second_result_path.exists()); - second_child.kill().await?; - timeout(CHILD_WAIT_TIMEOUT, second_child.wait()) - .await - .context("timed out stopping blocked OAuth client child")??; fs::set_permissions(&credentials_path, fs::Permissions::from_mode(0o600))?; fs::write(first_release_path, "release")?; wait_for_children(vec![first_child]).await?; + wait_for_children(vec![second_child]).await?; + assert_eq!( + fs::read_to_string(&second_result_path)?, + "error:Auth required: OAuth refresh token was rejected; reauthentication required" + ); + assert_eq!(refresh_attempts.load(Ordering::SeqCst), 3); + + server.verify().await; + Ok(()) +} + +#[cfg(unix)] +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn concurrent_waiter_observes_recorded_persistence_failure() -> anyhow::Result<()> { + use std::os::unix::fs::PermissionsExt; + + let server = MockServer::start().await; + mount_oauth_metadata(&server, /*expected_requests*/ 1).await; + let refresh_attempts = Arc::new(AtomicUsize::new(0)); + let responder_attempts = Arc::clone(&refresh_attempts); + Mock::given(method("POST")) + .and(path("/oauth/token")) + .respond_with(move |_request: &Request| { + match responder_attempts.fetch_add(1, Ordering::SeqCst) { + 0 => ResponseTemplate::new(200).set_body_json(json!({ + "access_token": NEW_ACCESS_TOKEN, + "token_type": "Bearer", + "expires_in": 31, + "refresh_token": NEW_REFRESH_TOKEN, + })), + _ => ResponseTemplate::new(200).set_body_json(json!({ + "access_token": FINAL_ACCESS_TOKEN, + "token_type": "Bearer", + "expires_in": 7200, + "refresh_token": FINAL_REFRESH_TOKEN, + })), + } + }) + .expect(2) + .mount(&server) + .await; + mount_mcp_server(&server, /*expected_requests*/ 2).await; + + let codex_home = TempDir::new()?; + let control_dir = TempDir::new()?; + let server_url = format!("{}/mcp", server.uri()); + seed_tokens(&codex_home, &server_url, /*expires_at*/ 0).await?; + + let ready_path = control_dir.path().join("client.ready"); + let go_path = control_dir.path().join("client.go"); + let result_path = control_dir.path().join("client.result"); + let child = spawn_concurrent_waiter_child( + &codex_home, + &server_url, + &ready_path, + &go_path, + &result_path, + )?; + wait_for_paths(std::slice::from_ref(&ready_path)).await?; + + let credentials_path = codex_home.path().join(".credentials.json"); + fs::set_permissions(&credentials_path, fs::Permissions::from_mode(0o400))?; + sleep(Duration::from_secs(2)).await; + fs::write(&go_path, "go")?; + wait_for_children(vec![child]).await?; + fs::set_permissions(&credentials_path, fs::Permissions::from_mode(0o600))?; + + let outcomes = fs::read_to_string(result_path)? + .lines() + .map(str::to_string) + .collect::>(); + assert_eq!(outcomes.len(), 2); + assert_eq!(outcomes[0], outcomes[1]); + assert!( + outcomes[0].starts_with(&format!( + "error:OAuth refresh succeeded but failed to persist rotated credentials for server {SERVER_NAME}:" + )), + "unexpected concurrent waiter outcome: {}", + outcomes[0] + ); + assert_eq!(refresh_attempts.load(Ordering::SeqCst), 2); server.verify().await; Ok(()) @@ -1239,6 +1415,9 @@ async fn oauth_refresh_race_client_child() -> anyhow::Result<()> { fs::write(ready_path, "ready")?; wait_for_paths(std::slice::from_ref(&go_path)).await?; + if let Ok(attempt_path) = std::env::var(ATTEMPT_PATH_ENV) { + fs::write(attempt_path, "attempt")?; + } let outcome = match initialize_client(&client).await { Ok(()) => "success".to_string(), Err(error) => format!("error:{error:#}"), @@ -1283,6 +1462,51 @@ async fn oauth_refresh_operation_timeout_child() -> anyhow::Result<()> { Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +#[ignore = "spawned by the OAuth refresh race integration tests"] +async fn oauth_refresh_concurrent_waiter_child() -> anyhow::Result<()> { + let server_url = std::env::var(SERVER_URL_ENV)?; + let ready_path = PathBuf::from(std::env::var(READY_PATH_ENV)?); + let go_path = PathBuf::from(std::env::var(GO_PATH_ENV)?); + let result_path = PathBuf::from(std::env::var(RESULT_PATH_ENV)?); + let client = RmcpClient::new_streamable_http_client( + SERVER_NAME, + &server_url, + /*bearer_token*/ None, + /*http_headers*/ None, + /*env_http_headers*/ None, + OAuthCredentialsStoreMode::File, + Environment::default_for_tests().get_http_client(), + /*auth_provider*/ None, + ) + .await?; + initialize_client(&client).await?; + + fs::write(ready_path, "ready")?; + wait_for_paths(std::slice::from_ref(&go_path)).await?; + let outcomes = timeout(Duration::from_secs(3), async { + tokio::join!( + client.list_tools(None, /*timeout*/ None), + client.list_tools(None, /*timeout*/ None), + ) + }) + .await + .map(|(first, second)| { + [first, second].map(|outcome| match outcome { + Ok(_) => "success".to_string(), + Err(error) => format!("error:{error:#}"), + }) + }) + .unwrap_or_else(|_| { + [ + "error:concurrent refresh waiters timed out".to_string(), + "error:concurrent refresh waiters timed out".to_string(), + ] + }); + fs::write(result_path, outcomes.join("\n"))?; + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 1)] #[ignore = "spawned by the OAuth refresh race integration tests"] async fn oauth_refresh_read_only_child() -> anyhow::Result<()> { @@ -1441,6 +1665,32 @@ fn spawn_client_child_held_open( Ok(command.spawn()?) } +fn spawn_client_child_with_attempt( + codex_home: &TempDir, + server_url: &str, + ready_path: &Path, + go_path: &Path, + attempt_path: &Path, + result_path: &Path, +) -> anyhow::Result { + let mut command = Command::new(std::env::current_exe()?); + command + .args([ + "oauth_refresh_race_client_child", + "--exact", + "--ignored", + "--nocapture", + ]) + .env("CODEX_HOME", codex_home.path()) + .env(SERVER_URL_ENV, server_url) + .env(READY_PATH_ENV, ready_path) + .env(GO_PATH_ENV, go_path) + .env(ATTEMPT_PATH_ENV, attempt_path) + .env(RESULT_PATH_ENV, result_path) + .kill_on_drop(true); + Ok(command.spawn()?) +} + fn spawn_operation_timeout_child( codex_home: &TempDir, server_url: &str, @@ -1491,6 +1741,30 @@ fn spawn_operation_timeout_child_held_open( Ok(command.spawn()?) } +fn spawn_concurrent_waiter_child( + codex_home: &TempDir, + server_url: &str, + ready_path: &Path, + go_path: &Path, + result_path: &Path, +) -> anyhow::Result { + let mut command = Command::new(std::env::current_exe()?); + command + .args([ + "oauth_refresh_concurrent_waiter_child", + "--exact", + "--ignored", + "--nocapture", + ]) + .env("CODEX_HOME", codex_home.path()) + .env(SERVER_URL_ENV, server_url) + .env(READY_PATH_ENV, ready_path) + .env(GO_PATH_ENV, go_path) + .env(RESULT_PATH_ENV, result_path) + .kill_on_drop(true); + Ok(command.spawn()?) +} + fn spawn_read_only_refresh_child( codex_home: &TempDir, server_url: &str,