From 6b13dffc5e4c334fe87f823b214ed15a4bf37a11 Mon Sep 17 00:00:00 2001 From: jif Date: Tue, 20 Oct 2026 23:20:41 +0000 Subject: [PATCH] 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 | 55 +++++ codex-rs/codex-mcp/src/ema_auth.rs | 116 +-------- codex-rs/codex-mcp/src/ema_token_source.rs | 265 +++++++++++++++++++++ codex-rs/codex-mcp/src/lib.rs | 2 + codex-rs/rmcp-client/src/ema_identity.rs | 3 +- 5 files changed, 327 insertions(+), 114 deletions(-) create mode 100644 codex-rs/codex-mcp/src/ema_account.rs create mode 100644 codex-rs/codex-mcp/src/ema_token_source.rs diff --git a/codex-rs/codex-mcp/src/ema_account.rs b/codex-rs/codex-mcp/src/ema_account.rs new file mode 100644 index 0000000000..c1ad3c1f7e --- /dev/null +++ b/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 +} diff --git a/codex-rs/codex-mcp/src/ema_auth.rs b/codex-rs/codex-mcp/src/ema_auth.rs index 503d2d056b..beb2b0258d 100644 --- a/codex-rs/codex-mcp/src/ema_auth.rs +++ b/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, 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 = 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, - http_client: Arc, -} - -impl EmaTokenSource for EnterpriseTokenSource { - fn endpoint(&self) -> &Url { - &self.endpoint - } - - async fn access_token(&self) -> Result { - 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::(), - 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::() { - 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()) - } - } - } -} diff --git a/codex-rs/codex-mcp/src/ema_token_source.rs b/codex-rs/codex-mcp/src/ema_token_source.rs new file mode 100644 index 0000000000..af48a1d851 --- /dev/null +++ b/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), +} + +enum IdpLease { + Account(AccountOAuthGrantLease), + Keyring(EmaCredentialLease), +} + +enum ExchangeReceipt { + Account(AccountOAuthGrant), + Keyring(EmaCredentialLease), +} + +pub(crate) struct EnterpriseTokenSource { + endpoint: Url, + registration: McpEmaRegistration, + credentials: CredentialSource, + auth: Arc, + http_client: Arc, +} + +impl EnterpriseTokenSource { + pub(crate) async fn new( + registration: McpEmaRegistration, + auth: Arc, + 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()) + .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 { + 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 { + 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 { + 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::(), + 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 { + 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::() { + (*error).into() + } else { + match error.downcast_ref::() { + 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() + } + }) + } +} diff --git a/codex-rs/codex-mcp/src/lib.rs b/codex-rs/codex-mcp/src/lib.rs index ac348c7439..8ba8a1d220 100644 --- a/codex-rs/codex-mcp/src/lib.rs +++ b/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; diff --git a/codex-rs/rmcp-client/src/ema_identity.rs b/codex-rs/rmcp-client/src/ema_identity.rs index 937dcb3093..8bd7fdb51a 100644 --- a/codex-rs/rmcp-client/src/ema_identity.rs +++ b/codex-rs/rmcp-client/src/ema_identity.rs @@ -59,7 +59,8 @@ impl EmaIdpIdentity { /// 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(