diff --git a/codex-rs/core/src/session/mcp_projection.rs b/codex-rs/core/src/session/mcp_projection.rs index 30713e745a..7094e76b91 100644 --- a/codex-rs/core/src/session/mcp_projection.rs +++ b/codex-rs/core/src/session/mcp_projection.rs @@ -49,8 +49,10 @@ impl Session { return base; } let bindings = selected_bindings(selected_roots, resolved_roots); - let mut plugins = cache.plugins(&bindings).unwrap_or_default(); - if let Some(runtime) = cache.runtime(&bindings) { + let cached_runtime = cache.runtime_for_bindings(&bindings); + let mut plugins = cache.plugins_for_bindings(&bindings).unwrap_or_default(); + let mut discovered_plugin = false; + if cached_runtime.is_some() { let unresolved = bindings .iter() .filter(|(order, _)| { @@ -61,11 +63,7 @@ impl Session { .cloned() .collect::>(); let discovered = project_executor_plugins(&unresolved).await; - if discovered.is_empty() { - self.services - .publish_existing_mcp_runtime(Arc::clone(&runtime)); - return runtime; - } + discovered_plugin = !discovered.is_empty(); plugins.extend(discovered); plugins.sort_unstable_by_key(|(order, _)| *order); } else { @@ -85,6 +83,20 @@ impl Session { .collect::>(); let runtime_context = self.mcp_runtime_context(&turn_context.config, environments, &pinned_roots); + if !discovered_plugin + && let Some(runtime) = cached_runtime + && runtime.matches_projection(mcp_config.as_ref(), &runtime_context) + { + runtime + .manager() + .set_approval_policy(&turn_context.approval_policy); + runtime + .manager() + .set_permission_profile(turn_context.permission_profile()); + self.services + .publish_existing_mcp_runtime(Arc::clone(&runtime)); + return runtime; + } let runtime = if plugins.is_empty() { base } else { @@ -97,7 +109,7 @@ impl Session { ) .await }; - cache.insert_runtime(bindings, plugins, Arc::clone(&runtime)); + cache.replace_selected_runtime(bindings, plugins, Arc::clone(&runtime)); self.services .publish_existing_mcp_runtime(Arc::clone(&runtime)); runtime @@ -146,7 +158,7 @@ impl Session { let mut cache = self.services.selected_mcp_runtime.lock().await; let projection = self.project_mcp_config_inner(config).await; let current = cache - .runtime(&projection.bindings) + .runtime_for_bindings(&projection.bindings) .unwrap_or_else(|| self.services.latest_mcp_runtime()); if current.matches_projection(&projection.config, &projection.runtime_context) { current @@ -169,9 +181,9 @@ impl Session { ) .await; if projection.bindings.is_empty() { - cache.replace_base(Arc::clone(&runtime)); + cache.replace_base_and_invalidate_selected(Arc::clone(&runtime)); } else { - cache.insert_runtime( + cache.replace_selected_runtime( projection.bindings, projection.plugins, Arc::clone(&runtime), diff --git a/codex-rs/core/src/session/mcp_runtime.rs b/codex-rs/core/src/session/mcp_runtime.rs index 959607bb55..a58d526c7b 100644 --- a/codex-rs/core/src/session/mcp_runtime.rs +++ b/codex-rs/core/src/session/mcp_runtime.rs @@ -4,6 +4,23 @@ use codex_core_plugins::ExecutorPluginRuntime; use codex_exec_server::ResolvedSelectedCapabilityRoot; use codex_mcp::McpRuntimeSnapshot; +/// One live selected-plugin MCP runtime retained between model steps. +/// +/// A cached runtime is a reuse candidate only for the same ordered selected roots and the same +/// process-local environment instances. The caller additionally compares the effective MCP config +/// and runtime context before reuse. Selected environment contents are treated as stable, so +/// manifest and MCP config file changes do not invalidate this cache. +/// +/// Within a live session, the selected runtime is invalidated in exactly two ways: +/// +/// 1. [`Self::replace_base_and_invalidate_selected`] installs a new base MCP runtime. +/// 2. [`Self::replace_selected_runtime`] stores a newly projected runtime. This happens when the +/// bindings change, when a previously unavailable plugin appears, or when the effective config +/// or runtime context changes. An unavailable environment disappears from the binding list and +/// therefore follows this path; returning with a new environment instance rebuilds the live +/// runtime even when the stable environment ID is unchanged. +/// +/// In-flight [`McpRuntimeSnapshot`] values retain their manager until their model step finishes. #[derive(Default)] pub(crate) struct SelectedMcpRuntimeCache { base_runtime: Option>, @@ -17,7 +34,10 @@ struct CachedSelectedRuntime { } impl SelectedMcpRuntimeCache { - pub(crate) fn replace_base(&mut self, runtime: Arc) { + pub(crate) fn replace_base_and_invalidate_selected( + &mut self, + runtime: Arc, + ) { self.base_runtime = Some(runtime); self.runtime = None; } @@ -29,7 +49,7 @@ impl SelectedMcpRuntimeCache { .expect("base MCP runtime must be installed before capturing a step") } - pub(crate) fn runtime( + pub(crate) fn runtime_for_bindings( &self, bindings: &[(usize, ResolvedSelectedCapabilityRoot)], ) -> Option> { @@ -39,7 +59,7 @@ impl SelectedMcpRuntimeCache { .map(|cached| Arc::clone(&cached.runtime)) } - pub(crate) fn plugins( + pub(crate) fn plugins_for_bindings( &self, bindings: &[(usize, ResolvedSelectedCapabilityRoot)], ) -> Option> { @@ -49,7 +69,7 @@ impl SelectedMcpRuntimeCache { .map(|cached| cached.plugins.clone()) } - pub(crate) fn insert_runtime( + pub(crate) fn replace_selected_runtime( &mut self, bindings: Vec<(usize, ResolvedSelectedCapabilityRoot)>, plugins: Vec<(usize, ExecutorPluginRuntime)>, @@ -67,6 +87,9 @@ fn same_bindings( left: &[(usize, ResolvedSelectedCapabilityRoot)], right: &[(usize, ResolvedSelectedCapabilityRoot)], ) -> bool { + // Order is part of the key because later selected roots can be renamed when MCP server names + // collide. Arc identity is part of the key because live processes and connections belong to + // one exact environment instance, even when a replacement reuses the same stable ID. left.len() == right.len() && left .iter() diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index d931216a92..9179b96705 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -5268,6 +5268,7 @@ async fn session_new_fails_when_zsh_fork_enabled_without_packaged_zsh() { plugins_manager, mcp_manager, Arc::new(codex_extension_api::ExtensionRegistryBuilder::new().build()), + /*selected_capability_roots*/ Vec::new(), codex_extension_api::ExtensionDataInit::default(), /*supports_openai_form_elicitation*/ false, AgentControl::default(), @@ -5411,7 +5412,7 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) { let network_approval = Arc::new(NetworkApprovalService::default()); let mcp_runtime = uninitialized_mcp_runtime(config.as_ref()); let mut selected_mcp_runtime = crate::session::SelectedMcpRuntimeCache::default(); - selected_mcp_runtime.replace_base(Arc::clone(&mcp_runtime)); + selected_mcp_runtime.replace_base_and_invalidate_selected(Arc::clone(&mcp_runtime)); let services = SessionServices { mcp_runtime: Arc::new(arc_swap::ArcSwapOption::from(Some(mcp_runtime))), mcp_elicitation_managers: std::sync::Mutex::new(Vec::new()), @@ -5646,6 +5647,7 @@ async fn make_session_with_config_and_rx( plugins_manager, mcp_manager, Arc::new(codex_extension_api::ExtensionRegistryBuilder::new().build()), + /*selected_capability_roots*/ Vec::new(), codex_extension_api::ExtensionDataInit::default(), /*supports_openai_form_elicitation*/ false, AgentControl::default(), @@ -5752,6 +5754,7 @@ async fn make_session_with_history_source_and_agent_control_and_rx( plugins_manager, mcp_manager, Arc::new(codex_extension_api::ExtensionRegistryBuilder::new().build()), + /*selected_capability_roots*/ Vec::new(), codex_extension_api::ExtensionDataInit::default(), /*supports_openai_form_elicitation*/ false, agent_control, @@ -7487,7 +7490,7 @@ where let network_approval = Arc::new(NetworkApprovalService::default()); let mcp_runtime = uninitialized_mcp_runtime(config.as_ref()); let mut selected_mcp_runtime = crate::session::SelectedMcpRuntimeCache::default(); - selected_mcp_runtime.replace_base(Arc::clone(&mcp_runtime)); + selected_mcp_runtime.replace_base_and_invalidate_selected(Arc::clone(&mcp_runtime)); let services = SessionServices { mcp_runtime: Arc::new(arc_swap::ArcSwapOption::from(Some(mcp_runtime))), mcp_elicitation_managers: std::sync::Mutex::new(Vec::new()), diff --git a/codex-rs/core/src/state/service.rs b/codex-rs/core/src/state/service.rs index 39ddc54427..313c447a12 100644 --- a/codex-rs/core/src/state/service.rs +++ b/codex-rs/core/src/state/service.rs @@ -140,7 +140,7 @@ impl SessionServices { self.selected_mcp_runtime .lock() .await - .replace_base(Arc::clone(&runtime)); + .replace_base_and_invalidate_selected(Arc::clone(&runtime)); self.publish_existing_mcp_runtime(runtime); } diff --git a/codex-rs/core/tests/common/test_codex.rs b/codex-rs/core/tests/common/test_codex.rs index c62fe44c63..dd7adf8add 100644 --- a/codex-rs/core/tests/common/test_codex.rs +++ b/codex-rs/core/tests/common/test_codex.rs @@ -388,14 +388,6 @@ impl TestCodexBuilder { self } - pub fn with_selected_capability_roots( - mut self, - selected_capability_roots: Vec, - ) -> Self { - self.thread_extension_init.insert(selected_capability_roots); - self - } - pub fn with_user_instructions_provider( mut self, provider: Arc,