Add account credential support to EMA token exchange helpers (#55685)

## What changed

- Add an account IdP identity resolver that checks ID-token bindings against trusted issuer and client provenance before discovering the refresh-token endpoint. An expired ID token does not by itself invalidate the refresh grant.
- Generalize `EmaIdpIdentity` and token exchange requests to carry caller-owned credential leases. Invoke a `finish_idp` callback after the IdP exchange and before resource authorization server I/O so callers can release their refresh guards at that boundary.
- Reject empty expected OIDC issuers and client IDs, and expose an unverified issuer hint solely for explaining legacy login migration.

GitOrigin-RevId: 34473fed82a9a6e58997664cd5e6e67049b5ee7e

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

@@ -301,28 +301,33 @@ impl EmaTokenSource for EnterpriseTokenSource {
     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),
-        });
+        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() => {

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

@@ -41,7 +41,7 @@ struct EmaAuthorizationMetadata {
 }
 
 /// Inputs for non-interactive enterprise-managed MCP authorization.
-pub struct EmaAuthTokenExchangeRequest<'a> {
+pub struct EmaAuthTokenExchangeRequest<'a, L = EmaCredentialLease> {
     pub server_url: &'a str,
     pub resource: Option<&'a str>,
     pub scopes: &'a [String],
@@ -49,13 +49,13 @@ pub struct EmaAuthTokenExchangeRequest<'a> {
     pub expected_authorization_server_issuer: Option<&'a str>,
     pub idp_issuer: &'a str,
     pub idp_client_id: &'a str,
-    pub idp_identity: BoxFuture<'a, Result<EmaIdpIdentity>>,
+    pub idp_identity: BoxFuture<'a, Result<EmaIdpIdentity<L>>>,
     pub idp_http_client: Arc<dyn HttpClient>,
     pub resource_http_client: Arc<dyn HttpClient>,
 }
 
-impl<'a> From<&EmaAuthTokenExchangeRequest<'a>> for EmaDiscoveryRequest<'a> {
-    fn from(request: &EmaAuthTokenExchangeRequest<'a>) -> Self {
+impl<'a, L> From<&EmaAuthTokenExchangeRequest<'a, L>> for EmaDiscoveryRequest<'a> {
+    fn from(request: &EmaAuthTokenExchangeRequest<'a, L>) -> Self {
         Self {
             server_url: request.server_url,
             resource: request.resource,
@@ -66,9 +66,10 @@ impl<'a> From<&EmaAuthTokenExchangeRequest<'a>> for EmaDiscoveryRequest<'a> {
 }
 
 /// Discover the resource's authorization server, then perform the ID-JAG flow.
-pub async fn exchange_ema_auth_token(
-    request: EmaAuthTokenExchangeRequest<'_>,
-) -> Result<(EmaAccessToken, EmaCredentialLease)> {
+pub async fn exchange_ema_auth_token<L: Send, R: Send>(
+    request: EmaAuthTokenExchangeRequest<'_, L>,
+    finish_idp: impl FnOnce(L) -> R + Send,
+) -> Result<(EmaAccessToken, R)> {
     validate_ema_auth_resource(request.server_url, request.resource)?;
     let client_id = request
         .mcp_client_id
@@ -89,21 +90,24 @@ pub async fn exchange_ema_auth_token(
         refresh_token,
         credentials,
     } = request.idp_identity.await?;
-    let token = exchange_id_jag(EmaIdJagExchangeRequest {
-        resource: &resource,
-        scopes: request.scopes,
-        mcp_client_id: client_id,
-        authorization_server_issuer: issuer,
-        authorization_server_token_endpoint: &metadata.token_endpoint,
-        idp_token_endpoint: &token_endpoint,
-        idp_issuer: request.idp_issuer,
-        idp_client_id: request.idp_client_id,
-        refresh_token,
-        idp_http_client: request.idp_http_client,
-        resource_http_client: request.resource_http_client,
-    })
-    .await?;
-    Ok((token, credentials))
+    exchange_id_jag(
+        EmaIdJagExchangeRequest {
+            resource: &resource,
+            scopes: request.scopes,
+            mcp_client_id: client_id,
+            authorization_server_issuer: issuer,
+            authorization_server_token_endpoint: &metadata.token_endpoint,
+            idp_token_endpoint: &token_endpoint,
+            idp_issuer: request.idp_issuer,
+            idp_client_id: request.idp_client_id,
+            refresh_token,
+            idp_http_client: request.idp_http_client,
+            resource_http_client: request.resource_http_client,
+        },
+        credentials,
+        finish_idp,
+    )
+    .await
 }
 
 async fn discover_ema_authorization_metadata(

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

@@ -133,7 +133,7 @@ async fn discovery_pins_resource_issuer_and_public_jwt_bearer_before_identity()
             polled.store(true, Ordering::SeqCst);
             bail!("identity resolver reached")
         });
-        exchange_ema_auth_token(exchange)
+        exchange_ema_auth_token(exchange, std::convert::identity)
             .await
             .expect_err("stop before token exchange");
         assert_eq!(polled.load(Ordering::SeqCst), valid, "{additions}");
@@ -452,7 +452,7 @@ async fn rejected_grants_preserve_their_source_and_release_the_credential_lease(
                 credentials: crate::EmaCredentialLease { credentials },
             })
         });
-        let error = exchange_ema_auth_token(exchange)
+        let error = exchange_ema_auth_token(exchange, std::convert::identity)
             .await
             .expect_err("invalid grant");
         assert_eq!(

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

@@ -59,6 +59,9 @@ pub(crate) fn oidc_identity(
     expected_issuer: &str,
     expected_audience: &str,
 ) -> Result<OidcClaims> {
+    if expected_issuer.trim().is_empty() || expected_audience.trim().is_empty() {
+        bail!("OIDC identity assertion requires a non-empty expected issuer and client");
+    }
     let claims: OidcClaims = signed_jwt(assertion)?;
     if claims.iss != expected_issuer || claims.sub.trim().is_empty() {
         bail!("OIDC identity assertion issuer or subject does not match the enterprise IdP");
@@ -82,6 +85,12 @@ pub(crate) fn oidc_identity(
     Ok(claims)
 }
 
+/// Unverified issuer claim used only to explain a legacy login migration.
+/// Never use this hint to choose where a refresh token is sent.
+pub(crate) fn oidc_issuer_hint(assertion: &str, expected_issuer: &str) -> bool {
+    signed_jwt::<OidcClaims>(assertion).is_ok_and(|claims| claims.iss == expected_issuer)
+}
+
 pub(crate) fn validate_oidc_identity_assertion(
     assertion: &str,
     expected_issuer: &str,

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

@@ -58,9 +58,11 @@ pub(crate) struct EmaIdJagExchangeRequest<'a> {
 }
 
 /// Exchanges an enterprise IdP credential for a resource-bound MCP bearer token.
-pub(crate) async fn exchange_id_jag(
+pub(crate) async fn exchange_id_jag<L: Send, R: Send>(
     request: EmaIdJagExchangeRequest<'_>,
-) -> Result<EmaAccessToken> {
+    credentials: L,
+    finish_idp: impl FnOnce(L) -> R + Send,
+) -> Result<(EmaAccessToken, R)> {
     let idp_http = OAuthHttpClientAdapter::new_with_redirect_mode(
         request.idp_http_client,
         build_default_headers(/*http_headers*/ None, /*env_http_headers*/ None)?,
@@ -78,7 +80,7 @@ pub(crate) async fn exchange_id_jag(
     let refresh_token = RefreshToken::new(request.refresh_token);
     // Discovery has already verified both public-client registrations. The SDK owns
     // token validation and the two exchanges; Codex retains routing and credential policy.
-    let token = EmaExchangeRequest::new(
+    let assertion = EmaExchangeRequest::new(
         EmaAuthorizationServer::new(
             request.idp_issuer,
             request.idp_token_endpoint,
@@ -93,36 +95,47 @@ pub(crate) async fn exchange_id_jag(
         &refresh_token,
     )
     .with_scopes(request.scopes.iter().cloned())
-    .exchange(&idp_http, &resource_http)
+    .exchange_id_jag(&idp_http)
     .await
-    .map_err(|error| {
-        let failure = match &error {
-            EmaError::InvalidGrant(EmaExchangeStage::IdentityProvider) => {
-                Some(EmaAuthFailure::InvalidGrant {
-                    grant_source: EmaInvalidGrantSource::EnterpriseIdentity,
-                })
-            }
-            EmaError::InvalidGrant(EmaExchangeStage::ResourceAuthorizationServer) => {
-                Some(EmaAuthFailure::InvalidGrant {
-                    grant_source: EmaInvalidGrantSource::ResourceAuthorization,
-                })
-            }
-            EmaError::InsufficientUserAuthentication(_) => {
-                Some(EmaAuthFailure::InsufficientUserAuthentication)
-            }
-            // Future SDK stages must not invalidate the shared enterprise credential.
-            _ => None,
-        };
-        match failure {
-            Some(failure) => anyhow::Error::new(failure).context(error),
-            None => anyhow::Error::new(error),
+    .map_err(exchange_error)?;
+    drop(refresh_token);
+    // Only the IdP uses the refresh token. Release its lease before resource-AS I/O.
+    let receipt = finish_idp(credentials);
+    let token = assertion
+        .exchange(&resource_http)
+        .await
+        .map_err(exchange_error)?;
+    Ok((
+        EmaAccessToken {
+            access_token: token.access_token.secret().to_owned(),
+            expires_in: token.expires_in,
+            received_at: Instant::now(),
+        },
+        receipt,
+    ))
+}
+
+fn exchange_error(error: EmaError) -> anyhow::Error {
+    let failure = match &error {
+        EmaError::InvalidGrant(EmaExchangeStage::IdentityProvider) => {
+            Some(EmaAuthFailure::InvalidGrant {
+                grant_source: EmaInvalidGrantSource::EnterpriseIdentity,
+            })
         }
-    })?;
-    Ok(EmaAccessToken {
-        access_token: token.access_token.secret().to_owned(),
-        expires_in: token.expires_in,
-        received_at: Instant::now(),
-    })
+        EmaError::InvalidGrant(EmaExchangeStage::ResourceAuthorizationServer) => {
+            Some(EmaAuthFailure::InvalidGrant {
+                grant_source: EmaInvalidGrantSource::ResourceAuthorization,
+            })
+        }
+        EmaError::InsufficientUserAuthentication(_) => {
+            Some(EmaAuthFailure::InsufficientUserAuthentication)
+        }
+        _ => None,
+    };
+    match failure {
+        Some(failure) => anyhow::Error::new(failure).context(error),
+        None => anyhow::Error::new(error),
+    }
 }
 
 #[cfg(test)]

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

@@ -75,20 +75,25 @@ fn token_response() -> Value {
 async fn exchange(server: &MockServer, refresh_token: &str) -> Result<EmaAccessToken> {
     let issuer = format!("{}/idp", server.uri());
     let audience = format!("{}/as", server.uri());
-    exchange_id_jag(EmaIdJagExchangeRequest {
-        resource: &format!("{}/mcp", server.uri()),
-        scopes: &[],
-        mcp_client_id: "mcp-client",
-        authorization_server_issuer: &audience,
-        authorization_server_token_endpoint: &format!("{audience}/token"),
-        idp_token_endpoint: &format!("{issuer}/token"),
-        idp_issuer: &issuer,
-        idp_client_id: "idp-client",
-        refresh_token: refresh_token.to_string(),
-        idp_http_client: http_client(),
-        resource_http_client: http_client(),
-    })
+    exchange_id_jag(
+        EmaIdJagExchangeRequest {
+            resource: &format!("{}/mcp", server.uri()),
+            scopes: &[],
+            mcp_client_id: "mcp-client",
+            authorization_server_issuer: &audience,
+            authorization_server_token_endpoint: &format!("{audience}/token"),
+            idp_token_endpoint: &format!("{issuer}/token"),
+            idp_issuer: &issuer,
+            idp_client_id: "idp-client",
+            refresh_token: refresh_token.to_string(),
+            idp_http_client: http_client(),
+            resource_http_client: http_client(),
+        },
+        (),
+        std::convert::identity,
+    )
     .await
+    .map(|(token, ())| token)
 }
 
 #[tokio::test]
@@ -122,21 +127,26 @@ async fn public_client_round_trip_preserves_signed_narrowing() -> Result<()> {
                 .await;
         }
         let expected_subject = refresh_token.to_string();
-        let result = exchange_id_jag(EmaIdJagExchangeRequest {
-            resource: &resource,
-            scopes,
-            mcp_client_id: "mcp-client",
-            authorization_server_issuer: &audience,
-            authorization_server_token_endpoint: &format!("{audience}/token"),
-            idp_token_endpoint: &format!("{issuer}/token"),
-            idp_issuer: &issuer,
-            idp_client_id: "idp-client",
-            refresh_token: refresh_token.to_string(),
-            idp_http_client: Arc::clone(&client),
-            resource_http_client: Arc::clone(&client),
-        })
+        let result = exchange_id_jag(
+            EmaIdJagExchangeRequest {
+                resource: &resource,
+                scopes,
+                mcp_client_id: "mcp-client",
+                authorization_server_issuer: &audience,
+                authorization_server_token_endpoint: &format!("{audience}/token"),
+                idp_token_endpoint: &format!("{issuer}/token"),
+                idp_issuer: &issuer,
+                idp_client_id: "idp-client",
+                refresh_token: refresh_token.to_string(),
+                idp_http_client: Arc::clone(&client),
+                resource_http_client: Arc::clone(&client),
+            },
+            (),
+            std::convert::identity,
+        )
         .boxed()
-        .await;
+        .await
+        .map(|(token, ())| token);
         if !valid {
             assert!(result.is_err(), "invalid subject must fail before HTTP");
             assert!(

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

@@ -39,10 +39,60 @@ pub struct EmaIdpIdentityRequest<'a> {
 }
 
 /// An opaque IdP refresh token paired with the credentials admitted for its exchange.
-pub struct EmaIdpIdentity {
+pub struct EmaIdpIdentity<L = EmaCredentialLease> {
     pub(crate) token_endpoint: String,
     pub(crate) refresh_token: String,
-    pub(crate) credentials: EmaCredentialLease,
+    pub(crate) credentials: L,
+}
+
+impl<L> EmaIdpIdentity<L> {
+    /// Preserve an opaque credential lease when selecting an EMA credential source.
+    /// The mapping must retain its owned refresh guard until the IdP exchange finishes.
+    pub fn map_credentials<U>(self, map: impl FnOnce(L) -> U) -> EmaIdpIdentity<U> {
+        EmaIdpIdentity {
+            token_endpoint: self.token_endpoint,
+            refresh_token: self.refresh_token,
+            credentials: map(self.credentials),
+        }
+    }
+}
+
+/// 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`.
+/// 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>(
+    issuer: &str,
+    client_id: &str,
+    id_token: &str,
+    refresh_token: String,
+    credentials: L,
+    http_client: Arc<dyn HttpClient>,
+    redirect_mode: StreamableHttpRedirectMode,
+) -> Result<EmaIdpIdentity<L>> {
+    if refresh_token.trim().is_empty() {
+        return Err(ema_reauthentication_required(
+            "Codex account has no refresh token; sign in again",
+        ));
+    }
+    crate::ema_claims::oidc_identity(id_token, issuer, client_id)?;
+    let token_endpoint = resolve_ema_idp_token_endpoint(issuer, http_client, redirect_mode).await?;
+    Ok(EmaIdpIdentity {
+        token_endpoint,
+        refresh_token,
+        credentials,
+    })
+}
+
+/// Validate account ID-token bindings without treating an expired ID token as an expired grant.
+pub fn account_oidc_identity_is_usable(id_token: &str, issuer: &str, client_id: &str) -> bool {
+    crate::ema_claims::oidc_identity(id_token, issuer, client_id).is_ok()
+}
+
+/// Only an explanatory legacy-login hint, never authority for RT exchange.
+pub fn account_oidc_issuer_hint(id_token: &str, issuer: &str) -> bool {
+    crate::ema_claims::oidc_issuer_hint(id_token, issuer)
 }
 
 /// Returns whether stored credentials contain a refresh token bound to the configured login.
@@ -110,30 +160,9 @@ async fn resolve_ema_idp_identity_in<K: KeyringStore + Clone + 'static>(
         bail!("stored enterprise IdP credentials do not match the configured issuer and client");
     }
     stored_oidc_identity(credentials)?;
-    let client = OAuthHttpClientAdapter::new_with_redirect_mode(
-        request.http_client,
-        build_default_headers(/*http_headers*/ None, /*env_http_headers*/ None)?,
-        request.issuer,
-        /*has_configured_headers*/ false,
-        request.redirect_mode,
-    )?;
-    let (_, metadata) =
-        resolve_ema_idp_authorization_manager(request.issuer, Arc::new(client)).await?;
-    for (name, expected) in [
-        (
-            "identity_chaining_requested_token_types_supported",
-            ID_JAG_TOKEN_TYPE,
-        ),
-        ("grant_types_supported", TOKEN_EXCHANGE_GRANT_TYPE),
-    ] {
-        if advertised_capability(metadata.additional_fields.get(name), expected, name)?
-            == Some(false)
-        {
-            bail!(
-                "enterprise IdP does not advertise the required ID-JAG token exchange capability"
-            );
-        }
-    }
+    let token_endpoint =
+        resolve_ema_idp_token_endpoint(request.issuer, request.http_client, request.redirect_mode)
+            .await?;
     let credentials = request.credentials;
     let credential_lock = RefreshCredentialLock::acquire_for_server(
         &credentials.credentials.credentials().server_name,
@@ -160,7 +189,7 @@ async fn resolve_ema_idp_identity_in<K: KeyringStore + Clone + 'static>(
                 )
             })?;
         Ok(EmaIdpIdentity {
-            token_endpoint: metadata.token_endpoint,
+            token_endpoint,
             refresh_token: refresh_token.secret().to_string(),
             credentials,
         })
@@ -169,6 +198,37 @@ async fn resolve_ema_idp_identity_in<K: KeyringStore + Clone + 'static>(
     .map_err(|_| anyhow!("enterprise IdP credential reread task failed"))?
 }
 
+async fn resolve_ema_idp_token_endpoint(
+    issuer: &str,
+    http_client: Arc<dyn HttpClient>,
+    redirect_mode: StreamableHttpRedirectMode,
+) -> Result<String> {
+    let client = OAuthHttpClientAdapter::new_with_redirect_mode(
+        http_client,
+        build_default_headers(/*http_headers*/ None, /*env_http_headers*/ None)?,
+        issuer,
+        /*has_configured_headers*/ false,
+        redirect_mode,
+    )?;
+    let (_, metadata) = resolve_ema_idp_authorization_manager(issuer, Arc::new(client)).await?;
+    for (name, expected) in [
+        (
+            "identity_chaining_requested_token_types_supported",
+            ID_JAG_TOKEN_TYPE,
+        ),
+        ("grant_types_supported", TOKEN_EXCHANGE_GRANT_TYPE),
+    ] {
+        if advertised_capability(metadata.additional_fields.get(name), expected, name)?
+            == Some(false)
+        {
+            bail!(
+                "enterprise IdP does not advertise the required ID-JAG token exchange capability"
+            );
+        }
+    }
+    Ok(metadata.token_endpoint)
+}
+
 #[cfg(test)]
 #[path = "ema_identity_tests.rs"]
 mod tests;

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

@@ -56,6 +56,9 @@ pub use ema_http_client::EmaRequestAuthority;
 pub use ema_http_client::EmaTokenSource;
 pub use ema_identity::EmaIdpIdentity;
 pub use ema_identity::EmaIdpIdentityRequest;
+pub use ema_identity::account_oidc_identity_is_usable;
+pub use ema_identity::account_oidc_issuer_hint;
+pub use ema_identity::resolve_ema_account_idp_identity;
 pub use ema_identity::resolve_ema_idp_identity;
 pub use ema_identity::stored_ema_identity_is_usable;
 pub use enterprise_oauth_login::EnterpriseOAuthCredentialGuard;