From cb1a7928ada6b90a2035b6b8d19cb476f3abdaa8 Mon Sep 17 00:00:00 2001 From: Adam Perry Date: Fri, 5 Jun 2026 07:28:47 +0000 Subject: [PATCH] fix(rmcp): persist rotated tokens after timeout --- codex-rs/rmcp-client/src/oauth.rs | 24 +- codex-rs/rmcp-client/src/rmcp_client.rs | 47 +-- .../streamable_http_oauth_refresh_race.rs | 270 +++++++++++++++++- 3 files changed, 304 insertions(+), 37 deletions(-) diff --git a/codex-rs/rmcp-client/src/oauth.rs b/codex-rs/rmcp-client/src/oauth.rs index 5c25893f6a..1fbf7aaf0c 100644 --- a/codex-rs/rmcp-client/src/oauth.rs +++ b/codex-rs/rmcp-client/src/oauth.rs @@ -424,11 +424,28 @@ 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")? + } + #[expect( clippy::await_holding_invalid_type, reason = "AuthorizationManager async access must be serialized through its mutex" )] - pub(crate) async fn persist_if_needed(&self) -> Result<()> { + async fn persist_if_needed_inner(&self) -> Result<()> { let (client_id, maybe_credentials) = { let manager = self.inner.authorization_manager.clone(); let guard = manager.lock().await; @@ -502,7 +519,7 @@ impl OAuthPersistor { // A different process may have rotated the one-time refresh token. // Reload its durable result while all refreshers share the same lock. - let (_refresh_locks, stored) = lock_and_load_oauth_tokens( + let (refresh_locks, stored) = lock_and_load_oauth_tokens( &self.inner.server_name, &self.inner.url, self.inner.store_mode, @@ -572,7 +589,8 @@ impl OAuthPersistor { } } - self.persist_if_needed().await + self.persist_if_needed_with_refresh_locks(refresh_locks) + .await } } diff --git a/codex-rs/rmcp-client/src/rmcp_client.rs b/codex-rs/rmcp-client/src/rmcp_client.rs index 006375a17d..a47e3c24f4 100644 --- a/codex-rs/rmcp-client/src/rmcp_client.rs +++ b/codex-rs/rmcp-client/src/rmcp_client.rs @@ -432,12 +432,6 @@ impl RmcpClient { }; } - if let Some(runtime) = oauth_persistor - && let Err(error) = runtime.persist_if_needed().await - { - warn!("failed to persist OAuth tokens after initialize: {error}"); - } - Ok(initialize_result) } @@ -452,7 +446,6 @@ impl RmcpClient { async move { service.list_tools(params).await }.boxed() }) .await?; - self.persist_oauth_tokens().await; Ok(result) } @@ -485,7 +478,6 @@ impl RmcpClient { }) }) .collect::>>()?; - self.persist_oauth_tokens().await; Ok(ListToolsWithConnectorIdResult { next_cursor: result.next_cursor, tools, @@ -511,7 +503,6 @@ impl RmcpClient { async move { service.list_resources(params).await }.boxed() }) .await?; - self.persist_oauth_tokens().await; Ok(result) } @@ -526,7 +517,6 @@ impl RmcpClient { async move { service.list_resource_templates(params).await }.boxed() }) .await?; - self.persist_oauth_tokens().await; Ok(result) } @@ -541,7 +531,6 @@ impl RmcpClient { async move { service.read_resource(params).await }.boxed() }) .await?; - self.persist_oauth_tokens().await; Ok(result) } @@ -598,7 +587,6 @@ impl RmcpClient { .boxed() }) .await?; - self.persist_oauth_tokens().await; Ok(result) } @@ -627,7 +615,6 @@ impl RmcpClient { }, ) .await?; - self.persist_oauth_tokens().await; Ok(()) } @@ -649,7 +636,6 @@ impl RmcpClient { .boxed() }) .await?; - self.persist_oauth_tokens().await; Ok(response) } @@ -689,16 +675,6 @@ impl RmcpClient { drop(previous_state); } - /// This should be called after every tool call so that if a given tool call triggered - /// a refresh of the OAuth tokens, they are persisted. - async fn persist_oauth_tokens(&self) { - if let Some(runtime) = self.oauth_persistor().await - && let Err(error) = runtime.persist_if_needed().await - { - warn!("failed to persist OAuth tokens: {error}"); - } - } - async fn create_pending_transport( transport_recipe: &TransportRecipe, ) -> Result { @@ -887,6 +863,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}"); + } Ok((Arc::new(service), oauth_persistor)) }; @@ -953,12 +934,20 @@ impl RmcpClient { Fut: std::future::Future>, { let refresh_and_operation = async move { - if let Some(runtime) = oauth_persistor + if let Some(runtime) = &oauth_persistor && let Err(error) = runtime.refresh_if_needed().await { warn!("failed to refresh OAuth tokens: {error}"); } - operation(service).await.map_err(ClientOperationError::from) + 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}"); + } + Ok(result) }; match timeout { @@ -1045,12 +1034,6 @@ impl RmcpClient { }; } - if let Some(runtime) = oauth_persistor - && let Err(error) = runtime.persist_if_needed().await - { - warn!("failed to persist OAuth tokens after session recovery: {error}"); - } - Ok(()) } } diff --git a/codex-rs/rmcp-client/tests/streamable_http_oauth_refresh_race.rs b/codex-rs/rmcp-client/tests/streamable_http_oauth_refresh_race.rs index a615bd883b..b431510c8e 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 @@ -39,6 +39,7 @@ use wiremock::Mock; use wiremock::MockServer; use wiremock::Request; use wiremock::ResponseTemplate; +use wiremock::Times; use wiremock::matchers::method; use wiremock::matchers::path; @@ -50,6 +51,8 @@ const OLD_ACCESS_TOKEN: &str = "old-access-token"; const OLD_REFRESH_TOKEN: &str = "old-one-time-refresh-token"; const NEW_ACCESS_TOKEN: &str = "new-access-token"; const NEW_REFRESH_TOKEN: &str = "new-one-time-refresh-token"; +const FINAL_ACCESS_TOKEN: &str = "final-access-token"; +const FINAL_REFRESH_TOKEN: &str = "final-one-time-refresh-token"; const FALLBACK_LOCK_PREFIX: &str = "codex-mcp-oauth-fallback"; const FILE_OAUTH_REFRESH_LOCK_PREFIX: &str = "codex-mcp-oauth-refresh-file"; const KEYRING_SERVICE: &str = "Codex MCP Credentials"; @@ -64,6 +67,7 @@ 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"; const STORE_MODE_ENV: &str = "MCP_TEST_OAUTH_RACE_STORE_MODE"; +const RELEASE_PATH_ENV: &str = "MCP_TEST_OAUTH_RACE_RELEASE_PATH"; #[tokio::test(flavor = "multi_thread", worker_threads = 1)] async fn concurrent_processes_coordinate_one_time_refresh_and_reload_rotated_tokens() @@ -327,6 +331,104 @@ async fn fallback_write_lock_wait_uses_startup_timeout_budget() -> anyhow::Resul Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn rotated_credentials_persist_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| { + if 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": 7200, + "refresh_token": NEW_REFRESH_TOKEN, + })) + } else { + ResponseTemplate::new(400).set_body_json(json!({ + "error": "invalid_grant", + "error_description": "refresh token already used", + })) + } + }) + .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 (lock_path, _) = fallback_lock_path_and_user_namespace(&codex_home)?; + let fallback_lock = OpenOptions::new() + .read(true) + .write(true) + .create(true) + .truncate(false) + .open(lock_path)?; + fallback_lock.lock_shared()?; + + 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); + + 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")?; + sleep(Duration::from_millis(250)).await; + assert_eq!(refresh_attempts.load(Ordering::SeqCst), 1); + assert!(!second_result_path.exists()); + + drop(fallback_lock); + 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); + assert_eq!( + persisted_token_snapshot(&codex_home)?, + PersistedTokenSnapshot { + access_token: NEW_ACCESS_TOKEN.to_string(), + refresh_token: Some(NEW_REFRESH_TOKEN.to_string()), + } + ); + + 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; @@ -439,6 +541,111 @@ async fn fallback_write_lock_wait_uses_operation_timeout_budget() -> anyhow::Res Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn rotated_credentials_persist_after_operation_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| { + 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*/ 4..=5).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?; + sleep(Duration::from_secs(2)).await; + + let (lock_path, _) = fallback_lock_path_and_user_namespace(&codex_home)?; + let fallback_lock = OpenOptions::new() + .read(true) + .write(true) + .create(true) + .truncate(false) + .open(lock_path)?; + fallback_lock.lock_shared()?; + + 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)?, + format!("error:timed out awaiting tools/list after {OPERATION_TIMEOUT:?}") + ); + 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 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()); + + drop(fallback_lock); + wait_for_children(vec![second_child]).await?; + assert_eq!(fs::read_to_string(&second_result_path)?, "success"); + assert_eq!(refresh_attempts.load(Ordering::SeqCst), 2); + assert_eq!( + persisted_token_snapshot(&codex_home)?, + PersistedTokenSnapshot { + access_token: FINAL_ACCESS_TOKEN.to_string(), + refresh_token: Some(FINAL_REFRESH_TOKEN.to_string()), + } + ); + + 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 refresh_provider_call_uses_operation_timeout_budget() -> anyhow::Result<()> { let server = MockServer::start().await; @@ -451,7 +658,7 @@ async fn refresh_provider_call_uses_operation_timeout_budget() -> anyhow::Result let response = ResponseTemplate::new(200).set_body_json(json!({ "access_token": NEW_ACCESS_TOKEN, "token_type": "Bearer", - "expires_in": 30, + "expires_in": 31, "refresh_token": NEW_REFRESH_TOKEN, })); if responder_attempts.fetch_add(1, Ordering::SeqCst) == 0 { @@ -481,6 +688,7 @@ async fn refresh_provider_call_uses_operation_timeout_budget() -> anyhow::Result &result_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?; @@ -841,6 +1049,9 @@ async fn oauth_refresh_race_client_child() -> anyhow::Result<()> { Err(error) => format!("error:{error:#}"), }; fs::write(result_path, outcome)?; + if let Ok(release_path) = std::env::var(RELEASE_PATH_ENV) { + wait_for_paths(&[PathBuf::from(release_path)]).await?; + } Ok(()) } @@ -871,6 +1082,9 @@ async fn oauth_refresh_operation_timeout_child() -> anyhow::Result<()> { Err(error) => format!("error:{error:#}"), }; fs::write(result_path, outcome)?; + if let Ok(release_path) = std::env::var(RELEASE_PATH_ENV) { + wait_for_paths(&[PathBuf::from(release_path)]).await?; + } Ok(()) } @@ -923,7 +1137,7 @@ async fn mount_oauth_metadata(server: &MockServer, expected_requests: u64) { .await; } -async fn mount_mcp_server(server: &MockServer, expected_requests: u64) { +async fn mount_mcp_server(server: &MockServer, expected_requests: impl Into) { Mock::given(method("POST")) .and(path("/mcp")) .respond_with(|request: &Request| { @@ -1006,6 +1220,32 @@ fn spawn_client_child( Ok(command.spawn()?) } +fn spawn_client_child_held_open( + codex_home: &TempDir, + server_url: &str, + ready_path: &Path, + go_path: &Path, + result_path: &Path, + release_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(RESULT_PATH_ENV, result_path) + .env(RELEASE_PATH_ENV, release_path) + .kill_on_drop(true); + Ok(command.spawn()?) +} + fn spawn_operation_timeout_child( codex_home: &TempDir, server_url: &str, @@ -1030,6 +1270,32 @@ fn spawn_operation_timeout_child( Ok(command.spawn()?) } +fn spawn_operation_timeout_child_held_open( + codex_home: &TempDir, + server_url: &str, + ready_path: &Path, + go_path: &Path, + result_path: &Path, + release_path: &Path, +) -> anyhow::Result { + let mut command = Command::new(std::env::current_exe()?); + command + .args([ + "oauth_refresh_operation_timeout_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) + .env(RELEASE_PATH_ENV, release_path) + .kill_on_drop(true); + Ok(command.spawn()?) +} + fn spawn_read_only_refresh_child( codex_home: &TempDir, server_url: &str,