From 20e4e38d52dbb0080a3c6a1eb75ad4eefefb8617 Mon Sep 17 00:00:00 2001 From: Casey Chow Date: Thu, 4 Jun 2026 21:32:10 +0000 Subject: [PATCH] refactor(rmcp-client): simplify oauth refresh coordination --- codex-rs/rmcp-client/src/lib.rs | 1 + codex-rs/rmcp-client/src/oauth.rs | 161 ++++-------------- codex-rs/rmcp-client/src/oauth_lock.rs | 103 +++++++++++ .../tests/streamable_http_recovery.rs | 140 +++------------ .../tests/streamable_http_test_support.rs | 6 +- 5 files changed, 165 insertions(+), 246 deletions(-) create mode 100644 codex-rs/rmcp-client/src/oauth_lock.rs diff --git a/codex-rs/rmcp-client/src/lib.rs b/codex-rs/rmcp-client/src/lib.rs index 7bb461f241..677cc96fea 100644 --- a/codex-rs/rmcp-client/src/lib.rs +++ b/codex-rs/rmcp-client/src/lib.rs @@ -5,6 +5,7 @@ mod http_client_adapter; mod in_process_transport; mod logging_client_handler; mod oauth; +mod oauth_lock; mod perform_oauth_login; mod program_resolver; mod rmcp_client; diff --git a/codex-rs/rmcp-client/src/oauth.rs b/codex-rs/rmcp-client/src/oauth.rs index 38ab2c6734..d049c47b24 100644 --- a/codex-rs/rmcp-client/src/oauth.rs +++ b/codex-rs/rmcp-client/src/oauth.rs @@ -35,11 +35,9 @@ use sha2::Digest; use sha2::Sha256; use std::collections::BTreeMap; use std::fs; -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; @@ -47,6 +45,7 @@ use tracing::warn; use codex_keyring_store::DefaultKeyringStore; use codex_keyring_store::KeyringStore; +use codex_utils_home_dir::find_codex_home; use codex_utils_path::write_atomically; use rmcp::transport::auth::AuthError; use rmcp::transport::auth::AuthorizationManager; @@ -54,9 +53,10 @@ use rmcp::transport::auth::CredentialStore; use rmcp::transport::auth::InMemoryCredentialStore; use rmcp::transport::auth::StoredCredentials; use tokio::sync::Mutex; -use tokio::sync::OwnedMutexGuard; -use codex_utils_home_dir::find_codex_home; +use crate::oauth_lock::acquire_fallback_store_lock; +use crate::oauth_lock::acquire_oauth_server_lock; +use crate::oauth_lock::acquire_oauth_server_lock_async; const KEYRING_SERVICE: &str = "Codex MCP Credentials"; const REFRESH_SKEW_MILLIS: u64 = 30_000; @@ -407,16 +407,11 @@ impl OAuthPersistor { } let new_token_response = WrappedOAuthTokenResponse(credentials.clone()); - let same_token = current_credentials - .as_ref() - .map(|prev| prev.token_response == new_token_response) - .unwrap_or(false); - let expires_at = if same_token { - current_credentials - .as_ref() - .and_then(|prev| prev.expires_at) - } else { - compute_expires_at_millis(&credentials) + let expires_at = match current_credentials.as_ref() { + Some(previous) if previous.token_response == new_token_response => { + previous.expires_at + } + _ => compute_expires_at_millis(&credentials), }; let stored = StoredOAuthTokens { server_name: self.inner.server_name.clone(), @@ -470,12 +465,7 @@ impl OAuthPersistor { reason = "AuthorizationManager async access must be serialized through its mutex" )] pub(crate) async fn refresh_if_needed(&self) -> Result<()> { - let expires_at = { - let guard = self.inner.current_credentials.lock().await; - guard.as_ref().and_then(|tokens| tokens.expires_at) - }; - - if !token_needs_refresh(expires_at) { + if !self.current_token_needs_refresh().await { return Ok(()); } @@ -489,20 +479,11 @@ impl OAuthPersistor { return Err(anyhow::anyhow!("Auth required for server")); } - let expires_at = { - let guard = self.inner.current_credentials.lock().await; - guard.as_ref().and_then(|tokens| tokens.expires_at) - }; - - if !token_needs_refresh(expires_at) { + if !self.current_token_needs_refresh().await { return Ok(()); } let previous_credentials = self.inner.current_credentials.lock().await.clone(); - let previous_refresh_token = previous_credentials - .as_ref() - .and_then(|tokens| tokens.token_response.0.refresh_token()) - .map(|token| token.secret().clone()); let mut refreshed_credentials = { let manager = self.inner.authorization_manager.clone(); let guard = manager.lock().await; @@ -526,10 +507,11 @@ impl OAuthPersistor { }; if refreshed_credentials.refresh_token().is_none() - && let Some(refresh_token) = previous_refresh_token && let Some(previous_credentials) = previous_credentials.as_ref() + && let Some(refresh_token) = previous_credentials.token_response.0.refresh_token() { - refreshed_credentials.set_refresh_token(Some(RefreshToken::new(refresh_token))); + refreshed_credentials + .set_refresh_token(Some(RefreshToken::new(refresh_token.secret().clone()))); self.replace_manager_credentials( &previous_credentials.client_id, refreshed_credentials.clone(), @@ -542,6 +524,11 @@ impl OAuthPersistor { Ok(()) } + async fn current_token_needs_refresh(&self) -> bool { + let guard = self.inner.current_credentials.lock().await; + token_needs_refresh(guard.as_ref().and_then(|tokens| tokens.expires_at)) + } + async fn persist_refreshed_credentials_with_retry(&self) { for attempt in 1..=PERSIST_RETRY_ATTEMPTS { match self.persist_if_needed_locked().await { @@ -596,7 +583,11 @@ impl OAuthPersistor { Ok(CredentialReload::Replaced) } None => { - self.clear_manager_credentials().await; + { + let manager = self.inner.authorization_manager.clone(); + let mut guard = manager.lock().await; + guard.set_credential_store(InMemoryCredentialStore::new()); + } { let mut current_credentials = self.inner.current_credentials.lock().await; *current_credentials = None; @@ -613,15 +604,16 @@ impl OAuthPersistor { client_id: &str, token_response: OAuthTokenResponse, ) -> Result<()> { + let scopes = token_response + .scopes() + .map(|scopes| scopes.iter().map(|scope| scope.to_string()).collect()) + .unwrap_or_default(); let store = InMemoryCredentialStore::new(); store .save(StoredCredentials::new( client_id.to_string(), - Some(token_response.clone()), - token_response - .scopes() - .map(|scopes| scopes.iter().map(|scope| scope.to_string()).collect()) - .unwrap_or_default(), + Some(token_response), + scopes, Some( SystemTime::now() .duration_since(UNIX_EPOCH) @@ -637,97 +629,6 @@ 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 OAuthFileLock { - _file: fs::File, - _in_process_guard: Option>, -} - -fn open_oauth_lock_file(path: PathBuf) -> Result { - if let Some(parent) = path.parent() { - fs::create_dir_all(parent)?; - } - Ok(OpenOptions::new() - .read(true) - .write(true) - .create(true) - .truncate(false) - .open(path)?) -} - -fn acquire_oauth_server_lock(server_name: &str, url: &str) -> Result { - let file = open_oauth_lock_file(oauth_server_lock_path(server_name, url)?)?; - file.lock()?; - Ok(OAuthFileLock { - _file: file, - _in_process_guard: None, - }) -} - -async fn acquire_oauth_server_lock_async(server_name: &str, url: &str) -> Result { - let in_process_lock = oauth_server_lock_for(server_name, url); - let in_process_guard = in_process_lock.lock_owned().await; - let path = oauth_server_lock_path(server_name, url)?; - let file_lock = tokio::task::spawn_blocking(move || { - let file = open_oauth_lock_file(path)?; - file.lock()?; - Ok::<_, anyhow::Error>(file) - }) - .await - .context("OAuth credential lock task failed")??; - Ok(OAuthFileLock { - _file: file_lock, - _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 file = open_oauth_lock_file(oauth_lock_dir()?.join("fallback-store.lock"))?; - file.lock()?; - Ok(FallbackStoreLock { - _in_process_guard: in_process_guard, - _file: file, - }) -} - -fn oauth_server_lock_path(server_name: &str, url: &str) -> Result { - let digest = sha_256_prefix(&Value::String(format!("{server_name}\n{url}")))?; - Ok(oauth_lock_dir()?.join(format!("server-{digest}.lock"))) -} - -fn oauth_lock_dir() -> Result { - Ok(find_codex_home()?.join(".mcp-oauth-locks").to_path_buf()) } fn refresh_failure_requires_reauth(message: &str) -> bool { @@ -947,7 +848,7 @@ fn write_fallback_file(store: &FallbackFile) -> Result<()> { Ok(()) } -fn sha_256_prefix(value: &Value) -> Result { +pub(super) fn sha_256_prefix(value: &Value) -> Result { let serialized = serde_json::to_string(&value).context("failed to serialize MCP OAuth key payload")?; let mut hasher = Sha256::new(); diff --git a/codex-rs/rmcp-client/src/oauth_lock.rs b/codex-rs/rmcp-client/src/oauth_lock.rs new file mode 100644 index 0000000000..838af85141 --- /dev/null +++ b/codex-rs/rmcp-client/src/oauth_lock.rs @@ -0,0 +1,103 @@ +use std::collections::BTreeMap; +use std::fs; +use std::fs::OpenOptions; +use std::path::PathBuf; +use std::sync::Arc; +use std::sync::OnceLock; + +use anyhow::Context; +use anyhow::Result; +use codex_utils_home_dir::find_codex_home; +use serde_json::Value; +use tokio::sync::Mutex; +use tokio::sync::OwnedMutexGuard; + +use crate::oauth::sha_256_prefix; + +pub(super) struct OAuthFileLock { + _file: fs::File, + _in_process_guard: Option>, +} + +pub(super) struct FallbackStoreLock { + _in_process_guard: std::sync::MutexGuard<'static, ()>, + _file: fs::File, +} + +pub(super) fn acquire_oauth_server_lock(server_name: &str, url: &str) -> Result { + let file = lock_oauth_file(oauth_server_lock_path(server_name, url)?)?; + Ok(OAuthFileLock { + _file: file, + _in_process_guard: None, + }) +} + +pub(super) async fn acquire_oauth_server_lock_async( + server_name: &str, + url: &str, +) -> Result { + let in_process_lock = oauth_server_lock_for(server_name, url); + let in_process_guard = in_process_lock.lock_owned().await; + let path = oauth_server_lock_path(server_name, url)?; + let file_lock = tokio::task::spawn_blocking(move || lock_oauth_file(path)) + .await + .context("OAuth credential lock task failed")??; + Ok(OAuthFileLock { + _file: file_lock, + _in_process_guard: Some(in_process_guard), + }) +} + +pub(super) 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 file = lock_oauth_file(oauth_lock_dir()?.join("fallback-store.lock"))?; + Ok(FallbackStoreLock { + _in_process_guard: in_process_guard, + _file: file, + }) +} + +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() +} + +fn lock_oauth_file(path: PathBuf) -> Result { + let file = open_oauth_lock_file(path)?; + file.lock()?; + Ok(file) +} + +fn open_oauth_lock_file(path: PathBuf) -> Result { + if let Some(parent) = path.parent() { + fs::create_dir_all(parent)?; + } + Ok(OpenOptions::new() + .read(true) + .write(true) + .create(true) + .truncate(false) + .open(path)?) +} + +fn oauth_server_lock_path(server_name: &str, url: &str) -> Result { + let digest = sha_256_prefix(&Value::String(format!("{server_name}\n{url}")))?; + Ok(oauth_lock_dir()?.join(format!("server-{digest}.lock"))) +} + +fn oauth_lock_dir() -> Result { + Ok(find_codex_home()?.join(".mcp-oauth-locks").to_path_buf()) +} diff --git a/codex-rs/rmcp-client/tests/streamable_http_recovery.rs b/codex-rs/rmcp-client/tests/streamable_http_recovery.rs index cfea26af24..5f26ff8a94 100644 --- a/codex-rs/rmcp-client/tests/streamable_http_recovery.rs +++ b/codex-rs/rmcp-client/tests/streamable_http_recovery.rs @@ -193,17 +193,7 @@ async fn streamable_http_oauth_refreshes_expired_token_before_initialize() -> an let server_url = format!("{base_url}/mcp"); save_expired_oauth_tokens(&server_url).await?; - let client = RmcpClient::new_streamable_http_client( - OAUTH_TEST_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?; + let client = create_oauth_file_client(&server_url).await?; initialize_client(&client).await?; let result = call_echo_tool(&client, "after-refresh").await?; @@ -233,17 +223,7 @@ async fn streamable_http_oauth_preserves_refresh_token_when_refresh_response_omi let server_url = format!("{base_url}/mcp"); save_expired_oauth_tokens(&server_url).await?; - let client = RmcpClient::new_streamable_http_client( - OAUTH_TEST_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?; + let client = create_oauth_file_client(&server_url).await?; initialize_client(&client).await?; let credentials = std::fs::read_to_string(codex_home.dir.path().join(".credentials.json"))?; @@ -272,28 +252,8 @@ async fn streamable_http_oauth_concurrent_initializes_share_refreshed_credential let server_url = format!("{base_url}/mcp"); save_expired_oauth_tokens(&server_url).await?; - let client_a = RmcpClient::new_streamable_http_client( - OAUTH_TEST_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?; - let client_b = RmcpClient::new_streamable_http_client( - OAUTH_TEST_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?; + let client_a = create_oauth_file_client(&server_url).await?; + let client_b = create_oauth_file_client(&server_url).await?; let (initialized_a, initialized_b) = tokio::join!(initialize_client(&client_a), initialize_client(&client_b)); @@ -411,17 +371,7 @@ async fn streamable_http_oauth_unexpired_token_does_not_require_writable_codex_h )?; let result = async { - let client = RmcpClient::new_streamable_http_client( - OAUTH_TEST_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?; + let client = create_oauth_file_client(&server_url).await?; initialize_client(&client).await } .await; @@ -447,19 +397,9 @@ async fn streamable_http_oauth_refresh_timeout_keeps_refresh_running() -> anyhow let server_url = format!("{base_url}/mcp"); save_expired_oauth_tokens(&server_url).await?; - let client = RmcpClient::new_streamable_http_client( - OAUTH_TEST_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?; + let client = create_oauth_file_client(&server_url).await?; - let error = initialize_client_with_timeout(&client, Some(Duration::from_millis(50))) + let error = initialize_client_with_timeout(&client, Duration::from_millis(50)) .await .unwrap_err(); assert!( @@ -511,18 +451,8 @@ async fn streamable_http_oauth_logout_wins_against_detached_refresh() -> anyhow: let server_url = format!("{base_url}/mcp"); save_expired_oauth_tokens(&server_url).await?; - let client = RmcpClient::new_streamable_http_client( - OAUTH_TEST_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_with_timeout(&client, Some(Duration::from_millis(50))) + let client = create_oauth_file_client(&server_url).await?; + initialize_client_with_timeout(&client, Duration::from_millis(50)) .await .unwrap_err(); @@ -556,19 +486,9 @@ async fn streamable_http_oauth_refresh_and_initialize_share_timeout_budget() -> let server_url = format!("{base_url}/mcp"); save_expired_oauth_tokens(&server_url).await?; - let client = RmcpClient::new_streamable_http_client( - OAUTH_TEST_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?; + let client = create_oauth_file_client(&server_url).await?; - let error = initialize_client_with_timeout(&client, Some(Duration::from_millis(250))) + let error = initialize_client_with_timeout(&client, Duration::from_millis(250)) .await .unwrap_err(); assert!( @@ -597,17 +517,7 @@ async fn streamable_http_oauth_transient_refresh_failure_does_not_require_login( let server_url = format!("{base_url}/mcp"); save_expired_oauth_tokens(&server_url).await?; - let client = RmcpClient::new_streamable_http_client( - OAUTH_TEST_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?; + let client = create_oauth_file_client(&server_url).await?; let error = initialize_client(&client).await.unwrap_err(); assert!(!error.to_string().contains("Auth required")); @@ -630,17 +540,7 @@ async fn streamable_http_oauth_refresh_failure_reports_auth_required() -> anyhow let server_url = format!("{base_url}/mcp"); save_expired_oauth_tokens(&server_url).await?; - let client = RmcpClient::new_streamable_http_client( - OAUTH_TEST_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?; + let client = create_oauth_file_client(&server_url).await?; let error = initialize_client(&client).await.unwrap_err(); assert!( @@ -718,6 +618,20 @@ async fn save_expired_oauth_tokens(server_url: &str) -> anyhow::Result<()> { .await } +async fn create_oauth_file_client(server_url: &str) -> anyhow::Result { + RmcpClient::new_streamable_http_client( + OAUTH_TEST_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 +} + async fn save_test_oauth_tokens( server_name: &str, server_url: &str, diff --git a/codex-rs/rmcp-client/tests/streamable_http_test_support.rs b/codex-rs/rmcp-client/tests/streamable_http_test_support.rs index aa86b97779..e88fee4088 100644 --- a/codex-rs/rmcp-client/tests/streamable_http_test_support.rs +++ b/codex-rs/rmcp-client/tests/streamable_http_test_support.rs @@ -92,17 +92,17 @@ pub(crate) async fn create_client(base_url: &str) -> anyhow::Result } pub(crate) async fn initialize_client(client: &RmcpClient) -> anyhow::Result<()> { - initialize_client_with_timeout(client, Some(Duration::from_secs(5))).await + initialize_client_with_timeout(client, Duration::from_secs(5)).await } pub(crate) async fn initialize_client_with_timeout( client: &RmcpClient, - timeout: Option, + timeout: Duration, ) -> anyhow::Result<()> { client .initialize( init_params(), - timeout, + Some(timeout), Box::new(|_, _| { async { Ok(ElicitationResponse {