diff --git a/codex-rs/cli/src/mcp_cmd.rs b/codex-rs/cli/src/mcp_cmd.rs index 38a7f323a8..eb2b70c13d 100644 --- a/codex-rs/cli/src/mcp_cmd.rs +++ b/codex-rs/cli/src/mcp_cmd.rs @@ -660,7 +660,9 @@ async fn run_logout(config: &Config, logout_args: LogoutArgs) -> Result<()> { &url, config.mcp_oauth_credentials_store_mode, config.auth_keyring_backend_kind(), - ) { + ) + .await + { Ok(true) => println!("Removed OAuth credentials for '{name}'."), Ok(false) => println!("No OAuth credentials stored for '{name}'."), Err(err) => return Err(anyhow!("failed to delete OAuth credentials: {err}")), diff --git a/codex-rs/codex-mcp/src/rmcp_client.rs b/codex-rs/codex-mcp/src/rmcp_client.rs index e469b865a2..49010cc98b 100644 --- a/codex-rs/codex-mcp/src/rmcp_client.rs +++ b/codex-rs/codex-mcp/src/rmcp_client.rs @@ -1088,11 +1088,6 @@ pub(crate) async fn make_rmcp_client( runtime_auth_provider: Option, protocol_mode: McpProtocolMode, ) -> Result { - if oauth_refresh_mode == McpOAuthRefreshMode::Coordinated { - warn!( - "MCP OAuth refresh coordination is not available in this build; using legacy refresh" - ); - } let config = server.config().clone(); if matches!(config.auth, McpServerAuth::ChatGpt) && !config.is_local_environment() @@ -1215,6 +1210,7 @@ pub(crate) async fn make_rmcp_client( runtime_auth_provider, protocol_mode, redirect_mode, + oauth_refresh_mode, ) .await .map_err(StartupOutcomeError::from) diff --git a/codex-rs/core/tests/suite/mcp_oauth_refresh_tests.rs b/codex-rs/core/tests/suite/mcp_oauth_refresh_tests.rs index 36a7139352..309370f9c7 100644 --- a/codex-rs/core/tests/suite/mcp_oauth_refresh_tests.rs +++ b/codex-rs/core/tests/suite/mcp_oauth_refresh_tests.rs @@ -66,7 +66,8 @@ async fn oauth_mode_refresh_replaces_the_live_connection( "synthetic-access-token", "synthetic-refresh-token", OAuthCredentialExpiry::Valid, - )?; + ) + .await?; let fixture = test_codex() .with_home(home) .with_config(move |config| { diff --git a/codex-rs/core/tests/suite/rmcp_client.rs b/codex-rs/core/tests/suite/rmcp_client.rs index ebdeaf4823..87f47b497e 100644 --- a/codex-rs/core/tests/suite/rmcp_client.rs +++ b/codex-rs/core/tests/suite/rmcp_client.rs @@ -3774,7 +3774,8 @@ async fn streamable_http_with_oauth_round_trip_impl() -> anyhow::Result<()> { "expired-access-token", refresh_token, OAuthCredentialExpiry::Expired, - )?; + ) + .await?; let discovered_credential_name = credential_config.oauth_credential_name(discovered_server_name); write_fallback_oauth_tokens( @@ -3784,7 +3785,8 @@ async fn streamable_http_with_oauth_round_trip_impl() -> anyhow::Result<()> { expected_token, refresh_token, OAuthCredentialExpiry::Valid, - )?; + ) + .await?; // Phase 4: configure Codex with the OAuth-backed Streamable HTTP MCP // server and build the fixture in the active local or remote-aware mode. @@ -3792,6 +3794,10 @@ async fn streamable_http_with_oauth_round_trip_impl() -> anyhow::Result<()> { .with_model_info_override("gpt-5.4", |model| model.supports_search_tool = false) .with_home(temp_home.clone()) .with_config(move |config| { + config + .features + .enable(Feature::McpOAuthRefreshCoordination) + .expect("test config should allow coordinated MCP OAuth refresh"); config.mcp_oauth_credentials_store_mode = OAuthCredentialsStoreMode::Auto; insert_mcp_server( config, @@ -3957,12 +3963,15 @@ async fn streamable_http_with_oauth_round_trip_impl() -> anyhow::Result<()> { .await .context("the newly discovered OAuth server did not recover after its store was unlocked")??; - assert!(codex_rmcp_client::delete_oauth_tokens( - discovered_credential_name.as_ref(), - http_server.url(), - OAuthCredentialsStoreMode::File, - codex_config::types::AuthKeyringBackendKind::default(), - )?); + assert!( + codex_rmcp_client::delete_oauth_tokens( + discovered_credential_name.as_ref(), + http_server.url(), + OAuthCredentialsStoreMode::File, + codex_config::types::AuthKeyringBackendKind::default(), + ) + .await? + ); fixture.codex.refresh_runtime_config(refreshed_config).await; let logged_out_startup = tokio::time::timeout( Duration::from_secs(5), @@ -3990,7 +3999,8 @@ async fn streamable_http_with_oauth_round_trip_impl() -> anyhow::Result<()> { expected_token, refresh_token, OAuthCredentialExpiry::Valid, - )?; + ) + .await?; // Phase 6: submit the user turn that should invoke the OAuth-backed tool. fixture @@ -4431,7 +4441,7 @@ enum OAuthCredentialExpiry { Expired, } -fn write_fallback_oauth_tokens( +async fn write_fallback_oauth_tokens( server_name: &str, server_url: &str, client_id: &str, @@ -4468,6 +4478,7 @@ fn write_fallback_oauth_tokens( OAuthCredentialsStoreMode::File, codex_config::types::AuthKeyringBackendKind::default(), ) + .await } struct EnvVarGuard { diff --git a/codex-rs/rmcp-client/src/oauth.rs b/codex-rs/rmcp-client/src/oauth.rs index 3b28dc2273..8d2882acc0 100644 --- a/codex-rs/rmcp-client/src/oauth.rs +++ b/codex-rs/rmcp-client/src/oauth.rs @@ -16,13 +16,13 @@ //! //! If the keyring is not available or fails, we fall back to CODEX_HOME/.credentials.json which is consistent with other coding CLI agents. -#[cfg(test)] mod credential_store; mod ema_identity; mod issuer_binding; mod refresh_lock; mod refresh_transaction; mod resolved_store; +mod runtime; mod store_lock; #[cfg(test)] @@ -74,7 +74,6 @@ use tokio::sync::Mutex; use codex_utils_home_dir::find_codex_home; -#[cfg(test)] pub(crate) use self::credential_store::OAuthCredentialStore; pub(crate) use self::ema_identity::stored_oidc_identity; pub(crate) use self::issuer_binding::validate_authorization_server_endpoints; @@ -85,6 +84,7 @@ pub(crate) use self::resolved_store::ResolvedOAuthCredentialStore; pub(crate) use self::resolved_store::ResolvedOAuthTokens; pub(crate) use self::resolved_store::resolve_oauth_tokens_from_store_policy; use self::resolved_store::try_resolve_oauth_tokens_from_store_policy; +pub(crate) use self::runtime::OAuthRuntime; const KEYRING_SERVICE: &str = "Codex MCP Credentials"; const MCP_OAUTH_SECRET_PREFIX: &str = "MCP_OAUTH"; @@ -440,12 +440,13 @@ enum OAuthKeyringLoadError { Backend(#[from] anyhow::Error), } -pub fn save_oauth_tokens( +pub async fn save_oauth_tokens( server_name: &str, tokens: &StoredOAuthTokens, store_mode: OAuthCredentialsStoreMode, keyring_backend_kind: AuthKeyringBackendKind, ) -> Result<()> { + let _lock = RefreshCredentialLock::acquire_for_server(server_name, &tokens.url).await?; let keyring_store = DefaultKeyringStore; match store_mode { OAuthCredentialsStoreMode::Auto => save_oauth_tokens_with_keyring_with_fallback_to_file( @@ -585,12 +586,13 @@ fn save_oauth_tokens_with_keyring_with_fallback_to_file Result { + let _lock = RefreshCredentialLock::acquire_for_server(server_name, url).await?; let keyring_store = DefaultKeyringStore; delete_oauth_tokens_from_keyring_and_file( &keyring_store, diff --git a/codex-rs/rmcp-client/src/oauth/credential_store.rs b/codex-rs/rmcp-client/src/oauth/credential_store.rs index 01c0114332..0429efde42 100644 --- a/codex-rs/rmcp-client/src/oauth/credential_store.rs +++ b/codex-rs/rmcp-client/src/oauth/credential_store.rs @@ -4,8 +4,10 @@ //! against an in-memory store. Initialization may save tokenless client credentials, which //! this refresh adapter does not support. //! -//! RMCP must hold the credential lock across its authoritative reread, token exchange, and -//! save. Saves and clears require an active guard; no store operation reacquires the lock. +//! Ordinary token reads use the cached credentials. Refresh-guard acquisition rereads +//! the pinned store before RMCP exchanges the token and saves the result. Codex +//! preparation rechecks freshness under that guard before asking RMCP to refresh. +//! Saves and clears require an active guard; no store operation reacquires the lock. //! The runtime snapshot advances only for credentials compatible with this connection, //! so replacement logins still cause a rebuild. //! Synchronous store operations run off the Tokio workers so caller deadlines remain pollable. @@ -22,21 +24,29 @@ use futures::future::BoxFuture; use oauth2::Scope; use oauth2::TokenResponse; use rmcp::transport::auth::AuthError; +use rmcp::transport::auth::AuthorizationManager; +use rmcp::transport::auth::CredentialRefreshGuard; use rmcp::transport::auth::CredentialStore; use rmcp::transport::auth::StoredCredentials; use tokio::sync::Mutex; use tracing::warn; +use crate::oauth_http_client::PROACTIVE_REFRESH_TIMEOUT; + use super::RefreshCredentialLock; use super::ResolvedOAuthCredentialStore; use super::StoredOAuthTokens; use super::WrappedOAuthTokenResponse; use super::normalized_oauth_credentials; use super::refresh_expires_in_from_timestamp; +use super::refresh_transaction::REFRESH_REQUEST_TIMEOUT; +use super::token_needs_refresh; +use super::validate_refresh_token_issuer; #[derive(Clone)] pub(crate) struct OAuthCredentialStore { inner: Arc>, + held_refresh_guard: Option>, } struct OAuthCredentialStoreInner { @@ -67,6 +77,39 @@ impl OAuthCredentialStore { last_credentials: Mutex::new(Some(tokens)), refresh_guard: Mutex::new(Weak::new()), }), + held_refresh_guard: None, + } + } + + pub(crate) async fn refresh_if_needed(&self, manager: &mut AuthorizationManager) -> Result<()> { + let guard = self.acquire_transaction_guard().await?; + let tokens = self + .inner + .last_credentials + .lock() + .await + .clone() + .ok_or(AuthError::AuthorizationRequired)?; + if !token_needs_refresh(tokens.expires_at) { + return Ok(()); + } + let metadata = manager.resolve_metadata().await?.metadata; + validate_refresh_token_issuer(&metadata, &tokens)?; + manager.set_metadata(metadata); + manager.configure_client_id(&tokens.client_id)?; + // Reuse the guard for RMCP's exchange and save after the locked freshness check. + manager.set_credential_store(Self { + inner: Arc::clone(&self.inner), + held_refresh_guard: Some(guard), + }); + let result = PROACTIVE_REFRESH_TIMEOUT + .scope(REFRESH_REQUEST_TIMEOUT, manager.refresh_token()) + .await; + manager.set_credential_store(self.clone()); + match result { + Ok(_) => Ok(()), + Err(AuthError::TokenRefreshRejected(_)) => Err(AuthError::AuthorizationRequired.into()), + Err(error) => Err(error.into()), } } @@ -74,11 +117,27 @@ impl OAuthCredentialStore { normalized_oauth_credentials(self.inner.last_credentials.lock().await.as_ref()) } - pub(crate) async fn acquire_transaction_guard(&self) -> Result> { + pub(crate) async fn acquire_transaction_guard( + &self, + ) -> Result, AuthError> { let guard = Arc::new( RefreshCredentialLock::acquire_for_server(&self.inner.server_name, &self.inner.url) - .await?, + .await + .map_err(credential_store_error)?, ); + let inner = Arc::clone(&self.inner); + let tokens = tokio::task::spawn_blocking(move || { + inner + .store + .load(&inner.keyring, &inner.server_name, &inner.url) + }) + .await + .context("OAuth credential load task failed") + .map_err(credential_store_error)? + .map_err(credential_store_error)? + .ok_or(AuthError::AuthorizationRequired)?; + self.validate_connection(&tokens)?; + *self.inner.last_credentials.lock().await = Some(tokens); *self.inner.refresh_guard.lock().await = Arc::downgrade(&guard); Ok(guard) } @@ -111,25 +170,13 @@ impl CredentialStore for OAuthCredentialStore Self: 'async_trait, { Box::pin(async move { - let inner = Arc::clone(&self.inner); - let Some(tokens) = tokio::task::spawn_blocking(move || { - inner - .store - .load(&inner.keyring, &inner.server_name, &inner.url) - }) - .await - .context("OAuth credential load task failed") - .map_err(credential_store_error)? - .map_err(credential_store_error)? - else { - // Do not acknowledge removal as a live connection update: reconciliation - // must see that this connection still belongs to the previous credentials. - return Ok(None); - }; - self.validate_connection(&tokens)?; - let credentials = rmcp_credentials(&tokens); - *self.inner.last_credentials.lock().await = Some(tokens); - Ok(Some(credentials)) + Ok(self + .inner + .last_credentials + .lock() + .await + .as_ref() + .map(rmcp_credentials)) }) } @@ -228,10 +275,26 @@ impl CredentialStore for OAuthCredentialStore .map_err(credential_store_error) }) } + + fn acquire_refresh_guard<'life0, 'async_trait>( + &'life0 self, + ) -> BoxFuture<'async_trait, Result, AuthError>> + where + 'life0: 'async_trait, + Self: 'async_trait, + { + Box::pin(async move { + let guard = match &self.held_refresh_guard { + Some(guard) => Arc::clone(guard), + None => self.acquire_transaction_guard().await?, + }; + Ok(Some(CredentialRefreshGuard::new(guard))) + }) + } } fn credential_store_error(error: anyhow::Error) -> AuthError { - AuthError::InternalError(format!("{error:#}")) + AuthError::CredentialStoreError(format!("{error:#}")) } fn rmcp_credentials(tokens: &StoredOAuthTokens) -> StoredCredentials { diff --git a/codex-rs/rmcp-client/src/oauth/refresh_transaction.rs b/codex-rs/rmcp-client/src/oauth/refresh_transaction.rs index 3728ef330f..2da7ec2adb 100644 --- a/codex-rs/rmcp-client/src/oauth/refresh_transaction.rs +++ b/codex-rs/rmcp-client/src/oauth/refresh_transaction.rs @@ -28,7 +28,7 @@ use super::refresh_lock::RefreshCredentialLock; use super::token_needs_refresh; use super::validate_refresh_token_issuer; -const REFRESH_REQUEST_TIMEOUT: Duration = Duration::from_secs(45); +pub(super) const REFRESH_REQUEST_TIMEOUT: Duration = Duration::from_secs(45); impl OAuthPersistor { pub(crate) async fn refresh_if_needed(&self) -> Result<()> { diff --git a/codex-rs/rmcp-client/src/oauth/runtime.rs b/codex-rs/rmcp-client/src/oauth/runtime.rs new file mode 100644 index 0000000000..fbe43c09fc --- /dev/null +++ b/codex-rs/rmcp-client/src/oauth/runtime.rs @@ -0,0 +1,64 @@ +//! Keeps one OAuth refresh and persistence owner for the lifetime of a connection. +//! +//! Both modes use cached expiry to prepare credentials before the MCP request budget. +//! Coordinated preparation owns a task so caller cancellation cannot interrupt persistence. +//! It locks the manager before the credential store, matching RMCP's lock order. +//! RMCP commits through the pinned store while Legacy retains Codex's persistor. + +use std::sync::Arc; + +use anyhow::Context; +use anyhow::Result; +use rmcp::transport::auth::AuthorizationManager; +use tokio::sync::Mutex; +use tracing::warn; + +use super::OAuthPersistor; +use super::credential_store::OAuthCredentialStore; +use super::token_needs_refresh; + +#[derive(Clone)] +pub(crate) enum OAuthRuntime { + Legacy(OAuthPersistor), + Coordinated { + auth_manager: Arc>, + store: OAuthCredentialStore, + }, +} + +impl OAuthRuntime { + #[expect( + clippy::await_holding_invalid_type, + reason = "AuthorizationManager access follows RMCP's manager-before-credential lock order" + )] + pub(crate) async fn refresh_if_needed(&self) -> Result<()> { + match self { + Self::Legacy(persistor) => persistor.refresh_if_needed().await, + Self::Coordinated { + auth_manager, + store, + } => { + let expires_at = store + .stored_credentials() + .await + .and_then(|tokens| tokens.expires_at); + if !token_needs_refresh(expires_at) { + return Ok(()); + } + let auth_manager = Arc::clone(auth_manager); + let store = store.clone(); + tokio::spawn(async move { + let mut manager = auth_manager.lock().await; + let result = store.refresh_if_needed(&mut manager).await; + if result.is_err() { + // Keep the summary in the owned task without logging credential/provider data. + warn!("MCP OAuth preparation failed"); + } + result + }) + .await + .context("OAuth refresh task failed")? + } + } + } +} diff --git a/codex-rs/rmcp-client/src/oauth/tests/credential_store_tests.rs b/codex-rs/rmcp-client/src/oauth/tests/credential_store_tests.rs index bd5d2659cd..d8b4e3bdb1 100644 --- a/codex-rs/rmcp-client/src/oauth/tests/credential_store_tests.rs +++ b/codex-rs/rmcp-client/src/oauth/tests/credential_store_tests.rs @@ -51,14 +51,22 @@ async fn mutations_require_and_retain_the_transaction_guard() -> Result<()> { store.save(credentials.clone()).await }; assert!( - matches!(result, Err(AuthError::InternalError(error)) if error.contains("active refresh guard")) + matches!(result, Err(AuthError::CredentialStoreError(error)) if error.contains("active refresh guard")) ); assert_eq!( load_oauth_tokens_from_file(&initial.server_name, &initial.url)?.is_none(), !clear, ); + // Guard acquisition now rereads storage. Restore the deleted credential first, + // then delete it again under the guard so the queued save must recreate it. + save_oauth_tokens_to_file(&initial)?; + let guard = CredentialStore::acquire_refresh_guard(&store) + .await? + .expect("coordinated refresh guard"); + if !clear { + store.clear().await?; + } let aggregate_lock = OAuthStoreLock::acquire_for_write(OAuthStore::File)?; - let guard = store.acquire_transaction_guard().await?; let mut mutation = if clear { store.clear() } else { @@ -117,7 +125,7 @@ async fn save_publishes_only_persisted_credentials() -> Result<()> { .expect("durable keyring credentials"); let mut expected = initial.clone(); if fail_save { - assert!(matches!(result, Err(AuthError::InternalError(_)))); + assert!(matches!(result, Err(AuthError::CredentialStoreError(_)))); } else { result?; expected @@ -156,9 +164,14 @@ async fn pinned_read_failure_does_not_adopt_fallback_credentials() -> Result<()> ResolvedOAuthCredentialStore::Keyring(AuthKeyringBackendKind::Direct), keyring, ); + let cached = store.load().await?.expect("cached credentials"); + assert_eq!( + cached.token_response.unwrap().access_token().secret(), + initial.token_response.0.access_token().secret() + ); assert!(matches!( - store.load().await, - Err(AuthError::InternalError(_)) + store.acquire_transaction_guard().await, + Err(AuthError::CredentialStoreError(_)) )); assert_eq!( store.stored_credentials().await, @@ -191,14 +204,21 @@ async fn replacement_or_removal_does_not_acknowledge_a_new_runtime_snapshot() -> replacement.issuer = issuer; save_oauth_tokens_to_file(&replacement)?; assert!(matches!( - store.load().await, + store.acquire_transaction_guard().await, Err(AuthError::AuthorizationRequired) )); assert_eq!(store.stored_credentials().await, original_snapshot); } - let _guard = store.acquire_transaction_guard().await?; + save_oauth_tokens_to_file(&initial)?; + let guard = store.acquire_transaction_guard().await?; store.clear().await?; - assert!(store.load().await?.is_none()); + assert!(load_oauth_tokens_from_file(&initial.server_name, &initial.url)?.is_none()); + drop(guard); + assert!(matches!( + store.acquire_transaction_guard().await, + Err(AuthError::AuthorizationRequired) + )); + assert!(store.load().await?.is_some()); assert_eq!(store.stored_credentials().await, original_snapshot); Ok(()) } @@ -219,7 +239,7 @@ async fn storage_roundtrip_preserves_absolute_and_unknown_expiry() -> Result<()> (Some(future_received_at), Some(future_deadline)), (None, None), ] { - let _guard = store.acquire_transaction_guard().await?; + let guard = store.acquire_transaction_guard().await?; let mut credentials = store.load().await?.unwrap(); credentials.token_received_at = received_at; credentials @@ -236,6 +256,8 @@ async fn storage_roundtrip_preserves_absolute_and_unknown_expiry() -> Result<()> .0 .set_expires_in(expected_deadline.map(|_| Duration::ZERO).as_ref()); assert_tokens_match_without_expiry(&durable, &expected); + drop(guard); + let _guard = store.acquire_transaction_guard().await?; let reloaded = store.load().await?.unwrap(); let expires_in = reloaded.token_response.unwrap().expires_in(); if let Some(expected_deadline) = expected_deadline { diff --git a/codex-rs/rmcp-client/src/oauth/tests/persistor_tests.rs b/codex-rs/rmcp-client/src/oauth/tests/persistor_tests.rs index b285fb6590..aa8155043f 100644 --- a/codex-rs/rmcp-client/src/oauth/tests/persistor_tests.rs +++ b/codex-rs/rmcp-client/src/oauth/tests/persistor_tests.rs @@ -9,14 +9,18 @@ use codex_config::types::OAuthCredentialsStoreMode; use codex_exec_server::RouteAwareHttpClient; use codex_http_client::HttpClientFactory; use codex_http_client::OutboundProxyPolicy; +use codex_keyring_store::DefaultKeyringStore; use http::HeaderMap; use keyring::Error as KeyringError; use oauth2::AccessToken; +use oauth2::RefreshToken; use oauth2::TokenResponse; use pretty_assertions::assert_eq; +use rmcp::transport::auth::AuthClient; use rmcp::transport::auth::AuthError; use rmcp::transport::auth::AuthorizationManager; use rmcp::transport::auth::OAuthState; +use rmcp::transport::streamable_http_client::StreamableHttpClient; use tokio::sync::Mutex as TokioMutex; use tracing::Event; use tracing::Id; @@ -29,6 +33,7 @@ use wiremock::Mock; use wiremock::MockServer; use wiremock::ResponseTemplate; use wiremock::matchers::body_string_contains; +use wiremock::matchers::header; use wiremock::matchers::method; use wiremock::matchers::path; @@ -36,21 +41,221 @@ use super::MockKeyringStore; use super::TempCodexHome; use super::assert_tokens_match_without_expiry; use super::sample_tokens; +use crate::http_client_adapter::StreamableHttpClientAdapter; +use crate::http_client_adapter::StreamableHttpRedirectMode; +use crate::oauth::OAuthCredentialStore; use crate::oauth::OAuthPersistor; +use crate::oauth::OAuthRuntime; use crate::oauth::ResolvedOAuthCredentialStore; use crate::oauth::StoredOAuthTokens; use crate::oauth::WrappedOAuthTokenResponse; use crate::oauth::compute_store_key; +use crate::oauth::delete_oauth_tokens; use crate::oauth::load_oauth_tokens_from_file; use crate::oauth::refresh_lock::RefreshCredentialLock; +use crate::oauth::save_oauth_tokens; use crate::oauth::save_oauth_tokens_to_file; use crate::oauth::stored_oauth_credentials; use crate::oauth_http_client::OAuthHttpClientAdapter; +use crate::oauth_http_client::PROACTIVE_REFRESH_TIMEOUT; use crate::startup_error::is_authentication_required_error; const REFRESH_LOCK_CONTENTION_EVENT_TARGET: &str = "codex_rmcp_client::oauth::refresh_lock::contention"; +#[tokio::test(flavor = "current_thread")] +async fn login_and_logout_follow_the_completed_refresh() -> Result<()> { + let (_env, _server, initial) = test_context().await?; + let mut replacement = initial.clone(); + replacement + .token_response + .0 + .set_access_token(AccessToken::new("replacement-access-token".into())); + replacement + .token_response + .0 + .set_refresh_token(Some(RefreshToken::new("replacement-refresh-token".into()))); + let mut refreshed = initial.clone(); + refreshed + .token_response + .0 + .set_refresh_token(Some(RefreshToken::new("rotated-refresh-token".into()))); + + for expected in [Some(replacement), None] { + save_oauth_tokens_to_file(&initial)?; + 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 mutation = tokio::spawn({ + let initial = initial.clone(); + let expected = expected.clone(); + async move { + match expected { + Some(tokens) => { + save_oauth_tokens( + &tokens.server_name, + &tokens, + OAuthCredentialsStoreMode::File, + AuthKeyringBackendKind::Direct, + ) + .await + } + None => { + assert!( + delete_oauth_tokens( + &initial.server_name, + &initial.url, + OAuthCredentialsStoreMode::File, + AuthKeyringBackendKind::Direct, + ) + .await? + ); + Ok(()) + } + } + } + }); + wait_for_lock_contention(contended_rx, /*expected_count*/ 1).await?; + assert!(!mutation.is_finished()); + assert_eq!( + load_oauth_tokens_from_file(&initial.server_name, &initial.url)?, + Some(initial.clone()) + ); + // Complete the in-flight refresh before allowing the login or logout to write. + save_oauth_tokens_to_file(&refreshed)?; + drop(held_lock); + mutation.await??; + assert_eq!( + load_oauth_tokens_from_file(&initial.server_name, &initial.url)?, + expected + ); + } + Ok(()) +} + +#[tokio::test(flavor = "current_thread")] +async fn fresh_preparation_skips_storage_and_refresh_lock() -> Result<()> { + let (env, server, mut initial) = test_context().await?; + let fresh = sample_tokens(); + initial.token_response = fresh.token_response; + initial.expires_at = fresh.expires_at; + std::fs::create_dir(env.path().join(".credentials.json"))?; + Mock::given(method("POST")) + .and(path("/oauth/token")) + .respond_with(ResponseTemplate::new(500)) + .expect(0) + .mount(&server) + .await; + + for mode in [ + crate::McpOAuthRefreshMode::Legacy, + crate::McpOAuthRefreshMode::Coordinated, + ] { + let runtime = runtime_for(&initial, mode).await?; + let held_lock = + RefreshCredentialLock::acquire_for_server(&initial.server_name, &initial.url).await?; + tokio::time::timeout( + Duration::from_millis(/*millis*/ 100), + runtime.refresh_if_needed(), + ) + .await + .context("fresh preparation must not wait for the refresh lock")??; + drop(held_lock); + } + server.verify().await; + Ok(()) +} + +#[tokio::test(flavor = "current_thread")] +async fn coordinated_401_refresh_rereads_and_persists_before_retry() -> Result<()> { + let (_env, server, mut initial) = test_context().await?; + let fresh = sample_tokens(); + initial.token_response = fresh.token_response; + initial.expires_at = fresh.expires_at; + save_oauth_tokens_to_file(&initial)?; + let mut manager = authorization_manager_for(&initial).await?; + manager.set_credential_store(OAuthCredentialStore::new( + initial.clone(), + ResolvedOAuthCredentialStore::File, + DefaultKeyringStore, + )); + let client = AuthClient::new( + StreamableHttpClientAdapter::new( + Arc::new(RouteAwareHttpClient::new(HttpClientFactory::new( + OutboundProxyPolicy::ReqwestDefault, + ))), + HeaderMap::new(), + /*auth_provider*/ None, + /*has_configured_headers*/ false, + StreamableHttpRedirectMode::Legacy, + Arc::default(), + ), + manager, + ); + let mut durable = initial.clone(); + durable + .token_response + .0 + .set_access_token(AccessToken::new("durable-access-token".into())); + durable + .token_response + .0 + .set_refresh_token(Some(RefreshToken::new("durable-refresh-token".into()))); + save_oauth_tokens_to_file(&durable)?; + + Mock::given(method("POST")) + .and(path("/mcp")) + .and(header("authorization", "Bearer access-token")) + .respond_with(ResponseTemplate::new(401).insert_header("www-authenticate", "Bearer")) + .expect(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/oauth/token")) + .and(body_string_contains("refresh_token=durable-refresh-token")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "access_token": "refreshed-access-token", + "refresh_token": "rotated-refresh-token", + "token_type": "Bearer", + "expires_in": 3600, + }))) + .expect(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/mcp")) + .and(header("authorization", "Bearer refreshed-access-token")) + .respond_with(move |_: &wiremock::Request| { + let saved = load_oauth_tokens_from_file(&durable.server_name, &durable.url) + .expect("read persisted refresh before retry") + .expect("refresh must be persisted before retry"); + let mut expected = durable.token_response.0.clone(); + expected.set_access_token(AccessToken::new("refreshed-access-token".into())); + expected.set_refresh_token(Some(RefreshToken::new("rotated-refresh-token".into()))); + expected.set_expires_in(saved.token_response.0.expires_in().as_ref()); + assert_eq!(saved.token_response, WrappedOAuthTokenResponse(expected)); + ResponseTemplate::new(202) + }) + .expect(1) + .mount(&server) + .await; + client + .post_message( + initial.url.into(), + serde_json::from_value( + serde_json::json!({"jsonrpc": "2.0", "id": 1, "method": "ping"}), + )?, + /*session_id*/ None, + /*auth_token*/ None, + Default::default(), + ) + .await?; + server.verify().await; + Ok(()) +} + struct LockContentionSubscriber { contended_tx: mpsc::Sender<()>, } @@ -95,6 +300,11 @@ impl Subscriber for LockContentionSubscriber { #[tokio::test(flavor = "current_thread")] async fn concurrent_refreshes_call_provider_once_and_carry_omitted_fields() -> Result<()> { + assert_concurrent_refreshes(crate::McpOAuthRefreshMode::Legacy).await?; + assert_concurrent_refreshes(crate::McpOAuthRefreshMode::Coordinated).await +} + +async fn assert_concurrent_refreshes(mode: crate::McpOAuthRefreshMode) -> Result<()> { let (_env, server, initial) = test_context().await?; Mock::given(method("POST")) .and(path("/oauth/token")) @@ -119,9 +329,8 @@ async fn concurrent_refreshes_call_provider_once_and_carry_omitted_fields() -> R let _subscriber_guard = tracing::subscriber::set_default(LockContentionSubscriber { contended_tx }); - let first = persistor_for(&initial).await?; - let second = persistor_for(&initial).await?; - let initial_credentials = first.stored_credentials().await; + let first = runtime_for(&initial, mode).await?; + let second = runtime_for(&initial, mode).await?; let first_task = tokio::spawn({ let first = first.clone(); async move { first.refresh_if_needed().await } @@ -140,11 +349,16 @@ async fn concurrent_refreshes_call_provider_once_and_carry_omitted_fields() -> R // 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?; + if let OAuthRuntime::Legacy(persistor) = &first { + persistor.persist_if_needed().await?; + } let stored = load_oauth_tokens_from_file(&initial.server_name, &initial.url)? .expect("refreshed credentials should be stored"); assert_eq!(stored.issuer, initial.issuer); - let live_credentials = first.stored_credentials().await; + let live_credentials = match &first { + OAuthRuntime::Legacy(persistor) => persistor.stored_credentials().await, + OAuthRuntime::Coordinated { store, .. } => store.stored_credentials().await, + }; let disk_credentials = stored_oauth_credentials( &initial.server_name, &initial.url, @@ -152,7 +366,6 @@ async fn concurrent_refreshes_call_provider_once_and_carry_omitted_fields() -> R AuthKeyringBackendKind::Direct, )?; assert_eq!(live_credentials, disk_credentials); - assert_ne!(live_credentials, initial_credentials); 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 @@ -176,7 +389,7 @@ async fn resolved_keyring_read_error_preserves_in_memory_credentials() -> Result 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 manager = Arc::new(TokioMutex::new(authorization_manager_for(&initial).await?)); let persistor = OAuthPersistor::new( initial.server_name.clone(), initial.url.clone(), @@ -354,6 +567,11 @@ async fn transient_refresh_failure_does_not_require_reauthorization() -> Result< #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn caller_cancellation_does_not_cancel_refresh_persistence() -> Result<()> { + assert_caller_cancellation(crate::McpOAuthRefreshMode::Legacy).await?; + assert_caller_cancellation(crate::McpOAuthRefreshMode::Coordinated).await +} + +async fn assert_caller_cancellation(mode: crate::McpOAuthRefreshMode) -> 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(); @@ -381,11 +599,8 @@ async fn caller_cancellation_does_not_cancel_refresh_persistence() -> Result<()> .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 } - }); + let runtime = runtime_for(&initial, mode).await?; + let caller = tokio::spawn(async move { runtime.refresh_if_needed().await }); tokio::task::spawn_blocking(move || { request_received_rx @@ -452,11 +667,83 @@ async fn provider_timeout_releases_lock_and_preserves_durable_credentials() -> R 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 coordinated_provider_timeout_excludes_lock_wait() -> Result<()> { + let (_env, server, initial) = test_context().await?; + mount_delayed_refresh(&server, "late-access-token").await; + save_oauth_tokens_to_file(&initial)?; + let (auth_manager, _) = coordinated_manager_for(&initial).await?; + 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 mut refresh = tokio::spawn( + PROACTIVE_REFRESH_TIMEOUT.scope(Duration::from_millis(/*millis*/ 50), async move { + auth_manager.lock().await.get_access_token().await + }), + ); + wait_for_lock_contention(contended_rx, /*expected_count*/ 1).await?; + assert!( + tokio::time::timeout(Duration::from_millis(/*millis*/ 100), &mut refresh) + .await + .is_err(), + "lock wait must not consume the provider timeout" + ); + drop(held_lock); + let error = refresh + .await? + .expect_err("the scoped provider request must time out"); + assert!(matches!(error, AuthError::TokenRefreshFailed(_))); + let _lock = tokio::time::timeout( + Duration::from_millis(/*millis*/ 100), + RefreshCredentialLock::acquire_for_server(&initial.server_name, &initial.url), + ) + .await??; + let stored = load_oauth_tokens_from_file(&initial.server_name, &initial.url)?.unwrap(); + assert_tokens_match_without_expiry(&stored, &initial); + server.verify().await; + Ok(()) +} + +async fn runtime_for( + tokens: &StoredOAuthTokens, + mode: crate::McpOAuthRefreshMode, +) -> Result { + Ok(match mode { + crate::McpOAuthRefreshMode::Legacy => OAuthRuntime::Legacy(persistor_for(tokens).await?), + crate::McpOAuthRefreshMode::Coordinated => { + let (auth_manager, store) = coordinated_manager_for(tokens).await?; + OAuthRuntime::Coordinated { + auth_manager, + store, + } + } + }) +} + +async fn coordinated_manager_for( + tokens: &StoredOAuthTokens, +) -> Result<(Arc>, OAuthCredentialStore)> { + let mut manager = authorization_manager_for(tokens).await?; + let store = OAuthCredentialStore::new( + tokens.clone(), + ResolvedOAuthCredentialStore::File, + DefaultKeyringStore, + ); + manager.set_credential_store(store.clone()); + Ok((Arc::new(TokioMutex::new(manager)), store)) +} + async fn persistor_for(tokens: &StoredOAuthTokens) -> Result { Ok(OAuthPersistor::new( tokens.server_name.clone(), tokens.url.clone(), - authorization_manager_for(tokens).await?, + Arc::new(TokioMutex::new(authorization_manager_for(tokens).await?)), ResolvedOAuthCredentialStore::File, Some(tokens.clone()), )) @@ -470,9 +757,7 @@ async fn test_context() -> Result<(TempCodexHome, MockServer, StoredOAuthTokens) Ok((env, server, tokens)) } -async fn authorization_manager_for( - tokens: &StoredOAuthTokens, -) -> Result>> { +async fn authorization_manager_for(tokens: &StoredOAuthTokens) -> Result { let oauth_http_client = Arc::new(OAuthHttpClientAdapter::new( Arc::new(RouteAwareHttpClient::new(HttpClientFactory::new( OutboundProxyPolicy::ReqwestDefault, @@ -492,7 +777,7 @@ async fn authorization_manager_for( } _ => anyhow::bail!("unexpected OAuth state"), }; - Ok(Arc::new(TokioMutex::new(manager))) + Ok(manager) } async fn mount_oauth_metadata(server: &MockServer) { diff --git a/codex-rs/rmcp-client/src/oauth_http_client.rs b/codex-rs/rmcp-client/src/oauth_http_client.rs index 14c70aba1b..9e4e21e93d 100644 --- a/codex-rs/rmcp-client/src/oauth_http_client.rs +++ b/codex-rs/rmcp-client/src/oauth_http_client.rs @@ -38,6 +38,11 @@ const MAX_OAUTH_HTTP_RESPONSE_BODY_BYTES: usize = 1024 * 1024; const MAX_OAUTH_HTTP_REDIRECTS: usize = 10; static NEXT_OAUTH_REQUEST_ID: AtomicU64 = AtomicU64::new(0); +tokio::task_local! { + /// Bounds provider HTTP work during preparation, excluding credential lock waits and saves. + pub(crate) static PROACTIVE_REFRESH_TIMEOUT: Duration; +} + #[derive(Debug, thiserror::Error)] enum OAuthHttpClientAdapterError { #[error("unsupported OAuth HTTP redirect policy")] @@ -324,7 +329,16 @@ fn oauth_redirect_policy( impl OAuthHttpClient for OAuthHttpClientAdapter { fn execute(&self, request: OAuthHttpRequest) -> OAuthHttpClientFuture<'_> { - Box::pin(self.execute_request(request.request, request.redirect_policy, request.timeout)) + Box::pin(async move { + let operation = + self.execute_request(request.request, request.redirect_policy, request.timeout); + match PROACTIVE_REFRESH_TIMEOUT.try_with(|duration| *duration) { + Ok(duration) => tokio::time::timeout(duration, operation) + .await + .map_err(|_| oauth_http_client_error(OAuthHttpClientAdapterError::TimedOut))?, + Err(_) => operation.await, + } + }) } } diff --git a/codex-rs/rmcp-client/src/perform_oauth_login.rs b/codex-rs/rmcp-client/src/perform_oauth_login.rs index b2ac1bc7c4..67280978e2 100644 --- a/codex-rs/rmcp-client/src/perform_oauth_login.rs +++ b/codex-rs/rmcp-client/src/perform_oauth_login.rs @@ -721,7 +721,8 @@ impl OauthLoginFlow { &stored, self.store_mode, self.keyring_backend_kind, - )?; + ) + .await?; Ok(()) } diff --git a/codex-rs/rmcp-client/src/rmcp_client.rs b/codex-rs/rmcp-client/src/rmcp_client.rs index b549a0621d..36cc5e88d9 100644 --- a/codex-rs/rmcp-client/src/rmcp_client.rs +++ b/codex-rs/rmcp-client/src/rmcp_client.rs @@ -77,7 +77,9 @@ use crate::http_client_adapter::StreamableHttpClientAdapter; use crate::http_client_adapter::StreamableHttpClientAdapterError; use crate::http_client_adapter::StreamableHttpRedirectMode; use crate::in_process_transport::InProcessTransportFactory; +use crate::oauth::OAuthCredentialStore; use crate::oauth::OAuthPersistor; +use crate::oauth::OAuthRuntime; use crate::oauth::ResolvedOAuthCredentialStore; use crate::oauth::ResolvedOAuthTokens; use crate::oauth::StoredOAuthTokens; @@ -85,6 +87,7 @@ use crate::oauth::install_tokens_in_manager; use crate::oauth::resolve_oauth_tokens_from_store_policy; use crate::oauth::validate_refresh_token_issuer; use crate::oauth_http_client::OAuthHttpClientAdapter; +use crate::oauth_refresh_mode::McpOAuthRefreshMode; use crate::protocol_mode::McpProtocolMode; use crate::stdio_server_launcher::StdioServerCommand; use crate::stdio_server_launcher::StdioServerLauncher; @@ -112,7 +115,7 @@ enum PendingTransport { }, StreamableHttpWithOAuth { transport: StreamableHttpClientTransport>, - oauth_persistor: OAuthPersistor, + oauth_runtime: OAuthRuntime, }, StreamableHttpWithAccessTokenOnly { transport: StreamableHttpClientTransport>, @@ -125,7 +128,7 @@ enum ClientState { }, Ready { service: Arc>, - oauth: Option, + oauth: Option, }, Closed, } @@ -160,6 +163,7 @@ enum TransportRecipe { http_client: Arc, auth_provider: Option, redirect_mode: StreamableHttpRedirectMode, + oauth_refresh_mode: McpOAuthRefreshMode, initialize_deadline: Arc>>, }, } @@ -539,6 +543,7 @@ impl RmcpClient { auth_provider, protocol_mode, StreamableHttpRedirectMode::Legacy, + McpOAuthRefreshMode::Legacy, ) .await } @@ -556,6 +561,7 @@ impl RmcpClient { auth_provider: Option, protocol_mode: McpProtocolMode, redirect_mode: StreamableHttpRedirectMode, + oauth_refresh_mode: McpOAuthRefreshMode, ) -> Result { let transport_recipe = TransportRecipe::StreamableHttp { server_name: server_name.to_string(), @@ -569,6 +575,7 @@ impl RmcpClient { http_client, auth_provider, redirect_mode, + oauth_refresh_mode, initialize_deadline: Arc::new(StdMutex::new(None)), }; let transport = Self::create_pending_transport(&transport_recipe).await?; @@ -611,7 +618,7 @@ impl RmcpClient { } }; - let (service, oauth_persistor) = self + let (service, oauth_runtime) = self .connect_pending_transport_with_initialize_retries( pending_transport, client_service.clone(), @@ -640,11 +647,11 @@ impl RmcpClient { } *guard = ClientState::Ready { service, - oauth: oauth_persistor.clone(), + oauth: oauth_runtime.clone(), }; } - if let Some(runtime) = oauth_persistor + if let Some(OAuthRuntime::Legacy(runtime)) = oauth_runtime && let Err(error) = runtime.persist_if_needed().await { warn!("failed to persist OAuth tokens after initialize: {error}"); @@ -937,7 +944,7 @@ impl RmcpClient { } } - async fn oauth_persistor(&self) -> Option { + async fn oauth_runtime(&self) -> Option { let guard = self.state.lock().await; match &*guard { ClientState::Ready { @@ -950,8 +957,11 @@ impl RmcpClient { /// Returns `None` when this client does not manage stored OAuth credentials. pub async fn managed_oauth_credentials(&self) -> Option> { - let persistor = self.oauth_persistor().await?; - Some(persistor.stored_credentials().await) + let runtime = self.oauth_runtime().await?; + Some(match runtime { + OAuthRuntime::Legacy(persistor) => persistor.stored_credentials().await, + OAuthRuntime::Coordinated { store, .. } => store.stored_credentials().await, + }) } /// Returns whether an initialized transport or its underlying service has stopped. @@ -985,16 +995,16 @@ impl RmcpClient { /// This should be called after every tool call so that if a given tool call triggered /// a refresh of the OAuth tokens, they are persisted. async fn persist_oauth_tokens(&self) { - if let Some(runtime) = self.oauth_persistor().await + if let Some(OAuthRuntime::Legacy(runtime)) = self.oauth_runtime().await && let Err(error) = runtime.persist_if_needed().await { warn!("failed to persist OAuth tokens: {error}"); } } - /// OAuth uses independent lock/request bounds and completes before the operation timeout starts. + /// Prepares managed OAuth before the MCP operation timeout starts. async fn refresh_oauth_if_needed(&self) -> Result<()> { - if let Some(runtime) = self.oauth_persistor().await { + if let Some(runtime) = self.oauth_runtime().await { runtime.refresh_if_needed().await?; } Ok(()) @@ -1026,6 +1036,7 @@ impl RmcpClient { http_client, auth_provider, redirect_mode, + oauth_refresh_mode, initialize_deadline, } => { let has_configured_headers = matches!( @@ -1112,6 +1123,7 @@ impl RmcpClient { Arc::clone(http_client), has_configured_headers, *redirect_mode, + *oauth_refresh_mode, Arc::clone(initialize_deadline), ) .await @@ -1180,7 +1192,7 @@ impl RmcpClient { timeout: Option, ) -> Result<( Arc>, - Option, + Option, )> { let _initialize_deadline = match &self.transport_recipe { TransportRecipe::StreamableHttp { @@ -1198,7 +1210,7 @@ impl RmcpClient { TransportRecipe::InProcess { .. } | TransportRecipe::Stdio { .. } => None, }; let lifecycle = self.protocol_mode.client_lifecycle(); - let (transport, oauth_persistor) = match pending_transport { + let (transport, oauth_runtime) = match pending_transport { PendingTransport::InProcess { transport } => ( client_service .serve_with_lifecycle(transport, lifecycle) @@ -1219,12 +1231,12 @@ impl RmcpClient { ), PendingTransport::StreamableHttpWithOAuth { transport, - oauth_persistor, + oauth_runtime, } => ( client_service .serve_with_lifecycle(transport, lifecycle) .boxed(), - Some(oauth_persistor), + Some(oauth_runtime), ), PendingTransport::StreamableHttpWithAccessTokenOnly { transport } => ( client_service @@ -1250,7 +1262,7 @@ impl RmcpClient { let service = match service_result { Ok(service) => service, Err(error) => { - if let Some(runtime) = oauth_persistor.as_ref() + if let Some(OAuthRuntime::Legacy(runtime)) = oauth_runtime.as_ref() && let Err(persist_error) = runtime.persist_if_needed().await { warn!( @@ -1268,7 +1280,7 @@ impl RmcpClient { .set_response_cache_config(ClientCacheConfig::disabled()) .await; - Ok((Arc::new(service), oauth_persistor)) + Ok((Arc::new(service), oauth_runtime)) } async fn run_service_operation( @@ -1459,7 +1471,7 @@ impl RmcpClient { .clone() .ok_or_else(|| anyhow!("MCP client cannot recover before initialize succeeds"))?; let pending_transport = Self::create_pending_transport(&self.transport_recipe).await?; - let (service, oauth_persistor) = self + let (service, oauth_runtime) = self .connect_pending_transport_with_initialize_retries( pending_transport, initialize_context.client_service, @@ -1478,11 +1490,11 @@ impl RmcpClient { } *guard = ClientState::Ready { service, - oauth: oauth_persistor.clone(), + oauth: oauth_runtime.clone(), }; } - if let Some(runtime) = oauth_persistor + if let Some(OAuthRuntime::Legacy(runtime)) = oauth_runtime && let Err(error) = runtime.persist_if_needed().await { warn!("failed to persist OAuth tokens after session recovery: {error}"); @@ -1502,6 +1514,7 @@ async fn create_oauth_transport_and_runtime( http_client: Arc, has_configured_headers: bool, redirect_mode: StreamableHttpRedirectMode, + oauth_refresh_mode: McpOAuthRefreshMode, initialize_deadline: Arc>>, ) -> Result { let oauth_http_client = Arc::new(OAuthHttpClientAdapter::new_with_redirect_mode( @@ -1533,6 +1546,18 @@ async fn create_oauth_transport_and_runtime( runtime_tokens.issuer = None; } install_tokens_in_manager(&mut manager, &runtime_tokens).await?; + let coordinated_store = match oauth_refresh_mode { + McpOAuthRefreshMode::Coordinated if !use_stored_access_token_only => { + let store = OAuthCredentialStore::new( + initial_tokens.clone(), + credential_store, + DefaultKeyringStore, + ); + manager.set_credential_store(store.clone()); + Some(store) + } + McpOAuthRefreshMode::Legacy | McpOAuthRefreshMode::Coordinated => None, + }; let auth_client = AuthClient::new( StreamableHttpClientAdapter::new( @@ -1559,17 +1584,23 @@ async fn create_oauth_transport_and_runtime( return Ok(PendingTransport::StreamableHttpWithAccessTokenOnly { transport }); } - let runtime = OAuthPersistor::new( - server_name.to_string(), - url.to_string(), - auth_manager, - credential_store, - Some(initial_tokens), - ); + let runtime = match coordinated_store { + Some(store) => OAuthRuntime::Coordinated { + auth_manager, + store, + }, + None => OAuthRuntime::Legacy(OAuthPersistor::new( + server_name.to_string(), + url.to_string(), + auth_manager, + credential_store, + Some(initial_tokens), + )), + }; Ok(PendingTransport::StreamableHttpWithOAuth { transport, - oauth_persistor: runtime, + oauth_runtime: runtime, }) } diff --git a/codex-rs/rmcp-client/src/streamable_http_retry.rs b/codex-rs/rmcp-client/src/streamable_http_retry.rs index 2a71a7defc..3f564630eb 100644 --- a/codex-rs/rmcp-client/src/streamable_http_retry.rs +++ b/codex-rs/rmcp-client/src/streamable_http_retry.rs @@ -14,7 +14,7 @@ use tracing::warn; use crate::elicitation_client_service::ElicitationClientService; use crate::http_client_adapter::StreamableHttpClientAdapterError; -use crate::oauth::OAuthPersistor; +use crate::oauth::OAuthRuntime; use super::PendingTransport; use super::RmcpClient; @@ -30,7 +30,7 @@ impl RmcpClient { timeout: Option, ) -> Result<( Arc>, - Option, + Option, )> { let should_retry = match &initial_transport { PendingTransport::InProcess { .. } | PendingTransport::Stdio { .. } => false, @@ -63,14 +63,11 @@ impl RmcpClient { } } }; - if let PendingTransport::StreamableHttpWithOAuth { - oauth_persistor, .. - } = &transport - { + if let PendingTransport::StreamableHttpWithOAuth { oauth_runtime, .. } = &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?; + oauth_runtime.refresh_if_needed().await?; if let Some(deadline) = retry_deadline.as_mut() { *deadline += refresh_started_at.elapsed(); } 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 b1855e8a1e..9b8162d077 100644 --- a/codex-rs/rmcp-client/tests/streamable_http_oauth_startup.rs +++ b/codex-rs/rmcp-client/tests/streamable_http_oauth_startup.rs @@ -13,6 +13,8 @@ use codex_config::types::OAuthCredentialsStoreMode; use codex_exec_server::Environment; use codex_rmcp_client::McpAuthState; use codex_rmcp_client::McpLoginRequirement; +use codex_rmcp_client::McpOAuthRefreshMode; +use codex_rmcp_client::McpProtocolMode; use codex_rmcp_client::OAuthDiscoveryTimeout; use codex_rmcp_client::RmcpClient; use codex_rmcp_client::StoredOAuthTokens; @@ -43,6 +45,7 @@ use wiremock::matchers::method; use wiremock::matchers::path; use streamable_http_test_support::initialize_client; +use streamable_http_test_support::initialize_client_with_timeout; const SERVER_NAME: &str = "test-streamable-http-oauth-startup"; const EXPIRED_ACCESS_TOKEN: &str = "expired-access-token"; @@ -52,6 +55,7 @@ const RESOURCE_API_KEY: &str = "resource-api-key-secret"; const RESOURCE_USER_AGENT: &str = "resource-only-user-agent"; const MCP_USER_AGENT: &str = concat!("codex-mcp-client/", env!("CARGO_PKG_VERSION")); const CHILD_SERVER_URL_ENV: &str = "MCP_TEST_OAUTH_STARTUP_SERVER_URL"; +const CHILD_REFRESH_MODE_ENV: &str = "MCP_TEST_OAUTH_REFRESH_MODE"; const CHILD_HELPER_COMMAND_ENV: &str = "MCP_TEST_OAUTH_STARTUP_HELPER_COMMAND"; const CHILD_RESOURCE_API_KEY_ENV: &str = "MCP_TEST_OAUTH_STARTUP_RESOURCE_API_KEY"; const CHILD_STORED_ISSUER_ENV: &str = "MCP_TEST_OAUTH_STARTUP_STORED_ISSUER"; @@ -76,25 +80,44 @@ enum IssuerMismatchAccessToken { #[tokio::test(flavor = "multi_thread", worker_threads = 1)] async fn refreshes_expired_persisted_token_before_initialize() -> anyhow::Result<()> { - assert_expired_token_refresh(OAuthStartupScenario::DirectAuthorizationMetadata).await + assert_expired_token_refresh( + OAuthStartupScenario::DirectAuthorizationMetadata, + McpOAuthRefreshMode::Coordinated, + ) + .await } #[tokio::test(flavor = "multi_thread", worker_threads = 1)] async fn refreshes_oauth_with_gateway_headers_helper() -> anyhow::Result<()> { - assert_expired_token_refresh(OAuthStartupScenario::GatewayHeadersHelper).await + assert_expired_token_refresh( + OAuthStartupScenario::GatewayHeadersHelper, + McpOAuthRefreshMode::Legacy, + ) + .await } #[tokio::test(flavor = "multi_thread", worker_threads = 1)] async fn refreshes_oauth_after_gateway_rejects_token_request() -> anyhow::Result<()> { - assert_expired_token_refresh(OAuthStartupScenario::SameOriginGatewayHeadersHelper).await + assert_expired_token_refresh( + OAuthStartupScenario::SameOriginGatewayHeadersHelper, + McpOAuthRefreshMode::Legacy, + ) + .await } #[tokio::test(flavor = "multi_thread", worker_threads = 1)] async fn refresh_uses_discovered_protected_resource_audience() -> anyhow::Result<()> { - assert_expired_token_refresh(OAuthStartupScenario::ProtectedResourceMetadata).await + assert_expired_token_refresh( + OAuthStartupScenario::ProtectedResourceMetadata, + McpOAuthRefreshMode::Legacy, + ) + .await } -async fn assert_expired_token_refresh(scenario: OAuthStartupScenario) -> anyhow::Result<()> { +async fn assert_expired_token_refresh( + scenario: OAuthStartupScenario, + refresh_mode: McpOAuthRefreshMode, +) -> anyhow::Result<()> { let server = MockServer::start().await; let authorization_server = MockServer::start().await; let same_origin_gateway = matches!( @@ -198,12 +221,18 @@ async fn assert_expired_token_refresh(scenario: OAuthStartupScenario) -> anyhow: { ResponseTemplate::new(/*s*/ 401) } else { - ResponseTemplate::new(/*s*/ 200).set_body_json(json!({ + let response = ResponseTemplate::new(/*s*/ 200).set_body_json(json!({ "access_token": REFRESHED_ACCESS_TOKEN, "token_type": "Bearer", "expires_in": 7200, "refresh_token": REFRESH_TOKEN, - })) + })); + if refresh_mode == McpOAuthRefreshMode::Coordinated { + // Longer than the child's handshake timeout: refresh must finish first. + response.set_delay(Duration::from_secs(/*secs*/ 2)) + } else { + response + } } }) .expect(if same_origin_gateway { 2 } else { 1 }) @@ -266,6 +295,14 @@ async fn assert_expired_token_refresh(scenario: OAuthStartupScenario) -> anyhow: .env(CHILD_SERVER_URL_ENV, server_url) .env(CHILD_STORED_ISSUER_ENV, authorization_server_issuer) .env(CHILD_RESOURCE_API_KEY_ENV, RESOURCE_API_KEY) + .env( + CHILD_REFRESH_MODE_ENV, + if refresh_mode == McpOAuthRefreshMode::Coordinated { + "coordinated" + } else { + "legacy" + }, + ) .env("MCP_TEST_AMBIENT_SECRET", "must-not-reach-helper"); if with_headers_helper { command.env( @@ -632,7 +669,8 @@ async fn persisted_credentials_auth_status_child() -> anyhow::Result<()> { &tokens, OAuthCredentialsStoreMode::File, AuthKeyringBackendKind::default(), - )?; + ) + .await?; let status = auth_status(LEGACY_REFRESHABLE_SERVER_URL).await?; assert_eq!( @@ -662,7 +700,8 @@ async fn persisted_credentials_auth_status_child() -> anyhow::Result<()> { &tokens, OAuthCredentialsStoreMode::File, AuthKeyringBackendKind::default(), - )?; + ) + .await?; let status = auth_status(UNEXPIRED_SERVER_URL).await?; assert_eq!(status, McpAuthState::OAuth); @@ -686,7 +725,8 @@ async fn persisted_credentials_auth_status_child() -> anyhow::Result<()> { &tokens, OAuthCredentialsStoreMode::File, AuthKeyringBackendKind::default(), - )?; + ) + .await?; let status = auth_status(REFRESHABLE_SERVER_URL).await?; assert_eq!(status, McpAuthState::OAuth); @@ -713,6 +753,11 @@ async fn auth_status(server_url: &str) -> anyhow::Result { #[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)?; + let refresh_mode = if std::env::var(CHILD_REFRESH_MODE_ENV).as_deref() == Ok("coordinated") { + McpOAuthRefreshMode::Coordinated + } else { + McpOAuthRefreshMode::Legacy + }; // Save an expired access token with a valid refresh token so startup must // refresh before sending the initialize request. @@ -736,7 +781,8 @@ async fn oauth_startup_child() -> anyhow::Result<()> { &tokens, OAuthCredentialsStoreMode::File, AuthKeyringBackendKind::default(), - )?; + ) + .await?; // This mirrors create_client's transport and initialization setup, except // it omits the direct bearer token. Supplying that token would bypass the @@ -750,7 +796,7 @@ async fn oauth_startup_child() -> anyhow::Result<()> { std::env::current_dir()?, )?; } - let client = RmcpClient::new_streamable_http_client( + let client = RmcpClient::new_streamable_http_client_with_protocol_mode_and_redirect_mode( SERVER_NAME, &server_url, /*bearer_token*/ None, @@ -766,10 +812,17 @@ async fn oauth_startup_child() -> anyhow::Result<()> { AuthKeyringBackendKind::default(), http_client, /*auth_provider*/ None, + McpProtocolMode::Legacy, + StreamableHttpRedirectMode::Legacy, + refresh_mode, ) .await?; - initialize_client(&client).await?; + if refresh_mode == McpOAuthRefreshMode::Coordinated { + initialize_client_with_timeout(&client, Duration::from_secs(/*secs*/ 1)).await?; + } else { + initialize_client(&client).await?; + } Ok(()) } @@ -795,7 +848,8 @@ async fn expired_unrefreshable_startup_child() -> anyhow::Result<()> { &tokens, OAuthCredentialsStoreMode::File, AuthKeyringBackendKind::default(), - )?; + ) + .await?; let client = RmcpClient::new_streamable_http_client( SERVER_NAME, @@ -852,7 +906,8 @@ async fn issuer_startup_child() -> anyhow::Result<()> { &tokens, OAuthCredentialsStoreMode::File, AuthKeyringBackendKind::default(), - )?; + ) + .await?; let client = RmcpClient::new_streamable_http_client( SERVER_NAME, diff --git a/codex-rs/rmcp-client/tests/streamable_http_oauth_store_pinning.rs b/codex-rs/rmcp-client/tests/streamable_http_oauth_store_pinning.rs index 3846fe4f81..1656774650 100644 --- a/codex-rs/rmcp-client/tests/streamable_http_oauth_store_pinning.rs +++ b/codex-rs/rmcp-client/tests/streamable_http_oauth_store_pinning.rs @@ -203,6 +203,7 @@ async fn transport_provided_bearer_token_avoids_placeholder_headers_and_redirect /*auth_provider*/ None, McpProtocolMode::Legacy, StreamableHttpRedirectMode::AgentPluginV1, + codex_rmcp_client::McpOAuthRefreshMode::Legacy, ) .await?; @@ -267,7 +268,8 @@ async fn auto_store_remains_pinned_across_session_recovery_child() -> anyhow::Re &file_tokens, OAuthCredentialsStoreMode::File, AuthKeyringBackendKind::Direct, - )?; + ) + .await?; let file_snapshot = stored_oauth_credential_snapshot( SERVER_NAME, &server_url, @@ -281,13 +283,15 @@ async fn auto_store_remains_pinned_across_session_recovery_child() -> anyhow::Re &keyring_tokens, OAuthCredentialsStoreMode::Auto, AuthKeyringBackendKind::Direct, - )?; + ) + .await?; save_oauth_tokens( SERVER_NAME, &file_tokens, OAuthCredentialsStoreMode::File, AuthKeyringBackendKind::Direct, - )?; + ) + .await?; assert_eq!( file_snapshot.reload( SERVER_NAME, 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 43bd3141bf..d289d6f83d 100644 --- a/codex-rs/rmcp-client/tests/streamable_http_test_support.rs +++ b/codex-rs/rmcp-client/tests/streamable_http_test_support.rs @@ -105,10 +105,17 @@ pub(crate) async fn create_client_with_http_client( } pub(crate) async fn initialize_client(client: &RmcpClient) -> anyhow::Result<()> { + initialize_client_with_timeout(client, Duration::from_secs(/*secs*/ 5)).await +} + +pub(crate) async fn initialize_client_with_timeout( + client: &RmcpClient, + timeout: Duration, +) -> anyhow::Result<()> { client .initialize( init_params(), - Some(Duration::from_secs(5)), + Some(timeout), Box::new(|_, _| { async { Ok(ElicitationResponse {