diff --git a/codex-rs/rmcp-client/src/oauth.rs b/codex-rs/rmcp-client/src/oauth.rs index e9a9316a40..5c25893f6a 100644 --- a/codex-rs/rmcp-client/src/oauth.rs +++ b/codex-rs/rmcp-client/src/oauth.rs @@ -247,16 +247,23 @@ fn save_oauth_tokens_with_keyring( server_name: &str, tokens: &StoredOAuthTokens, ) -> Result<()> { + let key = save_oauth_tokens_to_keyring(keyring_store, server_name, tokens)?; + if let Err(error) = delete_oauth_tokens_from_file(&key) { + warn!("failed to remove OAuth tokens from fallback storage: {error:?}"); + } + Ok(()) +} + +fn save_oauth_tokens_to_keyring( + keyring_store: &K, + server_name: &str, + tokens: &StoredOAuthTokens, +) -> Result { let serialized = serde_json::to_string(tokens).context("failed to serialize OAuth tokens")?; let key = compute_store_key(server_name, &tokens.url)?; match keyring_store.save(KEYRING_SERVICE, &key, &serialized) { - Ok(()) => { - if let Err(error) = delete_oauth_tokens_from_file(&key) { - warn!("failed to remove OAuth tokens from fallback storage: {error:?}"); - } - Ok(()) - } + Ok(()) => Ok(key), Err(error) => { let message = format!( "failed to write OAuth tokens to keyring: {}", @@ -284,6 +291,46 @@ fn save_oauth_tokens_with_keyring_with_fallback_to_file( } } +async fn save_oauth_tokens_for_runtime( + server_name: &str, + tokens: &StoredOAuthTokens, + store_mode: OAuthCredentialsStoreMode, +) -> Result<()> { + let keyring_store = DefaultKeyringStore; + match store_mode { + OAuthCredentialsStoreMode::Auto => { + match save_oauth_tokens_with_keyring_async(&keyring_store, server_name, tokens).await { + Ok(()) => Ok(()), + Err(error) => { + let message = error.to_string(); + warn!("falling back to file storage for OAuth tokens: {message}"); + save_oauth_tokens_to_file_async(tokens) + .await + .with_context(|| { + format!("failed to write OAuth tokens to keyring: {message}") + }) + } + } + } + OAuthCredentialsStoreMode::File => save_oauth_tokens_to_file_async(tokens).await, + OAuthCredentialsStoreMode::Keyring => { + save_oauth_tokens_with_keyring_async(&keyring_store, server_name, tokens).await + } + } +} + +async fn save_oauth_tokens_with_keyring_async( + keyring_store: &K, + server_name: &str, + tokens: &StoredOAuthTokens, +) -> Result<()> { + let key = save_oauth_tokens_to_keyring(keyring_store, server_name, tokens)?; + if let Err(error) = delete_oauth_tokens_from_file_async(&key).await { + warn!("failed to remove OAuth tokens from fallback storage: {error:?}"); + } + Ok(()) +} + pub fn delete_oauth_tokens( server_name: &str, url: &str, @@ -300,7 +347,29 @@ fn delete_oauth_tokens_from_keyring_and_file( url: &str, ) -> Result { let key = compute_store_key(server_name, url)?; - let keyring_result = keyring_store.delete(KEYRING_SERVICE, &key); + let keyring_removed = delete_oauth_tokens_from_keyring(keyring_store, store_mode, &key)?; + let file_removed = delete_oauth_tokens_from_file(&key)?; + Ok(keyring_removed || file_removed) +} + +async fn delete_oauth_tokens_for_runtime( + server_name: &str, + url: &str, + store_mode: OAuthCredentialsStoreMode, +) -> Result { + let keyring_store = DefaultKeyringStore; + let key = compute_store_key(server_name, url)?; + let keyring_removed = delete_oauth_tokens_from_keyring(&keyring_store, store_mode, &key)?; + let file_removed = delete_oauth_tokens_from_file_async(&key).await?; + Ok(keyring_removed || file_removed) +} + +fn delete_oauth_tokens_from_keyring( + keyring_store: &K, + store_mode: OAuthCredentialsStoreMode, + key: &str, +) -> Result { + let keyring_result = keyring_store.delete(KEYRING_SERVICE, key); let keyring_removed = match keyring_result { Ok(removed) => removed, Err(error) => { @@ -315,9 +384,7 @@ fn delete_oauth_tokens_from_keyring_and_file( } } }; - - let file_removed = delete_oauth_tokens_from_file(&key)?; - Ok(keyring_removed || file_removed) + Ok(keyring_removed) } #[derive(Clone)] @@ -389,18 +456,24 @@ impl OAuthPersistor { expires_at, }; if last_credentials.as_ref() != Some(&stored) { - save_oauth_tokens(&self.inner.server_name, &stored, self.inner.store_mode)?; + save_oauth_tokens_for_runtime( + &self.inner.server_name, + &stored, + self.inner.store_mode, + ) + .await?; *last_credentials = 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_for_runtime( &self.inner.server_name, &self.inner.url, self.inner.store_mode, ) + .await { warn!( "failed to remove OAuth tokens for server {}: {error}", @@ -744,6 +817,27 @@ fn acquire_fallback_write_lock() -> Result { Ok(file) } +async fn acquire_fallback_write_lock_async() -> Result { + let path = fallback_lock_path()?; + let file = open_oauth_lock_file(&path)?; + loop { + match file.try_lock() { + Ok(()) => return Ok(file), + Err(fs::TryLockError::WouldBlock) => { + tokio::time::sleep(Duration::from_millis(10)).await; + } + Err(error) => { + return Err(error).with_context(|| { + format!( + "failed to acquire OAuth fallback write lock at {}", + path.display() + ) + }); + } + } + } +} + fn acquire_fallback_read_lock() -> Result { let path = fallback_lock_path()?; let file = open_oauth_lock_file(&path)?; @@ -931,6 +1025,15 @@ fn load_oauth_tokens_from_file_unlocked( fn save_oauth_tokens_to_file(tokens: &StoredOAuthTokens) -> Result<()> { let _write_lock = acquire_fallback_write_lock()?; + save_oauth_tokens_to_file_unlocked(tokens) +} + +async fn save_oauth_tokens_to_file_async(tokens: &StoredOAuthTokens) -> Result<()> { + let _write_lock = acquire_fallback_write_lock_async().await?; + save_oauth_tokens_to_file_unlocked(tokens) +} + +fn save_oauth_tokens_to_file_unlocked(tokens: &StoredOAuthTokens) -> Result<()> { let key = compute_store_key(&tokens.server_name, &tokens.url)?; let mut store = read_fallback_file_unlocked()?.unwrap_or_default(); @@ -964,6 +1067,11 @@ fn delete_oauth_tokens_from_file(key: &str) -> Result { delete_oauth_tokens_from_file_unlocked(key) } +async fn delete_oauth_tokens_from_file_async(key: &str) -> Result { + let _write_lock = acquire_fallback_write_lock_async().await?; + delete_oauth_tokens_from_file_unlocked(key) +} + fn delete_oauth_tokens_from_file_unlocked(key: &str) -> Result { let mut store = match read_fallback_file_unlocked()? { Some(store) => store, 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 4765771b67..a615bd883b 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 @@ -272,6 +272,61 @@ async fn fallback_read_lock_wait_uses_startup_timeout_budget() -> anyhow::Result Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn fallback_write_lock_wait_uses_startup_timeout_budget() -> anyhow::Result<()> { + let server = MockServer::start().await; + mount_oauth_metadata(&server, /*expected_requests*/ 1).await; + Mock::given(method("POST")) + .and(path("/oauth/token")) + .respond_with(ResponseTemplate::new(200).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*/ 0).await; + + let codex_home = TempDir::new()?; + let control_dir = TempDir::new()?; + let server_url = format!("{}/mcp", server.uri()); + seed_tokens(&codex_home, &server_url, /*expires_at*/ 0).await?; + + let (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 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_client_child( + &codex_home, + &server_url, + &ready_path, + &go_path, + &result_path, + )?; + wait_for_paths(std::slice::from_ref(&ready_path)).await?; + fs::write(&go_path, "go")?; + wait_for_paths(std::slice::from_ref(&result_path)).await?; + assert_eq!( + fs::read_to_string(&result_path)?, + "error:timed out handshaking with MCP server after 5s" + ); + + drop(fallback_lock); + wait_for_children(vec![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; @@ -328,6 +383,62 @@ async fn fallback_read_lock_wait_uses_operation_timeout_budget() -> anyhow::Resu Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn fallback_write_lock_wait_uses_operation_timeout_budget() -> anyhow::Result<()> { + let server = MockServer::start().await; + mount_oauth_metadata(&server, /*expected_requests*/ 1).await; + Mock::given(method("POST")) + .and(path("/oauth/token")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "access_token": NEW_ACCESS_TOKEN, + "token_type": "Bearer", + "expires_in": 31, + "refresh_token": NEW_REFRESH_TOKEN, + }))) + .expect(2) + .mount(&server) + .await; + mount_mcp_server(&server, /*expected_requests*/ 2).await; + + let codex_home = TempDir::new()?; + let control_dir = TempDir::new()?; + let server_url = format!("{}/mcp", server.uri()); + seed_tokens(&codex_home, &server_url, /*expires_at*/ 0).await?; + + let ready_path = control_dir.path().join("client.ready"); + let go_path = control_dir.path().join("client.go"); + let result_path = control_dir.path().join("client.result"); + let child = spawn_operation_timeout_child( + &codex_home, + &server_url, + &ready_path, + &go_path, + &result_path, + )?; + wait_for_paths(std::slice::from_ref(&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(&go_path, "go")?; + wait_for_children(vec![child]).await?; + assert_eq!( + fs::read_to_string(result_path)?, + format!("error:timed out awaiting tools/list after {OPERATION_TIMEOUT:?}") + ); + + drop(fallback_lock); + 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;