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

Preserve enterprise MCP authorization across runtime refreshes (#55696)

Browse Source 
## Why

Configuration reloads must cancel stale enterprise MCP startups while preserving access for running threads that still pass validation. Replacing a runtime during `thread/revert` must also leave the replacement able to use enterprise tools.

## What changed

- Separate pending startup cancellation from published thread lifetimes, with independent startup state for forks and replacement runtimes.
- Retain eligible enterprise MCP connections across refreshes and explicit reconnects while applying updated tool restrictions and approvals. Do not create replacement enterprise connections after initial startup.
- Revoke enterprise authorization when a thread is removed, drain requests when access is disabled, and revoke a connection's authority when token renewal requires authentication.

## Testing

Extend regression coverage to verify enterprise tool calls survive `thread/revert`, inherited fork configurations are cancelled after a reload, and enterprise calls fail after logout while ordinary MCP calls remain available.

GitOrigin-RevId: 2d634be3011854b57f73c741bd1ec6b5aed94ee5
This commit is contained in:
jif
2026-10-20 23:03:56 +00:00
committed by copyberry
parent 6fb1c3903d
commit feadd9441e
20 changed files with 468 additions and 85 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/app-server/src/config_manager.rs
+26 -9
View File
@@ -57,7 +57,13 @@ pub(crate) struct ConfigManager {
network_policy_reload: Arc<tokio::sync::Semaphore>,
network_policy_snapshot: Arc<RwLock<Option<Arc<ApplicationPolicySnapshot>>>>,
/// Captured before config loading so a reload also revokes pending startups.
mcp_ema_cancellation: Arc<Mutex<CancellationToken>>,
mcp_ema_cancellation: Arc<Mutex<EmaCancellationRoots>>,
}
#[derive(Default)]
struct EmaCancellationRoots {
startup: CancellationToken,
lifetime: CancellationToken,
}
/// Startup configuration, including whether loading had to fall back to defaults.
@@ -117,22 +123,33 @@ impl ConfigManager {
self.codex_home.as_path()
}
/// Retires EMA access for every config already loaded or being loaded.
/// New loads resolve their own settings under the next cancellation scope.
/// Explicit account, policy-owner, and enterprise-login changes revoke cached access.
pub(crate) fn revoke_mcp_ema(&self) {
let mut current = self
.mcp_ema_cancellation
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
current.cancel();
*current = CancellationToken::new();
current.startup.cancel();
current.lifetime.cancel();
*current = EmaCancellationRoots::default();
}
fn capture_mcp_ema_cancellation(&self) -> CancellationToken {
self.mcp_ema_cancellation
/// Published threads are validated individually by the existing refresh loop.
pub(crate) fn revoke_mcp_ema_startups(&self) {
let mut current = self
.mcp_ema_cancellation
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.child_token()
.unwrap_or_else(std::sync::PoisonError::into_inner);
current.startup.cancel();
current.startup = CancellationToken::new();
}
fn capture_mcp_ema_cancellation(&self) -> codex_config::McpEmaCancellation {
let current = self
.mcp_ema_cancellation
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
codex_config::McpEmaCancellation::pending(current.startup.clone(), current.lifetime.clone())
}
pub(crate) fn with_embedded_network_policy(
codex-rs/app-server/src/mcp_refresh.rs
+20 -1
View File
@@ -46,7 +46,7 @@ async fn reload_mcp_config_with_policy(
config_manager: &ConfigManager,
policy: ReloadPolicy,
) -> io::Result<()> {
config_manager.revoke_mcp_ema();
config_manager.revoke_mcp_ema_startups();
let mut rejected = false;
let mut load_error = None;
let mut completed: Vec<Arc<CodexThread>> = Vec::new();
@@ -416,6 +416,14 @@ enabled = false
.mcp_ema_cancellation()
.is_cancelled()
);
// Prediction forks inherit a published config and may outlive their parent.
let parent = thread_manager
.start_thread(codex_core::StartThreadOptions::new(pending_startup.clone()))
.await?;
let (inherited_fork, _) = thread_manager
.fork_options_from_parent(parent.thread_id)
.await?;
thread_manager.remove_thread(&parent.thread_id).await;
std::fs::write(
temp_dir.path().join(codex_config::CONFIG_TOML_FILE),
r#"
@@ -427,6 +435,17 @@ enabled = false
reload_mcp_config(&thread_manager, &config_manager).await?;
let fork = thread_manager.start_thread(inherited_fork).await?;
assert!(
fork.thread
.config()
.await
.config_layer_stack
.mcp_ema_cancellation()
.is_cancelled()
);
fork.thread.shutdown_and_wait().await?;
parent.thread.shutdown_and_wait().await?;
assert!(
pending_startup
.config_layer_stack
codex-rs/app-server/src/request_processors/config_processor.rs
+1 -1
View File
@@ -374,7 +374,7 @@ pub(super) async fn reload_user_config(
}
async fn reload_user_config_inner(config_manager: &ConfigManager, thread_manager: &ThreadManager) {
config_manager.revoke_mcp_ema();
config_manager.revoke_mcp_ema_startups();
if let Err(err) = config_manager.load_config_layers(/*cwd*/ None).await {
tracing::warn!("failed to rebuild user config for runtime refresh: {err}");
for thread_id in thread_manager.list_thread_ids().await {
codex-rs/app-server/src/request_processors/thread_processor.rs
+10 -1
View File
@@ -2211,8 +2211,17 @@ impl ThreadRequestProcessor {
"thread/revert only supports paginated threads",
));
}
let mut config = thread.config().await.as_ref().clone();
// Give the replacement its own pending lifetime before removal revokes the old runtime.
let ema_cancellation = config
.config_layer_stack
.mcp_ema_cancellation()
.for_startup();
config.config_layer_stack = config
.config_layer_stack
.with_mcp_ema_cancellation(ema_cancellation);
let runtime_snapshot = ThreadRevertRuntimeSnapshot {
config: thread.config().await.as_ref().clone(),
config,
settings: thread.restorable_thread_settings().await,
client_mcp_extensions: thread.client_mcp_extensions(),
};
codex-rs/app-server/tests/suite/v2/account_enterprise_tests.rs
+53 -6
View File
@@ -31,6 +31,7 @@ use codex_feedback::CodexFeedback;
use codex_login::CODEX_ACCESS_TOKEN_ENV_VAR;
use codex_protocol::protocol::SessionSource;
use codex_rmcp_client::stored_oauth_credentials;
use core_test_support::responses;
use core_test_support::skip_if_no_network;
use keyring::credential::Credential;
use keyring::credential::CredentialApi;
@@ -411,6 +412,7 @@ oauth = {{ client_id = "mcp-client", authorization_server_issuer = "{origin}/as"
.await?;
// Development builds otherwise force file storage, bypassing the mock keyring.
config.cli_auth_credentials_store_mode = AuthCredentialsStoreMode::Auto;
let state_db = codex_rollout::state_db::try_init(&config).await?;
let config = Arc::new(config);
let mut client = in_process::start(InProcessStartArgs {
arg0_paths: Arg0DispatchPaths::default(),
@@ -422,7 +424,7 @@ oauth = {{ client_id = "mcp-client", authorization_server_issuer = "{origin}/as"
thread_config_loader: Arc::new(NoopThreadConfigLoader),
feedback: CodexFeedback::new(),
log_db: None,
state_db: None,
state_db: Some(state_db),
environment_manager: Arc::new(EnvironmentManager::default_for_tests()),
config_warnings: Vec::new(),
embedded_network_policy: Default::default(),
@@ -578,7 +580,8 @@ oauth = {{ client_id = "mcp-client", authorization_server_issuer = "{origin}/as"
}
// A fresh thread uses the saved enterprise grant and renews it for MCP requests.
let connected = enterprise_rpc(&client, "thread/start", json!({})).await?;
let connected =
enterprise_rpc(&client, "thread/start", json!({"historyMode":"paginated"})).await?;
let connected_id = connected["thread"]["id"]
.as_str()
.context("connected thread ID")?;
@@ -597,6 +600,44 @@ oauth = {{ client_id = "mcp-client", authorization_server_issuer = "{origin}/as"
enterprise_rpc(&client, "mcpServer/tool/call", call("enterprise")).await?;
let ordinary_result = enterprise_rpc(&client, "mcpServer/tool/call", call("ordinary")).await?;
assert_eq!(enterprise_result, ordinary_result);
// Revert replaces the runtime without revoking its replacement's EMA access.
let _response =
responses::mount_sse_once(&server, responses::sse_completed("revert-turn")).await;
let turn = enterprise_rpc(
&client,
"turn/start",
json!({"threadId":connected_id,"input":[{"type":"text","text":"Before revert"}]}),
)
.await?;
timeout(DEFAULT_READ_TIMEOUT, async {
loop {
if let InProcessServerEvent::ServerNotification(notification) =
client.next_event().await.context("turn notification")?
&& let ServerNotification::TurnCompleted(completed) = *notification
&& completed.thread_id == connected_id
{
assert_eq!(
completed.turn.status,
codex_app_server_protocol::TurnStatus::Completed
);
return anyhow::Ok(());
}
}
})
.await??;
enterprise_rpc(
&client,
"thread/revert",
json!({"threadId":connected_id,"beforeTurnId":turn["turn"]["id"]}),
)
.await?;
assert_eq!(
enterprise_rpc(&client, "mcpServer/tool/call", call("enterprise"))
.await
.context("enterprise tools must survive revert")?,
enterprise_result
);
let bulk_status = enterprise_rpc(
&client,
"mcpServerStatus/list",
@@ -707,12 +748,13 @@ oauth = {{ client_id = "mcp-client", authorization_server_issuer = "{origin}/as"
[
"Bearer enterprise-resource-1",
"Bearer enterprise-resource-2",
"Bearer ordinary-access"
"Bearer ordinary-access",
"Bearer enterprise-resource-3"
]
);
assert_eq!(
resource_exchanges.load(std::sync::atomic::Ordering::SeqCst),
2
3
);
// Replacing the enterprise grant fences old threads, including cached bearers.
@@ -766,7 +808,7 @@ oauth = {{ client_id = "mcp-client", authorization_server_issuer = "{origin}/as"
);
assert_eq!(
resource_exchanges.load(std::sync::atomic::Ordering::SeqCst),
2
3
);
std::fs::write(&auth_path, saved_auth)?;
@@ -863,7 +905,7 @@ oauth = {{ client_id = "mcp-client", authorization_server_issuer = "{origin}/as"
);
assert_eq!(
resource_exchanges.load(std::sync::atomic::Ordering::SeqCst),
2
3
);
assert!(
enterprise_rpc(
@@ -878,6 +920,11 @@ oauth = {{ client_id = "mcp-client", authorization_server_issuer = "{origin}/as"
.is_err()
);
assert!(!home.path().join("auth.json").exists());
assert!(
enterprise_rpc(&client, "mcpServer/tool/call", call("enterprise"))
.await
.is_err()
);
assert_eq!(
enterprise_rpc(&client, "mcpServer/tool/call", call("ordinary")).await?,
ordinary_result
codex-rs/codex-mcp/src/connection_manager.rs
+83 -9
View File
@@ -283,6 +283,7 @@ impl McpConnectionSet {
/// notifications are then skipped and interactive elicitations are declined.
pub async fn new(
ema_auth: Option<crate::ema_auth::EmaAuthContext>,
existing_ema: Option<&Self>,
previous: Option<&Self>,
publication_gate: McpPublicationGate,
input: McpRuntimeInput,
@@ -337,15 +338,27 @@ impl McpConnectionSet {
// Explicit reconnects have no previous set and must replace their clients eagerly.
let allow_deferred_startup =
startup_policy == McpStartupPolicy::LazyWhenCached && previous.is_some();
let reusable_previous = previous.filter(|previous| {
!previous.servers.is_empty()
&& previous.elicitation_requests.update(
Arc::clone(&config),
elicitation_reviewer.clone(),
elicitation_lifecycle.clone(),
)
});
let elicitation_requests = if let Some(previous) = reusable_previous {
// Retained EMA clients keep callbacks attached to this manager even when
// an explicit reconnect replaces the ordinary clients.
let elicitation_previous = previous
.or_else(|| {
existing_ema.filter(|previous| {
previous
.servers
.values()
.any(|view| view.connection.client.ema_authority.is_some())
})
})
.filter(|previous| {
!previous.servers.is_empty()
&& previous.elicitation_requests.update(
Arc::clone(&config),
elicitation_reviewer.clone(),
elicitation_lifecycle.clone(),
)
});
let reusable_previous = previous.filter(|_| elicitation_previous.is_some());
let elicitation_requests = if let Some(previous) = elicitation_previous {
previous.elicitation_requests.clone()
} else {
ElicitationRequestManager::new(
@@ -557,6 +570,67 @@ impl McpConnectionSet {
Some(_) => None,
},
};
if matches!(configured_config.auth, McpServerAuth::EmaAuth) && existing_ema.is_some() {
if let Some(previous_view) = existing_ema
.filter(|_| elicitation_previous.is_some())
.and_then(|previous| previous.servers.get(&server_name))
{
let connection = &previous_view.connection;
let authority_is_current = || {
connection
.client
.ema_authority
.as_ref()
.is_some_and(|authority| authority.check_new_request().is_ok())
};
let same_identity = connection.identity.as_ref() == Some(&connection_identity)
&& previous_view.catalog_item_limit == catalog_item_limit
&& previous_view.protocol_mode == protocol_mode;
let pending = same_identity
&& !connection.client.startup_complete.load(Ordering::Acquire)
&& connection.startup_timeout == startup_timeout
&& !connection.startup_is_dormant()
&& !connection.client.cancel_token.is_cancelled();
let ready = same_identity
&& !pending
&& connection
.reusable_client(&connection_identity)
.await
.is_some_and(|client| {
expected_protocol_mode == Some(client.client.protocol_mode())
});
if authority_is_current() && (pending || ready) {
servers.insert(
server_name.clone(),
McpServerView {
connection: Arc::clone(connection),
protocol_mode,
startup_readiness: configured_config.startup_readiness,
refresh_tool_list: config.mcp_tool_list_changed,
refresh_timeout: startup_timeout,
metadata,
tool_filter: configured_tool_filter,
tool_timeout: configured_tool_timeout,
catalog_item_limit,
},
);
if pending {
let client = connection.client.clone();
let publication_gate = publication_gate.clone();
join_set.spawn(async move {
if !publication_gate.wait().await {
return (server_name, Err(StartupOutcomeError::Cancelled));
}
(server_name, client.client().await)
});
} else {
reused_ready.push(server_name);
}
}
}
// Post-start publications may only retain the original EMA connection.
continue;
}
if let Some(previous_view) =
reusable_previous.and_then(|previous| previous.servers.get(&server_name))
&& previous_view.refresh_tool_list == config.mcp_tool_list_changed
codex-rs/codex-mcp/src/connection_manager_tests.rs
+25
View File
@@ -579,6 +579,7 @@ async fn prepared_call_timeout_includes_trusted_access_lookup() {
pub(crate) async fn create_ready_async_managed_client(tools: Vec<ToolInfo>) -> AsyncManagedClient {
AsyncManagedClient {
ema_authority: None,
client: futures::future::ready::<Result<ManagedClient, StartupOutcomeError>>(Ok(
create_test_managed_client(tools).await,
))
@@ -749,6 +750,7 @@ fn create_gated_async_managed_client(
(
AsyncManagedClient {
ema_authority: None,
client,
is_codex_apps_mcp_server: false,
server_capabilities: Arc::new(std::sync::Mutex::new(None)),
@@ -817,6 +819,7 @@ pub(crate) async fn create_test_manager_with_ready_apps_client(
manager.insert_test_client(
CODEX_APPS_MCP_SERVER_NAME.to_string(),
AsyncManagedClient {
ema_authority: None,
client: futures::future::ready::<Result<ManagedClient, StartupOutcomeError>>(Ok(
managed_client,
))
@@ -873,6 +876,7 @@ fn create_test_manager_with_failed_apps_startup(
manager.insert_test_client(
CODEX_APPS_MCP_SERVER_NAME.to_string(),
AsyncManagedClient {
ema_authority: None,
client,
is_codex_apps_mcp_server: true,
server_capabilities: Arc::new(std::sync::Mutex::new(None)),
@@ -2060,6 +2064,7 @@ async fn read_only_apps_discovery_never_uses_a_shared_writable_catalog() -> anyh
cancellation.cancel();
let manager = McpConnectionSet::new(
/*ema_auth*/ None,
/*existing_ema*/ None,
/*previous*/ None,
McpPublicationGate::already_published(),
McpRuntimeInput {
@@ -2134,6 +2139,7 @@ async fn hosted_apps_protocol_mode_is_independent_of_generic_mode() -> anyhow::R
startup_cancellation_token.cancel();
let manager = McpConnectionSet::new(
/*ema_auth*/ None,
/*existing_ema*/ None,
/*previous*/ None,
McpPublicationGate::already_published(),
McpRuntimeInput {
@@ -2241,6 +2247,7 @@ async fn codex_apps_extension_does_not_share_host_owned_tools_cache() -> anyhow:
startup_cancellation_token.cancel();
let manager = McpConnectionSet::new(
/*ema_auth*/ None,
/*existing_ema*/ None,
/*previous*/ None,
McpPublicationGate::already_published(),
McpRuntimeInput {
@@ -2318,6 +2325,7 @@ async fn list_all_tools_uses_shared_codex_apps_cache_while_client_is_pending() {
manager.insert_test_client(
CODEX_APPS_MCP_SERVER_NAME.to_string(),
AsyncManagedClient {
ema_authority: None,
client: pending_client,
is_codex_apps_mcp_server: true,
server_capabilities: Arc::new(std::sync::Mutex::new(None)),
@@ -2390,6 +2398,7 @@ async fn capture_binding_preserves_ready_client_catalog_and_shares_prepared_meta
manager.insert_test_client(
CODEX_APPS_MCP_SERVER_NAME.to_string(),
AsyncManagedClient {
ema_authority: None,
client: futures::future::ready(Ok(ready_client)).boxed().shared(),
is_codex_apps_mcp_server: true,
server_capabilities: Arc::new(std::sync::Mutex::new(None)),
@@ -2786,6 +2795,7 @@ async fn list_available_server_infos_uses_cache_while_client_is_pending() {
manager.insert_test_client(
CODEX_APPS_MCP_SERVER_NAME.to_string(),
AsyncManagedClient {
ema_authority: None,
client: pending_client,
is_codex_apps_mcp_server: true,
server_capabilities: Arc::new(std::sync::Mutex::new(None)),
@@ -2883,6 +2893,7 @@ async fn capture_binding_exposes_cached_tools_before_startup() {
manager.insert_test_client(
CODEX_APPS_MCP_SERVER_NAME.to_string(),
AsyncManagedClient {
ema_authority: None,
client: pending_client,
is_codex_apps_mcp_server: true,
server_capabilities: Arc::new(std::sync::Mutex::new(None)),
@@ -2983,6 +2994,7 @@ async fn capture_binding_skips_pending_optional_servers_after_configured_shared_
manager.insert_test_client(
server_name.to_string(),
AsyncManagedClient {
ema_authority: None,
client: futures::future::pending::<Result<ManagedClient, StartupOutcomeError>>()
.boxed()
.shared(),
@@ -3235,6 +3247,7 @@ async fn capture_binding_shares_optional_startup_grace_across_connection_sets()
manager.insert_test_client(
"pending",
AsyncManagedClient {
ema_authority: None,
client: futures::future::pending::<Result<ManagedClient, StartupOutcomeError>>()
.boxed()
.shared(),
@@ -3916,6 +3929,7 @@ async fn list_all_tools_blocks_while_client_is_pending_without_cached_tools() {
manager.insert_test_client(
CODEX_APPS_MCP_SERVER_NAME.to_string(),
AsyncManagedClient {
ema_authority: None,
client: pending_client,
is_codex_apps_mcp_server: true,
server_capabilities: Arc::new(std::sync::Mutex::new(None)),
@@ -3974,6 +3988,7 @@ async fn shutdown_cancels_pending_tool_listing() {
manager.insert_test_client(
CODEX_APPS_MCP_SERVER_NAME.to_string(),
AsyncManagedClient {
ema_authority: None,
client: pending_client,
is_codex_apps_mcp_server: true,
server_capabilities: Arc::new(std::sync::Mutex::new(None)),
@@ -4021,6 +4036,7 @@ async fn shutdown_continues_after_caller_is_aborted() {
manager.insert_test_client(
CODEX_APPS_MCP_SERVER_NAME.to_string(),
AsyncManagedClient {
ema_authority: None,
client: blocking_client,
is_codex_apps_mcp_server: true,
server_capabilities: Arc::new(std::sync::Mutex::new(None)),
@@ -4074,6 +4090,7 @@ async fn list_all_tools_does_not_block_when_shared_codex_apps_cache_is_empty() {
manager.insert_test_client(
CODEX_APPS_MCP_SERVER_NAME.to_string(),
AsyncManagedClient {
ema_authority: None,
client: pending_client,
is_codex_apps_mcp_server: true,
server_capabilities: Arc::new(std::sync::Mutex::new(None)),
@@ -4127,6 +4144,7 @@ async fn list_all_tools_uses_shared_codex_apps_cache_when_client_startup_fails()
manager.insert_test_client(
CODEX_APPS_MCP_SERVER_NAME.to_string(),
AsyncManagedClient {
ema_authority: None,
client: failed_client,
is_codex_apps_mcp_server: true,
server_capabilities: Arc::new(std::sync::Mutex::new(None)),
@@ -4814,6 +4832,7 @@ async fn executor_owned_chatgpt_mcp_accepts_only_safe_explicit_authorization() -
let manager = McpConnectionSet::new(
/*ema_auth*/ None,
/*existing_ema*/ None,
/*previous*/ None,
McpPublicationGate::already_published(),
McpRuntimeInput {
@@ -4933,6 +4952,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,
/*existing_ema*/ None,
/*previous*/ None,
McpPublicationGate::already_published(),
McpRuntimeInput {
@@ -5382,6 +5402,7 @@ async fn reconcile_reusable_server_with_mcp_config(
let (tx_event, _rx_event) = async_channel::unbounded();
McpConnectionSet::new(
/*ema_auth*/ None,
/*existing_ema*/ None,
Some(previous),
McpPublicationGate::already_published(),
McpRuntimeInput {
@@ -5778,6 +5799,7 @@ async fn reconciliation_reuses_connection_without_relisting_regular_tools() -> a
connection: Arc::new(McpServerConnection {
identity: Some(reusable_server_identity("docs", &config, &runtime_context)),
client: AsyncManagedClient {
ema_authority: None,
client: futures::future::ready(Ok(managed_client)).boxed().shared(),
is_codex_apps_mcp_server: false,
server_capabilities: Arc::new(std::sync::Mutex::new(None)),
@@ -6339,6 +6361,7 @@ async fn reconciliation_replaces_connection_when_protocol_mode_changes() {
let reconciled = McpConnectionSet::new(
/*ema_auth*/ None,
/*existing_ema*/ None,
Some(&previous),
McpPublicationGate::already_published(),
McpRuntimeInput {
@@ -6398,6 +6421,7 @@ async fn reconciliation_reuses_legacy_stdio_server_when_modern_protocol_is_enabl
let reconciled = McpConnectionSet::new(
/*ema_auth*/ None,
/*existing_ema*/ None,
Some(&previous),
McpPublicationGate::already_published(),
McpRuntimeInput {
@@ -6597,6 +6621,7 @@ async fn reconciliation_replaces_closed_connections() -> anyhow::Result<()> {
view.connection = Arc::new(McpServerConnection {
identity: Some(reusable_server_identity("docs", &config, &runtime_context)),
client: AsyncManagedClient {
ema_authority: None,
client: futures::future::ready(Ok(connected_client))
.boxed()
.shared(),
codex-rs/codex-mcp/src/ema_auth.rs
+29 -23
View File
@@ -1,5 +1,6 @@
//! Dedicated enterprise grants and one thread's EMA authorization lifetime.
//! Account credentials identify the owner; only the separate keyring grant is exchanged.
//! One thread's EMA authorization lifetime and guarded HTTP dispatch.
//! Account and delivered-policy changes cancel the same owner and its admitted requests.
//! Credential selection and recovery belong to each connection's token source.
use std::sync::Arc;
@@ -7,7 +8,6 @@ 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;
@@ -51,6 +51,12 @@ use crate::rmcp_client::StartupOutcomeError;
pub(crate) type EmaAuthContext = Result<Arc<EmaAuth>, StartupOutcomeError>;
#[derive(Clone)]
pub(crate) struct EmaConnectionAuth {
pub(crate) owner: Arc<EmaAuth>,
pub(crate) authority: EmaRequestAuthority,
}
pub fn ema_auth_scope(auth: Option<&CodexAuth>) -> Option<McpEmaAuthScope> {
let auth = auth?;
McpEmaAuthScope::new(auth.get_chatgpt_user_id()?, auth.get_account_id()?)
@@ -58,11 +64,11 @@ pub fn ema_auth_scope(auth: Option<&CodexAuth>) -> Option<McpEmaAuthScope> {
/// Retained by the thread across MCP refreshes; a revoked owner cannot be replaced.
pub(crate) struct EmaAuth {
scope: McpEmaAuthScope,
config: Arc<McpConfig>,
manager: Arc<AuthManager>,
account: CodexAuth,
authority: EmaRequestAuthority,
pub(crate) scope: McpEmaAuthScope,
pub(crate) config: Arc<McpConfig>,
pub(crate) manager: Arc<AuthManager>,
pub(crate) account: CodexAuth,
pub(crate) authority: EmaRequestAuthority,
_revocation: AbortOnDropHandle<()>,
}
@@ -84,24 +90,23 @@ impl EmaAuth {
}
let cancellation = config.mcp_ema_cancellation.clone();
let policy = config.config_layer_stack.cloud_config_binding().cloned();
let authority = EmaRequestAuthority::new(cancellation.clone(), {
let authority = EmaRequestAuthority::new(cancellation.token(), {
let cancellation = cancellation.clone();
let changes = changes.clone();
let policy = policy.clone();
move || {
changes.borrow().login_generation == login_generation
&& policy.as_ref().is_none_or(|binding| {
binding.read().status == CloudConfigBundleBindingStatus::Current
})
!cancellation.is_cancelled()
&& changes.borrow().login_generation == login_generation
}
});
authority.check()?;
})
.with_policy(policy.clone());
authority.check_new_request()?;
let revocation = tokio::spawn(async move {
tokio::select! {
biased;
() = cancellation.cancelled() => {}
() = async {
match policy {
Some(binding) => binding.invalidated().await,
Some(policy) => policy.invalidated().await,
None => std::future::pending().await,
}
} => {}
@@ -142,12 +147,13 @@ impl EmaAuth {
pub(crate) async fn connect(
self: &Arc<Self>,
authority: EmaRequestAuthority,
server_name: &str,
server: &EffectiveMcpServer,
http_client: Arc<dyn HttpClient>,
protocol_mode: McpProtocolMode,
) -> Result<RmcpClient> {
self.check_account().await?;
self.authority.check()?;
let config = server.config();
config
.validate_ema_auth_transport()
@@ -176,8 +182,8 @@ impl EmaAuth {
.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(),
authority: authority.clone(),
cancellation: self.config.mcp_ema_cancellation.token(),
});
let url = registration.server_url().to_owned();
let source = EnterpriseTokenSource {
@@ -192,7 +198,7 @@ impl EmaAuth {
http_client,
source,
token,
self.authority.clone(),
authority,
));
RmcpClient::new_streamable_http_client_with_protocol_mode_and_redirect_mode(
server_name,
@@ -258,7 +264,7 @@ impl HttpClient for EnterpriseExchangeHttpClient {
params: HttpRequestParams,
) -> BoxFuture<'_, Result<HttpRequestResponse, ExecServerError>> {
Box::pin(async move {
self.authority.check()?;
self.authority.check_new_request()?;
tokio::select! {
biased;
() = self.cancellation.cancelled() => Err(ExecServerError::HttpRequest("enterprise MCP authorization was revoked".into())),
@@ -275,7 +281,7 @@ impl HttpClient for EnterpriseExchangeHttpClient {
params: HttpRequestParams,
) -> BoxFuture<'_, Result<(HttpRequestResponse, HttpResponseBodyStream), ExecServerError>> {
Box::pin(async move {
self.authority.check()?;
self.authority.check_new_request()?;
tokio::select! {
biased;
() = self.cancellation.cancelled() => Err(ExecServerError::HttpRequest("enterprise MCP authorization was revoked".into())),
codex-rs/codex-mcp/src/mcp/mod.rs
+3 -1
View File
@@ -133,7 +133,7 @@ pub struct McpConfig {
/// Trusted enterprise IdP inherited after normal catalog and policy resolution.
pub mcp_enterprise_managed_auth: Option<McpEnterpriseManagedAuthConfig>,
/// Retires only EMA connections and pending exchanges for this config snapshot.
pub mcp_ema_cancellation: tokio_util::sync::CancellationToken,
pub mcp_ema_cancellation: codex_config::McpEmaCancellation,
pub xaa_enabled: bool,
/// Preferred credential store for MCP OAuth tokens.
pub mcp_oauth_credentials_store_mode: OAuthCredentialsStoreMode,
@@ -464,6 +464,7 @@ pub async fn read_mcp_resource(
let runtime_config = config.for_threadless_operations(&mcp_servers);
let manager = McpConnectionSet::new(
/*ema_auth*/ None,
/*existing_ema*/ None,
/*previous*/ None,
McpPublicationGate::already_published(),
McpRuntimeInput {
@@ -547,6 +548,7 @@ pub async fn collect_mcp_server_status_snapshot_with_detail(
let runtime_config = config.for_threadless_operations(&mcp_servers);
let mcp_connection_manager = McpConnectionSet::new(
/*ema_auth*/ None,
/*existing_ema*/ None,
/*previous*/ None,
McpPublicationGate::already_published(),
McpRuntimeInput {
codex-rs/codex-mcp/src/mcp/mod_tests.rs
+1 -1
View File
@@ -90,7 +90,7 @@ pub(crate) fn test_mcp_config(codex_home: PathBuf) -> McpConfig {
requires_read_only_mcp_tools: false,
codex_home,
mcp_enterprise_managed_auth: None,
mcp_ema_cancellation: tokio_util::sync::CancellationToken::new(),
mcp_ema_cancellation: codex_config::McpEmaCancellation::default(),
xaa_enabled: false,
mcp_oauth_credentials_store_mode: OAuthCredentialsStoreMode::default(),
oauth_refresh_mode: McpOAuthRefreshMode::Legacy,
codex-rs/codex-mcp/src/rmcp_client.rs
+18 -2
View File
@@ -297,7 +297,7 @@ fn codex_apps_reconnect_backoff(consecutive_failures: u32) -> Duration {
#[derive(Clone)]
struct ManagedClientStartup {
ema_auth: Option<crate::ema_auth::EmaAuthContext>,
ema_auth: Option<Result<crate::ema_auth::EmaConnectionAuth, StartupOutcomeError>>,
originator: AuthStorageOriginator,
server_name: String,
server: EffectiveMcpServer,
@@ -447,6 +447,7 @@ impl ManagedClientStartup {
#[derive(Clone)]
pub(crate) struct AsyncManagedClient {
pub(crate) ema_authority: Option<codex_rmcp_client::EmaRequestAuthority>,
pub(crate) client: ManagedClientFuture,
pub(crate) is_codex_apps_mcp_server: bool,
pub(crate) cached_server_info: Option<McpServerInfo>,
@@ -498,6 +499,18 @@ impl AsyncManagedClient {
};
let startup_complete = Arc::new(AtomicBool::new(false));
let server_capabilities = Arc::new(StdMutex::new(None));
let ema_auth = ema_auth
.filter(|_| matches!(server.config().auth, McpServerAuth::EmaAuth))
.map(|auth| {
auth.map(|owner| crate::ema_auth::EmaConnectionAuth {
authority: owner.authority.child(),
owner,
})
});
let ema_authority = ema_auth
.as_ref()
.and_then(|auth| auth.as_ref().ok())
.map(|auth| auth.authority.clone());
let startup = Arc::new(ManagedClientStartup {
ema_auth,
originator: AuthStorageOriginator::current(),
@@ -536,6 +549,7 @@ impl AsyncManagedClient {
});
Self {
client,
ema_authority,
is_codex_apps_mcp_server,
cached_server_info,
server_capabilities,
@@ -1241,7 +1255,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>,
ema_auth: Option<Result<crate::ema_auth::EmaConnectionAuth, StartupOutcomeError>>,
server_name: &str,
server: EffectiveMcpServer,
store_mode: OAuthCredentialsStoreMode,
@@ -1277,7 +1291,9 @@ pub(crate) async fn make_rmcp_client(
))
})??;
return auth
.owner
.connect(
auth.authority,
server_name,
&server,
runtime_context.local_http_client(),
codex-rs/codex-mcp/src/runtime.rs
+20 -2
View File
@@ -3,7 +3,8 @@
//! This module contains the thread-owned MCP runtime and data that describes the
//! environment in which MCP servers execute. Transport startup lives in
//! [`crate::rmcp_client`] and connection-set behavior lives in
//! [`crate::connection_manager`].
//! [`crate::connection_manager`]. Registry removal revokes EMA authority before
//! the thread becomes untracked.
mod status;
@@ -337,7 +338,24 @@ impl McpRuntime {
self.latest_hard_refresh_codex_apps_tools_cache().await
}
/// Revoke a removed thread's EMA access before releasing the registry lock.
pub fn revoke_ema(&self) {
if let Some(Ok(auth)) = self.ema_auth.get() {
auth.authority.revoke();
}
}
/// Drain EMA responses after its configuration binding has been revoked.
pub async fn quiesce_ema(&self) {
if let Some(Ok(auth)) = self.ema_auth.get() {
auth.authority.revoke();
auth.authority.quiesce().await;
}
}
async fn publish(&self, input: McpRuntimeInput, previous: Option<&McpConnectionSet>) {
let current = self.current.load_full();
let retain_ema = current.config.is_some();
let ema_auth = (input.config.xaa_enabled
&& input.mcp_servers.values().any(|server| {
server.config().enabled
@@ -359,7 +377,6 @@ impl McpRuntime {
let config = Arc::clone(&input.config);
let auth = input.auth.clone();
let auth_token = auth.as_ref().and_then(|auth| auth.get_token().ok());
let current = self.current.load_full();
let auth_generation = if current.auth_matches(auth.as_ref()) {
Arc::clone(&current.auth_generation)
} else {
@@ -372,6 +389,7 @@ impl McpRuntime {
let connections = Arc::new(
McpConnectionSet::new(
ema_auth,
retain_ema.then_some(current.connections.as_ref()),
previous,
publication_gate,
input,
codex-rs/config/src/cloud_config_bundle.rs
+5 -2
View File
@@ -343,10 +343,13 @@ pub enum CloudConfigBundleBindingStatus {
}
impl CloudConfigBundleBinding {
/// Completes when this captured policy can no longer authorize EMA requests.
/// Wait until this delivered policy no longer authorizes its EMA owner.
pub async fn invalidated(&self) {
let mut changes = self.policy.0.subscribe();
while self.read().status == CloudConfigBundleBindingStatus::Current {
loop {
if self.read().status != CloudConfigBundleBindingStatus::Current {
return;
}
if changes.changed().await.is_err() {
return;
}
codex-rs/config/src/lib.rs
+2
View File
@@ -23,6 +23,8 @@ pub mod loader;
mod marketplace_edit;
mod mcp_edit;
mod mcp_ema;
mod mcp_ema_cancellation;
pub use mcp_ema_cancellation::McpEmaCancellation;
mod mcp_requirements;
mod mcp_types;
mod merge;
codex-rs/config/src/mcp_ema_cancellation.rs
+84
View File
@@ -0,0 +1,84 @@
//! Carries EMA startup cancellation across configuration loading and thread publication.
//! Once published, a thread's existing config refresh gate owns selective revocation.
use std::sync::Arc;
use std::sync::atomic::AtomicBool;
use std::sync::atomic::Ordering;
use tokio_util::sync::CancellationToken;
#[derive(Clone, Debug, Default)]
pub struct McpEmaCancellation(Arc<State>);
#[derive(Debug, Default)]
struct State {
startup: CancellationToken,
owner: CancellationToken,
lifetime: CancellationToken,
published: AtomicBool,
}
impl McpEmaCancellation {
/// Capture before loading configuration so reload can fence a delayed startup.
pub fn pending(startup: CancellationToken, owner: CancellationToken) -> Self {
let lifetime = owner.child_token();
Self(Arc::new(State {
startup,
owner,
lifetime,
published: AtomicBool::new(false),
}))
}
/// A fork gets its own publication state, even when it inherits loaded config.
pub fn for_startup(&self) -> Self {
let pending = Self::pending(self.0.startup.clone(), self.0.owner.clone());
if self.is_cancelled() {
pending.cancel();
}
pending
}
/// A validated session keeps its authority and records the current startup fence
/// for later forks. The session publication gate must protect this replacement.
pub fn refreshed(&self, incoming: &Self) -> Self {
Self(Arc::new(State {
startup: incoming.0.startup.clone(),
owner: self.0.owner.clone(),
lifetime: self.0.lifetime.clone(),
published: AtomicBool::new(true),
}))
}
/// Call while holding the thread registry's publication lock. A reload either
/// cancels this pending snapshot or discovers the published thread for validation.
pub fn publish(&self) {
if !self.0.published.swap(true, Ordering::AcqRel) && self.0.startup.is_cancelled() {
self.cancel();
}
}
pub fn cancel(&self) {
self.0.lifetime.cancel();
}
pub fn is_cancelled(&self) -> bool {
self.0.lifetime.is_cancelled()
|| (!self.0.published.load(Ordering::Acquire) && self.0.startup.is_cancelled())
}
pub fn token(&self) -> CancellationToken {
self.0.lifetime.clone()
}
pub async fn cancelled(&self) {
tokio::select! {
biased;
() = self.0.lifetime.cancelled() => return,
() = self.0.startup.cancelled() => {}
}
if self.0.published.load(Ordering::Acquire) {
self.0.lifetime.cancelled().await;
}
}
}
codex-rs/config/src/state.rs
+4 -7
View File
@@ -248,7 +248,7 @@ pub struct ConfigLayerStack {
/// Runtime-only EMA authority; it does not change the configuration's contents.
cloud_config_binding: Option<crate::CloudConfigBundleBinding>,
/// Runtime-only cancellation inherited by config snapshots and pending EMA startups.
mcp_ema_cancellation: tokio_util::sync::CancellationToken,
mcp_ema_cancellation: crate::McpEmaCancellation,
/// Cached TOML projection derived only from `requirements_toml`.
/// Construction validates provider definitions and reports serialization errors,
/// so `effective_config()` can replace complete entries without a fallible conversion.
@@ -314,7 +314,7 @@ impl ConfigLayerStack {
)?);
Ok(Self {
cloud_config_binding: None,
mcp_ema_cancellation: tokio_util::sync::CancellationToken::new(),
mcp_ema_cancellation: crate::McpEmaCancellation::default(),
model_provider_requirements,
layers,
requirements,
@@ -349,15 +349,12 @@ impl ConfigLayerStack {
self.cloud_config_binding.as_ref()
}
pub fn with_mcp_ema_cancellation(
mut self,
cancellation: tokio_util::sync::CancellationToken,
) -> Self {
pub fn with_mcp_ema_cancellation(mut self, cancellation: crate::McpEmaCancellation) -> Self {
self.mcp_ema_cancellation = cancellation;
self
}
pub fn mcp_ema_cancellation(&self) -> &tokio_util::sync::CancellationToken {
pub fn mcp_ema_cancellation(&self) -> &crate::McpEmaCancellation {
&self.mcp_ema_cancellation
}
codex-rs/core/src/config/runtime_refresh.rs
+13 -3
View File
@@ -92,8 +92,8 @@ impl Config {
let is_enterprise = |server: &McpServerConfig| {
matches!(server.auth, McpServerAuth::EmaAuth) || server.ema_registration().is_some()
};
// Refresh ordinary servers from the complete host-owned set. Keep the
// initially admitted EMA configuration until a fresh session is constructed.
// Refresh ordinary servers from the complete host-owned set. Keep EMA
// admission pinned, while applying current tool restrictions and approvals.
let mut servers: HashMap<_, _> = incoming
.mcp_servers
.get()
@@ -109,7 +109,17 @@ impl Config {
.get()
.iter()
.filter(|(_, server)| is_enterprise(server))
.map(|(name, server)| (name.clone(), server.clone())),
.map(|(name, server)| {
let mut server = server.clone();
if let Some(incoming) = incoming.mcp_servers.get().get(name) {
server.enabled_tools.clone_from(&incoming.enabled_tools);
server.disabled_tools.clone_from(&incoming.disabled_tools);
server.omit_tools_from.clone_from(&incoming.omit_tools_from);
server.default_tools_approval_mode = incoming.default_tools_approval_mode;
server.tools.clone_from(&incoming.tools);
}
(name.clone(), server)
}),
);
let mut config = self.clone();
if !matches!(scope, RuntimeConfigRefresh::UserFiles) {
codex-rs/core/src/session/config_refresh.rs
+32 -12
View File
@@ -14,18 +14,20 @@ use tracing::warn;
impl Session {
pub(crate) async fn disable_mcp_enterprise_auth(&self) {
let mut state = self.state.lock().await;
let mut config = state
.session_configuration
.original_config_do_not_use
.as_ref()
.clone();
config.config_layer_stack.mcp_ema_cancellation().cancel();
config.disable_mcp_enterprise_auth();
state.session_configuration.original_config_do_not_use = Arc::new(config);
self.services.mcp_runtime.invalidate_resource_caches();
self.mark_mcp_runtime_dirty();
drop(state);
{
let mut state = self.state.lock().await;
let mut config = state
.session_configuration
.original_config_do_not_use
.as_ref()
.clone();
config.config_layer_stack.mcp_ema_cancellation().cancel();
config.disable_mcp_enterprise_auth();
state.session_configuration.original_config_do_not_use = Arc::new(config);
self.services.mcp_runtime.invalidate_resource_caches();
self.mark_mcp_runtime_dirty();
}
self.services.mcp_runtime.quiesce_ema().await;
self.schedule_mcp_prewarm();
}
@@ -164,6 +166,17 @@ impl Session {
if config.mcp_enterprise_managed_auth.is_none() {
config.config_layer_stack.mcp_ema_cancellation().cancel();
}
if !matches!(scope, RuntimeConfigRefresh::UserFiles) {
let cancellation = config
.config_layer_stack
.mcp_ema_cancellation()
.refreshed(next_config.config_layer_stack.mcp_ema_cancellation());
let config = Arc::make_mut(&mut config);
config.config_layer_stack = config
.config_layer_stack
.clone()
.with_mcp_ema_cancellation(cancellation);
}
state.session_configuration.original_config_do_not_use = Arc::clone(&config);
drop(policy_guard);
if matches!(scope, RuntimeConfigRefresh::Mcp) {
@@ -174,6 +187,13 @@ impl Session {
.then(|| self.build_effective_session_config(&state.session_configuration));
(previous_config, new_config)
};
if config
.config_layer_stack
.mcp_ema_cancellation()
.is_cancelled()
{
self.services.mcp_runtime.quiesce_ema().await;
}
self.emit_config_changed_contributors(previous_config.as_ref(), new_config.as_ref());
self.schedule_mcp_prewarm();
self.refresh_hooks(config).await;
codex-rs/core/src/thread_manager.rs
+36 -5
View File
@@ -1408,7 +1408,7 @@ impl ThreadManager {
/// as `Arc<CodexThread>`, it is possible that other references to it exist elsewhere.
/// Returns the thread if the thread was found and removed.
pub async fn remove_thread(&self, thread_id: &ThreadId) -> Option<Arc<CodexThread>> {
self.state.threads.write().await.remove(thread_id)
self.state.remove_thread(thread_id).await
}
/// Removes a runtime for a client request, leaving internal workers with their owner.
@@ -1426,7 +1426,9 @@ impl ThreadManager {
"live internal threads can only be removed by their owner".to_owned(),
));
}
Ok(threads.remove(thread_id))
Ok(threads.remove(thread_id).inspect(|thread| {
thread.session.services.mcp_runtime.revoke_ema();
}))
}
/// Removes a thread only if `thread_id` still maps to `expected`.
@@ -1443,7 +1445,9 @@ impl ThreadManager {
.get(thread_id)
.is_some_and(|thread| Arc::ptr_eq(thread, expected))
{
threads.remove(thread_id)
threads.remove(thread_id).inspect(|thread| {
thread.session.services.mcp_runtime.revoke_ema();
})
} else {
None
}
@@ -1485,7 +1489,9 @@ impl ThreadManager {
let mut tracked_threads = self.state.threads.write().await;
for thread_id in &report.completed {
tracked_threads.remove(thread_id);
if let Some(thread) = tracked_threads.remove(thread_id) {
thread.session.services.mcp_runtime.revoke_ema();
}
}
report
@@ -1850,7 +1856,13 @@ impl ThreadManagerState {
/// Remove a thread from the manager by ID, returning it when present.
pub(crate) async fn remove_thread(&self, thread_id: &ThreadId) -> Option<Arc<CodexThread>> {
self.threads.write().await.remove(thread_id)
self.threads
.write()
.await
.remove(thread_id)
.inspect(|thread| {
thread.session.services.mcp_runtime.revoke_ema();
})
}
pub(crate) async fn effective_multi_agent_version_for_spawn(
@@ -2094,6 +2106,17 @@ impl ThreadManagerState {
/// Spawn a new thread with optional history and register it with the manager.
async fn spawn_thread(&self, mut request: ThreadSpawnRequest) -> CodexResult<NewThread> {
let cancellation = request
.options
.config
.config_layer_stack
.mcp_ema_cancellation()
.for_startup();
request.options.config.config_layer_stack = request
.options
.config
.config_layer_stack
.with_mcp_ema_cancellation(cancellation);
let runtime = request.agent_control.runtime().clone();
let membership = runtime.admit_start()?;
let owns_startup = request.startup.is_none();
@@ -2241,6 +2264,7 @@ impl ThreadManagerState {
});
}
threads.remove(&resumed.conversation_id);
thread.session.services.mcp_runtime.revoke_ema();
}
}
let agent_control = match (&self.agent_control_factory, agent_control) {
@@ -2511,6 +2535,12 @@ impl ThreadManagerState {
}
};
let ema_cancellation = session
.get_config()
.await
.config_layer_stack
.mcp_ema_cancellation()
.clone();
{
let mut threads = self.threads.write().await;
if let std::collections::hash_map::Entry::Vacant(e) = threads.entry(thread_id) {
@@ -2520,6 +2550,7 @@ impl ThreadManagerState {
.agent_control
.propagate_config_update(update);
}
ema_cancellation.publish();
let thread = Arc::new(CodexThread::new(
session,
io,
codex-rs/rmcp-client/src/ema_http_client.rs
+3
View File
@@ -216,6 +216,9 @@ impl<S: EmaTokenSource> EmaAuthenticatedHttpClient<S> {
if let Err(refusal) = self.authority.check_new_request() {
return refusal;
}
if is_authentication_required_error(&error) {
self.authority.revoke();
}
auth_error(error)
})?;
let token = CachedToken::new(renewed);