Reuse Codex account grants for enterprise MCP authentication (#55698)

## What changed

- Prefer stored enterprise credentials when available; otherwise use a validated Codex account OAuth grant whose issuer matches the enterprise MCP configuration.
- Pin each connection to its credential source and OAuth client, reject stale account grants, and hold the account refresh lease only during the IdP exchange.
- Report actionable sign-in and issuer configuration errors, and classify token failures by credential source.

GitOrigin-RevId: 26673f276a31e4ca5290f6b120f0ac44c77d7239

codex-rs/codex-mcp/src/ema_account.rs

@@ -0,0 +1,55 @@
+//! Validates persisted account grants against the configured enterprise issuer.
+
+use codex_config::McpServerIdpOAuthConfig;
+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;
+
+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
+            && tokens
+                .login_owner_id
+                .as_ref()
+                .is_some_and(|id| !id.is_empty())
+            && !tokens.refresh_token.trim().is_empty()
+            && tokens
+                .id_token
+                .chatgpt_user_id
+                .as_ref()
+                .is_some_and(|id| !id.is_empty())
+            && tokens.account_id.as_ref().is_some_and(|id| !id.is_empty())
+            && tokens.id_token.chatgpt_account_id == tokens.account_id
+            && account_oidc_identity_is_usable(
+                &tokens.id_token.raw_jwt,
+                &provenance.issuer,
+                &provenance.client_id,
+            )
+    })
+}
+
+pub(crate) fn account_configuration_issue(
+    auth: &CodexAuth,
+    idp: &McpServerIdpOAuthConfig,
+) -> Option<&'static str> {
+    let tokens = auth.get_token_data().ok()?;
+    if tokens.login_owner_id.is_none()
+        && (tokens
+            .oauth_provenance
+            .as_ref()
+            .is_some_and(|provenance| provenance.issuer == idp.issuer)
+            || account_oidc_issuer_hint(&tokens.id_token.raw_jwt, &idp.issuer))
+    {
+        return Some("Sign in to your Codex account again to enable this enterprise MCP server");
+    }
+    if tokens.oauth_provenance.as_ref().is_some_and(|provenance| {
+        provenance.issuer != idp.issuer
+            && provenance.issuer.trim_end_matches('/') == idp.issuer.trim_end_matches('/')
+    }) {
+        return Some(
+            "Codex sign-in and enterprise MCP issuer strings differ by a trailing slash; align their OIDC issuer configuration",
+        );
+    }
+    None
+}

codex-rs/codex-mcp/src/ema_auth.rs

@@ -9,7 +9,6 @@ use anyhow::Result;
 use anyhow::anyhow;
 use anyhow::bail;
 use codex_config::McpEmaAuthScope;
-use codex_config::McpEmaRegistration;
 use codex_config::McpServerIdpOAuthConfig;
 use codex_config::types::AuthKeyringBackendKind;
 use codex_config::types::OAuthCredentialsStoreMode;
@@ -20,13 +19,8 @@ use codex_exec_server::HttpRequestResponse;
 use codex_exec_server::HttpResponseBodyStream;
 use codex_login::AuthManager;
 use codex_login::CodexAuth;
-use codex_rmcp_client::EmaAccessToken;
 use codex_rmcp_client::EmaAuthFailure;
-use codex_rmcp_client::EmaAuthTokenExchangeRequest;
 use codex_rmcp_client::EmaAuthenticatedHttpClient;
-use codex_rmcp_client::EmaCredentialSnapshot;
-use codex_rmcp_client::EmaIdpIdentityRequest;
-use codex_rmcp_client::EmaInvalidGrantSource;
 use codex_rmcp_client::EmaRequestAuthority;
 use codex_rmcp_client::EmaTokenSource;
 use codex_rmcp_client::McpOAuthRefreshMode;
@@ -35,18 +29,16 @@ use codex_rmcp_client::RmcpClient;
 use codex_rmcp_client::StoredOAuthCredentialSnapshot;
 use codex_rmcp_client::StreamableHttpBearerToken;
 use codex_rmcp_client::StreamableHttpRedirectMode;
-use codex_rmcp_client::exchange_ema_auth_token;
-use codex_rmcp_client::resolve_ema_idp_identity;
 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;
 use tokio_util::task::AbortOnDropHandle;
-use url::Url;
 
 use crate::EffectiveMcpServer;
 use crate::McpConfig;
 use crate::McpServerSource;
+use crate::ema_token_source::EnterpriseTokenSource;
 use crate::rmcp_client::StartupOutcomeError;
 
 pub(crate) type EmaAuthContext = Result<Arc<EmaAuth>, StartupOutcomeError>;
@@ -171,28 +163,14 @@ impl EmaAuth {
             .ema_registration()
             .context("enterprise MCP registration is missing")?
             .clone();
-        let credentials = stored_ema_credentials(
-            registration.idp(),
-            &self.scope,
-            self.config.auth_keyring_backend_kind,
-        )
-        .await?
-        .ok_or(EmaAuthFailure::ReauthenticationRequired)?
-        .pin_ema_credentials()
-        .await?;
         let exchange_client: Arc<dyn HttpClient> = Arc::new(EnterpriseExchangeHttpClient {
             inner: Arc::clone(&http_client),
             authority: authority.clone(),
             cancellation: self.config.mcp_ema_cancellation.token(),
         });
         let url = registration.server_url().to_owned();
-        let source = EnterpriseTokenSource {
-            endpoint: Url::parse(registration.server_url())?,
-            registration,
-            credentials,
-            auth: Arc::clone(self),
-            http_client: exchange_client,
-        };
+        let source =
+            EnterpriseTokenSource::new(registration, Arc::clone(self), exchange_client).await?;
         let token = source.access_token().await?;
         let http_client = Arc::new(EmaAuthenticatedHttpClient::new(
             http_client,
@@ -293,91 +271,3 @@ impl HttpClient for EnterpriseExchangeHttpClient {
         })
     }
 }
-
-struct EnterpriseTokenSource {
-    endpoint: Url,
-    registration: McpEmaRegistration,
-    credentials: EmaCredentialSnapshot,
-    auth: Arc<EmaAuth>,
-    http_client: Arc<dyn HttpClient>,
-}
-
-impl EmaTokenSource for EnterpriseTokenSource {
-    fn endpoint(&self) -> &Url {
-        &self.endpoint
-    }
-
-    async fn access_token(&self) -> Result<EmaAccessToken> {
-        self.auth.check_account().await?;
-        let idp = self.registration.idp();
-        let exchange = exchange_ema_auth_token(
-            EmaAuthTokenExchangeRequest {
-                server_url: self.registration.server_url(),
-                resource: self.registration.resource(),
-                scopes: self.registration.scopes(),
-                mcp_client_id: self.registration.client_id(),
-                expected_authorization_server_issuer: self
-                    .registration
-                    .authorization_server_issuer(),
-                idp_issuer: &idp.issuer,
-                idp_client_id: &idp.client_id,
-                idp_identity: Box::pin(async {
-                    self.auth.check_account().await?;
-                    resolve_ema_idp_identity(EmaIdpIdentityRequest {
-                        issuer: &idp.issuer,
-                        client_id: &idp.client_id,
-                        credentials: self.credentials.acquire_current_ema_credentials().await?,
-                        http_client: Arc::clone(&self.http_client),
-                        redirect_mode: StreamableHttpRedirectMode::Legacy,
-                    })
-                    .await
-                }),
-                idp_http_client: Arc::clone(&self.http_client),
-                resource_http_client: Arc::clone(&self.http_client),
-            },
-            std::convert::identity,
-        );
-        let result = tokio::select! {
-            biased;
-            () = self.auth.config.mcp_ema_cancellation.cancelled() => {
-                bail!(EmaAuthFailure::ReauthenticationRequired);
-            }
-            result = exchange => result,
-        };
-        self.auth.check_account().await?;
-        match result {
-            Ok((token, _credentials)) => {
-                self.credentials.acquire_current_ema_credentials().await?;
-                self.auth.check_account().await?;
-                Ok(token)
-            }
-            Err(error) => {
-                if matches!(
-                    error.downcast_ref::<EmaAuthFailure>(),
-                    Some(EmaAuthFailure::InvalidGrant {
-                        grant_source: EmaInvalidGrantSource::EnterpriseIdentity,
-                    })
-                ) && self
-                    .credentials
-                    .invalidate_ema_credentials_if_current(async {
-                        self.auth.check_account().await.is_ok()
-                    })
-                    .await
-                    .is_err()
-                {
-                    tracing::warn!("failed to clear rejected enterprise credentials");
-                }
-                Err(match error.downcast_ref::<EmaAuthFailure>() {
-                    Some(failure) => codex_rmcp_client::EmaTokenSourceError::new(
-                        codex_rmcp_client::EmaCredentialKind::Keyring,
-                        *failure,
-                    ),
-                    None => codex_rmcp_client::EmaTokenSourceError::unexpected(
-                        codex_rmcp_client::EmaCredentialKind::Keyring,
-                    ),
-                }
-                .into())
-            }
-        }
-    }
-}

codex-rs/codex-mcp/src/ema_token_source.rs

@@ -0,0 +1,265 @@
+//! Pins one EMA connection to an account login or a dedicated enterprise grant.
+//! Account refresh exclusion ends with the IdP exchange, before resource-server I/O.
+
+use std::sync::Arc;
+
+use anyhow::Result;
+use anyhow::bail;
+use codex_config::McpEmaRegistration;
+use codex_exec_server::HttpClient;
+use codex_login::AccountOAuthGrant;
+use codex_login::AccountOAuthGrantLease;
+use codex_rmcp_client::EmaAccessToken;
+use codex_rmcp_client::EmaAuthFailure;
+use codex_rmcp_client::EmaAuthTokenExchangeRequest;
+use codex_rmcp_client::EmaCredentialKind;
+use codex_rmcp_client::EmaCredentialLease;
+use codex_rmcp_client::EmaCredentialSnapshot;
+use codex_rmcp_client::EmaIdpIdentityRequest;
+use codex_rmcp_client::EmaInvalidGrantSource;
+use codex_rmcp_client::EmaTokenSource;
+use codex_rmcp_client::EmaTokenSourceError;
+use codex_rmcp_client::StreamableHttpRedirectMode;
+use codex_rmcp_client::exchange_ema_auth_token;
+use codex_rmcp_client::resolve_ema_account_idp_identity;
+use codex_rmcp_client::resolve_ema_idp_identity;
+use url::Url;
+
+use crate::ema_account::account_configuration_issue;
+use crate::ema_account::account_grant_is_usable;
+use crate::ema_auth::EmaAuth;
+use crate::ema_auth::stored_ema_credentials;
+
+enum CredentialSource {
+    Account { client_id: String },
+    Keyring(Box<EmaCredentialSnapshot>),
+}
+
+enum IdpLease {
+    Account(AccountOAuthGrantLease),
+    Keyring(EmaCredentialLease),
+}
+
+enum ExchangeReceipt {
+    Account(AccountOAuthGrant),
+    Keyring(EmaCredentialLease),
+}
+
+pub(crate) struct EnterpriseTokenSource {
+    endpoint: Url,
+    registration: McpEmaRegistration,
+    credentials: CredentialSource,
+    auth: Arc<EmaAuth>,
+    http_client: Arc<dyn HttpClient>,
+}
+
+impl EnterpriseTokenSource {
+    pub(crate) async fn new(
+        registration: McpEmaRegistration,
+        auth: Arc<EmaAuth>,
+        http_client: Arc<dyn HttpClient>,
+    ) -> Result<Self> {
+        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())
+                .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 }
+        };
+        Ok(Self {
+            endpoint: Url::parse(registration.server_url())?,
+            registration,
+            credentials,
+            auth,
+            http_client,
+        })
+    }
+
+    async fn account_grant(&self, client_id: &str) -> Result<AccountOAuthGrant> {
+        self.auth.check_account().await?;
+        let lease = self
+            .auth
+            .manager
+            .lease_account_oauth_grant()
+            .await?
+            .ok_or(EmaAuthFailure::ReauthenticationRequired)?;
+        if !account_grant_is_usable(lease.tokens(), self.registration.idp())
+            || lease
+                .tokens()
+                .oauth_provenance
+                .as_ref()
+                .is_none_or(|provenance| provenance.client_id != client_id)
+        {
+            bail!(EmaAuthFailure::ReauthenticationRequired);
+        }
+        Ok(lease.release())
+    }
+
+    async fn exchange(&self, account: Option<&AccountOAuthGrant>) -> Result<EmaAccessToken> {
+        let idp = self.registration.idp();
+        let client_id = match &self.credentials {
+            CredentialSource::Account { client_id } => client_id,
+            CredentialSource::Keyring(_) => &idp.client_id,
+        };
+        let identity = Box::pin(async {
+            self.auth.check_account().await?;
+            match &self.credentials {
+                CredentialSource::Account { .. } => {
+                    let grant = account.ok_or(EmaAuthFailure::ReauthenticationRequired)?;
+                    let tokens = grant.tokens();
+                    // Metadata discovery must not block ordinary account refresh.
+                    let identity = resolve_ema_account_idp_identity(
+                        &idp.issuer,
+                        client_id,
+                        &tokens.id_token.raw_jwt,
+                        tokens.refresh_token.clone(),
+                        /*credentials*/ (),
+                        Arc::clone(&self.http_client),
+                        StreamableHttpRedirectMode::Legacy,
+                    )
+                    .await?;
+                    let lease = self
+                        .auth
+                        .manager
+                        .lease_account_oauth_grant()
+                        .await?
+                        .ok_or(EmaAuthFailure::ReauthenticationRequired)?;
+                    if lease.tokens() != tokens {
+                        bail!(EmaAuthFailure::ReauthenticationRequired);
+                    }
+                    Ok(identity.map_credentials(|()| IdpLease::Account(lease)))
+                }
+                CredentialSource::Keyring(snapshot) => {
+                    resolve_ema_idp_identity(EmaIdpIdentityRequest {
+                        issuer: &idp.issuer,
+                        client_id,
+                        credentials: snapshot.acquire_current_ema_credentials().await?,
+                        http_client: Arc::clone(&self.http_client),
+                        redirect_mode: StreamableHttpRedirectMode::Legacy,
+                    })
+                    .await
+                    .map(|identity| identity.map_credentials(IdpLease::Keyring))
+                }
+            }
+        });
+        let (token, receipt) = exchange_ema_auth_token(
+            EmaAuthTokenExchangeRequest {
+                server_url: self.registration.server_url(),
+                resource: self.registration.resource(),
+                scopes: self.registration.scopes(),
+                mcp_client_id: self.registration.client_id(),
+                expected_authorization_server_issuer: self
+                    .registration
+                    .authorization_server_issuer(),
+                idp_issuer: &idp.issuer,
+                idp_client_id: client_id,
+                idp_identity: identity,
+                idp_http_client: Arc::clone(&self.http_client),
+                resource_http_client: Arc::clone(&self.http_client),
+            },
+            |lease| match lease {
+                IdpLease::Account(lease) => ExchangeReceipt::Account(lease.release()),
+                IdpLease::Keyring(lease) => ExchangeReceipt::Keyring(lease),
+            },
+        )
+        .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
+                    ));
+                }
+            }
+            ExchangeReceipt::Keyring(lease) => {
+                drop(lease);
+                if let CredentialSource::Keyring(snapshot) = &self.credentials {
+                    snapshot.acquire_current_ema_credentials().await?;
+                }
+            }
+        }
+        self.auth.check_account().await?;
+        Ok(token)
+    }
+
+    async fn acquire_token(&self) -> Result<EmaAccessToken> {
+        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
+            }
+            CredentialSource::Keyring(snapshot) => {
+                let result = self.exchange(/*account*/ None).await;
+                if result.as_ref().is_err_and(|error| {
+                    matches!(
+                        error.downcast_ref::<EmaAuthFailure>(),
+                        Some(EmaAuthFailure::InvalidGrant {
+                            grant_source: EmaInvalidGrantSource::EnterpriseIdentity
+                        })
+                    )
+                }) && snapshot
+                    .invalidate_ema_credentials_if_current(async {
+                        self.auth.check_account().await.is_ok()
+                    })
+                    .await
+                    .is_err()
+                {
+                    tracing::warn!("failed to clear rejected enterprise credentials");
+                }
+                result
+            }
+        }
+    }
+}
+
+impl EmaTokenSource for EnterpriseTokenSource {
+    fn endpoint(&self) -> &Url {
+        &self.endpoint
+    }
+
+    async fn access_token(&self) -> Result<EmaAccessToken> {
+        let kind = match self.credentials {
+            CredentialSource::Account { .. } => EmaCredentialKind::Account,
+            CredentialSource::Keyring(_) => EmaCredentialKind::Keyring,
+        };
+        let result = tokio::select! {
+            biased;
+            () = self.auth.config.mcp_ema_cancellation.cancelled() => Err(EmaAuthFailure::ReauthenticationRequired.into()),
+            result = self.acquire_token() => result,
+        };
+        result.map_err(|error| {
+            if let Some(error) = error.downcast_ref::<EmaTokenSourceError>() {
+                (*error).into()
+            } else {
+                match error.downcast_ref::<EmaAuthFailure>() {
+                    Some(failure) => EmaTokenSourceError::new(kind, *failure),
+                    None if codex_rmcp_client::is_authentication_required_error(&error) => {
+                        EmaTokenSourceError::new(kind, EmaAuthFailure::ReauthenticationRequired)
+                    }
+                    None => EmaTokenSourceError::unexpected(kind),
+                }
+                .into()
+            }
+        })
+    }
+}

codex-rs/codex-mcp/src/lib.rs

@@ -116,7 +116,9 @@ mod client_tool_catalog;
 pub(crate) mod codex_apps;
 pub(crate) mod connection_manager;
 pub(crate) mod elicitation;
+mod ema_account;
 mod ema_auth;
+mod ema_token_source;
 mod event_stream;
 mod executor_environment_http_client;
 pub(crate) mod mcp;

codex-rs/rmcp-client/src/ema_identity.rs

@@ -59,7 +59,8 @@ impl<L> EmaIdpIdentity<L> {
 
 /// Resolve a refresh token issued by the account's validated OAuth authority.
 /// The caller supplies trusted persisted issuer/client provenance and an ID/refresh-token
-/// pair from one admitted account record, retaining its refresh guard in `credentials`.
+/// pair from one admitted account record. Discovery may run before acquiring a refresh
+/// guard; before exchange, attach a lease that still matches that record with map_credentials.
 /// ID-token claims corroborate that record; this does not verify a JWT signature or
 /// authenticate an expired ID token. Only the refresh token is exchanged with the IdP.
 pub async fn resolve_ema_account_idp_identity<L: Send>(