diff --git a/codex-rs/app-server/src/request_processors/mcp_processor.rs b/codex-rs/app-server/src/request_processors/mcp_processor.rs index 4164a96c5d..5553aca947 100644 --- a/codex-rs/app-server/src/request_processors/mcp_processor.rs +++ b/codex-rs/app-server/src/request_processors/mcp_processor.rs @@ -435,6 +435,7 @@ impl McpRequestProcessor { let mut snapshot = collect_mcp_server_status_snapshot_with_detail( &mcp_config, auth.as_ref(), + Some(&self.auth_manager), runtime_context, mcp_manager.codex_apps_tools_cache(), mcp_manager.tool_catalog_cache(), diff --git a/codex-rs/app-server/tests/suite/v2/account_enterprise_tests.rs b/codex-rs/app-server/tests/suite/v2/account_enterprise_tests.rs index 4946223158..6808e75883 100644 --- a/codex-rs/app-server/tests/suite/v2/account_enterprise_tests.rs +++ b/codex-rs/app-server/tests/suite/v2/account_enterprise_tests.rs @@ -320,7 +320,7 @@ async fn enterprise_login_attempt_lifecycle() -> Result<()> { "token_type": "N_A", "resource": token_resource, "scope": "files.read", })) }) - .with_priority(1) + .with_priority(2) .mount(&idp_server) .await; Mock::given(method("GET")) @@ -659,6 +659,7 @@ oauth = {{ client_id = "mcp-client", authorization_server_issuer = "{origin}/as" assert_eq!(status, &selected_status["data"][0]); assert!(!status["tools"]["calendar_list_events"].is_null()); if name == "enterprise" { + assert_eq!(status["authStatus"], "oAuth"); assert_eq!(status["toolListMeta"], catalog_meta); } } @@ -844,6 +845,73 @@ oauth = {{ client_id = "mcp-client", authorization_server_issuer = "{origin}/as" ordinary_result ); + // Rejected and missing grants need sign-in; administrator configuration errors do not. + let reauthentication = + Some(codex_app_server_protocol::McpServerStartupFailureReason::ReauthenticationRequired); + for (error, reason) in [ + (Some("invalid_client"), None), + (Some("insufficient_user_authentication"), reauthentication), + (Some("invalid_grant"), reauthentication), + (None, reauthentication), + ] { + if let Some(error) = error { + Mock::given(method("POST")) + .and(path("/idp/token")) + .and(wiremock::matchers::body_string_contains( + "grant-type%3Atoken-exchange", + )) + .respond_with(ResponseTemplate::new(400).set_body_json(json!({"error": error}))) + .with_priority(1) + .up_to_n_times(1) + .expect(1) + .mount(&idp_server) + .await; + } + let rejected = enterprise_rpc(&client, "thread/start", json!({})).await?; + let rejected_id = rejected["thread"]["id"].as_str().context("thread ID")?; + let failure = timeout(DEFAULT_READ_TIMEOUT, async { + loop { + if let InProcessServerEvent::ServerNotification(notification) = + client.next_event().await.context("startup notification")? + && let ServerNotification::McpServerStatusUpdated(notification) = *notification + && notification.thread_id.as_deref() == Some(rejected_id) + && notification.name == "enterprise" + && notification.status + == codex_app_server_protocol::McpServerStartupState::Failed + { + return anyhow::Ok(notification); + } + } + }) + .await??; + assert_eq!(failure.failure_reason, reason, "{error:?}: {failure:?}"); + if reason.is_some() { + let status = enterprise_rpc( + &client, + "mcpServerStatus/list", + json!({"threadId": rejected_id, "serverName": "enterprise"}), + ) + .await?; + assert_eq!(status["data"][0]["authStatus"], "notLoggedIn"); + } + } + + let recovery = enterprise_rpc( + &client, + "mcpServer/oauth/login", + json!({"name":"enterprise","threadId":thread_id}), + ) + .await?; + browser + .get(enterprise_callback(&recovery, &issuer, callback_port)?) + .header(http::header::CONNECTION, "close") + .send() + .await?; + assert_eq!( + enterprise_completion(&mut client).await?, + json!({"name":"enterprise","threadId":thread_id,"loginId":recovery["loginId"],"success":true}) + ); + let stalled = enterprise_rpc( &client, "mcpServer/oauth/login", diff --git a/codex-rs/cli/src/mcp_cmd.rs b/codex-rs/cli/src/mcp_cmd.rs index 02eb45c3e4..1915cb2abf 100644 --- a/codex-rs/cli/src/mcp_cmd.rs +++ b/codex-rs/cli/src/mcp_cmd.rs @@ -693,7 +693,8 @@ async fn run_list(config: &Config, list_args: ListArgs) -> Result<()> { AuthManager::shared_from_config(config, /*enable_codex_api_key_env*/ true).await?; let auth = auth_manager.auth().await; let mcp_servers = mcp_manager.configured_servers(config).await; - let effective_mcp_servers = mcp_manager.effective_servers(config, auth.as_ref()).await; + let runtime_config = mcp_manager.runtime_config(config).await; + let effective_mcp_servers = codex_mcp::effective_mcp_servers(&runtime_config, auth.as_ref()); let mut entries: Vec<_> = mcp_servers.iter().collect(); entries.sort_by_key(|(name, _)| *name); @@ -710,6 +711,8 @@ async fn run_list(config: &Config, list_args: ListArgs) -> Result<()> { config.mcp_oauth_credentials_store_mode, config.auth_keyring_backend_kind(), auth.as_ref(), + Some(&auth_manager), + &runtime_config.ema_state, &runtime_context, ) .await; diff --git a/codex-rs/codex-mcp/src/connection_manager/startup.rs b/codex-rs/codex-mcp/src/connection_manager/startup.rs index a4d6263b1f..71a8f4d979 100644 --- a/codex-rs/codex-mcp/src/connection_manager/startup.rs +++ b/codex-rs/codex-mcp/src/connection_manager/startup.rs @@ -59,6 +59,13 @@ pub(super) fn mcp_startup_failure_reason( if !error.is_authentication_required() { return None; } + if let StartupOutcomeError::Failed { + is_ema_authentication_required: true, + .. + } = error + { + return Some(McpStartupFailureReason::ReauthenticationRequired); + } match auth_state { Some( McpAuthState::LoggedOut(McpLoginRequirement::Reauthentication) | McpAuthState::OAuth, diff --git a/codex-rs/codex-mcp/src/connection_manager_tests.rs b/codex-rs/codex-mcp/src/connection_manager_tests.rs index 044d475b89..449fbb9e65 100644 --- a/codex-rs/codex-mcp/src/connection_manager_tests.rs +++ b/codex-rs/codex-mcp/src/connection_manager_tests.rs @@ -611,6 +611,7 @@ async fn connection_statuses_observe_clients_without_starting_them() { StartupOutcomeError::Failed { error: "broken".to_string(), is_authentication_required: false, + is_ema_authentication_required: false, }, ), ( @@ -618,6 +619,7 @@ async fn connection_statuses_observe_clients_without_starting_them() { StartupOutcomeError::Failed { error: "login".to_string(), is_authentication_required: true, + is_ema_authentication_required: false, }, ), ( @@ -692,6 +694,7 @@ async fn connection_statuses_follow_latest_reconnect_outcome() { 0 | 1 => Err(StartupOutcomeError::Failed { error: "retry failed".to_string(), is_authentication_required: attempt == 0, + is_ema_authentication_required: false, }), _ => Ok(recovered), } @@ -856,6 +859,7 @@ fn create_test_manager_with_failed_apps_startup( let client: ManagedClientFuture = futures::future::ready(Err(StartupOutcomeError::Failed { error: "startup failed".to_string(), is_authentication_required: false, + is_ema_authentication_required: false, })) .boxed() .shared(); @@ -3180,6 +3184,7 @@ async fn stable_catalog_revisions_ignore_terminal_optional_server_failures() { StartupOutcomeError::Failed { error: "optional startup failed".to_string(), is_authentication_required: false, + is_ema_authentication_required: false, }, )) .boxed() @@ -4129,6 +4134,7 @@ async fn list_all_tools_uses_shared_codex_apps_cache_when_client_startup_fails() StartupOutcomeError::Failed { error: "startup failed".to_string(), is_authentication_required: false, + is_ema_authentication_required: false, }, )) .boxed() @@ -4288,6 +4294,7 @@ async fn later_tool_list_retries_after_failed_reconnect_and_keeps_cached_tools() Err(StartupOutcomeError::Failed { error: "recreated startup failed".to_string(), is_authentication_required: false, + is_ema_authentication_required: false, }) } else { Ok(recovered_client) @@ -4817,6 +4824,8 @@ async fn executor_owned_chatgpt_mcp_accepts_only_safe_explicit_authorization() - OAuthCredentialsStoreMode::default(), AuthKeyringBackendKind::default(), Some(&hosted_auth), + /*auth_manager*/ None, + &crate::EmaState::default(), &runtime_context, ) .await; @@ -5089,6 +5098,7 @@ fn mcp_init_error_display_prompts_for_login_when_auth_required() { StartupOutcomeError::Failed { error: "OAuth refresh token was rejected: invalid_grant".to_string(), is_authentication_required: true, + is_ema_authentication_required: false, }, ] { let display = mcp_init_error_display( @@ -5120,6 +5130,7 @@ fn mcp_init_error_display_identifies_oauth_reauthentication() { let error = StartupOutcomeError::Failed { error: "authorization required: Bearer error=\"invalid_token\"".to_string(), is_authentication_required: true, + is_ema_authentication_required: false, }; let executor_config: McpServerConfig = serde_json::from_value(serde_json::json!({ "url": "https://example.com/mcp", @@ -5183,6 +5194,7 @@ fn mcp_startup_failure_reason_requires_existing_oauth_and_auth_failure() { let error = StartupOutcomeError::Failed { error: "startup failed".to_string(), is_authentication_required, + is_ema_authentication_required: false, }; assert_eq!( @@ -6103,6 +6115,7 @@ async fn reconciliation_retries_non_oauth_authentication_failures() { connection.client.client = futures::future::ready(Err(StartupOutcomeError::Failed { error: "bearer token rejected".to_string(), is_authentication_required: true, + is_ema_authentication_required: false, })) .boxed() .shared(); @@ -6328,6 +6341,7 @@ async fn reconciliation_replaces_connection_when_auth_mode_changes() -> anyhow:: StartupOutcomeError::Failed { error, is_authentication_required, + .. } => { assert_eq!( (error.as_str(), is_authentication_required), diff --git a/codex-rs/codex-mcp/src/ema_account.rs b/codex-rs/codex-mcp/src/ema_account.rs index c1ad3c1f7e..9abd5b1d07 100644 --- a/codex-rs/codex-mcp/src/ema_account.rs +++ b/codex-rs/codex-mcp/src/ema_account.rs @@ -1,11 +1,68 @@ -//! Validates persisted account grants against the configured enterprise issuer. +//! Account-grant eligibility and registration-specific sign-in requirements. +//! The MCP service owns this state across threads; it retains no credentials. +use std::sync::Mutex; + +use codex_config::McpEmaRegistration; use codex_config::McpServerIdpOAuthConfig; +use codex_login::AccountOAuthGrantRevision; use codex_login::CodexAuth; use codex_login::TokenData; use codex_rmcp_client::account_oidc_identity_is_usable; use codex_rmcp_client::account_oidc_issuer_hint; +/// Remembers a rejected account credential only for the affected EMA registration. +#[derive(Default)] +pub struct EmaState { + rejected: Mutex>, +} + +impl std::fmt::Debug for EmaState { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("EmaState").finish_non_exhaustive() + } +} + +impl EmaState { + pub(crate) fn requires_sign_in( + &self, + registration: &McpEmaRegistration, + tokens: &TokenData, + ) -> bool { + let revision = AccountOAuthGrantRevision::from(tokens); + self.rejected + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .iter() + .any(|(registered, rejected)| registered == registration && rejected == &revision) + } + + pub(crate) fn reject(&self, registration: &McpEmaRegistration, tokens: &TokenData) { + let mut rejected = self + .rejected + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + // Replacing a registration's marker also bounds repeated refresh/rejection cycles. + rejected.retain(|(registered, _)| registered != registration); + rejected.push(( + registration.clone(), + AccountOAuthGrantRevision::from(tokens), + )); + } + + pub(crate) fn account_is_usable( + &self, + auth: &CodexAuth, + registration: &McpEmaRegistration, + ) -> bool { + matches!(auth, CodexAuth::Chatgpt(_)) + && auth.get_token_data().is_ok_and(|tokens| { + account_grant_is_usable(&tokens, registration.idp()) + && !self.requires_sign_in(registration, &tokens) + }) + } +} + pub(crate) fn account_grant_is_usable(tokens: &TokenData, idp: &McpServerIdpOAuthConfig) -> bool { tokens.oauth_provenance.as_ref().is_some_and(|provenance| { provenance.issuer == idp.issuer diff --git a/codex-rs/codex-mcp/src/ema_auth.rs b/codex-rs/codex-mcp/src/ema_auth.rs index beb2b0258d..35d03f8cd0 100644 --- a/codex-rs/codex-mcp/src/ema_auth.rs +++ b/codex-rs/codex-mcp/src/ema_auth.rs @@ -22,14 +22,12 @@ use codex_login::CodexAuth; use codex_rmcp_client::EmaAuthFailure; use codex_rmcp_client::EmaAuthenticatedHttpClient; use codex_rmcp_client::EmaRequestAuthority; -use codex_rmcp_client::EmaTokenSource; use codex_rmcp_client::McpOAuthRefreshMode; use codex_rmcp_client::McpProtocolMode; use codex_rmcp_client::RmcpClient; use codex_rmcp_client::StoredOAuthCredentialSnapshot; use codex_rmcp_client::StreamableHttpBearerToken; use codex_rmcp_client::StreamableHttpRedirectMode; -use codex_rmcp_client::stored_ema_identity_is_usable; use codex_rmcp_client::stored_oauth_credential_snapshot; use futures::future::BoxFuture; use tokio_util::sync::CancellationToken; @@ -169,9 +167,9 @@ impl EmaAuth { cancellation: self.config.mcp_ema_cancellation.token(), }); let url = registration.server_url().to_owned(); - let source = + let mut source = EnterpriseTokenSource::new(registration, Arc::clone(self), exchange_client).await?; - let token = source.access_token().await?; + let token = source.initial_access_token().await?; let http_client = Arc::new(EmaAuthenticatedHttpClient::new( http_client, source, @@ -204,24 +202,18 @@ pub(crate) async fn stored_ema_credentials( ) -> Result> { let name = idp.credential_name(scope); let idp = idp.clone(); + let originator = codex_otel::auth_storage::AuthStorageOriginator::current(); tokio::task::spawn_blocking(move || { - tracing::subscriber::with_default(tracing::subscriber::NoSubscriber::default(), || { - stored_oauth_credential_snapshot( - &name, - &idp.issuer, - OAuthCredentialsStoreMode::Keyring, - keyring_backend, - ) - .map(|snapshot| { - snapshot.filter(|snapshot| { - stored_ema_identity_is_usable( - snapshot.credentials(), - &idp.issuer, - &idp.client_id, - ) - }) + originator.sync_scope(|| { + tracing::subscriber::with_default(tracing::subscriber::NoSubscriber::default(), || { + stored_oauth_credential_snapshot( + &name, + &idp.issuer, + OAuthCredentialsStoreMode::Keyring, + keyring_backend, + ) + .map_err(|_| anyhow!("failed to read enterprise credentials")) }) - .map_err(|_| anyhow!("failed to read enterprise credentials")) }) }) .await diff --git a/codex-rs/codex-mcp/src/ema_token_source.rs b/codex-rs/codex-mcp/src/ema_token_source.rs index af48a1d851..d9ef85eb1c 100644 --- a/codex-rs/codex-mcp/src/ema_token_source.rs +++ b/codex-rs/codex-mcp/src/ema_token_source.rs @@ -9,6 +9,7 @@ use codex_config::McpEmaRegistration; use codex_exec_server::HttpClient; use codex_login::AccountOAuthGrant; use codex_login::AccountOAuthGrantLease; +use codex_login::AccountOAuthGrantRevision; use codex_rmcp_client::EmaAccessToken; use codex_rmcp_client::EmaAuthFailure; use codex_rmcp_client::EmaAuthTokenExchangeRequest; @@ -60,28 +61,36 @@ impl EnterpriseTokenSource { http_client: Arc, ) -> Result { auth.check_account().await?; - let credentials = if let Some(snapshot) = stored_ema_credentials( - registration.idp(), - &auth.scope, - auth.config.auth_keyring_backend_kind, - ) - .await? - { - CredentialSource::Keyring(Box::new(snapshot.pin_ema_credentials().await?)) - } else { - let account = auth.manager.lease_account_oauth_grant().await?; - let client_id = account - .as_ref() - .filter(|lease| account_grant_is_usable(lease.tokens(), registration.idp())) - .and_then(|lease| lease.tokens().oauth_provenance.as_ref()) - .map(|provenance| provenance.client_id.clone()) + let account = auth.manager.lease_account_oauth_grant().await?; + let client_id = account + .as_ref() + .filter(|lease| { + account_grant_is_usable(lease.tokens(), registration.idp()) + && !auth + .config + .ema_state + .requires_sign_in(®istration, lease.tokens()) + }) + .and_then(|lease| lease.tokens().oauth_provenance.as_ref()) + .map(|provenance| provenance.client_id.clone()); + drop(account); + let credentials = match client_id { + Some(client_id) => CredentialSource::Account { client_id }, + None => { + let snapshot = stored_ema_credentials( + registration.idp(), + &auth.scope, + auth.config.auth_keyring_backend_kind, + ) + .await? .ok_or_else(|| { anyhow::Error::new(EmaAuthFailure::ReauthenticationRequired).context( account_configuration_issue(&auth.account, registration.idp()) .unwrap_or("Complete enterprise sign-in to use this MCP server"), ) })?; - CredentialSource::Account { client_id } + CredentialSource::Keyring(Box::new(snapshot.pin_ema_credentials().await?)) + } }; Ok(Self { endpoint: Url::parse(registration.server_url())?, @@ -92,6 +101,42 @@ impl EnterpriseTokenSource { }) } + pub(crate) async fn initial_access_token(&mut self) -> Result { + let error = match self.access_token().await { + Ok(token) => return Ok(token), + Err(error) => error, + }; + // A restart forgets the rejection marker. A saved enterprise login can still + // satisfy step-up, but a live connection never silently changes its source. + if matches!(self.credentials, CredentialSource::Account { .. }) + && error + .downcast_ref::() + .is_some_and(EmaTokenSourceError::is_authentication_required) + && let Some(lease) = self.auth.manager.lease_account_oauth_grant().await? + { + let rejected = self + .auth + .config + .ema_state + .requires_sign_in(&self.registration, lease.tokens()); + drop(lease); + if rejected + && let Ok(Some(snapshot)) = stored_ema_credentials( + self.registration.idp(), + &self.auth.scope, + self.auth.config.auth_keyring_backend_kind, + ) + .await + { + self.auth.check_account().await?; + self.credentials = + CredentialSource::Keyring(Box::new(snapshot.pin_ema_credentials().await?)); + return self.access_token().await; + } + } + Err(error) + } + async fn account_grant(&self, client_id: &str) -> Result { self.auth.check_account().await?; let lease = self @@ -109,6 +154,16 @@ impl EnterpriseTokenSource { { bail!(EmaAuthFailure::ReauthenticationRequired); } + if self + .auth + .config + .ema_state + .requires_sign_in(&self.registration, lease.tokens()) + { + bail!(EmaAuthFailure::InsufficientUserAuthentication { + grant_source: EmaInvalidGrantSource::EnterpriseIdentity + }); + } Ok(lease.release()) } @@ -141,9 +196,22 @@ impl EnterpriseTokenSource { .lease_account_oauth_grant() .await? .ok_or(EmaAuthFailure::ReauthenticationRequired)?; - if lease.tokens() != tokens { + if AccountOAuthGrantRevision::from(lease.tokens()) + != AccountOAuthGrantRevision::from(tokens) + || !account_grant_is_usable(lease.tokens(), idp) + { bail!(EmaAuthFailure::ReauthenticationRequired); } + if self + .auth + .config + .ema_state + .requires_sign_in(&self.registration, lease.tokens()) + { + bail!(EmaAuthFailure::InsufficientUserAuthentication { + grant_source: EmaInvalidGrantSource::EnterpriseIdentity, + }); + } Ok(identity.map_credentials(|()| IdpLease::Account(lease))) } CredentialSource::Keyring(snapshot) => { @@ -182,12 +250,17 @@ impl EnterpriseTokenSource { .await?; match receipt { ExchangeReceipt::Account(grant) => { - if !grant.check_current().await? { - // A successful exchange with a stale receipt must fail this call. - bail!(EmaTokenSourceError::new( - EmaCredentialKind::Account, - EmaAuthFailure::ReauthenticationRequired - )); + if grant.current_tokens().await?.as_ref().is_none_or(|tokens| { + !account_grant_is_usable(tokens, idp) + || self + .auth + .config + .ema_state + .requires_sign_in(&self.registration, tokens) + }) { + // Discard this bearer. The bounded retry below reacquires the grant + // only if its credential changed under the same admitted owner. + bail!(EmaAuthFailure::ReauthenticationRequired); } } ExchangeReceipt::Keyring(lease) => { @@ -205,8 +278,49 @@ impl EnterpriseTokenSource { self.auth.check_account().await?; match &self.credentials { CredentialSource::Account { client_id } => { - let grant = self.account_grant(client_id).await?; - self.exchange(Some(&grant)).await + let mut grant = self.account_grant(client_id).await?; + for attempt in 0..2 { + let error = match self.exchange(Some(&grant)).await { + Ok(token) => return Ok(token), + Err(error) => error, + }; + if idp_rejected(&error) + || matches!( + error.downcast_ref::(), + Some(EmaAuthFailure::ReauthenticationRequired) + ) + { + let current = self.account_grant(client_id).await?; + if !current.same_credential(&grant) { + if attempt == 0 { + grant = current; + continue; + } + // Repeated rotation is transient; a later call can renew + // without revoking this still-current account owner. + bail!(EmaTokenSourceError::unexpected(EmaCredentialKind::Account)); + } + if idp_rejected(&error) + && let Some(lease) = + self.auth.manager.lease_account_oauth_grant().await? + { + if AccountOAuthGrantRevision::from(lease.tokens()) + != AccountOAuthGrantRevision::from(grant.tokens()) + { + // Refresh may also rotate after the first credential check. + bail!(EmaTokenSourceError::unexpected(EmaCredentialKind::Account)); + } + // Serialize rejection publication with refresh and other leases. + // A late rejection cannot replace a newer credential's marker. + self.auth + .config + .ema_state + .reject(&self.registration, lease.tokens()); + } + } + return Err(error); + } + unreachable!("each exchange returns or retries once") } CredentialSource::Keyring(snapshot) => { let result = self.exchange(/*account*/ None).await; @@ -263,3 +377,16 @@ impl EmaTokenSource for EnterpriseTokenSource { }) } } + +fn idp_rejected(error: &anyhow::Error) -> bool { + matches!( + error.downcast_ref::(), + Some( + EmaAuthFailure::InvalidGrant { + grant_source: EmaInvalidGrantSource::EnterpriseIdentity + } | EmaAuthFailure::InsufficientUserAuthentication { + grant_source: EmaInvalidGrantSource::EnterpriseIdentity + } | EmaAuthFailure::IdentityClientRejected + ) + ) +} diff --git a/codex-rs/codex-mcp/src/lib.rs b/codex-rs/codex-mcp/src/lib.rs index 8ba8a1d220..5768aa71cd 100644 --- a/codex-rs/codex-mcp/src/lib.rs +++ b/codex-rs/codex-mcp/src/lib.rs @@ -119,6 +119,7 @@ pub(crate) mod elicitation; mod ema_account; mod ema_auth; mod ema_token_source; +pub use ema_account::EmaState; mod event_stream; mod executor_environment_http_client; pub(crate) mod mcp; diff --git a/codex-rs/codex-mcp/src/mcp/auth.rs b/codex-rs/codex-mcp/src/mcp/auth.rs index 9f745b3cf5..4abc7c7d5b 100644 --- a/codex-rs/codex-mcp/src/mcp/auth.rs +++ b/codex-rs/codex-mcp/src/mcp/auth.rs @@ -200,6 +200,8 @@ pub async fn compute_auth_statuses<'a, I>( store_mode: OAuthCredentialsStoreMode, keyring_backend_kind: AuthKeyringBackendKind, auth: Option<&CodexAuth>, + auth_manager: Option<&std::sync::Arc>, + ema_state: &crate::EmaState, runtime_context: &McpRuntimeContext, ) -> HashMap where @@ -221,6 +223,8 @@ where store_mode, keyring_backend_kind, auth, + auth_manager, + ema_state, &runtime_context, redirect_mode, ) @@ -243,12 +247,15 @@ where join_all(futures).await.into_iter().collect() } +#[allow(clippy::too_many_arguments)] async fn compute_auth_status( server_name: &str, config: &McpServerConfig, store_mode: OAuthCredentialsStoreMode, keyring_backend_kind: AuthKeyringBackendKind, auth: Option<&CodexAuth>, + auth_manager: Option<&std::sync::Arc>, + ema_state: &crate::EmaState, runtime_context: &McpRuntimeContext, redirect_mode: StreamableHttpRedirectMode, ) -> Result { @@ -265,21 +272,49 @@ async fn compute_auth_status( let Some(scope) = crate::ema_auth_scope(auth) else { return Ok(McpAuthState::LoggedOut(McpLoginRequirement::Login)); }; + if matches!(auth, Some(CodexAuth::Chatgpt(_))) + && auth + .and_then(|auth| auth.get_token_data().ok()) + .is_none_or(|tokens| tokens.login_owner_id.as_ref().is_none_or(String::is_empty)) + { + return Ok(McpAuthState::LoggedOut(McpLoginRequirement::Login)); + } + let account_usable = match auth_manager { + Some(manager) => match manager.lease_account_oauth_grant().await? { + Some(lease) => { + crate::ema_account::account_grant_is_usable(lease.tokens(), registration.idp()) + && !ema_state.requires_sign_in(registration, lease.tokens()) + } + None => false, + }, + None => auth.is_some_and(|auth| ema_state.account_is_usable(auth, registration)), + }; + if account_usable { + return Ok(McpAuthState::BearerToken); + } return Ok( - if crate::ema_auth::stored_ema_credentials( + match crate::ema_auth::stored_ema_credentials( registration.idp(), &scope, keyring_backend_kind, ) .await? - .is_some() { - McpAuthState::OAuth - } else { - McpAuthState::LoggedOut(McpLoginRequirement::Login) + Some(snapshot) + if codex_rmcp_client::stored_ema_identity_is_usable( + snapshot.credentials(), + ®istration.idp().issuer, + ®istration.idp().client_id, + ) => + { + McpAuthState::OAuth + } + Some(_) => McpAuthState::LoggedOut(McpLoginRequirement::Reauthentication), + None => McpAuthState::LoggedOut(McpLoginRequirement::Login), }, ); } + let has_runtime_auth = matches!(config.auth, McpServerAuth::ChatGpt) && auth.is_some_and(CodexAuth::uses_codex_backend) && matches!( diff --git a/codex-rs/codex-mcp/src/mcp/mod.rs b/codex-rs/codex-mcp/src/mcp/mod.rs index 56203cf4b2..d73825b6f6 100644 --- a/codex-rs/codex-mcp/src/mcp/mod.rs +++ b/codex-rs/codex-mcp/src/mcp/mod.rs @@ -135,6 +135,8 @@ pub struct McpConfig { /// Retires only EMA connections and pending exchanges for this config snapshot. pub mcp_ema_cancellation: codex_config::McpEmaCancellation, pub xaa_enabled: bool, + /// Shared sign-in requirements for account-backed enterprise MCP. + pub ema_state: Arc, /// Preferred credential store for MCP OAuth tokens. pub mcp_oauth_credentials_store_mode: OAuthCredentialsStoreMode, /// OAuth refresh ownership selected for new MCP connections. @@ -508,9 +510,11 @@ pub struct McpServerStatusSnapshot { pub server_names: Vec, } +#[allow(clippy::too_many_arguments)] pub async fn collect_mcp_server_status_snapshot_with_detail( config: &McpConfig, auth: Option<&CodexAuth>, + auth_manager: Option<&Arc>, runtime_context: McpRuntimeContext, codex_apps_tools_cache: ConnectorRuntimeManager, tool_catalog_cache: crate::McpToolCatalogCache, @@ -538,6 +542,8 @@ pub async fn collect_mcp_server_status_snapshot_with_detail( config.mcp_oauth_credentials_store_mode, config.auth_keyring_backend_kind, auth, + auth_manager, + &config.ema_state, &runtime_context, ) .await; @@ -859,11 +865,16 @@ pub(crate) async fn collect_mcp_server_status_snapshot_from_manager( .insert(tool_name, tool); } - // Status-only discovery has no event channel. Report OAuth failures from the completed + // Status-only discovery has no event channel. Report OAuth and EMA failures from the completed // connection attempt instead of retaining the credential-presence status read beforehand. let mut auth_statuses = auth_statuses_from_entries(&auth_status_entries); for server_name in mcp_connection_manager.authentication_failed_servers().await { - if auth_statuses.get(&server_name) == Some(&McpAuthStatus::OAuth) { + if auth_statuses.get(&server_name) == Some(&McpAuthStatus::OAuth) + || auth_status_entries + .get(&server_name) + .and_then(|entry| entry.config.as_ref()) + .is_some_and(|config| matches!(config.auth, McpServerAuth::EmaAuth)) + { auth_statuses.insert(server_name, McpAuthStatus::NotLoggedIn); } } diff --git a/codex-rs/codex-mcp/src/mcp/mod_tests.rs b/codex-rs/codex-mcp/src/mcp/mod_tests.rs index 9c1435e435..b464628d10 100644 --- a/codex-rs/codex-mcp/src/mcp/mod_tests.rs +++ b/codex-rs/codex-mcp/src/mcp/mod_tests.rs @@ -25,10 +25,11 @@ use std::path::PathBuf; use std::sync::Arc; #[tokio::test] -async fn status_snapshot_only_downgrades_oauth_authentication_failures() { +async fn status_snapshot_only_downgrades_oauth_and_ema_authentication_failures() { let auth_failure = StartupOutcomeError::Failed { error: "OAuth refresh token was rejected".to_string(), is_authentication_required: true, + is_ema_authentication_required: false, }; let mut manager = McpConnectionSet::empty(/*prefix_mcp_tool_names*/ true); let mut auth_status_entries = HashMap::new(); @@ -45,8 +46,20 @@ async fn status_snapshot_only_downgrades_oauth_authentication_failures() { Some(StartupOutcomeError::Failed { error: "provider temporarily unavailable".to_string(), is_authentication_required: false, + is_ema_authentication_required: false, }), ), + ( + "ema-failed", + McpAuthState::BearerToken, + Some(auth_failure.clone()), + ), + ("ema-ready", McpAuthState::BearerToken, None), + ( + "ema-provider-error", + McpAuthState::BearerToken, + Some(anyhow::anyhow!("provider temporarily unavailable").into()), + ), ("bearer", McpAuthState::BearerToken, Some(auth_failure)), ] { let mut client = create_ready_async_managed_client(Vec::new()).await; @@ -57,7 +70,12 @@ async fn status_snapshot_only_downgrades_oauth_authentication_failures() { auth_status_entries.insert( name.to_string(), McpAuthStatusEntry { - config: None, + config: name.starts_with("ema-").then(|| { + serde_json::from_value(serde_json::json!({ + "url": "https://example.com/mcp", "auth": "ema_auth", + })) + .expect("EMA config") + }), auth_state, }, ); @@ -78,6 +96,9 @@ async fn status_snapshot_only_downgrades_oauth_authentication_failures() { ("oauth-failed".to_string(), McpAuthStatus::NotLoggedIn), ("oauth-ready".to_string(), McpAuthStatus::OAuth), ("oauth-provider-error".to_string(), McpAuthStatus::OAuth), + ("ema-failed".to_string(), McpAuthStatus::NotLoggedIn), + ("ema-ready".to_string(), McpAuthStatus::BearerToken), + ("ema-provider-error".to_string(), McpAuthStatus::BearerToken), ("bearer".to_string(), McpAuthStatus::BearerToken), ]), ); @@ -90,6 +111,7 @@ pub(crate) fn test_mcp_config(codex_home: PathBuf) -> McpConfig { requires_read_only_mcp_tools: false, codex_home, mcp_enterprise_managed_auth: None, + ema_state: Default::default(), mcp_ema_cancellation: codex_config::McpEmaCancellation::default(), xaa_enabled: false, mcp_oauth_credentials_store_mode: OAuthCredentialsStoreMode::default(), diff --git a/codex-rs/codex-mcp/src/rmcp_client.rs b/codex-rs/codex-mcp/src/rmcp_client.rs index 24295d9bfe..b60e35438f 100644 --- a/codex-rs/codex-mcp/src/rmcp_client.rs +++ b/codex-rs/codex-mcp/src/rmcp_client.rs @@ -69,6 +69,8 @@ use codex_protocol::protocol::Event; use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::McpStartupStatus; use codex_protocol::protocol::McpStartupUpdateEvent; +use codex_rmcp_client::EmaAuthFailure; +use codex_rmcp_client::EmaTokenSourceError; use codex_rmcp_client::ExecutorStdioServerLauncher; use codex_rmcp_client::LocalStdioServerLauncher; use codex_rmcp_client::McpOAuthRefreshMode; @@ -678,6 +680,7 @@ pub(crate) enum StartupOutcomeError { Failed { error: String, is_authentication_required: bool, + is_ema_authentication_required: bool, }, } @@ -688,6 +691,7 @@ impl StartupOutcomeError { Self::Failed { error, is_authentication_required, + .. } => *is_authentication_required || error.contains("Auth required"), } } @@ -699,6 +703,9 @@ impl From for StartupOutcomeError { Self::Failed { error: format!("{error:#}"), is_authentication_required, + is_ema_authentication_required: is_authentication_required + && (error.downcast_ref::().is_some() + || error.downcast_ref::().is_some()), } } } diff --git a/codex-rs/codex-mcp/src/runtime.rs b/codex-rs/codex-mcp/src/runtime.rs index 6d57cd1684..4d46c67696 100644 --- a/codex-rs/codex-mcp/src/runtime.rs +++ b/codex-rs/codex-mcp/src/runtime.rs @@ -128,6 +128,7 @@ struct PublishedMcpRuntime { config: Option>, auth: Option, auth_token: Option, + auth_manager: Option>, auth_generation: Arc<()>, plugins_available: bool, ready_selected_capability_roots: Vec, @@ -238,6 +239,7 @@ impl McpRuntime { config: None, auth: None, auth_token: None, + auth_manager: None, auth_generation: Arc::new(()), plugins_available: false, ready_selected_capability_roots: Vec::new(), @@ -376,6 +378,7 @@ impl McpRuntime { let (publish, publication_gate) = McpPublicationGate::pending(); let config = Arc::clone(&input.config); let auth = input.auth.clone(); + let auth_manager = input.auth_manager.clone(); let auth_token = auth.as_ref().and_then(|auth| auth.get_token().ok()); let auth_generation = if current.auth_matches(auth.as_ref()) { Arc::clone(¤t.auth_generation) @@ -415,6 +418,7 @@ impl McpRuntime { config: Some(config), auth, auth_token, + auth_manager, auth_generation, plugins_available, ready_selected_capability_roots, @@ -1119,6 +1123,7 @@ mod tests { config: Some(Arc::new(config)), auth: None, auth_token: None, + auth_manager: None, auth_generation: Arc::new(()), plugins_available: false, ready_selected_capability_roots: Vec::new(), @@ -1193,6 +1198,7 @@ mod tests { ))), auth: None, auth_token: None, + auth_manager: None, auth_generation: Arc::new(()), plugins_available: false, ready_selected_capability_roots: Vec::new(), diff --git a/codex-rs/codex-mcp/src/runtime/status.rs b/codex-rs/codex-mcp/src/runtime/status.rs index bb6b0f2624..c183524d1c 100644 --- a/codex-rs/codex-mcp/src/runtime/status.rs +++ b/codex-rs/codex-mcp/src/runtime/status.rs @@ -40,11 +40,24 @@ impl McpRuntime { && expected_registration .is_none_or(|expected| config.mcp_server_catalog.server(name) == Some(expected)) }); + if config.xaa_enabled + && servers + .values() + .any(|server| matches!(server.config().auth, codex_config::McpServerAuth::EmaAuth)) + { + // Finish selected startup and any rejected-grant cleanup before reading auth. + current + .connections + .list_available_server_infos(|name| name == server) + .await; + } let auth_statuses = compute_auth_statuses( servers.iter(), config.mcp_oauth_credentials_store_mode, config.auth_keyring_backend_kind, current.auth.as_ref(), + current.auth_manager.as_ref(), + &config.ema_state, &runtime_context, ) .await; diff --git a/codex-rs/core/src/config/mod.rs b/codex-rs/core/src/config/mod.rs index d7cef148ba..31e7dc2c39 100644 --- a/codex-rs/core/src/config/mod.rs +++ b/codex-rs/core/src/config/mod.rs @@ -1878,6 +1878,7 @@ impl Config { requires_read_only_mcp_tools: false, codex_home: self.codex_home.to_path_buf(), mcp_enterprise_managed_auth: self.mcp_enterprise_managed_auth.clone(), + ema_state: Default::default(), mcp_ema_cancellation: self.config_layer_stack.mcp_ema_cancellation().clone(), xaa_enabled: self.features.enabled(Feature::UseXaa) && self.mcp_enterprise_managed_auth.is_some(), diff --git a/codex-rs/core/src/mcp.rs b/codex-rs/core/src/mcp.rs index a3046d892b..f2d4ea9f29 100644 --- a/codex-rs/core/src/mcp.rs +++ b/codex-rs/core/src/mcp.rs @@ -116,6 +116,7 @@ pub struct McpManager { codex_apps_tools_cache: ConnectorRuntimeManager, tool_catalog_cache: McpToolCatalogCache, shared_definitions: Arc, + ema_state: Arc, } impl McpManager { @@ -139,6 +140,7 @@ impl McpManager { codex_apps_tools_cache, tool_catalog_cache: McpToolCatalogCache::default(), shared_definitions: Arc::new(SharedMcpDefinitions::default()), + ema_state: Arc::default(), } } @@ -457,6 +459,7 @@ impl McpManager { !selected_plugins.is_empty() || !loaded_plugins.capability_summaries().is_empty(); let mut mcp_config = config .to_mcp_config_with_loaded_plugins(&loaded_plugins, selected_plugin_registrations); + mcp_config.ema_state = Arc::clone(&self.ema_state); let mut catalog = mcp_config.mcp_server_catalog.to_builder(); // Isolated sessions do not inherit the implicit Apps connection. if mcp_config.apps_enabled diff --git a/codex-rs/core/tests/suite/snapshots/all__suite__scenarios__mcp_ema__enterprise_sign_in_required.snap b/codex-rs/core/tests/suite/snapshots/all__suite__scenarios__mcp_ema__enterprise_sign_in_required.snap index 117c68691e..76b00adc6b 100644 --- a/codex-rs/core/tests/suite/snapshots/all__suite__scenarios__mcp_ema__enterprise_sign_in_required.snap +++ b/codex-rs/core/tests/suite/snapshots/all__suite__scenarios__mcp_ema__enterprise_sign_in_required.snap @@ -3,4 +3,4 @@ source: core/tests/suite/scenarios_mcp_ema.rs assertion_line: 65 expression: error --- -The enterprise MCP server is not logged in. Sign in to your enterprise provider in Codex, then start a new thread. +The enterprise MCP server requires OAuth reauthentication. Sign in to your enterprise provider in Codex, then start a new thread. diff --git a/codex-rs/login/src/auth/account_grant.rs b/codex-rs/login/src/auth/account_grant.rs index 217828f052..0ef95caf0a 100644 --- a/codex-rs/login/src/auth/account_grant.rs +++ b/codex-rs/login/src/auth/account_grant.rs @@ -2,6 +2,8 @@ //! exchange; callers must recheck ownership before publishing its result. Storage reads //! observe other processes at these boundaries, not continuously. +use sha2::Digest; +use sha2::Sha256; use std::fmt; use std::io; use std::sync::Arc; @@ -18,6 +20,38 @@ pub struct AccountOAuthGrant { pub(super) login_generation: u64, } +/// Identity of a refresh grant, excluding access and ID token rotation. +/// This opaque digest compares revisions; it does not authorize a destination. +#[derive(Clone, PartialEq, Eq)] +pub struct AccountOAuthGrantRevision([u8; 32]); + +impl From<&TokenData> for AccountOAuthGrantRevision { + fn from(tokens: &TokenData) -> Self { + let mut hash = Sha256::new(); + let provenance = tokens.oauth_provenance.as_ref(); + for value in [ + tokens.login_owner_id.as_deref(), + tokens.account_id.as_deref(), + tokens.id_token.chatgpt_user_id.as_deref(), + provenance.map(|provenance| provenance.issuer.as_str()), + provenance.map(|provenance| provenance.client_id.as_str()), + Some(tokens.refresh_token.as_str()), + ] { + let value = value.unwrap_or_default(); + hash.update(value.len().to_be_bytes()); + hash.update(value.as_bytes()); + } + Self(hash.finalize().into()) + } +} + +impl fmt::Debug for AccountOAuthGrantRevision { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("AccountOAuthGrantRevision") + .finish_non_exhaustive() + } +} + /// Keeps ordinary refresh from rotating the account token during an IdP exchange. /// Login and logout remain available and invalidate the grant independently of this lease. pub struct AccountOAuthGrantLease { @@ -57,7 +91,9 @@ impl AccountOAuthGrant { /// Compare credential revisions without exposing the refresh token to state holders. pub fn same_credential(&self, other: &Self) -> bool { - self.tokens == other.tokens + self.login_generation == other.login_generation + && AccountOAuthGrantRevision::from(&self.tokens) + == AccountOAuthGrantRevision::from(&other.tokens) } pub(super) fn owner_is_current(&self) -> bool { @@ -73,14 +109,19 @@ impl AccountOAuthGrant { .is_some_and(|tokens| same_account_grant(&self.tokens, &tokens)) } - /// Check after asynchronous work, before accepting a token or publishing a connection. + /// Read the still-current grant after asynchronous work. Consumers must validate + /// the latest ID token's eligibility before accepting an exchange result. /// A storage failure is an error, never proof that the saved grant is still valid. - pub async fn check_current(&self) -> io::Result { + pub async fn current_tokens(&self) -> io::Result> { if !self.owner_is_current() { - return Ok(false); + return Ok(None); } let stored = self.manager.read_account_grant_tokens().await?; - Ok(self.owner_is_current() && stored.as_ref().is_some_and(|tokens| self.tokens == *tokens)) + Ok(stored.filter(|tokens| { + self.owner_is_current() + && AccountOAuthGrantRevision::from(&self.tokens) + == AccountOAuthGrantRevision::from(tokens) + })) } } diff --git a/codex-rs/login/src/auth/manager.rs b/codex-rs/login/src/auth/manager.rs index f8a5d76036..db0d611914 100644 --- a/codex-rs/login/src/auth/manager.rs +++ b/codex-rs/login/src/auth/manager.rs @@ -4,6 +4,7 @@ mod workspace_routing; mod account_grant; pub use account_grant::AccountOAuthGrant; pub use account_grant::AccountOAuthGrantLease; +pub use account_grant::AccountOAuthGrantRevision; use account_grant::same_account_grant; use account_grant::same_account_login; diff --git a/codex-rs/login/src/lib.rs b/codex-rs/login/src/lib.rs index 4698c66ff3..b0863ee7c2 100644 --- a/codex-rs/login/src/lib.rs +++ b/codex-rs/login/src/lib.rs @@ -34,6 +34,7 @@ pub use success_page::LoginSuccessPageBrand; pub use auth::AccountOAuthGrant; pub use auth::AccountOAuthGrantLease; +pub use auth::AccountOAuthGrantRevision; pub use auth::AgentIdentityAuthPolicy; pub use auth::AuthChangeState; pub use auth::AuthConfig; diff --git a/codex-rs/login/tests/suite/auth_refresh.rs b/codex-rs/login/tests/suite/auth_refresh.rs index 303ac14c58..0ccd7d4ca1 100644 --- a/codex-rs/login/tests/suite/auth_refresh.rs +++ b/codex-rs/login/tests/suite/auth_refresh.rs @@ -288,12 +288,12 @@ async fn account_grant_tracks_refresh_replacement_and_logout() -> Result<()> { refresh_wake.0.load(Ordering::SeqCst), "releasing the account grant must wake the blocked refresh" ); - assert!(grant.check_current().await?); + assert!(grant.current_tokens().await?.is_some()); refresh .await .context("refresh should succeed once the IdP lease is released")?; assert!( - !grant.check_current().await?, + grant.current_tokens().await?.is_none(), "the old credential was rotated" ); let after = *changes.borrow(); @@ -345,6 +345,32 @@ async fn account_grant_tracks_refresh_replacement_and_logout() -> Result<()> { .context("rotated account grant")?; assert_eq!(lease.tokens(), &refreshed_tokens); let grant = lease.release(); + server.verify().await; + server.reset().await; + Mock::given(method("POST")) + .and(path("/oauth/token")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "access_token": "access-only-refresh", + "id_token": jwt_with_payload(json!({ + "sub": "user-123", "iat": 42, + "https://api.openai.com/auth": { "chatgpt_user_id": "user-123", "chatgpt_account_id": "account-id" } + })), + }))) + .expect(1) + .mount(&server) + .await; + ctx.auth_manager.refresh_token_from_authority().await?; + let current = ctx + .auth_manager + .lease_account_oauth_grant() + .await? + .context("same refresh grant")? + .release(); + assert!(grant.same_credential(¤t)); + assert_eq!( + grant.current_tokens().await?, + Some(current.tokens().clone()) + ); let mut replacement = stored.clone(); replacement .tokens @@ -358,7 +384,7 @@ async fn account_grant_tracks_refresh_replacement_and_logout() -> Result<()> { AuthCredentialsStoreMode::File, AuthKeyringBackendKind::default(), )?; - assert!(!grant.check_current().await?); + assert!(grant.current_tokens().await?.is_none()); assert!( !ctx.auth_manager .account_login_is_current(&cached_auth) @@ -369,7 +395,7 @@ async fn account_grant_tracks_refresh_replacement_and_logout() -> Result<()> { .auth_cached() .context("original cached auth")? .get_token_data()?, - cached, + current.tokens().clone(), ); ctx.auth_manager.reload().await; assert!( @@ -406,7 +432,7 @@ async fn account_grant_tracks_refresh_replacement_and_logout() -> Result<()> { assert_eq!(after.owner_generation, before.owner_generation); assert!(policy.acquire(&endpoint).is_ok()); assert_eq!(after.login_generation, before.login_generation + 1); - assert!(!grant.check_current().await?); + assert!(grant.current_tokens().await?.is_none()); assert!( ctx.auth_manager .lease_account_oauth_grant() @@ -420,7 +446,7 @@ async fn account_grant_tracks_refresh_replacement_and_logout() -> Result<()> { changes.borrow().login_generation, before.login_generation + 2 ); - assert!(!grant.check_current().await?); + assert!(grant.current_tokens().await?.is_none()); server.verify().await; server.reset().await; for login_owner_id in [Some("login-during-refresh"), None] { @@ -484,7 +510,7 @@ async fn account_grant_tracks_refresh_replacement_and_logout() -> Result<()> { latest_login.last_refresh = stored.last_refresh; } assert_eq!(stored, latest_login); - assert!(!grant.check_current().await?); + assert!(grant.current_tokens().await?.is_none()); ctx.auth_manager.reload().await; assert_eq!( ctx.auth_manager @@ -504,7 +530,7 @@ async fn account_grant_tracks_refresh_replacement_and_logout() -> Result<()> { .context("restored login grant")?; // Logout must complete even while an IdP exchange holds the refresh lease. ctx.auth_manager.logout().await?; - assert!(!lease.release().check_current().await?); + assert!(lease.release().current_tokens().await?.is_none()); assert!( ctx.auth_manager .lease_account_oauth_grant() diff --git a/codex-rs/rmcp-client/src/ema_http_client_regression_tests.rs b/codex-rs/rmcp-client/src/ema_http_client_regression_tests.rs index c6edacb4d7..8870831876 100644 --- a/codex-rs/rmcp-client/src/ema_http_client_regression_tests.rs +++ b/codex-rs/rmcp-client/src/ema_http_client_regression_tests.rs @@ -237,15 +237,26 @@ async fn credential_failures_are_classified_before_dispatch() -> anyhow::Result< }) .context("resource authorization failed without invalidating the enterprise identity"), anyhow::anyhow!("temporary network failure"), + anyhow::Error::new(crate::EmaTokenSourceError::unexpected( + crate::EmaCredentialKind::Account, + )), ] { let (client, source, params, mut requests) = controlled_fixture(); *client.cached_token() = None; *source.renewal_error.lock().unwrap() = Some(error); assert!(matches!( - mode.start(client, params).await?, + mode.start(Arc::clone(&client), params.clone()).await?, Err(ExecServerError::HttpRequest(_)) )); assert!(requests.try_recv().is_err()); + let retry = mode.start(client, params); + let request = requests + .recv() + .await + .expect("retry after transient failure"); + assert_eq!(bearer(&request), "Bearer renewed-2"); + request.response.send(/*t*/ 200).expect("complete retry"); + assert_eq!(retry.await??.status, 200); } } Ok(())