openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-10-20 17:38:45 +00:00
Code Issues Packages Projects Releases Wiki Activity

Enable enterprise-managed authorization for MCP connections (#55334)

Browse Source 
## What changed

- Connect directly configured MCP servers using account-scoped enterprise grants stored in the keyring, with token exchange and authenticated HTTP requests.
- Retain authorization for the thread across MCP refreshes and revoke it when the account owner changes, cloud configuration becomes invalid, or an enterprise login replaces the grant.
- Negotiate the enterprise-managed authorization extension with the `2026-07-28` protocol, bypass shared tool catalog caching for enterprise servers, and report credential status with enterprise sign-in guidance.

GitOrigin-RevId: c0f96d37194212758a211ae6a4107205df96df26
This commit is contained in:
jif
2026-10-19 17:57:35 +00:00
committed by copyberry
parent fbb2f9dbec
commit 1569abc5dc
10 changed files with 461 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/app-server/src/request_processors/account_processor/enterprise_login.rs
+9 -3
View File
@@ -240,11 +240,17 @@ impl EnterpriseLoginState {
// Keep cancellation responsive during the async checks above.
// Once admitted, retain this guard through the synchronous write.
let active = lifecycle.active.lock().await;
(active.as_ref().is_some_and(|attempt| {
let current = active.as_ref().is_some_and(|attempt| {
attempt.login_id == completion_id && !attempt.cancel.is_cancelled()
})
&& ema_auth_scope(lifecycle.auth_manager.auth_cached().as_ref()).as_ref() == Some(&scope))
.then_some(active)
&& ema_auth_scope(lifecycle.auth_manager.auth_cached().as_ref()).as_ref() == Some(&scope);
if !current {
return None;
}
// The credential lock is held: revoke cached bearers before replacing
// the grant, and let new threads pin only the committed credential.
lifecycle.config_manager.revoke_mcp_ema();
Some(active)
}).await?;
// A successful commit is no longer cancelable. Clear it while
// still holding the commit gate.
codex-rs/codex-mcp/src/connection_manager.rs
+31 -1
View File
@@ -26,6 +26,7 @@ use startup::should_share_codex_apps_tools_cache;
pub(crate) use tool_catalog::BindingCatalogRevision;
pub use tool_catalog::tool_is_model_visible;
use codex_config::McpServerAuth;
use std::collections::HashMap;
use std::sync::Arc;
use std::sync::OnceLock;
@@ -69,6 +70,8 @@ use codex_diagnostics::Gauge;
use codex_diagnostics::GaugeGuard;
use codex_otel::auth_storage::AuthStorageOriginator;
use codex_protocol::mcp::CallToolResult;
use codex_protocol::mcp::ClientMcpExtensions;
use codex_protocol::mcp::ENTERPRISE_MANAGED_AUTHORIZATION_EXTENSION_ID;
use codex_protocol::mcp::McpServerInfo;
use codex_protocol::protocol::Event;
use codex_protocol::protocol::EventMsg;
@@ -279,6 +282,7 @@ impl McpConnectionSet {
/// Creates an MCP connection manager. Threadless callers can pass no `tx_event`; startup
/// notifications are then skipped and interactive elicitations are declined.
pub async fn new(
ema_auth: Option<crate::ema_auth::EmaAuthContext>,
previous: Option<&Self>,
publication_gate: McpPublicationGate,
input: McpRuntimeInput,
@@ -408,6 +412,26 @@ impl McpConnectionSet {
} else {
default_protocol_mode
};
let uses_ema_auth =
config.xaa_enabled && matches!(configured_config.auth, McpServerAuth::EmaAuth);
let protocol_mode = if uses_ema_auth {
crate::McpProtocolMode::V20260728
} else {
protocol_mode
};
let client_mcp_extensions = if uses_ema_auth {
ClientMcpExtensions::new(
client_mcp_extensions
.iter()
.map(|(id, settings)| (id.to_string(), settings.clone()))
.chain(std::iter::once((
ENTERPRISE_MANAGED_AUTHORIZATION_EXTENSION_ID.to_string(),
serde_json::json!({}),
))),
)
} else {
client_mcp_extensions
};
let configured_tool_filter = ToolFilter::from_config(&configured_config);
let startup_timeout = configured_config
.startup_timeout_sec
@@ -643,6 +667,7 @@ impl McpConnectionSet {
let cancel_token = startup_cancellation_token.child_token();
let mut tool_catalog_cache_context = if server_name == CODEX_APPS_MCP_SERVER_NAME
|| server.requires_read_only_mcp_tools()
|| matches!(configured_config.auth, McpServerAuth::EmaAuth)
{
None
} else if let Ok(environment) = resolved_environment.as_ref() {
@@ -666,6 +691,7 @@ impl McpConnectionSet {
}
let has_runtime_auth = runtime_auth_provider.is_some();
let async_managed_client = AsyncManagedClient::new(
ema_auth.clone(),
server_name.clone(),
startup_submit_id.clone(),
server,
@@ -763,7 +789,11 @@ impl McpConnectionSet {
}
if let Some(tx_event) = tx_event.as_ref() {
let auth_state = match &outcome {
Err(error) if error.is_authentication_required() && !has_runtime_auth => {
Err(error)
if error.is_authentication_required()
&& !has_runtime_auth
&& !matches!(configured_config.auth, McpServerAuth::EmaAuth) =>
{
match &configured_config.transport {
McpServerTransportConfig::StreamableHttp {
url,
codex-rs/codex-mcp/src/connection_manager/startup.rs
+8 -5
View File
@@ -101,11 +101,14 @@ pub(super) fn mcp_init_error_display(
"GitHub MCP does not support OAuth. Log in by adding a personal access token (https://github.com/settings/personal-access-tokens) to your environment and config.toml:\n[mcp_servers.{server_key}]\nbearer_token_env_var = CODEX_GITHUB_PERSONAL_ACCESS_TOKEN"
)
} else if error.is_authentication_required() {
let recovery_hint = if config.is_some_and(|config| !config.is_local_environment()) {
"Use your client's MCP OAuth sign-in flow.".to_string()
} else {
format!("Run `codex mcp login {server_name}`.")
};
let recovery_hint =
if config.is_some_and(|config| matches!(config.auth, McpServerAuth::EmaAuth)) {
"Sign in to your enterprise provider in Codex, then start a new thread.".to_string()
} else if config.is_some_and(|config| !config.is_local_environment()) {
"Use your client's MCP OAuth sign-in flow.".to_string()
} else {
format!("Run `codex mcp login {server_name}`.")
};
let auth_status = match reason {
Some(McpStartupFailureReason::ReauthenticationRequired) => {
"requires OAuth reauthentication"
codex-rs/codex-mcp/src/connection_manager_tests.rs
+8
View File
@@ -2059,6 +2059,7 @@ async fn read_only_apps_discovery_never_uses_a_shared_writable_catalog() -> anyh
let cancellation = CancellationToken::new();
cancellation.cancel();
let manager = McpConnectionSet::new(
/*ema_auth*/ None,
/*previous*/ None,
McpPublicationGate::already_published(),
McpRuntimeInput {
@@ -2132,6 +2133,7 @@ async fn hosted_apps_protocol_mode_is_independent_of_generic_mode() -> anyhow::R
let startup_cancellation_token = CancellationToken::new();
startup_cancellation_token.cancel();
let manager = McpConnectionSet::new(
/*ema_auth*/ None,
/*previous*/ None,
McpPublicationGate::already_published(),
McpRuntimeInput {
@@ -2238,6 +2240,7 @@ async fn codex_apps_extension_does_not_share_host_owned_tools_cache() -> anyhow:
let startup_cancellation_token = CancellationToken::new();
startup_cancellation_token.cancel();
let manager = McpConnectionSet::new(
/*ema_auth*/ None,
/*previous*/ None,
McpPublicationGate::already_published(),
McpRuntimeInput {
@@ -4807,6 +4810,7 @@ async fn executor_owned_chatgpt_mcp_accepts_only_safe_explicit_authorization() -
);
let manager = McpConnectionSet::new(
/*ema_auth*/ None,
/*previous*/ None,
McpPublicationGate::already_published(),
McpRuntimeInput {
@@ -4925,6 +4929,7 @@ async fn no_local_runtime_fails_local_stdio_but_keeps_local_http_server() {
let cancel_token = CancellationToken::new();
let manager = McpConnectionSet::new(
/*ema_auth*/ None,
/*previous*/ None,
McpPublicationGate::already_published(),
McpRuntimeInput {
@@ -5373,6 +5378,7 @@ async fn reconcile_reusable_server_with_mcp_config(
) -> McpConnectionSet {
let (tx_event, _rx_event) = async_channel::unbounded();
McpConnectionSet::new(
/*ema_auth*/ None,
Some(previous),
McpPublicationGate::already_published(),
McpRuntimeInput {
@@ -6329,6 +6335,7 @@ async fn reconciliation_replaces_connection_when_protocol_mode_changes() {
mcp_config.protocol_mode = codex_rmcp_client::McpProtocolMode::V20260728;
let reconciled = McpConnectionSet::new(
/*ema_auth*/ None,
Some(&previous),
McpPublicationGate::already_published(),
McpRuntimeInput {
@@ -6387,6 +6394,7 @@ async fn reconciliation_reuses_legacy_stdio_server_when_modern_protocol_is_enabl
mcp_config.protocol_mode = codex_rmcp_client::McpProtocolMode::V20260728;
let reconciled = McpConnectionSet::new(
/*ema_auth*/ None,
Some(&previous),
McpPublicationGate::already_published(),
McpRuntimeInput {
codex-rs/codex-mcp/src/ema_auth.rs
+337 -1
View File
@@ -1,9 +1,345 @@
//! Derives the enterprise credential owner from the active Codex account.
//! Dedicated enterprise grants and one thread's EMA authorization lifetime.
//! Account credentials identify the owner; only the separate keyring grant is exchanged.
use std::sync::Arc;
use anyhow::Context;
use anyhow::Result;
use anyhow::anyhow;
use anyhow::bail;
use codex_config::CloudConfigBundleBindingStatus;
use codex_config::McpEmaAuthScope;
use codex_config::McpEmaRegistration;
use codex_config::McpServerIdpOAuthConfig;
use codex_config::types::AuthKeyringBackendKind;
use codex_config::types::OAuthCredentialsStoreMode;
use codex_exec_server::ExecServerError;
use codex_exec_server::HttpClient;
use codex_exec_server::HttpRequestParams;
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;
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::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::rmcp_client::StartupOutcomeError;
pub(crate) type EmaAuthContext = Result<Arc<EmaAuth>, StartupOutcomeError>;
pub fn ema_auth_scope(auth: Option<&CodexAuth>) -> Option<McpEmaAuthScope> {
let auth = auth?;
McpEmaAuthScope::new(auth.get_chatgpt_user_id()?, auth.get_account_id()?)
}
/// Retained by the thread across MCP refreshes; a revoked owner cannot be replaced.
pub(crate) struct EmaAuth {
scope: McpEmaAuthScope,
config: Arc<McpConfig>,
authority: EmaRequestAuthority,
_revocation: AbortOnDropHandle<()>,
}
impl EmaAuth {
pub(crate) fn new(
config: Arc<McpConfig>,
auth: Option<&CodexAuth>,
auth_manager: Option<&Arc<AuthManager>>,
) -> Result<Arc<Self>> {
let scope = ema_auth_scope(auth).ok_or(EmaAuthFailure::ReauthenticationRequired)?;
let manager = auth_manager.context("enterprise MCP requires an account owner")?;
let mut changes = manager.auth_change_state_receiver();
let owner_generation = changes.borrow().owner_generation;
// Never hold the watch guard while reading cached auth: writers take the auth lock first.
if ema_auth_scope(manager.auth_cached().as_ref()).as_ref() != Some(&scope)
|| changes.borrow().owner_generation != owner_generation
{
bail!(EmaAuthFailure::ReauthenticationRequired);
}
let cancellation = config.mcp_ema_cancellation.clone();
let policy = config.config_layer_stack.cloud_config_binding().cloned();
let authority = EmaRequestAuthority::new(cancellation.clone(), {
let changes = changes.clone();
let policy = policy.clone();
move || {
changes.borrow().owner_generation == owner_generation
&& policy.as_ref().is_none_or(|binding| {
binding.read().status == CloudConfigBundleBindingStatus::Current
})
}
});
authority.check()?;
let revocation = tokio::spawn(async move {
tokio::select! {
biased;
() = cancellation.cancelled() => {}
() = async {
match policy {
Some(binding) => binding.invalidated().await,
None => std::future::pending().await,
}
} => {}
() = async {
while changes.changed().await.is_ok() {
if changes.borrow().owner_generation != owner_generation {
break;
}
}
} => {}
}
cancellation.cancel();
});
Ok(Arc::new(Self {
scope,
config,
authority,
_revocation: AbortOnDropHandle::new(revocation),
}))
}
pub(crate) async fn connect(
self: &Arc<Self>,
server_name: &str,
server: &EffectiveMcpServer,
http_client: Arc<dyn HttpClient>,
protocol_mode: McpProtocolMode,
) -> Result<RmcpClient> {
self.authority.check()?;
let config = server.config();
config
.validate_ema_auth_transport()
.map_err(anyhow::Error::msg)?;
if !self.config.xaa_enabled
|| !self
.config
.mcp_server_catalog
.server(server_name)
.is_some_and(|server| matches!(server.source(), McpServerSource::Config))
{
bail!("enterprise MCP requires an enabled, directly configured server");
}
let registration = config
.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: self.authority.clone(),
cancellation: self.config.mcp_ema_cancellation.clone(),
});
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 token = source.access_token().await?;
let http_client = Arc::new(EmaAuthenticatedHttpClient::new(
http_client,
source,
token,
self.authority.clone(),
));
RmcpClient::new_streamable_http_client_with_protocol_mode_and_redirect_mode(
server_name,
&url,
Some(StreamableHttpBearerToken::ProvidedByHttpClient),
/*http_headers*/ None,
/*env_http_headers*/ None,
/*oauth_config*/ None,
OAuthCredentialsStoreMode::Keyring,
self.config.auth_keyring_backend_kind,
http_client,
/*auth_provider*/ None,
protocol_mode,
StreamableHttpRedirectMode::Legacy,
McpOAuthRefreshMode::Legacy,
)
.await
}
}
pub(crate) async fn stored_ema_credentials(
idp: &McpServerIdpOAuthConfig,
scope: &McpEmaAuthScope,
keyring_backend: AuthKeyringBackendKind,
) -> Result<Option<StoredOAuthCredentialSnapshot>> {
let name = idp.credential_name(scope);
let idp = idp.clone();
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,
)
})
})
.map_err(|_| anyhow!("failed to read enterprise credentials"))
})
})
.await
.map_err(|_| anyhow!("enterprise credential read task failed"))?
}
// Discovery and both credential exchanges must check the owner at each dispatch,
// including when the revocation watcher has not yet been scheduled.
struct EnterpriseExchangeHttpClient {
inner: Arc<dyn HttpClient>,
authority: EmaRequestAuthority,
cancellation: CancellationToken,
}
impl HttpClient for EnterpriseExchangeHttpClient {
fn http_request(
&self,
params: HttpRequestParams,
) -> BoxFuture<'_, Result<HttpRequestResponse, ExecServerError>> {
Box::pin(async move {
self.authority.check()?;
tokio::select! {
biased;
() = self.cancellation.cancelled() => Err(ExecServerError::HttpRequest("enterprise MCP authorization was revoked".into())),
response = self.inner.http_request(params) => {
self.authority.check()?;
response
}
}
})
}
fn http_request_stream(
&self,
params: HttpRequestParams,
) -> BoxFuture<'_, Result<(HttpRequestResponse, HttpResponseBodyStream), ExecServerError>> {
Box::pin(async move {
self.authority.check()?;
tokio::select! {
biased;
() = self.cancellation.cancelled() => Err(ExecServerError::HttpRequest("enterprise MCP authorization was revoked".into())),
response = self.inner.http_request_stream(params) => {
self.authority.check()?;
response.map(|(response, body)| (response, body.cancel_on(self.cancellation.clone())))
}
}
})
}
}
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.authority.check()?;
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.authority.check()?;
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 result = tokio::select! {
biased;
() = self.auth.config.mcp_ema_cancellation.cancelled() => {
bail!(EmaAuthFailure::ReauthenticationRequired);
}
result = exchange => result,
};
self.auth.authority.check()?;
match result {
Ok((token, _credentials)) => {
self.credentials.acquire_current_ema_credentials().await?;
self.auth.authority.check()?;
Ok(token)
}
Err(error) => {
if matches!(
error.downcast_ref::<EmaAuthFailure>(),
Some(EmaAuthFailure::InvalidGrant {
grant_source: EmaInvalidGrantSource::EnterpriseIdentity,
})
) && self
.credentials
.invalidate_ema_credentials_if_current()
.await
.is_err()
{
tracing::warn!("failed to clear rejected enterprise credentials");
}
Err(match error.downcast::<EmaAuthFailure>() {
Ok(failure) => failure.into(),
Err(_) => anyhow!("enterprise MCP token exchange failed"),
})
}
}
}
}
codex-rs/codex-mcp/src/event_stream.rs
+1
View File
@@ -79,6 +79,7 @@ impl McpEventStreamOpener {
let startup_timeout = connection.server.config().startup_timeout_sec
.unwrap_or(DEFAULT_STARTUP_TIMEOUT);
let client = Arc::new(tokio::time::timeout(startup_timeout, make_rmcp_client(
/*ema_auth*/ None,
CODEX_APPS_MCP_SERVER_NAME,
connection.server.clone(),
connection.store_mode,
codex-rs/codex-mcp/src/mcp/auth.rs
+24 -2
View File
@@ -10,6 +10,7 @@ use codex_config::types::OAuthCredentialsStoreMode;
use codex_exec_server::HttpClient;
use codex_login::CodexAuth;
use codex_rmcp_client::McpAuthState;
use codex_rmcp_client::McpLoginRequirement;
use codex_rmcp_client::McpOAuthCallbackMode;
use codex_rmcp_client::OAuthDiscoveryTimeout;
use codex_rmcp_client::OAuthProviderError;
@@ -255,8 +256,29 @@ async fn compute_auth_status(
return Ok(McpAuthState::Unsupported);
}
if matches!(config.auth, McpServerAuth::EmaAuth) {
// EMA connections are not enabled until the runtime stage of the stack.
return Ok(McpAuthState::Unsupported);
let Some(registration) = config.ema_registration() else {
return Ok(McpAuthState::Unsupported);
};
config
.validate_ema_auth_transport()
.map_err(anyhow::Error::msg)?;
let Some(scope) = crate::ema_auth_scope(auth) else {
return Ok(McpAuthState::LoggedOut(McpLoginRequirement::Login));
};
return Ok(
if crate::ema_auth::stored_ema_credentials(
registration.idp(),
&scope,
keyring_backend_kind,
)
.await?
.is_some()
{
McpAuthState::OAuth
} else {
McpAuthState::LoggedOut(McpLoginRequirement::Login)
},
);
}
let has_runtime_auth = matches!(config.auth, McpServerAuth::ChatGpt)
&& auth.is_some_and(CodexAuth::uses_codex_backend)
codex-rs/codex-mcp/src/mcp/mod.rs
+2
View File
@@ -463,6 +463,7 @@ pub async fn read_mcp_resource(
let cancel_token = CancellationToken::new();
let runtime_config = config.for_threadless_operations(&mcp_servers);
let manager = McpConnectionSet::new(
/*ema_auth*/ None,
/*previous*/ None,
McpPublicationGate::already_published(),
McpRuntimeInput {
@@ -545,6 +546,7 @@ pub async fn collect_mcp_server_status_snapshot_with_detail(
let cancel_token = CancellationToken::new();
let runtime_config = config.for_threadless_operations(&mcp_servers);
let mcp_connection_manager = McpConnectionSet::new(
/*ema_auth*/ None,
/*previous*/ None,
McpPublicationGate::already_published(),
McpRuntimeInput {
codex-rs/codex-mcp/src/rmcp_client.rs
+20 -3
View File
@@ -297,6 +297,7 @@ fn codex_apps_reconnect_backoff(consecutive_failures: u32) -> Duration {
#[derive(Clone)]
struct ManagedClientStartup {
ema_auth: Option<crate::ema_auth::EmaAuthContext>,
originator: AuthStorageOriginator,
server_name: String,
server: EffectiveMcpServer,
@@ -328,6 +329,7 @@ impl ManagedClientStartup {
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner) = None;
let Self {
ema_auth,
originator,
server_name,
server,
@@ -369,6 +371,7 @@ impl ManagedClientStartup {
let client = match tokio::time::timeout(
startup_timeout,
make_rmcp_client(
ema_auth,
&server_name,
server.clone(),
store_mode,
@@ -462,6 +465,7 @@ impl AsyncManagedClient {
#[instrument(level = "trace", skip_all, fields(server_name = %server_name))]
#[allow(clippy::too_many_arguments)]
pub(crate) fn new(
ema_auth: Option<crate::ema_auth::EmaAuthContext>,
server_name: String,
startup_submit_id: String,
server: EffectiveMcpServer,
@@ -495,6 +499,7 @@ impl AsyncManagedClient {
let startup_complete = Arc::new(AtomicBool::new(false));
let server_capabilities = Arc::new(StdMutex::new(None));
let startup = Arc::new(ManagedClientStartup {
ema_auth,
originator: AuthStorageOriginator::current(),
server_name,
server,
@@ -1236,6 +1241,7 @@ struct StartServerTaskParams {
#[allow(clippy::too_many_arguments)]
#[instrument(level = "trace", skip_all, fields(server_name = %server_name))]
pub(crate) async fn make_rmcp_client(
ema_auth: Option<crate::ema_auth::EmaAuthContext>,
server_name: &str,
server: EffectiveMcpServer,
store_mode: OAuthCredentialsStoreMode,
@@ -1265,9 +1271,20 @@ pub(crate) async fn make_rmcp_client(
}
}
if matches!(config.auth, McpServerAuth::EmaAuth) {
return Err(StartupOutcomeError::from(anyhow!(
"EMA MCP connections are not enabled in this version"
)));
let auth = ema_auth.ok_or_else(|| {
StartupOutcomeError::from(anyhow!(
"enterprise MCP requires a thread with an account owner"
))
})??;
return auth
.connect(
server_name,
&server,
runtime_context.local_http_client(),
protocol_mode,
)
.await
.map_err(Into::into);
}
if matches!(config.auth, McpServerAuth::ChatGpt)
&& !config.is_local_environment()
codex-rs/codex-mcp/src/runtime.rs
+21
View File
@@ -12,6 +12,7 @@ use std::collections::HashSet;
use std::path::PathBuf;
use std::sync::Arc;
use std::sync::Mutex;
use std::sync::OnceLock;
use std::sync::Weak;
use std::sync::atomic::AtomicBool;
use std::sync::atomic::AtomicU64;
@@ -112,6 +113,7 @@ pub struct McpRuntime {
resource_cache_generation: AtomicU64,
elicitation_router: ElicitationRequestRouter,
resource_origins: Mutex<ResourceOrigins>,
ema_auth: OnceLock<crate::ema_auth::EmaAuthContext>,
}
struct EventStreamCancellation {
@@ -251,6 +253,7 @@ impl McpRuntime {
resource_cache_generation: AtomicU64::new(0),
elicitation_router: ElicitationRequestRouter::default(),
resource_origins: Mutex::default(),
ema_auth: OnceLock::new(),
}
}
@@ -335,6 +338,23 @@ impl McpRuntime {
}
async fn publish(&self, input: McpRuntimeInput, previous: Option<&McpConnectionSet>) {
let ema_auth = (input.config.xaa_enabled
&& input.mcp_servers.values().any(|server| {
server.config().enabled
&& matches!(server.config().auth, codex_config::McpServerAuth::EmaAuth)
}))
.then(|| {
self.ema_auth
.get_or_init(|| {
crate::ema_auth::EmaAuth::new(
Arc::clone(&input.config),
input.auth.as_ref(),
input.auth_manager.as_ref(),
)
.map_err(Into::into)
})
.clone()
});
let (publish, publication_gate) = McpPublicationGate::pending();
let config = Arc::clone(&input.config);
let auth = input.auth.clone();
@@ -351,6 +371,7 @@ impl McpRuntime {
let ready_environments = input.runtime_context.ready_environments.clone();
let connections = Arc::new(
McpConnectionSet::new(
ema_auth,
previous,
publication_gate,
input,