diff --git a/codex-rs/app-server/src/config_manager.rs b/codex-rs/app-server/src/config_manager.rs index ee23da9ee5..4fe02c5264 100644 --- a/codex-rs/app-server/src/config_manager.rs +++ b/codex-rs/app-server/src/config_manager.rs @@ -57,7 +57,13 @@ pub(crate) struct ConfigManager { network_policy_reload: Arc, network_policy_snapshot: Arc>>>, /// Captured before config loading so a reload also revokes pending startups. - mcp_ema_cancellation: Arc>, + mcp_ema_cancellation: Arc>, +} + +#[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( diff --git a/codex-rs/app-server/src/mcp_refresh.rs b/codex-rs/app-server/src/mcp_refresh.rs index c0fad84f86..7fe1e32af4 100644 --- a/codex-rs/app-server/src/mcp_refresh.rs +++ b/codex-rs/app-server/src/mcp_refresh.rs @@ -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> = 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 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 2eaf2c5ffb..1c6b35dda6 100644 --- a/codex-rs/app-server/src/request_processors/config_processor.rs +++ b/codex-rs/app-server/src/request_processors/config_processor.rs @@ -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 { diff --git a/codex-rs/app-server/src/request_processors/thread_processor.rs b/codex-rs/app-server/src/request_processors/thread_processor.rs index de9b90925f..7c67ec17d4 100644 --- a/codex-rs/app-server/src/request_processors/thread_processor.rs +++ b/codex-rs/app-server/src/request_processors/thread_processor.rs @@ -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(), }; diff --git a/codex-rs/app-server/tests/suite/v2/account_enterprise_tests.rs b/codex-rs/app-server/tests/suite/v2/account_enterprise_tests.rs index 01dd50b903..4946223158 100644 --- a/codex-rs/app-server/tests/suite/v2/account_enterprise_tests.rs +++ b/codex-rs/app-server/tests/suite/v2/account_enterprise_tests.rs @@ -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 diff --git a/codex-rs/codex-mcp/src/connection_manager.rs b/codex-rs/codex-mcp/src/connection_manager.rs index ea7e36e5dc..3115ea9472 100644 --- a/codex-rs/codex-mcp/src/connection_manager.rs +++ b/codex-rs/codex-mcp/src/connection_manager.rs @@ -283,6 +283,7 @@ impl McpConnectionSet { /// notifications are then skipped and interactive elicitations are declined. pub async fn new( ema_auth: Option, + 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 diff --git a/codex-rs/codex-mcp/src/connection_manager_tests.rs b/codex-rs/codex-mcp/src/connection_manager_tests.rs index b2f35edb68..044d475b89 100644 --- a/codex-rs/codex-mcp/src/connection_manager_tests.rs +++ b/codex-rs/codex-mcp/src/connection_manager_tests.rs @@ -579,6 +579,7 @@ async fn prepared_call_timeout_includes_trusted_access_lookup() { pub(crate) async fn create_ready_async_managed_client(tools: Vec) -> AsyncManagedClient { AsyncManagedClient { + ema_authority: None, client: futures::future::ready::>(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::>(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::>() .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::>() .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(), diff --git a/codex-rs/codex-mcp/src/ema_auth.rs b/codex-rs/codex-mcp/src/ema_auth.rs index 82075fb93d..503d2d056b 100644 --- a/codex-rs/codex-mcp/src/ema_auth.rs +++ b/codex-rs/codex-mcp/src/ema_auth.rs @@ -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, StartupOutcomeError>; +#[derive(Clone)] +pub(crate) struct EmaConnectionAuth { + pub(crate) owner: Arc, + pub(crate) authority: EmaRequestAuthority, +} + pub fn ema_auth_scope(auth: Option<&CodexAuth>) -> Option { 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 { /// Retained by the thread across MCP refreshes; a revoked owner cannot be replaced. pub(crate) struct EmaAuth { - scope: McpEmaAuthScope, - config: Arc, - manager: Arc, - account: CodexAuth, - authority: EmaRequestAuthority, + pub(crate) scope: McpEmaAuthScope, + pub(crate) config: Arc, + pub(crate) manager: Arc, + 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, + authority: EmaRequestAuthority, server_name: &str, server: &EffectiveMcpServer, http_client: Arc, protocol_mode: McpProtocolMode, ) -> Result { - 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 = 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> { 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())), diff --git a/codex-rs/codex-mcp/src/mcp/mod.rs b/codex-rs/codex-mcp/src/mcp/mod.rs index 2a3c2dd83f..56203cf4b2 100644 --- a/codex-rs/codex-mcp/src/mcp/mod.rs +++ b/codex-rs/codex-mcp/src/mcp/mod.rs @@ -133,7 +133,7 @@ pub struct McpConfig { /// 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 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 { diff --git a/codex-rs/codex-mcp/src/mcp/mod_tests.rs b/codex-rs/codex-mcp/src/mcp/mod_tests.rs index 90db57a7f9..9c1435e435 100644 --- a/codex-rs/codex-mcp/src/mcp/mod_tests.rs +++ b/codex-rs/codex-mcp/src/mcp/mod_tests.rs @@ -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, diff --git a/codex-rs/codex-mcp/src/rmcp_client.rs b/codex-rs/codex-mcp/src/rmcp_client.rs index 86cbd1cedb..24295d9bfe 100644 --- a/codex-rs/codex-mcp/src/rmcp_client.rs +++ b/codex-rs/codex-mcp/src/rmcp_client.rs @@ -297,7 +297,7 @@ fn codex_apps_reconnect_backoff(consecutive_failures: u32) -> Duration { #[derive(Clone)] struct ManagedClientStartup { - ema_auth: Option, + ema_auth: Option>, originator: AuthStorageOriginator, server_name: String, server: EffectiveMcpServer, @@ -447,6 +447,7 @@ impl ManagedClientStartup { #[derive(Clone)] pub(crate) struct AsyncManagedClient { + pub(crate) ema_authority: Option, pub(crate) client: ManagedClientFuture, pub(crate) is_codex_apps_mcp_server: bool, pub(crate) cached_server_info: Option, @@ -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, + ema_auth: Option>, 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(), diff --git a/codex-rs/codex-mcp/src/runtime.rs b/codex-rs/codex-mcp/src/runtime.rs index 211a151e9c..6d57cd1684 100644 --- a/codex-rs/codex-mcp/src/runtime.rs +++ b/codex-rs/codex-mcp/src/runtime.rs @@ -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(¤t.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, diff --git a/codex-rs/config/src/cloud_config_bundle.rs b/codex-rs/config/src/cloud_config_bundle.rs index 0f6daac894..05d62fb25f 100644 --- a/codex-rs/config/src/cloud_config_bundle.rs +++ b/codex-rs/config/src/cloud_config_bundle.rs @@ -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; } diff --git a/codex-rs/config/src/lib.rs b/codex-rs/config/src/lib.rs index b51ad2f52f..52db805ce1 100644 --- a/codex-rs/config/src/lib.rs +++ b/codex-rs/config/src/lib.rs @@ -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; diff --git a/codex-rs/config/src/mcp_ema_cancellation.rs b/codex-rs/config/src/mcp_ema_cancellation.rs new file mode 100644 index 0000000000..209e9fd5ac --- /dev/null +++ b/codex-rs/config/src/mcp_ema_cancellation.rs @@ -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); + +#[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; + } + } +} diff --git a/codex-rs/config/src/state.rs b/codex-rs/config/src/state.rs index 8d8324e547..127f7d044a 100644 --- a/codex-rs/config/src/state.rs +++ b/codex-rs/config/src/state.rs @@ -248,7 +248,7 @@ 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, + 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 } diff --git a/codex-rs/core/src/config/runtime_refresh.rs b/codex-rs/core/src/config/runtime_refresh.rs index 5495ff631a..4428ed56e6 100644 --- a/codex-rs/core/src/config/runtime_refresh.rs +++ b/codex-rs/core/src/config/runtime_refresh.rs @@ -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) { diff --git a/codex-rs/core/src/session/config_refresh.rs b/codex-rs/core/src/session/config_refresh.rs index 99915a251b..08a99f81dc 100644 --- a/codex-rs/core/src/session/config_refresh.rs +++ b/codex-rs/core/src/session/config_refresh.rs @@ -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; diff --git a/codex-rs/core/src/thread_manager.rs b/codex-rs/core/src/thread_manager.rs index 1ac38d20fd..de3c37d322 100644 --- a/codex-rs/core/src/thread_manager.rs +++ b/codex-rs/core/src/thread_manager.rs @@ -1408,7 +1408,7 @@ impl ThreadManager { /// as `Arc`, 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> { - 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> { - 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 { + 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, diff --git a/codex-rs/rmcp-client/src/ema_http_client.rs b/codex-rs/rmcp-client/src/ema_http_client.rs index 29b3abab93..aace365851 100644 --- a/codex-rs/rmcp-client/src/ema_http_client.rs +++ b/codex-rs/rmcp-client/src/ema_http_client.rs @@ -216,6 +216,9 @@ impl EmaAuthenticatedHttpClient { 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);