diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 44ddf7b894..91f1a4f1cd 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -2877,6 +2877,7 @@ dependencies = [ "thiserror 2.0.18", "tokio", "tokio-stream", + "tokio-util", "toml 0.9.11+spec-1.1.0", "toml_edit 0.24.0+spec-1.1.0", "tonic", diff --git a/codex-rs/app-server/src/config_manager.rs b/codex-rs/app-server/src/config_manager.rs index d2c2db7481..199499826d 100644 --- a/codex-rs/app-server/src/config_manager.rs +++ b/codex-rs/app-server/src/config_manager.rs @@ -25,7 +25,9 @@ use std::collections::HashMap; use std::path::Path; use std::path::PathBuf; use std::sync::Arc; +use std::sync::Mutex; use std::sync::RwLock; +use tokio_util::sync::CancellationToken; use toml::Value as TomlValue; use tracing::instrument; use tracing::warn; @@ -54,6 +56,8 @@ pub(crate) struct ConfigManager { local_network_policy: codex_http_client::NetworkPolicyController, network_policy_reload: Arc, network_policy_snapshot: Arc>>>, + /// Captured before config loading so a reload also revokes pending startups. + mcp_ema_cancellation: Arc>, } /// Startup configuration, including whether loading had to fall back to defaults. @@ -99,6 +103,7 @@ impl ConfigManager { local_network_policy: Default::default(), network_policy_reload: Arc::new(tokio::sync::Semaphore::new(/*permits*/ 1)), network_policy_snapshot: Arc::default(), + mcp_ema_cancellation: Arc::default(), } } @@ -106,6 +111,24 @@ 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. + 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(); + } + + fn capture_mcp_ema_cancellation(&self) -> CancellationToken { + self.mcp_ema_cancellation + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .child_token() + } + pub(crate) fn with_embedded_network_policy( mut self, policy: crate::in_process::EmbeddedNetworkPolicy, @@ -171,6 +194,7 @@ impl ConfigManager { } else { warn!("failed to update cloud config bundle loader"); } + self.revoke_mcp_ema(); } pub(crate) fn clear_cloud_config_bundle_loader(&self) { @@ -180,6 +204,7 @@ impl ConfigManager { } else { warn!("failed to clear cloud config bundle loader"); } + self.revoke_mcp_ema(); } pub(crate) async fn sync_default_client_residency_requirement(&self) { @@ -213,6 +238,7 @@ impl ConfigManager { /// Loads system, user, and runtime settings without discovering a project /// from the app-server process's working directory. pub(crate) async fn load_non_project_config(&self) -> std::io::Result { + let mcp_ema_cancellation = self.capture_mcp_ema_cancellation(); let policy_load = self.refresh_application_network_policy().await?; let mut config = ConfigBuilder::default() .codex_home(self.codex_home.clone()) @@ -229,6 +255,9 @@ impl ConfigManager { self.apply_network_policy(&mut config); self.apply_runtime_feature_enablement(&mut config); self.apply_arg0_paths(&mut config); + config.config_layer_stack = config + .config_layer_stack + .with_mcp_ema_cancellation(mcp_ema_cancellation); Ok(config) } @@ -237,6 +266,7 @@ impl ConfigManager { session_layers: &ConfigLayerStack, cwd: &Path, ) -> std::io::Result { + let mcp_ema_cancellation = self.capture_mcp_ema_cancellation(); let codex_home = AbsolutePathBuf::from_absolute_path(&self.codex_home)?; let layers = self .load_config_layers_for_cwd(AbsolutePathBuf::from_absolute_path(cwd)?) @@ -251,6 +281,9 @@ impl ConfigManager { self.apply_network_policy(&mut config); self.apply_runtime_feature_enablement(&mut config); self.apply_arg0_paths(&mut config); + config.config_layer_stack = config + .config_layer_stack + .with_mcp_ema_cancellation(mcp_ema_cancellation); Ok(config) } @@ -260,6 +293,7 @@ impl ConfigManager { session_layers: &ConfigLayerStack, cwd: &Path, ) -> std::io::Result { + let mcp_ema_cancellation = self.capture_mcp_ema_cancellation(); let mut manager = self.clone(); manager.thread_config_loader = Arc::new(codex_config::NoopThreadConfigLoader); let refreshed_layers = manager @@ -276,6 +310,9 @@ impl ConfigManager { self.apply_network_policy(&mut config); self.apply_runtime_feature_enablement(&mut config); self.apply_arg0_paths(&mut config); + config.config_layer_stack = config + .config_layer_stack + .with_mcp_ema_cancellation(mcp_ema_cancellation); Ok(config) } @@ -365,6 +402,7 @@ impl ConfigManager { } pub(crate) async fn load_default_config(&self) -> std::io::Result { + let mcp_ema_cancellation = self.capture_mcp_ema_cancellation(); let mut loader_overrides = self.loader_overrides.clone(); loader_overrides.ignore_user_config = true; loader_overrides.ignore_project_config = true; @@ -381,6 +419,9 @@ impl ConfigManager { self.apply_network_policy(&mut config); self.apply_runtime_feature_enablement(&mut config); self.apply_arg0_paths(&mut config); + config.config_layer_stack = config + .config_layer_stack + .with_mcp_ema_cancellation(mcp_ema_cancellation); Ok(config) } @@ -467,6 +508,7 @@ impl ConfigManager { mut typesafe_overrides: ConfigOverrides, fallback_cwd: Option, ) -> std::io::Result { + let mcp_ema_cancellation = self.capture_mcp_ema_cancellation(); let policy_load = self.refresh_application_network_policy().await?; let mut request_overrides = request_overrides.unwrap_or_default(); if let Some(value) = request_overrides.remove("bypass_hook_trust") { @@ -502,6 +544,9 @@ impl ConfigManager { self.apply_network_policy(&mut config); self.apply_runtime_feature_enablement(&mut config); self.apply_arg0_paths(&mut config); + config.config_layer_stack = config + .config_layer_stack + .with_mcp_ema_cancellation(mcp_ema_cancellation); Ok(config) } diff --git a/codex-rs/app-server/src/mcp_refresh.rs b/codex-rs/app-server/src/mcp_refresh.rs index 3611ab3992..cf09564530 100644 --- a/codex-rs/app-server/src/mcp_refresh.rs +++ b/codex-rs/app-server/src/mcp_refresh.rs @@ -46,6 +46,7 @@ async fn reload_mcp_config_with_policy( config_manager: &ConfigManager, policy: ReloadPolicy, ) -> io::Result<()> { + config_manager.revoke_mcp_ema(); let mut rejected = false; let mut load_error = None; let mut completed: Vec> = Vec::new(); @@ -401,6 +402,20 @@ enabled = false } } let thread = good_thread.expect("good test thread should exist"); + // Config loading can finish before a startup is visible to ThreadManager. + let pending_startup = config_manager + .load_for_cwd( + /*request_overrides*/ None, + ConfigOverrides::default(), + Some(temp_dir.path().join("good")), + ) + .await?; + assert!( + !pending_startup + .config_layer_stack + .mcp_ema_cancellation() + .is_cancelled() + ); std::fs::write( temp_dir.path().join(codex_config::CONFIG_TOML_FILE), r#" @@ -412,6 +427,34 @@ enabled = false reload_mcp_config(&thread_manager, &config_manager).await?; + assert!( + pending_startup + .config_layer_stack + .mcp_ema_cancellation() + .is_cancelled() + ); + assert!( + thread + .config() + .await + .config_layer_stack + .mcp_ema_cancellation() + .is_cancelled() + ); + let fresh_startup = config_manager + .load_for_cwd( + /*request_overrides*/ None, + ConfigOverrides::default(), + Some(temp_dir.path().join("good")), + ) + .await?; + assert!( + !fresh_startup + .config_layer_stack + .mcp_ema_cancellation() + .is_cancelled() + ); + assert!( thread .config() diff --git a/codex-rs/app-server/src/request_processors/config_processor.rs b/codex-rs/app-server/src/request_processors/config_processor.rs index ff87f767a6..fc3e17e58e 100644 --- a/codex-rs/app-server/src/request_processors/config_processor.rs +++ b/codex-rs/app-server/src/request_processors/config_processor.rs @@ -374,6 +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(); 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 { diff --git a/codex-rs/codex-mcp/src/mcp/mod.rs b/codex-rs/codex-mcp/src/mcp/mod.rs index a8a323a5aa..09d016cb6c 100644 --- a/codex-rs/codex-mcp/src/mcp/mod.rs +++ b/codex-rs/codex-mcp/src/mcp/mod.rs @@ -132,6 +132,8 @@ pub struct McpConfig { pub codex_home: PathBuf, /// Trusted enterprise IdP inherited after normal catalog and policy resolution. pub mcp_enterprise_managed_auth: Option, + /// Retires only EMA connections and pending exchanges for this config snapshot. + pub mcp_ema_cancellation: tokio_util::sync::CancellationToken, pub xaa_enabled: bool, /// Preferred credential store for MCP OAuth tokens. pub mcp_oauth_credentials_store_mode: OAuthCredentialsStoreMode, diff --git a/codex-rs/codex-mcp/src/mcp/mod_tests.rs b/codex-rs/codex-mcp/src/mcp/mod_tests.rs index b0bea88b78..90db57a7f9 100644 --- a/codex-rs/codex-mcp/src/mcp/mod_tests.rs +++ b/codex-rs/codex-mcp/src/mcp/mod_tests.rs @@ -90,6 +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(), xaa_enabled: false, mcp_oauth_credentials_store_mode: OAuthCredentialsStoreMode::default(), oauth_refresh_mode: McpOAuthRefreshMode::Legacy, diff --git a/codex-rs/config/Cargo.toml b/codex-rs/config/Cargo.toml index b4be94bb58..0fc87168fa 100644 --- a/codex-rs/config/Cargo.toml +++ b/codex-rs/config/Cargo.toml @@ -40,6 +40,7 @@ serde_path_to_error = { workspace = true } sha2 = { workspace = true } thiserror = { workspace = true } tokio = { workspace = true, features = ["fs"] } +tokio-util = { workspace = true, features = ["rt"] } toml = { workspace = true, features = ["preserve_order"] } toml_edit = { workspace = true } tonic = { workspace = true, features = ["tls-native-roots", "tls-ring"] } diff --git a/codex-rs/config/src/state.rs b/codex-rs/config/src/state.rs index ed3d13ccce..8d8324e547 100644 --- a/codex-rs/config/src/state.rs +++ b/codex-rs/config/src/state.rs @@ -247,6 +247,8 @@ impl ConfigLayerEntry { pub struct ConfigLayerStack { /// Runtime-only EMA authority; it does not change the configuration's contents. cloud_config_binding: Option, + /// Runtime-only cancellation inherited by config snapshots and pending EMA startups. + mcp_ema_cancellation: tokio_util::sync::CancellationToken, /// 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. @@ -279,6 +281,7 @@ impl PartialEq for ConfigLayerStack { fn eq(&self, other: &Self) -> bool { let Self { cloud_config_binding: _, + mcp_ema_cancellation: _, model_provider_requirements, layers, requirements, @@ -311,6 +314,7 @@ impl ConfigLayerStack { )?); Ok(Self { cloud_config_binding: None, + mcp_ema_cancellation: tokio_util::sync::CancellationToken::new(), model_provider_requirements, layers, requirements, @@ -345,6 +349,18 @@ impl ConfigLayerStack { self.cloud_config_binding.as_ref() } + pub fn with_mcp_ema_cancellation( + mut self, + cancellation: tokio_util::sync::CancellationToken, + ) -> Self { + self.mcp_ema_cancellation = cancellation; + self + } + + pub fn mcp_ema_cancellation(&self) -> &tokio_util::sync::CancellationToken { + &self.mcp_ema_cancellation + } + /// Retains session layers while adopting current MCP, plugin, and feature restrictions. /// Rejected refreshes must not restore an earlier policy on the next user reload. pub fn with_mcp_requirements_from(&self, incoming: &Self) -> Self { @@ -476,6 +492,7 @@ impl ConfigLayerStack { Ok(Self { layers, cloud_config_binding: self.cloud_config_binding.clone(), + mcp_ema_cancellation: self.mcp_ema_cancellation.clone(), model_provider_requirements: self.model_provider_requirements.clone(), requirements: self.requirements.clone(), requirements_toml: self.requirements_toml.clone(), @@ -513,6 +530,7 @@ impl ConfigLayerStack { Self { layers, cloud_config_binding: self.cloud_config_binding.clone(), + mcp_ema_cancellation: self.mcp_ema_cancellation.clone(), model_provider_requirements: self.model_provider_requirements.clone(), requirements: self.requirements.clone(), requirements_toml: self.requirements_toml.clone(), diff --git a/codex-rs/core/src/agent/role.rs b/codex-rs/core/src/agent/role.rs index cf43be3b40..8a34de0108 100644 --- a/codex-rs/core/src/agent/role.rs +++ b/codex-rs/core/src/agent/role.rs @@ -258,6 +258,7 @@ mod role_overrides { config.config_layer_stack.requirements_toml().clone(), )? .with_cloud_config_binding(config.config_layer_stack.cloud_config_binding().cloned()) + .with_mcp_ema_cancellation(config.config_layer_stack.mcp_ema_cancellation().clone()) .with_user_and_project_exec_policy_rules_ignored( config .config_layer_stack diff --git a/codex-rs/core/src/config/mod.rs b/codex-rs/core/src/config/mod.rs index ff6c9eca1d..2f13f038ed 100644 --- a/codex-rs/core/src/config/mod.rs +++ b/codex-rs/core/src/config/mod.rs @@ -1878,6 +1878,7 @@ impl Config { requires_read_only_mcp_tools: false, codex_home: self.codex_home.to_path_buf(), mcp_enterprise_managed_auth: self.mcp_enterprise_managed_auth.clone(), + mcp_ema_cancellation: self.config_layer_stack.mcp_ema_cancellation().clone(), xaa_enabled: self.features.enabled(Feature::UseXaa) && self.mcp_enterprise_managed_auth.is_some(), mcp_oauth_credentials_store_mode: self.mcp_oauth_credentials_store_mode, diff --git a/codex-rs/core/src/config/runtime_refresh.rs b/codex-rs/core/src/config/runtime_refresh.rs index 05b017376a..5495ff631a 100644 --- a/codex-rs/core/src/config/runtime_refresh.rs +++ b/codex-rs/core/src/config/runtime_refresh.rs @@ -166,7 +166,8 @@ impl Config { if !matches!(scope, RuntimeConfigRefresh::Mcp) { config.tool_suggest = resolve_tool_suggest_config_from_layer_stack(&layers); } - config.config_layer_stack = layers; + config.config_layer_stack = layers + .with_mcp_ema_cancellation(self.config_layer_stack.mcp_ema_cancellation().clone()); config.plugins = cfg.plugins; if enterprise_retired { config.disable_mcp_enterprise_auth(); diff --git a/codex-rs/core/src/session/config_refresh.rs b/codex-rs/core/src/session/config_refresh.rs index 166c1f6045..99915a251b 100644 --- a/codex-rs/core/src/session/config_refresh.rs +++ b/codex-rs/core/src/session/config_refresh.rs @@ -20,6 +20,7 @@ impl Session { .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(); @@ -160,6 +161,9 @@ impl Session { } let previous_config = notify_contributors .then(|| self.build_effective_session_config(&state.session_configuration)); + if config.mcp_enterprise_managed_auth.is_none() { + config.config_layer_stack.mcp_ema_cancellation().cancel(); + } state.session_configuration.original_config_do_not_use = Arc::clone(&config); drop(policy_guard); if matches!(scope, RuntimeConfigRefresh::Mcp) {