From 04cfc0382512b429a53fe9a2f1f664afa8a4aef8 Mon Sep 17 00:00:00 2001 From: Steven Lee Date: Tue, 23 Jun 2026 23:27:17 +0000 Subject: [PATCH] Address OAuth refresh lock review feedback --- codex-rs/rmcp-client/src/oauth.rs | 108 +------ .../rmcp-client/src/oauth/refresh_lock.rs | 113 +++++++ .../src/oauth/refresh_lock_tests.rs | 42 +++ .../tests/streamable_http_oauth_startup.rs | 287 ++++++++++++++---- 4 files changed, 392 insertions(+), 158 deletions(-) create mode 100644 codex-rs/rmcp-client/src/oauth/refresh_lock.rs create mode 100644 codex-rs/rmcp-client/src/oauth/refresh_lock_tests.rs diff --git a/codex-rs/rmcp-client/src/oauth.rs b/codex-rs/rmcp-client/src/oauth.rs index d3e442de8b..1799f2d2e6 100644 --- a/codex-rs/rmcp-client/src/oauth.rs +++ b/codex-rs/rmcp-client/src/oauth.rs @@ -16,6 +16,8 @@ //! //! If the keyring is not available or fails, we fall back to CODEX_HOME/.credentials.json which is consistent with other coding CLI agents. +mod refresh_lock; + use anyhow::Context; use anyhow::Error; use anyhow::Result; @@ -41,8 +43,6 @@ use sha2::Digest; use sha2::Sha256; use std::collections::BTreeMap; use std::fs; -use std::fs::File; -use std::fs::OpenOptions; use std::io::ErrorKind; use std::path::PathBuf; use std::sync::Arc; @@ -58,17 +58,14 @@ use rmcp::transport::auth::CredentialStore as _; use rmcp::transport::auth::InMemoryCredentialStore; use rmcp::transport::auth::StoredCredentials; use tokio::sync::Mutex; -use tokio::time::sleep; use tokio::time::timeout; +use self::refresh_lock::RefreshCredentialLock; use codex_utils_home_dir::find_codex_home; const KEYRING_SERVICE: &str = "Codex MCP Credentials"; const MCP_OAUTH_SECRET_PREFIX: &str = "MCP_OAUTH"; const REFRESH_SKEW_MILLIS: u64 = 30_000; -const REFRESH_LOCK_DIR: &str = "mcp-oauth-refresh-locks"; -const REFRESH_LOCK_ACQUIRE_TIMEOUT: Duration = Duration::from_secs(60); -const REFRESH_LOCK_RETRY_SLEEP: Duration = Duration::from_millis(50); const REFRESH_REQUEST_TIMEOUT: Duration = Duration::from_secs(45); #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] @@ -832,61 +829,6 @@ impl OAuthPersistor { } } -struct RefreshCredentialLock { - _file: File, -} - -impl RefreshCredentialLock { - async fn acquire(store_key: &str) -> Result { - Self::acquire_with_timeout(store_key, REFRESH_LOCK_ACQUIRE_TIMEOUT).await - } - - async fn acquire_with_timeout(store_key: &str, acquire_timeout: Duration) -> Result { - let path = refresh_lock_path(store_key)?; - if let Some(parent) = path.parent() { - fs::create_dir_all(parent)?; - } - - let file = OpenOptions::new() - .read(true) - .write(true) - .create(true) - .truncate(false) - .open(&path) - .with_context(|| format!("failed to open OAuth refresh lock {}", path.display()))?; - - // Bound every contender, but keep the acquired lock for the full provider request and - // persistence transaction. Releasing it while awaiting the provider would allow concurrent - // use of a rotating refresh token. - match timeout(acquire_timeout, async { - loop { - match file.try_lock() { - Ok(()) => return Ok(()), - Err(std::fs::TryLockError::WouldBlock) => { - sleep(REFRESH_LOCK_RETRY_SLEEP).await; - } - Err(error) => return Err(std::io::Error::from(error)), - } - } - }) - .await - { - Ok(Ok(())) => {} - Ok(Err(error)) => { - return Err(error).with_context(|| { - format!("failed to lock OAuth refresh lock {}", path.display()) - }); - } - Err(_) => anyhow::bail!( - "timed out after {acquire_timeout:?} waiting for OAuth refresh lock {}", - path.display() - ), - } - - Ok(Self { _file: file }) - } -} - #[expect( clippy::await_holding_invalid_type, reason = "AuthorizationManager async access must be serialized through its mutex" @@ -1111,21 +1053,6 @@ fn fallback_file_path() -> Result { Ok(find_codex_home()?.join(FALLBACK_FILENAME).to_path_buf()) } -fn refresh_lock_path(store_key: &str) -> Result { - // Credential coordination is deliberately scoped to the active CODEX_HOME, alongside File - // and Secrets state. Coordinating the process-global Direct keyring across distinct homes - // would require a separately defined global lock namespace and is outside this transaction. - // TODO(stevenlee): define a safe per-user, cross-platform rendezvous before extending Direct - // keyring coordination across distinct CODEX_HOME values. - let mut hasher = Sha256::new(); - hasher.update(store_key.as_bytes()); - let digest = hasher.finalize(); - Ok(find_codex_home()? - .join(REFRESH_LOCK_DIR) - .join(format!("{digest:x}.lock")) - .to_path_buf()) -} - fn read_fallback_file() -> Result> { let path = fallback_file_path()?; let contents = match fs::read_to_string(&path) { @@ -1479,35 +1406,6 @@ mod tests { Ok(()) } - #[tokio::test] - async fn refresh_lock_acquisition_times_out_without_stealing() -> Result<()> { - let _env = TempCodexHome::new(); - let store_key = "test-store-key"; - let held_lock = - RefreshCredentialLock::acquire_with_timeout(store_key, Duration::from_millis(100)) - .await?; - - let error = - match RefreshCredentialLock::acquire_with_timeout(store_key, Duration::from_millis(50)) - .await - { - Ok(_) => panic!("contending lock acquisition should time out"), - Err(error) => error, - }; - assert!( - error - .to_string() - .contains("timed out after 50ms waiting for OAuth refresh lock"), - "unexpected error: {error:#}" - ); - - drop(held_lock); - let _reacquired = - RefreshCredentialLock::acquire_with_timeout(store_key, Duration::from_millis(100)) - .await?; - Ok(()) - } - #[tokio::test] async fn refresh_transaction_preserves_credentials_when_resolved_keyring_reread_fails() -> Result<()> { diff --git a/codex-rs/rmcp-client/src/oauth/refresh_lock.rs b/codex-rs/rmcp-client/src/oauth/refresh_lock.rs new file mode 100644 index 0000000000..b540066cd1 --- /dev/null +++ b/codex-rs/rmcp-client/src/oauth/refresh_lock.rs @@ -0,0 +1,113 @@ +//! Cross-process serialization for one MCP OAuth credential's refresh transaction. +//! +//! The guard is intentionally acquired before the authoritative credential reread and retained +//! through provider refresh and persistence. This prevents two processes from replaying the same +//! rotating refresh token or observing a partially persisted transaction. + +use anyhow::Context; +use anyhow::Result; +use codex_utils_home_dir::find_codex_home; +use sha2::Digest; +use sha2::Sha256; +use std::fs; +use std::fs::File; +use std::fs::OpenOptions; +use std::path::Path; +use std::path::PathBuf; +use std::time::Duration; +use tokio::time::sleep; +use tokio::time::timeout; + +const REFRESH_LOCK_DIR: &str = "mcp-oauth-refresh-locks"; +const REFRESH_LOCK_ACQUIRE_TIMEOUT: Duration = Duration::from_secs(/*secs*/ 60); +const REFRESH_LOCK_RETRY_SLEEP: Duration = Duration::from_millis(/*millis*/ 50); +// Keep this internal target stable so diagnostics and cross-process tests can distinguish actual +// WouldBlock contention from a contender that merely started late and observed persisted tokens. +const LOCK_CONTENTION_EVENT_TARGET: &str = "codex_rmcp_client::oauth::refresh_lock::contention"; + +pub(super) struct RefreshCredentialLock { + _file: File, +} + +impl RefreshCredentialLock { + pub(super) async fn acquire(store_key: &str) -> Result { + let codex_home = find_codex_home()?; + Self::acquire_in(&codex_home, store_key, REFRESH_LOCK_ACQUIRE_TIMEOUT).await + } + + async fn acquire_in( + codex_home: &Path, + store_key: &str, + acquire_timeout: Duration, + ) -> Result { + let path = refresh_lock_path(codex_home, store_key); + if let Some(parent) = path.parent() { + fs::create_dir_all(parent)?; + } + + let file = OpenOptions::new() + .read(true) + .write(true) + .create(true) + .truncate(false) + .open(&path) + .with_context(|| format!("failed to open OAuth refresh lock {}", path.display()))?; + + // Bound every contender, but keep the acquired lock for the full provider request and + // persistence transaction. Releasing it while awaiting the provider would allow concurrent + // use of a rotating refresh token. + let mut reported_contention = false; + match timeout(acquire_timeout, async { + loop { + match file.try_lock() { + Ok(()) => return Ok(()), + Err(std::fs::TryLockError::WouldBlock) => { + if !reported_contention { + tracing::debug!( + target: LOCK_CONTENTION_EVENT_TARGET, + lock_path = %path.display(), + "waiting for another process to finish refreshing MCP OAuth credentials" + ); + reported_contention = true; + } + sleep(REFRESH_LOCK_RETRY_SLEEP).await; + } + Err(error) => return Err(std::io::Error::from(error)), + } + } + }) + .await + { + Ok(Ok(())) => {} + Ok(Err(error)) => { + return Err(error).with_context(|| { + format!("failed to lock OAuth refresh lock {}", path.display()) + }); + } + Err(_) => anyhow::bail!( + "timed out after {acquire_timeout:?} waiting for OAuth refresh lock {}", + path.display() + ), + } + + Ok(Self { _file: file }) + } +} + +fn refresh_lock_path(codex_home: &Path, store_key: &str) -> PathBuf { + // Credential coordination is deliberately scoped to the active CODEX_HOME, alongside File + // and Secrets state. Coordinating the process-global Direct keyring across distinct homes + // would require a separately defined global lock namespace and is outside this transaction. + // TODO(stevenlee): define a safe per-user, cross-platform rendezvous before extending Direct + // keyring coordination across distinct CODEX_HOME values. + let mut hasher = Sha256::new(); + hasher.update(store_key.as_bytes()); + let digest = hasher.finalize(); + codex_home + .join(REFRESH_LOCK_DIR) + .join(format!("{digest:x}.lock")) +} + +#[cfg(test)] +#[path = "refresh_lock_tests.rs"] +mod tests; diff --git a/codex-rs/rmcp-client/src/oauth/refresh_lock_tests.rs b/codex-rs/rmcp-client/src/oauth/refresh_lock_tests.rs new file mode 100644 index 0000000000..2b90e73996 --- /dev/null +++ b/codex-rs/rmcp-client/src/oauth/refresh_lock_tests.rs @@ -0,0 +1,42 @@ +use super::RefreshCredentialLock; +use anyhow::Result; +use std::time::Duration; +use tempfile::tempdir; + +#[tokio::test] +async fn acquisition_times_out_without_stealing() -> Result<()> { + let codex_home = tempdir()?; + let store_key = "test-store-key"; + let held_lock = RefreshCredentialLock::acquire_in( + codex_home.path(), + store_key, + Duration::from_millis(/*millis*/ 100), + ) + .await?; + + let error = match RefreshCredentialLock::acquire_in( + codex_home.path(), + store_key, + Duration::from_millis(/*millis*/ 50), + ) + .await + { + Ok(_) => panic!("contending lock acquisition should time out"), + Err(error) => error, + }; + assert!( + error + .to_string() + .contains("timed out after 50ms waiting for OAuth refresh lock"), + "unexpected error: {error:#}" + ); + + drop(held_lock); + let _reacquired = RefreshCredentialLock::acquire_in( + codex_home.path(), + store_key, + Duration::from_millis(/*millis*/ 100), + ) + .await?; + 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 02866823f2..7adb150a60 100644 --- a/codex-rs/rmcp-client/tests/streamable_http_oauth_startup.rs +++ b/codex-rs/rmcp-client/tests/streamable_http_oauth_startup.rs @@ -1,11 +1,19 @@ 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::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; @@ -24,7 +32,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; @@ -42,13 +61,170 @@ 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 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 startup_refresh_does_not_consume_handshake_timeout() -> anyhow::Result<()> { let server = MockServer::start().await; @@ -133,39 +309,21 @@ async fn startup_refresh_does_not_consume_handshake_timeout() -> anyhow::Result< Ok(()) } -#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +#[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": format!("{}/oauth/token", server.uri()), + "token_endpoint": refresh_server.token_endpoint.clone(), "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_delay(Duration::from_millis(250)) - .set_body_json(json!({ - "access_token": REFRESHED_ACCESS_TOKEN, - "token_type": "Bearer", - "expires_in": 7200, - "refresh_token": ROTATED_REFRESH_TOKEN, - })), - ) - // This exact expectation proves that the two child processes issue one serialized provider - // refresh rather than both replaying the same rotating refresh token. - .expect(1) - .mount(&server) - .await; Mock::given(method("POST")) .and(path("/mcp")) .and(header( @@ -199,8 +357,6 @@ async fn concurrent_file_mode_startup_refreshes_once() -> anyhow::Result<()> { .mount(&server) .await; - let codex_home = TempDir::new()?; - let server_url = format!("{}/mcp", server.uri()); let seed_status = Command::new(std::env::current_exe()?) .args(["oauth_concurrency_seed_child", "--exact", "--ignored"]) .env("CODEX_HOME", codex_home.path()) @@ -212,34 +368,51 @@ async fn concurrent_file_mode_startup_refreshes_once() -> anyhow::Result<()> { "OAuth concurrency seed child failed: {seed_status}" ); - let release_file = codex_home.path().join("oauth-clients.release"); 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 mut first_child = Command::new(std::env::current_exe()?) - .args(["oauth_concurrency_client_child", "--exact", "--ignored"]) - .env("CODEX_HOME", codex_home.path()) - .env(CHILD_SERVER_URL_ENV, &server_url) - .env(CHILD_READY_FILE_ENV, &first_ready_file) - .env(CHILD_RELEASE_FILE_ENV, &release_file) - .kill_on_drop(true) - .spawn()?; - let mut second_child = Command::new(std::env::current_exe()?) - .args(["oauth_concurrency_client_child", "--exact", "--ignored"]) - .env("CODEX_HOME", codex_home.path()) - .env(CHILD_SERVER_URL_ENV, &server_url) - .env(CHILD_READY_FILE_ENV, &second_ready_file) - .env(CHILD_RELEASE_FILE_ENV, &release_file) - .kill_on_drop(true) - .spawn()?; + 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()?; - tokio::time::timeout(Duration::from_secs(10), async { - while !first_ready_file.exists() || !second_ready_file.exists() { - tokio::time::sleep(Duration::from_millis(10)).await; - } - }) - .await - .map_err(|_| anyhow::anyhow!("OAuth concurrency children did not become ready"))?; - fs::write(&release_file, b"release")?; + 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!( @@ -252,6 +425,8 @@ async fn concurrent_file_mode_startup_refreshes_once() -> anyhow::Result<()> { ); server.verify().await; + contention_result?; + assert_eq!(refresh_server.request_count(), 1); Ok(()) } @@ -428,6 +603,12 @@ 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, @@ -441,11 +622,11 @@ async fn oauth_concurrency_client_child() -> anyhow::Result<()> { ) .await?; - // Both processes must construct their OAuth client from the same expired snapshot before - // either begins refresh. This removes process scheduling from the contention setup. + // 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(10)).await; + tokio::time::sleep(Duration::from_millis(/*millis*/ 10)).await; } initialize_client(&client).await?; Ok(())