diff --git a/codex-rs/rmcp-client/Cargo.toml b/codex-rs/rmcp-client/Cargo.toml index 517f7270bbe0..a8464f3054d9 100644 --- a/codex-rs/rmcp-client/Cargo.toml +++ b/codex-rs/rmcp-client/Cargo.toml @@ -84,5 +84,10 @@ keyring = { workspace = true, features = ["windows-native"] } [target.'cfg(any(target_os = "freebsd", target_os = "openbsd"))'.dependencies] keyring = { workspace = true, features = ["sync-secret-service"] } +# This test is compiled through `#[path]` inside the inline `oauth::tests` module. Cargo-shear +# cannot resolve that nested module path and otherwise reports the linked file as unlinked. +[package.metadata.cargo-shear] +ignored-paths = ["src/oauth/tests/persistor_tests.rs"] + [lib] doctest = false diff --git a/codex-rs/rmcp-client/src/oauth.rs b/codex-rs/rmcp-client/src/oauth.rs index 601d0f45f4ee..b8736168856c 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; +mod refresh_transaction; mod resolved_store; mod store_lock; @@ -571,34 +573,6 @@ impl OAuthPersistor { Ok(()) } - - #[expect( - clippy::await_holding_invalid_type, - 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.last_credentials.lock().await; - guard.as_ref().and_then(|tokens| tokens.expires_at) - }; - - if !token_needs_refresh(expires_at) { - return Ok(()); - } - - { - let manager = self.inner.authorization_manager.clone(); - let guard = manager.lock().await; - guard.refresh_token().await.with_context(|| { - format!( - "failed to refresh OAuth tokens for server {}", - self.inner.server_name - ) - })?; - } - - self.persist_if_needed().await - } } const FALLBACK_FILENAME: &str = ".credentials.json"; @@ -858,6 +832,8 @@ mod tests { use keyring::Error as KeyringError; use pretty_assertions::assert_eq; use std::sync::Arc; + #[path = "persistor_tests.rs"] + mod persistor_tests; use super::test_support::TempCodexHome; 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 000000000000..9d1ecf9e60cc --- /dev/null +++ b/codex-rs/rmcp-client/src/oauth/refresh_lock.rs @@ -0,0 +1,104 @@ +//! 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 anyhow::anyhow; +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::time::Duration; +use tokio::time::sleep; +use tokio::time::timeout; + +const REFRESH_LOCK_DIR: &str = "mcp-oauth-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_for_server(server_name: &str, url: &str) -> Result { + let store_key = super::compute_store_key(server_name, url)?; + let codex_home = find_codex_home()?; + Self::acquire_in(&codex_home, &store_key, REFRESH_LOCK_ACQUIRE_TIMEOUT) + .await + .with_context(|| format!("failed to acquire OAuth credential lock for {server_name}")) + } + + async fn acquire_in( + codex_home: &Path, + store_key: &str, + acquire_timeout: Duration, + ) -> Result { + // Scope coordination to CODEX_HOME alongside File and Secrets state. Direct keyring + // coordination across homes needs a separate cross-platform rendezvous. + // TODO(stevenlee): define that rendezvous before expanding this lock's scope. + let mut hasher = Sha256::new(); + hasher.update(store_key.as_bytes()); + let path = codex_home + .join(REFRESH_LOCK_DIR) + .join(format!("{:x}.lock", hasher.finalize())); + 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; + 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 + .map_err(|_| { + anyhow!( + "timed out after {acquire_timeout:?} waiting for OAuth refresh lock {}", + path.display() + ) + })? + .with_context(|| format!("failed to lock OAuth refresh lock {}", path.display()))?; + + Ok(Self { _file: file }) + } +} + +#[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 000000000000..6344bd036cbd --- /dev/null +++ b/codex-rs/rmcp-client/src/oauth/refresh_lock_tests.rs @@ -0,0 +1,40 @@ +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 = RefreshCredentialLock::acquire_in( + codex_home.path(), + store_key, + Duration::from_millis(/*millis*/ 50), + ) + .await + .err() + .expect("contending lock acquisition should time out"); + 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/src/oauth/refresh_transaction.rs b/codex-rs/rmcp-client/src/oauth/refresh_transaction.rs new file mode 100644 index 000000000000..9a656ed3cf5f --- /dev/null +++ b/codex-rs/rmcp-client/src/oauth/refresh_transaction.rs @@ -0,0 +1,313 @@ +//! Serialized read-refresh-write transactions for MCP OAuth credentials. + +use std::time::Duration; +use std::time::SystemTime; +use std::time::UNIX_EPOCH; + +use anyhow::Context; +use anyhow::Result; +use codex_keyring_store::DefaultKeyringStore; +use codex_keyring_store::KeyringStore; +use oauth2::TokenResponse; +use rmcp::transport::auth::AuthError; +use rmcp::transport::auth::AuthorizationManager; +use rmcp::transport::auth::CredentialStore as _; +use rmcp::transport::auth::InMemoryCredentialStore; +use rmcp::transport::auth::OAuthTokenResponse; +use rmcp::transport::auth::StoredCredentials; +use tokio::time::timeout; +use tracing::debug; +use tracing::warn; + +use super::OAuthPersistor; +use super::OAuthPersistorInner; +use super::StoredOAuthTokens; +use super::WrappedOAuthTokenResponse; +use super::compute_expires_at_millis; +use super::refresh_lock::RefreshCredentialLock; +use super::token_needs_refresh; + +const REFRESH_REQUEST_TIMEOUT: Duration = Duration::from_secs(45); + +impl OAuthPersistor { + pub(crate) async fn refresh_if_needed(&self) -> Result<()> { + self.refresh_if_needed_in(&DefaultKeyringStore, REFRESH_REQUEST_TIMEOUT) + .await + } + + /// Injects the credential backend and provider timeout for deterministic failure-path tests. + pub(super) async fn refresh_if_needed_in( + &self, + keyring_store: &K, + refresh_request_timeout: Duration, + ) -> Result<()> { + let expires_at = { + let guard = self.inner.last_credentials.lock().await; + guard.as_ref().and_then(|tokens| tokens.expires_at) + }; + + if !token_needs_refresh(expires_at) { + return Ok(()); + } + + let persistor = self.clone(); + let keyring_store = keyring_store.clone(); + // Once the provider can consume a rotating token, caller cancellation must not cancel + // persistence. The owned task continues with independently bounded lock and request waits. + // A provider timeout leaves the outcome unknown and permits a later serialized retry: + // provider grace may recover, otherwise reauthorization is unavoidable. This residual + // risk is preferred to holding the credential lock indefinitely. + let transaction_task = tokio::spawn(async move { + let result = persistor + .refresh_transaction(&keyring_store, refresh_request_timeout) + .await; + + // Keep this summary inside the owned task so caller cancellation cannot suppress it. + if let Err(error) = &result { + warn!( + server_name = %persistor.inner.server_name, + refresh_reason = "expiry", + error = %error, + "MCP OAuth refresh transaction failed" + ); + } + + result + }); + transaction_task.await.with_context(|| { + format!( + "OAuth refresh task failed for server {}", + self.inner.server_name + ) + })? + } + + #[expect( + clippy::await_holding_invalid_type, + reason = "AuthorizationManager async access must be serialized through its Tokio mutex" + )] + #[tracing::instrument( + level = "debug", + skip_all, + fields( + server_name = %self.inner.server_name, + refresh_reason = "expiry", + ), + err + )] + async fn refresh_transaction( + &self, + keyring_store: &K, + refresh_request_timeout: Duration, + ) -> Result<()> { + debug!("waiting for the MCP OAuth credential transaction lock"); + let _lock = + RefreshCredentialLock::acquire_for_server(&self.inner.server_name, &self.inner.url) + .await?; + debug!("acquired the MCP OAuth credential transaction lock"); + + // Stay on the lifecycle-pinned store. A failure is surfaced rather than falling back and + // possibly replaying an older rotating refresh token from the other store. + debug!("rereading authoritative MCP OAuth credentials"); + let latest = self.inner.credential_store.load( + keyring_store, + &self.inner.server_name, + &self.inner.url, + )?; + + // The pre-lock snapshot is only a hint. This locked reread is authoritative, so adopt a + // winner from another process rather than refreshing its predecessor. + let Some(latest) = latest else { + let manager = self.inner.authorization_manager.clone(); + manager + .lock() + .await + .set_credential_store(InMemoryCredentialStore::new()); + *self.inner.last_credentials.lock().await = None; + return Err(AuthError::AuthorizationRequired).with_context(|| { + format!( + "OAuth tokens for server {} were removed before refresh; authorization required", + self.inner.server_name + ) + }); + }; + + if !token_needs_refresh(latest.expires_at) { + debug!("adopting newer MCP OAuth credentials without contacting the provider"); + let manager = self.inner.authorization_manager.clone(); + let mut guard = manager.lock().await; + install_tokens_in_manager_guard(&mut guard, &latest).await?; + *self.inner.last_credentials.lock().await = Some(latest); + return Ok(()); + } + + // Preserve RMCP's `AuthorizationRequired` marker only for credentials known to be + // unrefreshable. Network and provider failures below remain ordinary errors. + if latest + .token_response + .0 + .refresh_token() + .is_none_or(|refresh_token| refresh_token.secret().trim().is_empty()) + { + return Err(AuthError::AuthorizationRequired).with_context(|| { + format!( + "OAuth tokens for server {} cannot be refreshed; authorization required", + self.inner.server_name + ) + }); + } + + let manager = self.inner.authorization_manager.clone(); + // The provider uses a separate HTTP client and cannot re-enter `AuthClient`. Retain this + // async guard so requests cannot observe credentials while they are staged and committed. + let mut guard = manager.lock().await; + install_tokens_in_manager_guard(&mut guard, &latest) + .await + .context("failed to stage OAuth credentials for refresh")?; + // The owned task prevents caller deadlines from canceling after possible token rotation; + // this timeout independently bounds the provider request. + debug!( + timeout_ms = refresh_request_timeout.as_millis(), + "requesting refreshed MCP OAuth credentials from the provider" + ); + let refreshed = match timeout(refresh_request_timeout, guard.refresh_token()).await { + Ok(Ok(token_response)) => { + debug!("received refreshed MCP OAuth credentials from the provider"); + refreshed_tokens(token_response, &latest, &self.inner) + } + Ok(Err(error @ AuthError::TokenRefreshFailed(_))) => { + // RMCP 1.8 collapses definitive OAuth rejection (for example, + // `invalid_grant`) and transient token-endpoint failures into this string + // variant. Match RMCP's own request path for now so rejected refresh tokens + // prompt reauthorization instead of surfacing as generic MCP startup failures. + // This can also prompt reauthorization after a transient failure. + // TODO: When RMCP exposes a typed distinction for refresh-token rejection, + // map only that definitive rejection to `AuthorizationRequired` here. + warn!( + error = %error, + "MCP OAuth refresh failed; reauthorization required by RMCP compatibility policy" + ); + return Err(AuthError::AuthorizationRequired).with_context(|| { + format!( + "failed to refresh OAuth tokens for server {}: {error}", + self.inner.server_name + ) + }); + } + Ok(Err(error)) => { + warn!( + error = %error, + "MCP OAuth provider refresh failed" + ); + return Err(error).with_context(|| { + format!( + "failed to refresh OAuth tokens for server {}", + self.inner.server_name + ) + }); + } + Err(_) => { + warn!( + timeout_ms = refresh_request_timeout.as_millis(), + "MCP OAuth provider refresh timed out; the outcome is unknown and a later serialized retry is permitted" + ); + anyhow::bail!( + "timed out after {refresh_request_timeout:?} refreshing OAuth tokens for server {}", + self.inner.server_name + ); + } + }; + + // Persist to the pinned source before exposing the refreshed token. On failure, restore + // the prior in-process credential and return the error; serving an unpersisted token would + // hide the root cause until a later process restart. If the provider already consumed the + // prior token, the next refresh may require reauthorization. That is the deliberate + // fail-closed policy. + // TODO: Add a bounded persistence retry only if telemetry shows this is common; never + // silently switch stores or continue with an unpersisted credential. + debug!("persisting refreshed MCP OAuth credentials to the resolved store"); + if let Err(error) = + self.inner + .credential_store + .save(keyring_store, &self.inner.server_name, &refreshed) + { + warn!( + error = %error, + "failed to persist refreshed MCP OAuth credentials; returning the error and restoring the previous in-process credentials" + ); + install_tokens_in_manager_guard(&mut guard, &latest) + .await + .context( + "failed to restore previous OAuth credentials after refresh persistence failed", + )?; + return Err(error); + } + + // This layer retains RMCP's legacy persistence hook. Install the same merged response + // (including carried-forward refresh token/scopes) so that hook cannot overwrite durable + // credentials with the provider's partial response. + install_tokens_in_manager_guard(&mut guard, &refreshed) + .await + .context( + "refreshed OAuth tokens were persisted but could not be installed in the authorization manager", + )?; + *self.inner.last_credentials.lock().await = Some(refreshed); + drop(guard); + debug!("persisted refreshed MCP OAuth credentials and completed the transaction"); + Ok(()) + } +} + +async fn install_tokens_in_manager_guard( + authorization_manager: &mut AuthorizationManager, + tokens: &StoredOAuthTokens, +) -> Result<()> { + let store = InMemoryCredentialStore::new(); + let token_response = tokens.token_response.0.clone(); + let granted_scopes = token_response + .scopes() + .map(|scopes| scopes.iter().map(|scope| scope.to_string()).collect()) + .unwrap_or_default(); + let token_received_at = SystemTime::now() + .duration_since(UNIX_EPOCH) + .ok() + .map(|duration| duration.as_secs()); + store + .save(StoredCredentials::new( + tokens.client_id.clone(), + Some(token_response), + granted_scopes, + token_received_at, + )) + .await + .context("failed to stage OAuth tokens for authorization manager")?; + + authorization_manager.set_credential_store(store); + // TODO(stevenlee): Add an RMCP adoption API that atomically updates credentials, client ID, + // and private `current_scopes`; this path cannot synchronize RMCP's scope-upgrade state. + authorization_manager + .initialize_from_store() + .await + .context("failed to adopt refreshed OAuth tokens")?; + Ok(()) +} + +fn refreshed_tokens( + mut token_response: OAuthTokenResponse, + previous: &StoredOAuthTokens, + inner: &OAuthPersistorInner, +) -> StoredOAuthTokens { + if token_response.refresh_token().is_none() { + token_response.set_refresh_token(previous.token_response.0.refresh_token().cloned()); + } + if token_response.scopes().is_none() { + token_response.set_scopes(previous.token_response.0.scopes().cloned()); + } + StoredOAuthTokens { + server_name: inner.server_name.clone(), + url: inner.url.clone(), + client_id: previous.client_id.clone(), + expires_at: compute_expires_at_millis(&token_response), + token_response: WrappedOAuthTokenResponse(token_response), + } +} diff --git a/codex-rs/rmcp-client/src/oauth/tests/persistor_tests.rs b/codex-rs/rmcp-client/src/oauth/tests/persistor_tests.rs new file mode 100644 index 000000000000..f32c2b691c2b --- /dev/null +++ b/codex-rs/rmcp-client/src/oauth/tests/persistor_tests.rs @@ -0,0 +1,420 @@ +use std::sync::Arc; +use std::sync::mpsc; +use std::time::Duration; + +use anyhow::Context; +use anyhow::Result; +use codex_config::types::AuthKeyringBackendKind; +use keyring::Error as KeyringError; +use oauth2::AccessToken; +use oauth2::TokenResponse; +use pretty_assertions::assert_eq; +use rmcp::transport::auth::AuthError; +use rmcp::transport::auth::AuthorizationManager; +use rmcp::transport::auth::OAuthState; +use tokio::sync::Mutex as TokioMutex; +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::ResponseTemplate; +use wiremock::matchers::body_string_contains; +use wiremock::matchers::method; +use wiremock::matchers::path; + +use super::MockKeyringStore; +use super::TempCodexHome; +use super::assert_tokens_match_without_expiry; +use super::sample_tokens; +use crate::oauth::OAuthPersistor; +use crate::oauth::ResolvedOAuthCredentialStore; +use crate::oauth::StoredOAuthTokens; +use crate::oauth::WrappedOAuthTokenResponse; +use crate::oauth::compute_store_key; +use crate::oauth::load_oauth_tokens_from_file; +use crate::oauth::refresh_lock::RefreshCredentialLock; +use crate::oauth::save_oauth_tokens_to_file; +use crate::startup_error::is_authentication_required_error; + +const REFRESH_LOCK_CONTENTION_EVENT_TARGET: &str = + "codex_rmcp_client::oauth::refresh_lock::contention"; + +struct LockContentionSubscriber { + contended_tx: mpsc::Sender<()>, +} + +impl Subscriber for LockContentionSubscriber { + fn enabled(&self, metadata: &Metadata<'_>) -> bool { + metadata.target() == REFRESH_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 { + Id::from_u64(/*u*/ 1) + } + + fn record(&self, _span: &Id, _values: &Record<'_>) {} + + fn record_follows_from(&self, _span: &Id, _follows_from: &Id) {} + + fn event(&self, event: &Event<'_>) { + if self.enabled(event.metadata()) { + self.contended_tx + .send(()) + .expect("signal actual OAuth credential-lock contention"); + } + } + + fn enter(&self, _span: &Id) {} + + fn exit(&self, _span: &Id) {} +} + +#[tokio::test(flavor = "current_thread")] +async fn concurrent_refreshes_call_provider_once_and_carry_omitted_fields() -> Result<()> { + let (_env, server, initial) = test_context().await?; + Mock::given(method("POST")) + .and(path("/oauth/token")) + .and(body_string_contains("grant_type=refresh_token")) + .and(body_string_contains("refresh_token=refresh-token")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "access_token": "refreshed-access-token", + "token_type": "Bearer", + "expires_in": 3600, + }))) + .expect(1) + .mount(&server) + .await; + save_oauth_tokens_to_file(&initial)?; + + // Hold the real credential lock until both refresh transactions report WouldBlock. This makes + // the lock assertion independent of task scheduling and ensures removing transaction locking + // makes the test fail before either request can reach the provider. + let held_lock = + RefreshCredentialLock::acquire_for_server(&initial.server_name, &initial.url).await?; + let (contended_tx, contended_rx) = mpsc::channel(); + let _subscriber_guard = + tracing::subscriber::set_default(LockContentionSubscriber { contended_tx }); + + let first = persistor_for(&initial).await?; + let second = persistor_for(&initial).await?; + let first_task = tokio::spawn({ + let first = first.clone(); + async move { first.refresh_if_needed().await } + }); + let second_task = tokio::spawn({ + let second = second.clone(); + async move { second.refresh_if_needed().await } + }); + + wait_for_lock_contention(contended_rx, /*expected_count*/ 2).await?; + drop(held_lock); + first_task.await??; + second_task.await??; + server.verify().await; + + // Layer 2 still invokes the legacy RMCP persistence hook after operations. Exercise that hook + // so a raw provider response that omitted refresh token/scopes cannot overwrite the merged + // authoritative credential. + first.persist_if_needed().await?; + let stored = load_oauth_tokens_from_file(&initial.server_name, &initial.url)? + .expect("refreshed credentials should be stored"); + let mut expected_response = initial.token_response.0.clone(); + expected_response.set_access_token(AccessToken::new("refreshed-access-token".to_string())); + // File loads derive `expires_in` from stable `expires_at`, so it may tick down before this + // assertion. Normalize only that derived field and compare the complete token response so + // omitted refresh-token and scope carry-forward remain covered. + expected_response.set_expires_in(stored.token_response.0.expires_in().as_ref()); + assert_eq!( + stored.token_response, + WrappedOAuthTokenResponse(expected_response) + ); + Ok(()) +} + +#[expect( + clippy::await_holding_invalid_type, + reason = "AuthorizationManager async access must be serialized through its Tokio mutex" +)] +#[tokio::test(flavor = "current_thread")] +async fn resolved_keyring_read_error_preserves_in_memory_credentials() -> Result<()> { + let (_env, _server, initial) = test_context().await?; + let keyring_store = MockKeyringStore::default(); + let key = compute_store_key(&initial.server_name, &initial.url)?; + keyring_store.set_error(&key, KeyringError::Invalid("error".into(), "load".into())); + let manager = authorization_manager_for(&initial).await?; + let persistor = OAuthPersistor::new( + initial.server_name.clone(), + initial.url.clone(), + Arc::clone(&manager), + ResolvedOAuthCredentialStore::Keyring(AuthKeyringBackendKind::Direct), + Some(initial.clone()), + ); + + let error = persistor + .refresh_if_needed_in(&keyring_store, Duration::from_secs(/*secs*/ 45)) + .await + .expect_err("the resolved keyring read error should abort refresh"); + assert!( + error + .to_string() + .contains("failed to reread OAuth tokens from resolved keyring storage"), + "unexpected error: {error:#}" + ); + let guard = manager.lock().await; + let (_client_id, token_response) = guard.get_credentials().await?; + assert_eq!( + WrappedOAuthTokenResponse(token_response.expect("manager should retain credentials")), + initial.token_response + ); + Ok(()) +} + +#[tokio::test(flavor = "current_thread")] +async fn missing_authoritative_credentials_require_reauthorization() -> Result<()> { + let (_env, _server, initial) = test_context().await?; + let persistor = persistor_for(&initial).await?; + + let error = persistor + .refresh_if_needed() + .await + .expect_err("a removed authoritative credential should abort refresh"); + assert!(error.chain().any(|source| matches!( + source.downcast_ref::(), + Some(AuthError::AuthorizationRequired) + ))); + Ok(()) +} + +#[tokio::test(flavor = "current_thread")] +async fn rejected_refresh_token_requires_reauthorization() -> Result<()> { + let (_env, server, initial) = test_context().await?; + Mock::given(method("POST")) + .and(path("/oauth/token")) + .and(body_string_contains("grant_type=refresh_token")) + .and(body_string_contains("refresh_token=refresh-token")) + .respond_with(ResponseTemplate::new(400).set_body_json(serde_json::json!({ + "error": "invalid_grant", + "error_description": "refresh token expired or revoked", + }))) + .expect(1) + .mount(&server) + .await; + save_oauth_tokens_to_file(&initial)?; + let persistor = persistor_for(&initial).await?; + + let error = persistor + .refresh_if_needed() + .await + .expect_err("a provider-rejected refresh token should require reauthorization"); + assert!(is_authentication_required_error(&error)); + let stored = load_oauth_tokens_from_file(&initial.server_name, &initial.url)? + .expect("rejected refresh must preserve the durable credentials"); + assert_tokens_match_without_expiry(&stored, &initial); + server.verify().await; + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn caller_cancellation_does_not_cancel_refresh_persistence() -> Result<()> { + let (_env, server, initial) = test_context().await?; + let (request_received_tx, request_received_rx) = mpsc::channel(); + let (release_response_tx, release_response_rx) = mpsc::channel(); + let release_response_rx = Arc::new(std::sync::Mutex::new(release_response_rx)); + Mock::given(method("POST")) + .and(path("/oauth/token")) + .and(body_string_contains("grant_type=refresh_token")) + .and(body_string_contains("refresh_token=refresh-token")) + .respond_with(move |_request: &wiremock::Request| { + request_received_tx + .send(()) + .expect("signal OAuth refresh request"); + release_response_rx + .lock() + .expect("lock OAuth refresh response gate") + .recv() + .expect("wait to release OAuth refresh response"); + ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "access_token": "cancel-safe-access-token", + "token_type": "Bearer", + "expires_in": 3600, + })) + }) + .expect(1) + .mount(&server) + .await; + save_oauth_tokens_to_file(&initial)?; + let persistor = persistor_for(&initial).await?; + let caller = tokio::spawn({ + let persistor = persistor.clone(); + async move { persistor.refresh_if_needed().await } + }); + + tokio::task::spawn_blocking(move || { + request_received_rx + .recv_timeout(Duration::from_secs(/*secs*/ 5)) + .context("timed out waiting for OAuth refresh request") + }) + .await??; + caller.abort(); + assert!( + caller + .await + .expect_err("caller should be cancelled") + .is_cancelled() + ); + + release_response_tx + .send(()) + .context("release OAuth refresh response")?; + + // Reacquiring the same credential lock waits for the detached refresh task to persist and + // release it, avoiding a scheduler-sensitive sleep after cancellation. + let _lock = tokio::time::timeout( + Duration::from_secs(/*secs*/ 2), + RefreshCredentialLock::acquire_for_server(&initial.server_name, &initial.url), + ) + .await + .context("detached refresh did not release its credential lock")??; + let stored = load_oauth_tokens_from_file(&initial.server_name, &initial.url)? + .expect("detached refresh should persist credentials"); + assert_eq!( + stored.token_response.0.access_token().secret(), + "cancel-safe-access-token" + ); + server.verify().await; + Ok(()) +} + +#[tokio::test(flavor = "current_thread")] +async fn provider_timeout_releases_lock_and_preserves_durable_credentials() -> Result<()> { + let (_env, server, initial) = test_context().await?; + mount_delayed_refresh(&server, "late-access-token").await; + save_oauth_tokens_to_file(&initial)?; + let persistor = persistor_for(&initial).await?; + + let error = persistor + .refresh_if_needed_in( + &MockKeyringStore::default(), + Duration::from_millis(/*millis*/ 50), + ) + .await + .expect_err("provider request should reach its explicit timeout"); + assert!(error.to_string().contains("timed out after 50ms")); + + let _lock = tokio::time::timeout( + Duration::from_millis(/*millis*/ 100), + RefreshCredentialLock::acquire_for_server(&initial.server_name, &initial.url), + ) + .await + .context("provider timeout did not release the credential lock")??; + let stored = load_oauth_tokens_from_file(&initial.server_name, &initial.url)? + .expect("timed-out refresh must leave durable credentials present"); + assert_tokens_match_without_expiry(&stored, &initial); + server.verify().await; + Ok(()) +} + +async fn persistor_for(tokens: &StoredOAuthTokens) -> Result { + Ok(OAuthPersistor::new( + tokens.server_name.clone(), + tokens.url.clone(), + authorization_manager_for(tokens).await?, + ResolvedOAuthCredentialStore::File, + Some(tokens.clone()), + )) +} + +async fn test_context() -> Result<(TempCodexHome, MockServer, StoredOAuthTokens)> { + let env = TempCodexHome::new(); + let server = MockServer::start().await; + mount_oauth_metadata(&server).await; + let tokens = expired_tokens(&format!("{}/mcp", server.uri())); + Ok((env, server, tokens)) +} + +async fn authorization_manager_for( + tokens: &StoredOAuthTokens, +) -> Result>> { + let mut state = OAuthState::new(tokens.url.clone(), Some(reqwest::Client::new())).await?; + state + .set_credentials(&tokens.client_id, tokens.token_response.0.clone()) + .await?; + let manager = match state { + OAuthState::Authorized(manager) | OAuthState::Unauthorized(manager) => manager, + OAuthState::Session(_) | OAuthState::AuthorizedHttpClient(_) => { + anyhow::bail!("unexpected OAuth state") + } + _ => anyhow::bail!("unexpected OAuth state"), + }; + Ok(Arc::new(TokioMutex::new(manager))) +} + +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(serde_json::json!({ + "authorization_endpoint": format!("{}/oauth/authorize", server.uri()), + "token_endpoint": format!("{}/oauth/token", server.uri()), + "scopes_supported": ["scope-a", "scope-b"], + }))) + .mount(server) + .await; +} + +async fn mount_delayed_refresh(server: &MockServer, response_access_token: &str) { + Mock::given(method("POST")) + .and(path("/oauth/token")) + .and(body_string_contains("grant_type=refresh_token")) + .and(body_string_contains("refresh_token=refresh-token")) + .respond_with( + ResponseTemplate::new(200) + .set_delay(Duration::from_millis(/*millis*/ 200)) + .set_body_json(serde_json::json!({ + "access_token": response_access_token, + "token_type": "Bearer", + "expires_in": 3600, + })), + ) + .expect(1) + .mount(server) + .await; +} + +async fn wait_for_lock_contention(rx: mpsc::Receiver<()>, expected_count: usize) -> Result<()> { + tokio::task::spawn_blocking(move || { + for _ in 0..expected_count { + rx.recv_timeout(Duration::from_secs(/*secs*/ 5)) + .context("timed out waiting for lock contention")?; + } + Ok(()) + }) + .await? +} + +fn expired_tokens(url: &str) -> StoredOAuthTokens { + let mut tokens = sample_tokens(); + tokens.url = url.to_string(); + tokens.expires_at = Some(0); + tokens + .token_response + .0 + .set_expires_in(Some(&Duration::ZERO)); + tokens +} diff --git a/codex-rs/rmcp-client/src/oauth/tests/store_lock_tests.rs b/codex-rs/rmcp-client/src/oauth/tests/store_lock_tests.rs index 9ecbf2fff6fd..6e8209c2dc93 100644 --- a/codex-rs/rmcp-client/src/oauth/tests/store_lock_tests.rs +++ b/codex-rs/rmcp-client/src/oauth/tests/store_lock_tests.rs @@ -41,6 +41,10 @@ use crate::oauth::test_support::TempCodexHome; use codex_config::types::OAuthCredentialsStoreMode; const STORE_LOCK_CONTENTION_EVENT_TARGET: &str = "codex_rmcp_client::oauth::store_lock::contention"; +// Contention is proven by the tracing event emitted after a real WouldBlock. Keep the timeout +// generous because it only bounds a failed test; it must not turn worker scheduling latency into +// a false failure on loaded CI hosts. +const STORE_LOCK_CONTENTION_EVENT_TIMEOUT: Duration = Duration::from_secs(/*secs*/ 10); fn assert_tokens_match_without_expiry(actual: &StoredOAuthTokens, expected: &StoredOAuthTokens) { assert_eq!(actual.server_name, expected.server_name); @@ -294,7 +298,9 @@ where // This event is emitted only after `try_lock()` returns WouldBlock, so the test fails if the // operation stops acquiring the aggregate-store lock. - contended_rx.recv_timeout(Duration::from_secs(/*secs*/ 1))?; + contended_rx + .recv_timeout(STORE_LOCK_CONTENTION_EVENT_TIMEOUT) + .context("timed out waiting for actual OAuth store lock contention")?; drop(held_lock); let result = result_rx.recv_timeout(Duration::from_secs(/*secs*/ 10))??; worker @@ -327,7 +333,9 @@ fn file_store_lock_preserves_updates_for_different_servers() -> Result<()> { }); }); - contended_rx.recv_timeout(Duration::from_secs(/*secs*/ 1))?; + contended_rx + .recv_timeout(STORE_LOCK_CONTENTION_EVENT_TIMEOUT) + .context("timed out waiting for actual OAuth store lock contention")?; save_oauth_tokens_to_file_with_lock_held(&first)?; drop(held_lock); result_rx.recv_timeout(Duration::from_secs(/*secs*/ 10))??; @@ -396,7 +404,9 @@ fn secrets_store_lock_preserves_updates_for_different_servers() -> Result<()> { }); }); - contended_rx.recv_timeout(Duration::from_secs(/*secs*/ 1))?; + contended_rx + .recv_timeout(STORE_LOCK_CONTENTION_EVENT_TIMEOUT) + .context("timed out waiting for actual OAuth store lock contention")?; let first_serialized = serde_json::to_string(&first)?; save_oauth_tokens_to_secrets_keyring_with_lock_held( &keyring_store, diff --git a/codex-rs/rmcp-client/src/rmcp_client.rs b/codex-rs/rmcp-client/src/rmcp_client.rs index e056e6338243..ed81eefba798 100644 --- a/codex-rs/rmcp-client/src/rmcp_client.rs +++ b/codex-rs/rmcp-client/src/rmcp_client.rs @@ -497,7 +497,7 @@ impl RmcpClient { params: Option, timeout: Option, ) -> Result { - self.refresh_oauth_if_needed().await; + self.refresh_oauth_if_needed().await?; let result = self .run_service_operation("tools/list", timeout, move |service| { let params = params.clone(); @@ -514,7 +514,7 @@ impl RmcpClient { params: Option, timeout: Option, ) -> Result { - self.refresh_oauth_if_needed().await; + self.refresh_oauth_if_needed().await?; let result = self .run_service_operation("tools/list", timeout, move |service| { let params = params.clone(); @@ -559,7 +559,7 @@ impl RmcpClient { params: Option, timeout: Option, ) -> Result { - self.refresh_oauth_if_needed().await; + self.refresh_oauth_if_needed().await?; let result = self .run_service_operation("resources/list", timeout, move |service| { let params = params.clone(); @@ -575,7 +575,7 @@ impl RmcpClient { params: Option, timeout: Option, ) -> Result { - self.refresh_oauth_if_needed().await; + self.refresh_oauth_if_needed().await?; let result = self .run_service_operation("resources/templates/list", timeout, move |service| { let params = params.clone(); @@ -591,7 +591,7 @@ impl RmcpClient { params: ReadResourceRequestParams, timeout: Option, ) -> Result { - self.refresh_oauth_if_needed().await; + self.refresh_oauth_if_needed().await?; let result = self .run_service_operation("resources/read", timeout, move |service| { let params = params.clone(); @@ -609,7 +609,7 @@ impl RmcpClient { meta: Option, timeout: Option, ) -> Result { - self.refresh_oauth_if_needed().await; + self.refresh_oauth_if_needed().await?; let arguments = match arguments { Some(Value::Object(map)) => Some(map), Some(other) => { @@ -665,7 +665,7 @@ impl RmcpClient { method: &str, params: Option, ) -> Result<()> { - self.refresh_oauth_if_needed().await; + self.refresh_oauth_if_needed().await?; self.run_service_operation( "notifications/custom", /*timeout*/ None, @@ -695,7 +695,7 @@ impl RmcpClient { method: &str, params: Option, ) -> Result { - self.refresh_oauth_if_needed().await; + self.refresh_oauth_if_needed().await?; let response = self .run_service_operation("requests/custom", /*timeout*/ None, move |service| { let params = params.clone(); @@ -759,12 +759,12 @@ impl RmcpClient { } } - async fn refresh_oauth_if_needed(&self) { - if let Some(runtime) = self.oauth_persistor().await - && let Err(error) = runtime.refresh_if_needed().await - { - warn!("failed to refresh OAuth tokens: {error}"); + /// OAuth uses independent lock/request bounds and completes before the operation timeout starts. + async fn refresh_oauth_if_needed(&self) -> Result<()> { + if let Some(runtime) = self.oauth_persistor().await { + runtime.refresh_if_needed().await?; } + Ok(()) } async fn create_pending_transport( diff --git a/codex-rs/rmcp-client/src/streamable_http_retry.rs b/codex-rs/rmcp-client/src/streamable_http_retry.rs index 73da95de58ec..5c4d28881533 100644 --- a/codex-rs/rmcp-client/src/streamable_http_retry.rs +++ b/codex-rs/rmcp-client/src/streamable_http_retry.rs @@ -37,7 +37,7 @@ impl RmcpClient { PendingTransport::StreamableHttp { .. } | PendingTransport::StreamableHttpWithOAuth { .. } => true, }; - let retry_deadline = timeout.map(|duration| Instant::now() + duration); + let mut retry_deadline = timeout.map(|duration| Instant::now() + duration); let mut pending_transport = Some(initial_transport); for (attempt, retry_delay_ms) in STREAMABLE_HTTP_RETRY_DELAYS_MS @@ -62,6 +62,18 @@ impl RmcpClient { } } }; + if let PendingTransport::StreamableHttpWithOAuth { + oauth_persistor, .. + } = &transport + { + // OAuth refresh has its own lock and provider request bounds. Exclude it from the + // MCP handshake budget, and finish persistence before attempting initialize. + let refresh_started_at = Instant::now(); + oauth_persistor.refresh_if_needed().await?; + if let Some(deadline) = retry_deadline.as_mut() { + *deadline += refresh_started_at.elapsed(); + } + } let attempt_timeout = remaining_initialize_timeout(timeout, retry_deadline)?; match Self::connect_pending_transport(