From a525c9cf4b4bdac2cc084868a9607e618a8923da Mon Sep 17 00:00:00 2001 From: Adam Perry Date: Fri, 5 Jun 2026 05:23:49 +0000 Subject: [PATCH] fix(rmcp-client): serialize OAuth credential generations --- codex-rs/rmcp-client/src/oauth.rs | 346 ++++++++--- .../rmcp-client/src/perform_oauth_login.rs | 4 +- .../tests/streamable_http_oauth_startup.rs | 561 +++++++++--------- 3 files changed, 554 insertions(+), 357 deletions(-) diff --git a/codex-rs/rmcp-client/src/oauth.rs b/codex-rs/rmcp-client/src/oauth.rs index bf906b2564..d4ef0be37b 100644 --- a/codex-rs/rmcp-client/src/oauth.rs +++ b/codex-rs/rmcp-client/src/oauth.rs @@ -39,6 +39,7 @@ use std::fs::OpenOptions; use std::io::ErrorKind; 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; @@ -51,7 +52,7 @@ use rmcp::transport::auth::CredentialStore; use rmcp::transport::auth::InMemoryCredentialStore; use rmcp::transport::auth::StoredCredentials; use tokio::sync::Mutex; -use tokio::sync::Semaphore; +use tokio::sync::OwnedMutexGuard; use codex_utils_home_dir::find_codex_home; @@ -172,6 +173,15 @@ pub fn save_oauth_tokens( save_oauth_tokens_locked(server_name, tokens, store_mode) } +pub(crate) async fn save_oauth_tokens_async( + server_name: &str, + tokens: &StoredOAuthTokens, + store_mode: OAuthCredentialsStoreMode, +) -> Result<()> { + let _lock = acquire_oauth_server_lock_async(server_name, &tokens.url).await?; + save_oauth_tokens_locked(server_name, tokens, store_mode) +} + fn save_oauth_tokens_locked( server_name: &str, tokens: &StoredOAuthTokens, @@ -237,6 +247,15 @@ pub fn delete_oauth_tokens( server_name: &str, url: &str, store_mode: OAuthCredentialsStoreMode, +) -> Result { + let _lock = acquire_oauth_server_lock(server_name, url)?; + delete_oauth_tokens_locked(server_name, url, store_mode) +} + +fn delete_oauth_tokens_locked( + server_name: &str, + url: &str, + store_mode: OAuthCredentialsStoreMode, ) -> Result { let keyring_store = DefaultKeyringStore; delete_oauth_tokens_from_keyring_and_file(&keyring_store, store_mode, server_name, url) @@ -280,7 +299,12 @@ struct OAuthPersistorInner { authorization_manager: Arc>, store_mode: OAuthCredentialsStoreMode, last_credentials: Mutex>, - update_lock: Semaphore, +} + +enum CredentialReload { + Unchanged, + Replaced, + Removed, } impl OAuthPersistor { @@ -298,7 +322,6 @@ impl OAuthPersistor { authorization_manager, store_mode, last_credentials: Mutex::new(initial_credentials), - update_lock: Semaphore::new(1), }), } } @@ -306,7 +329,14 @@ 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<()> { - let _permit = self.inner.update_lock.acquire().await?; + let _server_lock = + acquire_oauth_server_lock_async(&self.inner.server_name, &self.inner.url).await?; + if !matches!( + self.reload_persisted_credentials_locked().await?, + CredentialReload::Unchanged + ) { + return Ok(()); + } self.persist_if_needed_locked().await } @@ -341,36 +371,22 @@ impl OAuthPersistor { token_response: new_token_response, expires_at, }; - if last_credentials + if !last_credentials .as_ref() .is_some_and(|previous| same_persisted_generation(previous, &stored)) { - return Ok(()); + save_oauth_tokens_locked( + &self.inner.server_name, + &stored, + self.inner.store_mode, + )?; + *self.inner.last_credentials.lock().await = Some(stored); } - - let _server_lock = - acquire_oauth_server_lock_async(&self.inner.server_name, &self.inner.url) - .await?; - if let Some(persisted) = load_oauth_tokens( - &self.inner.server_name, - &self.inner.url, - self.inner.store_mode, - )? && !last_credentials - .as_ref() - .is_some_and(|previous| same_persisted_generation(previous, &persisted)) - { - self.replace_manager_credentials(&persisted).await?; - *self.inner.last_credentials.lock().await = Some(persisted); - return Ok(()); - } - - save_oauth_tokens_locked(&self.inner.server_name, &stored, self.inner.store_mode)?; - *self.inner.last_credentials.lock().await = Some(stored); } None => { let mut last_serialized = self.inner.last_credentials.lock().await; if last_serialized.take().is_some() - && let Err(error) = delete_oauth_tokens( + && let Err(error) = delete_oauth_tokens_locked( &self.inner.server_name, &self.inner.url, self.inner.store_mode, @@ -389,11 +405,20 @@ impl OAuthPersistor { #[expect( clippy::await_holding_invalid_type, - reason = "AuthorizationManager async access and persisted credential adoption must be serialized" + reason = "the server credential lock must span reload, refresh, and persistence" )] pub(crate) async fn refresh_if_needed(&self) -> Result<()> { - let _permit = self.inner.update_lock.acquire().await?; - self.adopt_persisted_credentials_if_changed().await?; + let _server_lock = + acquire_oauth_server_lock_async(&self.inner.server_name, &self.inner.url).await?; + if matches!( + self.reload_persisted_credentials_locked().await?, + CredentialReload::Removed + ) { + return Err(anyhow::anyhow!( + "OAuth credentials were removed for server {}", + self.inner.server_name + )); + } let expires_at = { let guard = self.inner.last_credentials.lock().await; @@ -418,29 +443,41 @@ impl OAuthPersistor { self.persist_if_needed_locked().await } - async fn adopt_persisted_credentials_if_changed(&self) -> Result<()> { - let _server_lock = - acquire_oauth_server_lock_async(&self.inner.server_name, &self.inner.url).await?; - let Some(persisted) = load_oauth_tokens( + async fn reload_persisted_credentials_locked(&self) -> Result { + let loaded = load_oauth_tokens( &self.inner.server_name, &self.inner.url, self.inner.store_mode, - )? - else { - return Ok(()); - }; + ) + .with_context(|| { + format!( + "failed to reload OAuth tokens for server {}", + self.inner.server_name + ) + })?; let last_credentials = self.inner.last_credentials.lock().await.clone(); - if last_credentials - .as_ref() - .is_some_and(|previous| same_persisted_generation(previous, &persisted)) - { - return Ok(()); + let unchanged = match (&loaded, &last_credentials) { + (Some(loaded), Some(previous)) => same_persisted_generation(loaded, previous), + (None, None) => true, + _ => false, + }; + if unchanged { + return Ok(CredentialReload::Unchanged); } - self.replace_manager_credentials(&persisted).await?; - *self.inner.last_credentials.lock().await = Some(persisted); - Ok(()) + match loaded { + Some(tokens) => { + self.replace_manager_credentials(&tokens).await?; + *self.inner.last_credentials.lock().await = Some(tokens); + Ok(CredentialReload::Replaced) + } + None => { + self.clear_manager_credentials().await; + *self.inner.last_credentials.lock().await = None; + Ok(CredentialReload::Removed) + } + } } async fn replace_manager_credentials(&self, tokens: &StoredOAuthTokens) -> Result<()> { @@ -469,33 +506,112 @@ impl OAuthPersistor { guard.set_credential_store(store); Ok(()) } + + async fn clear_manager_credentials(&self) { + let manager = self.inner.authorization_manager.clone(); + let mut guard = manager.lock().await; + guard.set_credential_store(InMemoryCredentialStore::new()); + } +} + +fn oauth_server_lock_for(server_name: &str, url: &str) -> Arc> { + static OAUTH_SERVER_LOCKS: OnceLock>>>> = + OnceLock::new(); + + let mut locks = OAUTH_SERVER_LOCKS + .get_or_init(std::sync::Mutex::default) + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + locks + .entry(format!("{server_name}\n{url}")) + .or_insert_with(|| Arc::new(Mutex::new(()))) + .clone() } struct OAuthServerLock { _file: fs::File, + _in_process_guard: Option>, } fn acquire_oauth_server_lock(server_name: &str, url: &str) -> Result { - let digest = sha_256_prefix(&Value::String(format!("{server_name}\n{url}")))?; - let path = find_codex_home()?.join(format!(".mcp-oauth-{digest}.lock")); - let file = OpenOptions::new() + let path = oauth_server_lock_path(server_name, url)?; + let file = open_oauth_lock_file(&path)?; + file.lock() + .with_context(|| format!("failed to lock OAuth credentials at {}", path.display()))?; + Ok(OAuthServerLock { + _file: file, + _in_process_guard: None, + }) +} + +async fn acquire_oauth_server_lock_async(server_name: &str, url: &str) -> Result { + let in_process_guard = oauth_server_lock_for(server_name, url).lock_owned().await; + let path = oauth_server_lock_path(server_name, url)?; + let file = tokio::task::spawn_blocking({ + let path = path.clone(); + move || { + let file = open_oauth_lock_file(&path)?; + file.lock().with_context(|| { + format!("failed to lock OAuth credentials at {}", path.display()) + })?; + Ok::<_, anyhow::Error>(file) + } + }) + .await + .context("OAuth credential lock task failed")??; + Ok(OAuthServerLock { + _file: file, + _in_process_guard: Some(in_process_guard), + }) +} + +struct FallbackStoreLock { + _in_process_guard: std::sync::MutexGuard<'static, ()>, + _file: fs::File, +} + +fn acquire_fallback_store_lock() -> Result { + static FALLBACK_STORE_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(()); + + let in_process_guard = FALLBACK_STORE_LOCK + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + let path = find_codex_home()?.join(".mcp-oauth-locks/fallback-store.lock"); + let file = open_oauth_lock_file(&path)?; + file.lock().with_context(|| { + format!( + "failed to lock fallback OAuth credentials at {}", + path.display() + ) + })?; + Ok(FallbackStoreLock { + _in_process_guard: in_process_guard, + _file: file, + }) +} + +fn open_oauth_lock_file(path: &std::path::Path) -> Result { + if let Some(parent) = path.parent() { + fs::create_dir_all(parent)?; + } + OpenOptions::new() .read(true) .write(true) .create(true) .truncate(false) - .open(&path) - .with_context(|| format!("failed to open OAuth credential lock at {}", path.display()))?; - file.lock() - .with_context(|| format!("failed to lock OAuth credentials at {}", path.display()))?; - Ok(OAuthServerLock { _file: file }) + .open(path) + .with_context(|| format!("failed to open OAuth credential lock at {}", path.display())) } -async fn acquire_oauth_server_lock_async(server_name: &str, url: &str) -> Result { - let server_name = server_name.to_string(); - let url = url.to_string(); - tokio::task::spawn_blocking(move || acquire_oauth_server_lock(&server_name, &url)) - .await - .context("OAuth credential lock task failed")? +fn oauth_server_lock_path(server_name: &str, url: &str) -> Result { + let server_digest = sha_256_prefix(&Value::String(format!("{server_name}\n{url}")))?; + let user_identity = std::env::var_os("HOME") + .or_else(|| std::env::var_os("USERPROFILE")) + .unwrap_or_else(|| std::env::temp_dir().into_os_string()); + let user_digest = sha_256_prefix(&Value::String(user_identity.to_string_lossy().into_owned()))?; + Ok(std::env::temp_dir() + .join(format!("codex-mcp-oauth-{user_digest}")) + .join(format!("server-{server_digest}.lock"))) } fn same_persisted_generation(left: &StoredOAuthTokens, right: &StoredOAuthTokens) -> bool { @@ -507,16 +623,15 @@ fn same_persisted_generation(left: &StoredOAuthTokens, right: &StoredOAuthTokens return false; } - let normalize = |response: &WrappedOAuthTokenResponse| { - let mut value = serde_json::to_value(response).ok()?; - value.as_object_mut()?.remove("expires_in"); - Some(value) - }; + let mut left_response = left.token_response.0.clone(); + let mut right_response = right.token_response.0.clone(); + left_response.set_expires_in(None); + right_response.set_expires_in(None); match ( - normalize(&left.token_response), - normalize(&right.token_response), + serde_json::to_string(&left_response), + serde_json::to_string(&right_response), ) { - (Some(left), Some(right)) => left == right, + (Ok(left), Ok(right)) => left == right, _ => false, } } @@ -541,6 +656,7 @@ struct FallbackTokenEntry { } fn load_oauth_tokens_from_file(server_name: &str, url: &str) -> Result> { + let _lock = acquire_fallback_store_lock()?; let Some(store) = read_fallback_file()? else { return Ok(None); }; @@ -584,6 +700,7 @@ fn load_oauth_tokens_from_file(server_name: &str, url: &str) -> Result Result<()> { + let _lock = acquire_fallback_store_lock()?; let key = compute_store_key(&tokens.server_name, &tokens.url)?; let mut store = read_fallback_file()?.unwrap_or_default(); @@ -613,6 +730,7 @@ fn save_oauth_tokens_to_file(tokens: &StoredOAuthTokens) -> Result<()> { } fn delete_oauth_tokens_from_file(key: &str) -> Result { + let _lock = acquire_fallback_store_lock()?; let mut store = match read_fallback_file()? { Some(store) => store, None => return Ok(false), @@ -750,10 +868,13 @@ mod tests { use anyhow::Result; use keyring::Error as KeyringError; use pretty_assertions::assert_eq; + use std::sync::Barrier; use std::sync::Mutex; use std::sync::MutexGuard; use std::sync::OnceLock; use std::sync::PoisonError; + use std::sync::mpsc; + use std::thread; use tempfile::tempdir; use codex_keyring_store::tests::MockKeyringStore; @@ -999,6 +1120,97 @@ mod tests { assert_eq!(tokens.token_response.0.expires_in(), Some(Duration::ZERO)); } + #[test] + fn persisted_generation_ignores_reload_expires_in_drift() -> Result<()> { + let _env = TempCodexHome::new(); + let mut tokens = sample_tokens(); + tokens + .token_response + .0 + .set_expires_in(Some(&Duration::from_secs(7200))); + + super::save_oauth_tokens_to_file(&tokens)?; + let loaded = super::load_oauth_tokens_from_file(&tokens.server_name, &tokens.url)? + .expect("saved tokens should reload"); + + assert_ne!( + loaded.token_response.0.expires_in(), + tokens.token_response.0.expires_in() + ); + assert_ne!(loaded, tokens); + assert!(super::same_persisted_generation(&loaded, &tokens)); + Ok(()) + } + + #[test] + fn fallback_store_serializes_writes_for_different_servers() -> Result<()> { + let _env = TempCodexHome::new(); + let first = sample_tokens(); + let mut second = sample_tokens(); + second.server_name = "other-server".to_string(); + second.url = "https://other.example.test".to_string(); + second.token_response.0 = OAuthTokenResponse::new( + AccessToken::new("other-access-token".to_string()), + BasicTokenType::Bearer, + VendorExtraTokenFields::default(), + ); + + let fallback_lock = super::acquire_fallback_store_lock()?; + let barrier = Arc::new(Barrier::new(3)); + let (tx, rx) = mpsc::channel(); + let handles = [first.clone(), second.clone()] + .into_iter() + .map(|tokens| { + let barrier = Arc::clone(&barrier); + let tx = tx.clone(); + thread::spawn(move || { + barrier.wait(); + tx.send(super::save_oauth_tokens_to_file(&tokens)) + .expect("send save result"); + }) + }) + .collect::>(); + drop(tx); + + barrier.wait(); + thread::sleep(Duration::from_millis(50)); + assert!(matches!(rx.try_recv(), Err(mpsc::TryRecvError::Empty))); + + drop(fallback_lock); + for _ in 0..2 { + rx.recv_timeout(Duration::from_secs(5))??; + } + for handle in handles { + handle.join().expect("fallback writer should not panic"); + } + + let store = super::read_fallback_file()?.expect("fallback file should load"); + let first_key = super::compute_store_key(&first.server_name, &first.url)?; + let second_key = super::compute_store_key(&second.server_name, &second.url)?; + assert_eq!( + store + .keys() + .cloned() + .collect::>(), + [first_key, second_key].into_iter().collect() + ); + Ok(()) + } + + #[test] + fn server_generation_lock_is_shared_across_codex_homes() -> Result<()> { + let _env = TempCodexHome::new(); + let first = super::oauth_server_lock_path("server", "https://example.test")?; + let other_codex_home = tempdir()?; + unsafe { + std::env::set_var("CODEX_HOME", other_codex_home.path()); + } + let second = super::oauth_server_lock_path("server", "https://example.test")?; + + assert_eq!(first, second); + Ok(()) + } + fn assert_tokens_match_without_expiry( actual: &StoredOAuthTokens, expected: &StoredOAuthTokens, diff --git a/codex-rs/rmcp-client/src/perform_oauth_login.rs b/codex-rs/rmcp-client/src/perform_oauth_login.rs index a416b3d9b0..c1e8fa11b5 100644 --- a/codex-rs/rmcp-client/src/perform_oauth_login.rs +++ b/codex-rs/rmcp-client/src/perform_oauth_login.rs @@ -26,7 +26,7 @@ use urlencoding::decode; use crate::StoredOAuthTokens; use crate::WrappedOAuthTokenResponse; use crate::oauth::compute_expires_at_millis; -use crate::save_oauth_tokens; +use crate::oauth::save_oauth_tokens_async; use crate::utils::apply_default_headers; use crate::utils::build_default_headers; use codex_config::types::OAuthCredentialsStoreMode; @@ -571,7 +571,7 @@ impl OauthLoginFlow { token_response: WrappedOAuthTokenResponse(credentials), expires_at, }; - save_oauth_tokens(&self.server_name, &stored, self.store_mode)?; + save_oauth_tokens_async(&self.server_name, &stored, self.store_mode).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 cc7c43907d..b32f7f9e74 100644 --- a/codex-rs/rmcp-client/tests/streamable_http_oauth_startup.rs +++ b/codex-rs/rmcp-client/tests/streamable_http_oauth_startup.rs @@ -9,6 +9,7 @@ use codex_exec_server::Environment; use codex_rmcp_client::RmcpClient; use codex_rmcp_client::StoredOAuthTokens; use codex_rmcp_client::WrappedOAuthTokenResponse; +use codex_rmcp_client::delete_oauth_tokens; use codex_rmcp_client::save_oauth_tokens; use oauth2::AccessToken; use oauth2::RefreshToken; @@ -24,18 +25,16 @@ 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-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 CHILD_SERVER_URL_ENV: &str = "MCP_TEST_OAUTH_STARTUP_SERVER_URL"; const EXTERNAL_UPDATE_CHILD_SERVER_URL_ENV: &str = "MCP_TEST_OAUTH_EXTERNAL_UPDATE_SERVER_URL"; +const REFRESH_CONTENTION_SERVER_URL_ENV: &str = "MCP_TEST_OAUTH_REFRESH_CONTENTION_SERVER_URL"; +const REFRESH_CONTENTION_ATTEMPT_ENV: &str = "MCP_TEST_OAUTH_REFRESH_CONTENTION_ATTEMPT"; +const REFRESH_CONTENTION_DONE_ENV: &str = "MCP_TEST_OAUTH_REFRESH_CONTENTION_DONE"; const OLD_ACCESS_TOKEN: &str = "old-access-token"; const OLD_REFRESH_TOKEN: &str = "old-refresh-token"; const EXTERNAL_ACCESS_TOKEN: &str = "external-access-token"; @@ -44,186 +43,22 @@ const STALE_REFRESHED_ACCESS_TOKEN: &str = "stale-refreshed-access-token"; const STALE_ROTATED_REFRESH_TOKEN: &str = "stale-rotated-refresh-token"; #[tokio::test(flavor = "multi_thread", worker_threads = 1)] -async fn refreshes_expired_persisted_token_before_initialize() -> anyhow::Result<()> { +async fn running_client_adopts_external_oauth_update_and_deletion() -> 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": [""], - }))) - .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={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, - }))) - .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-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(2) - .mount(&server) - .await; - - let codex_home = TempDir::new()?; - let server_url = format!("{}/mcp", server.uri()); - - // Credential storage resolves CODEX_HOME from the process environment. - // Run the client half of the test in an ignored helper test so it can use - // an isolated home without mutating the parent test runner's environment. - let status = Command::new(std::env::current_exe()?) - .args(["oauth_startup_child", "--exact", "--ignored", "--nocapture"]) - .env("CODEX_HOME", codex_home.path()) - .env(CHILD_SERVER_URL_ENV, server_url) - .status() - .await?; - assert!(status.success(), "OAuth startup child failed: {status}"); - server.verify().await; - Ok(()) -} - -#[tokio::test(flavor = "multi_thread", worker_threads = 1)] -#[ignore = "spawned by refreshes_expired_persisted_token_before_initialize"] -async fn oauth_startup_child() -> anyhow::Result<()> { - let server_url = std::env::var(CHILD_SERVER_URL_ENV)?; - - // Save an expired access token with a valid refresh token so startup must - // refresh before sending the initialize request. - 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.clone(), - client_id: "test-client-id".to_string(), - token_response: WrappedOAuthTokenResponse(response), - expires_at: Some(0), - }; - save_oauth_tokens(SERVER_NAME, &tokens, OAuthCredentialsStoreMode::File)?; - - // This mirrors create_client's transport and initialization setup, except - // it omits the direct bearer token. Supplying that token would bypass the - // persisted OAuth credentials and the startup refresh under test. - 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?; - Ok(()) -} - -#[tokio::test(flavor = "multi_thread", worker_threads = 1)] -async fn running_client_adopts_external_oauth_update_before_next_operation() -> 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": [""], - }))) - .expect(1) - .mount(&server) - .await; + mount_oauth_metadata(&server).await; Mock::given(method("POST")) .and(path("/mcp")) .respond_with(|request: &Request| { - let body: Value = request.body_json().expect("valid JSON-RPC request"); + let Ok(body) = request.body_json::() else { + return ResponseTemplate::new(400); + }; let request_method = body.get("method").and_then(Value::as_str); let expected_access_token = match request_method { Some("tools/call") => EXTERNAL_ACCESS_TOKEN, _ => OLD_ACCESS_TOKEN, }; - let authorization = request - .headers - .get("authorization") - .and_then(|value| value.to_str().ok()); - if authorization != Some(format!("Bearer {expected_access_token}").as_str()) { - return ResponseTemplate::new(401); - } - - match request_method { - 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": { - "tools": {}, - }, - "serverInfo": { - "name": "oauth-external-update-test", - "version": "0.0.0-test", - }, - }, - })), - Some("notifications/initialized") => ResponseTemplate::new(202), - Some("tools/call") => ResponseTemplate::new(200).set_body_json(json!({ - "jsonrpc": "2.0", - "id": body.get("id").cloned().unwrap_or(Value::Null), - "result": { - "content": [], - "isError": false, - }, - })), - method => ResponseTemplate::new(400) - .set_body_string(format!("unexpected JSON-RPC method: {method:?}")), - } + oauth_mcp_response(request, expected_access_token) }) .expect(3) .mount(&server) @@ -251,7 +86,7 @@ async fn running_client_adopts_external_oauth_update_before_next_operation() -> } #[tokio::test(flavor = "multi_thread", worker_threads = 1)] -#[ignore = "spawned by running_client_adopts_external_oauth_update_before_next_operation"] +#[ignore = "spawned by running_client_adopts_external_oauth_update_and_deletion"] async fn oauth_external_update_child() -> anyhow::Result<()> { let server_url = std::env::var(EXTERNAL_UPDATE_CHILD_SERVER_URL_ENV)?; save_test_tokens( @@ -261,18 +96,7 @@ async fn oauth_external_update_child() -> anyhow::Result<()> { future_expiry()?, )?; - 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?; + let client = initialized_oauth_client(&server_url).await?; save_test_tokens( &server_url, @@ -291,13 +115,105 @@ async fn oauth_external_update_child() -> anyhow::Result<()> { .await?; assert_persisted_tokens(EXTERNAL_ACCESS_TOKEN, EXTERNAL_REFRESH_TOKEN)?; + let status = Command::new(std::env::current_exe()?) + .args([ + "oauth_external_delete_writer_child", + "--exact", + "--ignored", + "--nocapture", + ]) + .env("CODEX_HOME", std::env::var("CODEX_HOME")?) + .env(EXTERNAL_UPDATE_CHILD_SERVER_URL_ENV, &server_url) + .status() + .await?; + assert!( + status.success(), + "OAuth external delete writer failed: {status}" + ); + + client + .call_tool( + "call-after-external-delete".to_string(), + /*arguments*/ None, + /*meta*/ None, + Some(Duration::from_secs(5)), + ) + .await + .expect_err("externally deleted credentials must not be reused"); + assert!( + !std::path::Path::new(&std::env::var("CODEX_HOME")?) + .join(".credentials.json") + .exists() + ); Ok(()) } #[tokio::test(flavor = "multi_thread", worker_threads = 1)] -async fn refresh_does_not_overwrite_external_oauth_update() -> anyhow::Result<()> { +#[ignore = "spawned by oauth_external_update_child"] +async fn oauth_external_delete_writer_child() -> anyhow::Result<()> { + let server_url = std::env::var(EXTERNAL_UPDATE_CHILD_SERVER_URL_ENV)?; + assert!(delete_oauth_tokens( + SERVER_NAME, + &server_url, + OAuthCredentialsStoreMode::File, + )?); + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn refresh_holds_generation_lock_against_external_oauth_update() -> anyhow::Result<()> { + let server = MockServer::start().await; + let server_url = format!("{}/mcp", server.uri()); + mount_oauth_metadata(&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={OLD_REFRESH_TOKEN}" + ))) + .respond_with( + ResponseTemplate::new(200) + .set_delay(Duration::from_millis(1500)) + .set_body_json(json!({ + "access_token": STALE_REFRESHED_ACCESS_TOKEN, + "token_type": "Bearer", + "expires_in": 7200, + "refresh_token": STALE_ROTATED_REFRESH_TOKEN, + })), + ) + .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); + }; + let request_method = body.get("method").and_then(Value::as_str); + let expected_access_token = match ( + request_method, + body.pointer("/params/name").and_then(Value::as_str), + ) { + (Some("initialize" | "notifications/initialized"), _) => OLD_ACCESS_TOKEN, + (Some("tools/call"), Some("trigger-refresh-contention")) => { + STALE_REFRESHED_ACCESS_TOKEN + } + (Some("tools/call"), Some("call-after-refresh-contention")) => { + EXTERNAL_ACCESS_TOKEN + } + _ => OLD_ACCESS_TOKEN, + }; + oauth_mcp_response(request, expected_access_token) + }) + .expect(4) + .mount(&server) + .await; + let codex_home = TempDir::new()?; - let status = Command::new(std::env::current_exe()?) + let attempt_path = codex_home.path().join("external-update-attempt"); + let done_path = codex_home.path().join("external-update-done"); + let mut live_client = Command::new(std::env::current_exe()?) .args([ "oauth_external_update_during_refresh_child", "--exact", @@ -305,20 +221,104 @@ async fn refresh_does_not_overwrite_external_oauth_update() -> anyhow::Result<() "--nocapture", ]) .env("CODEX_HOME", codex_home.path()) - .status() - .await?; + .env(REFRESH_CONTENTION_SERVER_URL_ENV, &server_url) + .env(REFRESH_CONTENTION_DONE_ENV, &done_path) + .spawn()?; + + wait_for_request_path(&server, "/oauth/token").await?; + let mut external_writer = Command::new(std::env::current_exe()?) + .args([ + "oauth_external_update_writer_child", + "--exact", + "--ignored", + "--nocapture", + ]) + .env("CODEX_HOME", codex_home.path()) + .env(REFRESH_CONTENTION_SERVER_URL_ENV, &server_url) + .env(REFRESH_CONTENTION_ATTEMPT_ENV, &attempt_path) + .env(REFRESH_CONTENTION_DONE_ENV, &done_path) + .spawn()?; + + wait_for_path(&attempt_path).await?; + tokio::time::sleep(Duration::from_millis(100)).await; assert!( - status.success(), - "OAuth external update during refresh child failed: {status}" + !done_path.exists(), + "external writer should block while refresh owns the generation lock" ); + assert!( + external_writer.try_wait()?.is_none(), + "external writer exited before the refresh released its generation lock" + ); + + let writer_status = external_writer.wait().await?; + assert!( + writer_status.success(), + "OAuth external update writer failed: {writer_status}" + ); + let live_status = live_client.wait().await?; + assert!( + live_status.success(), + "OAuth refresh contention child failed: {live_status}" + ); + + server.verify().await; Ok(()) } #[tokio::test(flavor = "multi_thread", worker_threads = 1)] -#[ignore = "spawned by refresh_does_not_overwrite_external_oauth_update"] +#[ignore = "spawned by refresh_holds_generation_lock_against_external_oauth_update"] async fn oauth_external_update_during_refresh_child() -> anyhow::Result<()> { - let server = MockServer::start().await; - let server_url = format!("{}/mcp", server.uri()); + let server_url = std::env::var(REFRESH_CONTENTION_SERVER_URL_ENV)?; + let done_path = std::path::PathBuf::from(std::env::var(REFRESH_CONTENTION_DONE_ENV)?); + save_test_tokens( + &server_url, + OLD_ACCESS_TOKEN, + OLD_REFRESH_TOKEN, + future_expiry()?, + )?; + + let client = initialized_oauth_client(&server_url).await?; + save_test_tokens(&server_url, OLD_ACCESS_TOKEN, OLD_REFRESH_TOKEN, 0)?; + client + .call_tool( + "trigger-refresh-contention".to_string(), + /*arguments*/ None, + /*meta*/ None, + Some(Duration::from_secs(5)), + ) + .await?; + wait_for_path(&done_path).await?; + client + .call_tool( + "call-after-refresh-contention".to_string(), + /*arguments*/ None, + /*meta*/ None, + Some(Duration::from_secs(5)), + ) + .await?; + + assert_persisted_tokens(EXTERNAL_ACCESS_TOKEN, EXTERNAL_REFRESH_TOKEN)?; + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +#[ignore = "spawned by refresh_holds_generation_lock_against_external_oauth_update"] +async fn oauth_external_update_writer_child() -> anyhow::Result<()> { + let server_url = std::env::var(REFRESH_CONTENTION_SERVER_URL_ENV)?; + let attempt_path = std::env::var(REFRESH_CONTENTION_ATTEMPT_ENV)?; + let done_path = std::env::var(REFRESH_CONTENTION_DONE_ENV)?; + std::fs::write(attempt_path, b"attempting")?; + save_test_tokens( + &server_url, + EXTERNAL_ACCESS_TOKEN, + EXTERNAL_REFRESH_TOKEN, + future_expiry()?, + )?; + std::fs::write(done_path, b"done")?; + Ok(()) +} + +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!({ @@ -327,91 +327,14 @@ async fn oauth_external_update_during_refresh_child() -> anyhow::Result<()> { "scopes_supported": [""], }))) .expect(1) - .mount(&server) + .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={OLD_REFRESH_TOKEN}" - ))) - .respond_with({ - let server_url = server_url.clone(); - move |_request: &Request| { - save_test_tokens( - &server_url, - EXTERNAL_ACCESS_TOKEN, - EXTERNAL_REFRESH_TOKEN, - future_expiry().expect("future expiry"), - ) - .expect("save externally updated credentials"); - ResponseTemplate::new(200).set_body_json(json!({ - "access_token": STALE_REFRESHED_ACCESS_TOKEN, - "token_type": "Bearer", - "expires_in": 7200, - "refresh_token": STALE_ROTATED_REFRESH_TOKEN, - })) - } - }) - .expect(1) - .mount(&server) - .await; - Mock::given(method("POST")) - .and(path("/mcp")) - .respond_with(|request: &Request| { - let body: Value = request.body_json().expect("valid JSON-RPC request"); - let request_method = body.get("method").and_then(Value::as_str); - let expected_access_token = match request_method { - Some("tools/call") => EXTERNAL_ACCESS_TOKEN, - _ => STALE_REFRESHED_ACCESS_TOKEN, - }; - let authorization = request - .headers - .get("authorization") - .and_then(|value| value.to_str().ok()); - if authorization != Some(format!("Bearer {expected_access_token}").as_str()) { - return ResponseTemplate::new(401); - } +} - match request_method { - 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": { - "tools": {}, - }, - "serverInfo": { - "name": "oauth-external-refresh-test", - "version": "0.0.0-test", - }, - }, - })), - Some("notifications/initialized") => ResponseTemplate::new(202), - Some("tools/call") => ResponseTemplate::new(200).set_body_json(json!({ - "jsonrpc": "2.0", - "id": body.get("id").cloned().unwrap_or(Value::Null), - "result": { - "content": [], - "isError": false, - }, - })), - method => ResponseTemplate::new(400) - .set_body_string(format!("unexpected JSON-RPC method: {method:?}")), - } - }) - .expect(3) - .mount(&server) - .await; - - save_test_tokens(&server_url, OLD_ACCESS_TOKEN, OLD_REFRESH_TOKEN, 0)?; +async fn initialized_oauth_client(server_url: &str) -> anyhow::Result { let client = RmcpClient::new_streamable_http_client( SERVER_NAME, - &server_url, + server_url, /*bearer_token*/ None, /*http_headers*/ None, /*env_http_headers*/ None, @@ -421,17 +344,79 @@ async fn oauth_external_update_during_refresh_child() -> anyhow::Result<()> { ) .await?; initialize_client(&client).await?; - client - .call_tool( - "call-after-refresh-race".to_string(), - /*arguments*/ None, - /*meta*/ None, - Some(Duration::from_secs(5)), - ) - .await?; + Ok(client) +} - assert_persisted_tokens(EXTERNAL_ACCESS_TOKEN, EXTERNAL_REFRESH_TOKEN)?; - server.verify().await; +fn oauth_mcp_response(request: &Request, expected_access_token: &str) -> ResponseTemplate { + let authorization = request + .headers + .get("authorization") + .and_then(|value| value.to_str().ok()); + if authorization != Some(format!("Bearer {expected_access_token}").as_str()) { + return ResponseTemplate::new(401); + } + + let Ok(body) = request.body_json::() else { + return ResponseTemplate::new(400); + }; + 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": { + "tools": {}, + }, + "serverInfo": { + "name": "oauth-external-credential-test", + "version": "0.0.0-test", + }, + }, + })), + Some("notifications/initialized") => ResponseTemplate::new(202), + Some("tools/call") => ResponseTemplate::new(200).set_body_json(json!({ + "jsonrpc": "2.0", + "id": body.get("id").cloned().unwrap_or(Value::Null), + "result": { + "content": [], + "isError": false, + }, + })), + method => ResponseTemplate::new(400) + .set_body_string(format!("unexpected JSON-RPC method: {method:?}")), + } +} + +async fn wait_for_request_path(server: &MockServer, expected_path: &str) -> anyhow::Result<()> { + tokio::time::timeout(Duration::from_secs(5), async { + loop { + if server + .received_requests() + .await + .unwrap_or_default() + .iter() + .any(|request| request.url.path() == expected_path) + { + return; + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await?; + Ok(()) +} + +async fn wait_for_path(path: &std::path::Path) -> anyhow::Result<()> { + tokio::time::timeout(Duration::from_secs(5), async { + while !path.exists() { + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await?; Ok(()) }