diff --git a/codex-rs/rmcp-client/src/oauth.rs b/codex-rs/rmcp-client/src/oauth.rs index 1fbf7aaf0c..75bebddea7 100644 --- a/codex-rs/rmcp-client/src/oauth.rs +++ b/codex-rs/rmcp-client/src/oauth.rs @@ -399,6 +399,12 @@ struct OAuthPersistorInner { credential_store: InMemoryCredentialStore, store_mode: OAuthCredentialsStoreMode, last_credentials: Mutex>, + refresh_persistence_failure: Mutex>, +} + +struct RefreshPersistenceFailure { + message: String, + _refresh_locks: Vec, } impl OAuthPersistor { @@ -418,6 +424,7 @@ impl OAuthPersistor { credential_store, store_mode, last_credentials: Mutex::new(initial_credentials), + refresh_persistence_failure: Mutex::new(None), }), } } @@ -425,20 +432,10 @@ impl OAuthPersistor { /// Persists the latest stored credentials if they have changed. /// Deletes the credentials if they are no longer present. pub(crate) async fn persist_if_needed(&self) -> Result<()> { - self.persist_if_needed_with_refresh_locks(Vec::new()).await - } - - async fn persist_if_needed_with_refresh_locks( - &self, - refresh_locks: Vec, - ) -> Result<()> { let persistor = self.clone(); - tokio::spawn(async move { - let _refresh_locks = refresh_locks; - persistor.persist_if_needed_inner().await - }) - .await - .context("OAuth token persistence task failed")? + tokio::spawn(async move { persistor.persist_if_needed_inner().await }) + .await + .context("OAuth token persistence task failed")? } #[expect( @@ -508,6 +505,17 @@ impl OAuthPersistor { reason = "AuthorizationManager async access must be serialized through its mutex" )] pub(crate) async fn refresh_if_needed(&self) -> Result<()> { + let previous_failure = self + .inner + .refresh_persistence_failure + .lock() + .await + .as_ref() + .map(|failure| failure.message.clone()); + if let Some(message) = previous_failure { + return Err(anyhow!(message)); + } + let initial_expires_at = { let guard = self.inner.last_credentials.lock().await; guard.as_ref().and_then(|tokens| tokens.expires_at) @@ -568,8 +576,9 @@ impl OAuthPersistor { return Ok(()); } - { - let manager = self.inner.authorization_manager.clone(); + let persistor = self.clone(); + tokio::spawn(async move { + let manager = persistor.inner.authorization_manager.clone(); let guard = manager.lock().await; match guard.refresh_token().await { Ok(_) => {} @@ -582,15 +591,31 @@ impl OAuthPersistor { return Err(error).with_context(|| { format!( "failed to refresh OAuth tokens for server {}", - self.inner.server_name + persistor.inner.server_name ) }); } } - } + drop(guard); - self.persist_if_needed_with_refresh_locks(refresh_locks) - .await + if let Err(error) = persistor.persist_if_needed_inner().await { + let message = format!( + "OAuth refresh succeeded but failed to persist rotated credentials for server {}: {error:#}", + persistor.inner.server_name + ); + warn!("{message}"); + *persistor.inner.refresh_persistence_failure.lock().await = + Some(RefreshPersistenceFailure { + message: message.clone(), + _refresh_locks: refresh_locks, + }); + return Err(anyhow!(message)); + } + + Ok(()) + }) + .await + .context("OAuth refresh task failed")? } } diff --git a/codex-rs/rmcp-client/src/rmcp_client.rs b/codex-rs/rmcp-client/src/rmcp_client.rs index a47e3c24f4..f846406b15 100644 --- a/codex-rs/rmcp-client/src/rmcp_client.rs +++ b/codex-rs/rmcp-client/src/rmcp_client.rs @@ -225,6 +225,8 @@ where enum ClientOperationError { #[error(transparent)] Service(#[from] rmcp::service::ServiceError), + #[error(transparent)] + OAuth(anyhow::Error), #[error("timed out awaiting {label} after {duration:?}")] Timeout { label: String, duration: Duration }, } @@ -863,10 +865,11 @@ impl RmcpClient { let service = transport .await .map_err(|err| anyhow!("handshaking with MCP server failed: {err}"))?; - if let Some(runtime) = &oauth_persistor - && let Err(error) = runtime.persist_if_needed().await - { - warn!("failed to persist OAuth tokens after initialize: {error}"); + if let Some(runtime) = oauth_persistor.clone() { + Self::persist_oauth_tokens_in_background( + runtime, + "failed to persist OAuth tokens after initialize", + ); } Ok((Arc::new(service), oauth_persistor)) }; @@ -934,18 +937,17 @@ impl RmcpClient { Fut: std::future::Future>, { let refresh_and_operation = async move { - if let Some(runtime) = &oauth_persistor - && let Err(error) = runtime.refresh_if_needed().await - { - warn!("failed to refresh OAuth tokens: {error}"); + if let Some(runtime) = &oauth_persistor { + runtime + .refresh_if_needed() + .await + .map_err(ClientOperationError::OAuth)?; } let result = operation(service) .await .map_err(ClientOperationError::from)?; - if let Some(runtime) = oauth_persistor - && let Err(error) = runtime.persist_if_needed().await - { - warn!("failed to persist OAuth tokens: {error}"); + if let Some(runtime) = oauth_persistor { + Self::persist_oauth_tokens_in_background(runtime, "failed to persist OAuth tokens"); } Ok(result) }; @@ -963,6 +965,24 @@ 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 { + warn!("{error_context}: {error}"); + } + }); + } + fn is_session_expired_404(error: &ClientOperationError) -> bool { let ClientOperationError::Service(rmcp::service::ServiceError::TransportSend(error)) = error @@ -1112,6 +1132,7 @@ mod tests { use std::time::Duration; use pretty_assertions::assert_eq; + use tokio::sync::oneshot; use tokio::time; use super::*; @@ -1134,4 +1155,37 @@ 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(); + + 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 completed_result = "completed tool response"; + + 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"); + } } 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 b431510c8e..69a7ce5619 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 @@ -429,6 +429,87 @@ async fn rotated_credentials_persist_after_startup_timeout_before_next_refresh() Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn provider_refresh_persists_after_startup_timeout_before_next_refresh() -> anyhow::Result<()> +{ + 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| { + responder_attempts.fetch_add(1, Ordering::SeqCst); + 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, + })) + }) + .expect(1) + .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 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" + ); + assert_eq!(refresh_attempts.load(Ordering::SeqCst), 1); + wait_for_persisted_token_snapshot( + &codex_home, + PersistedTokenSnapshot { + access_token: NEW_ACCESS_TOKEN.to_string(), + refresh_token: Some(NEW_REFRESH_TOKEN.to_string()), + }, + ) + .await?; + + let second_ready_path = control_dir.path().join("client-b.ready"); + let second_go_path = control_dir.path().join("client-b.go"); + let second_result_path = control_dir.path().join("client-b.result"); + let second_child = spawn_client_child( + &codex_home, + &server_url, + &second_ready_path, + &second_go_path, + &second_result_path, + )?; + wait_for_paths(std::slice::from_ref(&second_ready_path)).await?; + fs::write(&second_go_path, "go")?; + wait_for_children(vec![second_child]).await?; + assert_eq!(fs::read_to_string(second_result_path)?, "success"); + assert_eq!(refresh_attempts.load(Ordering::SeqCst), 1); + + fs::write(first_release_path, "release")?; + wait_for_children(vec![first_child]).await?; + 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; @@ -655,16 +736,21 @@ async fn refresh_provider_call_uses_operation_timeout_budget() -> anyhow::Result Mock::given(method("POST")) .and(path("/oauth/token")) .respond_with(move |_request: &Request| { - let response = ResponseTemplate::new(200).set_body_json(json!({ - "access_token": NEW_ACCESS_TOKEN, - "token_type": "Bearer", - "expires_in": 31, - "refresh_token": NEW_REFRESH_TOKEN, - })); - if responder_attempts.fetch_add(1, Ordering::SeqCst) == 0 { - response - } else { - response.set_delay(Duration::from_secs(1)) + 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_delay(Duration::from_secs(1)) + .set_body_json(json!({ + "access_token": FINAL_ACCESS_TOKEN, + "token_type": "Bearer", + "expires_in": 7200, + "refresh_token": FINAL_REFRESH_TOKEN, + })), } }) .expect(2) @@ -680,23 +766,132 @@ async fn refresh_provider_call_uses_operation_timeout_budget() -> anyhow::Result 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_operation_timeout_child( + let release_path = control_dir.path().join("client.release"); + let child = spawn_operation_timeout_child_held_open( &codex_home, &server_url, &ready_path, &go_path, &result_path, + &release_path, )?; wait_for_paths(std::slice::from_ref(&ready_path)).await?; sleep(Duration::from_secs(2)).await; fs::write(&go_path, "go")?; - wait_for_children(vec![child]).await?; + wait_for_paths(std::slice::from_ref(&result_path)).await?; assert_eq!( - fs::read_to_string(result_path)?, + fs::read_to_string(&result_path)?, format!("error:timed out awaiting tools/list after {OPERATION_TIMEOUT:?}") ); assert_eq!(refresh_attempts.load(Ordering::SeqCst), 2); + wait_for_persisted_token_snapshot( + &codex_home, + PersistedTokenSnapshot { + access_token: FINAL_ACCESS_TOKEN.to_string(), + refresh_token: Some(FINAL_REFRESH_TOKEN.to_string()), + }, + ) + .await?; + + fs::write(release_path, "release")?; + wait_for_children(vec![child]).await?; + server.verify().await; + Ok(()) +} + +#[cfg(unix)] +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn rotated_token_storage_failure_blocks_another_process_refresh() -> 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_body_json(json!({ + "access_token": NEW_ACCESS_TOKEN, + "token_type": "Bearer", + "expires_in": 31, + "refresh_token": NEW_REFRESH_TOKEN, + })), + 1 => ResponseTemplate::new(200).set_body_json(json!({ + "access_token": FINAL_ACCESS_TOKEN, + "token_type": "Bearer", + "expires_in": 7200, + "refresh_token": FINAL_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*/ 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 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_operation_timeout_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?; + + 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(&first_go_path, "go")?; + wait_for_paths(std::slice::from_ref(&first_result_path)).await?; + let first_outcome = fs::read_to_string(&first_result_path)?; + assert!( + first_outcome.starts_with(&format!( + "error:OAuth refresh succeeded but failed to persist rotated credentials for server {SERVER_NAME}:" + )), + "unexpected first-client outcome: {first_outcome}" + ); + assert_eq!(refresh_attempts.load(Ordering::SeqCst), 2); + + let second_ready_path = control_dir.path().join("client-b.ready"); + let second_go_path = control_dir.path().join("client-b.go"); + let second_result_path = control_dir.path().join("client-b.result"); + let mut second_child = spawn_client_child( + &codex_home, + &server_url, + &second_ready_path, + &second_go_path, + &second_result_path, + )?; + wait_for_paths(std::slice::from_ref(&second_ready_path)).await?; + fs::write(&second_go_path, "go")?; + 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?; server.verify().await; Ok(()) @@ -1348,6 +1543,22 @@ async fn wait_for_paths(paths: &[PathBuf]) -> anyhow::Result<()> { .context("timed out waiting for OAuth race process barrier") } +async fn wait_for_persisted_token_snapshot( + codex_home: &TempDir, + expected: PersistedTokenSnapshot, +) -> anyhow::Result<()> { + timeout(Duration::from_secs(10), async { + loop { + if persisted_token_snapshot(codex_home).ok().as_ref() == Some(&expected) { + return; + } + sleep(Duration::from_millis(10)).await; + } + }) + .await + .context("timed out waiting for rotated OAuth credentials to persist") +} + fn read_outcomes(paths: &[PathBuf]) -> anyhow::Result> { paths .iter()