From 7498521d288b9b3b96ffba4eedf089d8d6e06a84 Mon Sep 17 00:00:00 2001 From: sayan-oai Date: Fri, 18 Sep 2026 00:46:19 +0000 Subject: [PATCH] Keep MCP policy evaluation consistent with turn environments (#46335) ## Why Environment settings saved for the next turn must not change MCP tool availability during the active turn. ## What changed - Use one captured environment snapshot for MCP policy evaluation and runtime publication, including selections that are still starting or have failed. - Compare both captured selections and ready environment handles when deciding whether to refresh the runtime, and rebuild the startup configuration when resolved selections change. - Preserve the configuration origin of failed selections so MCP authority evaluation retains whether configuration comes from the thread. - Box the MCP refresh future to keep it off the sampling request's stack. ## Testing Add a regression test that updates the environment MCP policy while a turn waits for user input, verifies the tool remains visible when that turn resumes, and verifies it disappears on the next turn. Extend snapshot tests to cover starting and failed selections. GitOrigin-RevId: cc2744b5137be69f000d70b0c7009b14079f39ae --- codex-rs/codex-mcp/src/runtime.rs | 49 ++++--- codex-rs/core/src/environment_selection.rs | 42 +++++- codex-rs/core/src/mcp.rs | 69 +++++----- codex-rs/core/src/session/environment.rs | 2 - codex-rs/core/src/session/mcp.rs | 54 +++----- codex-rs/core/src/session/mcp_runtime.rs | 54 ++------ codex-rs/core/src/session/mod.rs | 5 +- codex-rs/core/src/session/session.rs | 8 +- codex-rs/core/src/thread_manager_tests.rs | 12 +- codex-rs/core/tests/suite/rmcp_client.rs | 147 +++++++++++++++++++++ 10 files changed, 295 insertions(+), 147 deletions(-) diff --git a/codex-rs/codex-mcp/src/runtime.rs b/codex-rs/codex-mcp/src/runtime.rs index 17aa9f557b..7cc11b1771 100644 --- a/codex-rs/codex-mcp/src/runtime.rs +++ b/codex-rs/codex-mcp/src/runtime.rs @@ -35,6 +35,7 @@ use codex_protocol::mcp::McpResourceOriginCheckpoint; use codex_protocol::models::PermissionProfile; use codex_protocol::protocol::Event; use codex_protocol::protocol::EventMsg; +use codex_protocol::protocol::TurnEnvironmentSelection; use codex_rmcp_client::ElicitationResponse; use codex_rmcp_client::with_http_headers_helper; use codex_utils_path_uri::PathUri; @@ -120,7 +121,8 @@ struct PublishedMcpRuntime { auth_token: Option, plugins_available: bool, ready_selected_capability_roots: Vec, - selected_environments: HashMap>, + environment_selections: Arc<[TurnEnvironmentSelection]>, + ready_environments: HashMap>, cached_binding: Mutex>, } @@ -211,7 +213,8 @@ impl McpRuntime { auth_token: None, plugins_available: false, ready_selected_capability_roots: Vec::new(), - selected_environments: HashMap::new(), + environment_selections: Arc::default(), + ready_environments: HashMap::new(), cached_binding: Mutex::new(None), }), event_stream_cancellation: Mutex::new(EventStreamCancellation { @@ -313,7 +316,8 @@ impl McpRuntime { let auth_token = auth.as_ref().and_then(|auth| auth.get_token().ok()); let plugins_available = input.plugins_available; let ready_selected_capability_roots = input.ready_selected_capability_roots.clone(); - let selected_environments = input.runtime_context.selected_environments.clone(); + let environment_selections = Arc::clone(&input.runtime_context.environment_selections); + let ready_environments = input.runtime_context.ready_environments.clone(); let connections = Arc::new( McpConnectionSet::new( previous, @@ -343,7 +347,8 @@ impl McpRuntime { auth_token, plugins_available, ready_selected_capability_roots, - selected_environments, + environment_selections, + ready_environments, cached_binding: Mutex::new(None), })); let _ = publish.send(true); @@ -528,17 +533,19 @@ impl McpRuntime { self.current.load().ready_selected_capability_roots.clone() } - /// Whether this publication uses the currently ready environment handles. + /// Whether this publication uses the same selections and currently ready environment handles. pub fn current_environments_match( &self, - environments: &HashMap>, + selections: &[TurnEnvironmentSelection], + ready_environments: &HashMap>, ) -> bool { let current = self.current.load(); current.config.is_some() - && current.selected_environments.len() == environments.len() - && environments.iter().all(|(id, environment)| { + && current.environment_selections.as_ref() == selections + && current.ready_environments.len() == ready_environments.len() + && ready_environments.iter().all(|(id, environment)| { current - .selected_environments + .ready_environments .get(id) .is_some_and(|published| Arc::ptr_eq(published, environment)) }) @@ -749,11 +756,12 @@ pub struct SandboxState { /// Runtime context used when resolving per-server MCP environments. /// /// `McpConfig` describes what servers exist. This value carries the canonical -/// environment registry plus the host-local cwd used by local MCP processes. +/// environment registry, captured thread selections, and the host-local cwd used by MCP processes. #[derive(Clone)] pub struct McpRuntimeContext { environment_manager: Arc, - selected_environments: HashMap>, + environment_selections: Arc<[TurnEnvironmentSelection]>, + ready_environments: HashMap>, local_process_cwd: PathBuf, local_http_client: Arc, } @@ -797,18 +805,21 @@ impl McpRuntimeContext { ); Self { environment_manager, - selected_environments: HashMap::new(), + environment_selections: Arc::default(), + ready_environments: HashMap::new(), local_process_cwd, local_http_client, } } - /// Pins the concrete environment handles captured for this thread or model step. + /// Pins all captured selections and the concrete handles for environments already ready. pub fn with_selected_environments( mut self, - selected_environments: HashMap>, + environment_selections: Arc<[TurnEnvironmentSelection]>, + ready_environments: HashMap>, ) -> Self { - self.selected_environments = selected_environments; + self.environment_selections = environment_selections; + self.ready_environments = ready_environments; self } @@ -829,7 +840,7 @@ impl McpRuntimeContext { // HTTP is the one current exception: it can use the ambient HTTP client // even when no local Environment is configured. if let Some(environment) = self - .selected_environments + .ready_environments .get(&config.environment_id) .cloned() .or_else(|| { @@ -973,7 +984,8 @@ mod tests { auth_token: None, plugins_available: false, ready_selected_capability_roots: Vec::new(), - selected_environments: HashMap::new(), + environment_selections: Arc::default(), + ready_environments: HashMap::new(), cached_binding: Mutex::new(None), }); let before = McpRuntime::binding_from_published_runtime( @@ -1045,7 +1057,8 @@ mod tests { auth_token: None, plugins_available: false, ready_selected_capability_roots: Vec::new(), - selected_environments: HashMap::new(), + environment_selections: Arc::default(), + ready_environments: HashMap::new(), cached_binding: Mutex::new(None), }); let first = McpRuntime::binding_from_published_runtime( diff --git a/codex-rs/core/src/environment_selection.rs b/codex-rs/core/src/environment_selection.rs index 1416e359c1..e162210afd 100644 --- a/codex-rs/core/src/environment_selection.rs +++ b/codex-rs/core/src/environment_selection.rs @@ -819,7 +819,9 @@ impl ThreadEnvironments { for environment in selected.iter() { if let EnvironmentConfigState::Failed(error) = &environment.selection.config { environments.push(TurnEnvironmentState::Failed { - selection: environment.selection.clone(), + selection: environment + .config_origin + .into_input_selection(environment.selection.clone()), error: error.clone(), }); continue; @@ -858,6 +860,7 @@ pub(crate) enum TurnEnvironmentState { Starting(StartingTurnEnvironment), /// Unavailable for execution, but still selected when evaluating permissions. Failed { + // Keep the input form so connection failure doesn't hide whether config comes from the thread. selection: TurnEnvironmentSelection, error: String, }, @@ -874,7 +877,7 @@ impl TurnEnvironmentState { if matches!(selection.config, EnvironmentConfigState::Pending) { let Some(config) = environment.installed_config else { return Self::Failed { - selection, + selection: starting.config_origin.into_input_selection(selection), error: "Environment configuration was not supplied.".to_string(), }; }; @@ -903,7 +906,9 @@ impl TurnEnvironmentState { "skipping failed turn environment: {err}" ); Self::Failed { - selection: starting.selection, + selection: starting + .config_origin + .into_input_selection(starting.selection), error: err.to_string(), } } @@ -980,6 +985,17 @@ impl TurnEnvironmentSnapshot { }) } + pub(crate) fn ready_environment_handles(&self) -> HashMap> { + self.turn_environments() + .map(|environment| { + ( + environment.selection.environment_id.clone(), + Arc::clone(&environment.environment), + ) + }) + .collect() + } + /// Maps each captured environment to its exact ready handle, or `None` when it was starting. pub(crate) fn captured_environments(&self) -> HashMap>> { self.turn_environments() @@ -1058,6 +1074,20 @@ impl TurnEnvironmentSnapshot { .collect() } + /// Returns every captured selection, including those still starting or unable to connect. + pub(crate) fn all_selections(&self) -> Vec { + self.environments + .iter() + .map(|environment| match environment { + TurnEnvironmentState::Ready(environment) => environment.selection(), + TurnEnvironmentState::Starting(environment) => environment + .config_origin + .into_input_selection(environment.selection.clone()), + TurnEnvironmentState::Failed { selection, .. } => selection.clone(), + }) + .collect() + } + pub(crate) fn primary_filesystem(&self) -> Option> { self.primary() .map(|environment| environment.environment.get_filesystem()) @@ -1640,6 +1670,10 @@ url = "ws://127.0.0.1:8765" vec![resolved_remote.clone()] ); assert_eq!(starting.to_selections(), vec![local.clone()]); + assert_eq!( + starting.all_selections(), + vec![remote.clone(), local.clone()] + ); assert!(starting.single_local_environment().is_none()); let next_config = EnvironmentConfig { @@ -1753,6 +1787,8 @@ url = "ws://127.0.0.1:8765" // Failed selections must not turn an empty executable snapshot into Full Access. for snapshot in [starting.refresh_readiness(), environments.snapshot().await] { assert!(!snapshot.has_full_access(AskForApproval::Never, &PermissionProfile::Disabled)); + assert_eq!(snapshot.all_selections(), vec![selection.clone()]); + assert!(snapshot.ready_environment_handles().is_empty()); } let selected_root = SelectedCapabilityRoot { id: "failed-root".to_string(), diff --git a/codex-rs/core/src/mcp.rs b/codex-rs/core/src/mcp.rs index 77c469f4da..f1679f333a 100644 --- a/codex-rs/core/src/mcp.rs +++ b/codex-rs/core/src/mcp.rs @@ -2,7 +2,6 @@ use std::collections::HashMap; use std::sync::Arc; use crate::config::Config; -use crate::environment_selection::ThreadEnvironments; use codex_config::DEFAULT_MCP_SERVER_ENVIRONMENT_ID; use codex_config::McpServerConfig; use codex_connectors::ConnectorRuntimeManager; @@ -48,10 +47,38 @@ pub(crate) struct McpRuntimeProjection { pub(crate) enum McpEnvironmentScope<'a> { /// Controller-level operations without an associated thread. HostOnly, - /// Initial thread selections before the live environment store exists. - Initial(&'a [TurnEnvironmentSelection]), - /// Current attachment state for an existing thread. - Live(&'a ThreadEnvironments), + /// Every thread selection captured for this operation, including unavailable configurations. + Selected(&'a [TurnEnvironmentSelection]), +} + +impl<'a> McpEnvironmentScope<'a> { + pub(crate) fn authority_for(&self, environment_id: &str) -> McpEnvironmentAuthority<'a> { + let Self::Selected(selections) = *self else { + return McpEnvironmentAuthority::Unrestricted; + }; + let Some(selection) = selections + .iter() + .find(|selection| selection.environment_id == environment_id) + else { + return if environment_id == DEFAULT_MCP_SERVER_ENVIRONMENT_ID { + McpEnvironmentAuthority::Unrestricted + } else { + McpEnvironmentAuthority::SelectedPluginsOnly + }; + }; + match &selection.config { + EnvironmentConfigState::FromThread => McpEnvironmentAuthority::Unrestricted, + EnvironmentConfigState::Pending | EnvironmentConfigState::Failed(_) => { + McpEnvironmentAuthority::Unavailable + } + EnvironmentConfigState::Ready(config) => config + .mcp_policy + .as_ref() + .map_or(McpEnvironmentAuthority::Unrestricted, |policy| { + McpEnvironmentAuthority::Restricted(policy) + }), + } + } } pub(crate) struct McpThreadIdentity<'a> { @@ -318,38 +345,8 @@ impl McpManager { } => catalog.remove_extension(name, contributor_id, contribution_order), } } - let selections = match environment_scope { - McpEnvironmentScope::HostOnly => None, - McpEnvironmentScope::Initial(selections) => Some(selections.to_vec()), - McpEnvironmentScope::Live(environments) => Some(environments.selections()), - }; let catalog = catalog.build_with_environment_authority(|environment_id| { - let Some(selections) = selections.as_ref() else { - return McpEnvironmentAuthority::Unrestricted; - }; - let Some(selection) = selections - .iter() - .find(|selection| selection.environment_id == environment_id) - else { - return if environment_id == DEFAULT_MCP_SERVER_ENVIRONMENT_ID { - McpEnvironmentAuthority::Unrestricted - } else { - McpEnvironmentAuthority::SelectedPluginsOnly - }; - }; - - match &selection.config { - EnvironmentConfigState::FromThread => McpEnvironmentAuthority::Unrestricted, - EnvironmentConfigState::Pending | EnvironmentConfigState::Failed(_) => { - McpEnvironmentAuthority::Unavailable - } - EnvironmentConfigState::Ready(config) => config - .mcp_policy - .as_ref() - .map_or(McpEnvironmentAuthority::Unrestricted, |policy| { - McpEnvironmentAuthority::Restricted(policy) - }), - } + environment_scope.authority_for(environment_id) }); for conflict in catalog.conflicts() { tracing::warn!( diff --git a/codex-rs/core/src/session/environment.rs b/codex-rs/core/src/session/environment.rs index 2093251643..77e3217ac3 100644 --- a/codex-rs/core/src/session/environment.rs +++ b/codex-rs/core/src/session/environment.rs @@ -157,12 +157,10 @@ impl Session { } } if environments != self.services.turn_environments.selections() { - self.mark_mcp_runtime_dirty(); self.services.turn_environments.update_selections( &environments, &configuration.inferred_environment_config(), ); - self.schedule_mcp_prewarm(); } } self.services.turn_environments.snapshot() diff --git a/codex-rs/core/src/session/mcp.rs b/codex-rs/core/src/session/mcp.rs index 11555def8a..e16f82c121 100644 --- a/codex-rs/core/src/session/mcp.rs +++ b/codex-rs/core/src/session/mcp.rs @@ -108,6 +108,7 @@ impl Session { ) }; let environments = self.services.turn_environments.snapshot().await; + let environment_selections = environments.all_selections(); let selected_capability_roots = self .resolve_selected_capability_roots_for_step(&environments) .await; @@ -131,7 +132,7 @@ impl Session { session_source: &session_source, originator: &originator, disabled_plugin_ids: &disabled_plugin_ids, - environments: McpEnvironmentScope::Live(&self.services.turn_environments), + environments: McpEnvironmentScope::Selected(&environment_selections), }, &ready_selected_capability_roots, executor_capability_discovery.as_deref(), @@ -150,15 +151,8 @@ impl Session { local_process_cwd, ) .with_selected_environments( - environments - .turn_environments() - .map(|environment| { - ( - environment.selection.environment_id.clone(), - Arc::clone(&environment.environment), - ) - }) - .collect(), + environment_selections.into(), + environments.ready_environment_handles(), ); (mcp_config, runtime_context) } @@ -178,22 +172,13 @@ impl Session { return; }; loop { + // Compare and rebuild from current environments, not choices saved for a future turn. let environments = self.services.turn_environments.snapshot().await; - let ready_environments = environments - .turn_environments() - .map(|environment| { - ( - environment.selection.environment_id.clone(), - Arc::clone(&environment.environment), - ) - }) - .collect(); - // Attachment resolution can finish after the owner's configuration callback. - if !self - .services - .mcp_runtime - .current_environments_match(&ready_environments) - { + let environment_selections = environments.all_selections(); + if !self.services.mcp_runtime.current_environments_match( + &environment_selections, + &environments.ready_environment_handles(), + ) { self.mark_mcp_runtime_dirty(); } let auth = self.services.auth_manager.auth_cached(); @@ -213,7 +198,7 @@ impl Session { published: false, }; let auth = self.services.auth_manager.auth().await; - let desired = self.latest_mcp_desired_state(auth).await; + let desired = self.latest_mcp_desired_state(auth, environments).await; let selected_capability_roots = self .resolve_selected_capability_roots_for_step(&desired.environments) .await; @@ -237,7 +222,7 @@ impl Session { session_source: &desired.session_source, originator: &desired.originator, disabled_plugin_ids: &desired.disabled_plugin_ids, - environments: McpEnvironmentScope::Live(&self.services.turn_environments), + environments: McpEnvironmentScope::Selected(&environment_selections), }, &ready_selected_capability_roots, executor_capability_discovery.as_deref(), @@ -280,7 +265,9 @@ impl Session { .await .map_err(|_| anyhow::anyhow!("MCP runtime refresh semaphore closed"))?; let auth = self.services.auth_manager.auth().await; - let desired = self.latest_mcp_desired_state(auth).await; + let environments = self.services.turn_environments.snapshot().await; + let environment_selections = environments.all_selections(); + let desired = self.latest_mcp_desired_state(auth, environments).await; let selected_capability_roots = self .resolve_selected_capability_roots_for_step(&desired.environments) .await; @@ -304,7 +291,7 @@ impl Session { session_source: &desired.session_source, originator: &desired.originator, disabled_plugin_ids: &desired.disabled_plugin_ids, - environments: McpEnvironmentScope::Live(&self.services.turn_environments), + environments: McpEnvironmentScope::Selected(&environment_selections), }, &ready_selected_capability_roots, executor_capability_discovery.as_deref(), @@ -683,11 +670,14 @@ impl Session { .mcp_runtime .current_ready_selected_capability_roots(); let environments = self.services.turn_environments.snapshot().await; + let environment_selections = environments.all_selections(); + let mut desired = self.latest_mcp_desired_state(auth, environments).await; + desired.config = Arc::new(refresh_config.clone()); let executor_capability_discovery = self .executor_capability_discovery_for_step( refresh_config, &ready_selected_capability_roots, - &environments, + &desired.environments, ) .await; let mcp_projection = self @@ -701,14 +691,12 @@ impl Session { session_source: &turn_context.session_source, originator: &turn_context.originator, disabled_plugin_ids: &disabled_plugin_ids, - environments: McpEnvironmentScope::Live(&self.services.turn_environments), + environments: McpEnvironmentScope::Selected(&environment_selections), }, &ready_selected_capability_roots, executor_capability_discovery.as_deref(), ) .await; - let mut desired = self.latest_mcp_desired_state(auth).await; - desired.config = Arc::new(refresh_config.clone()); self.publish_mcp_runtime( &desired, mcp_projection, diff --git a/codex-rs/core/src/session/mcp_runtime.rs b/codex-rs/core/src/session/mcp_runtime.rs index fe05899606..9c0d291cc5 100644 --- a/codex-rs/core/src/session/mcp_runtime.rs +++ b/codex-rs/core/src/session/mcp_runtime.rs @@ -11,13 +11,11 @@ use codex_config::McpServerDisabledReason; use codex_config::McpServerTransportConfig; use codex_mcp::CODEX_APPS_MCP_SERVER_NAME; use codex_mcp::ElicitationReviewerHandle; -use codex_mcp::McpEnvironmentAuthority; use codex_mcp::McpServerRegistration; use codex_mcp::McpServerSource; use codex_mcp::McpStartupPolicy; use codex_mcp::PreparedMcpCall; use codex_protocol::capabilities::SelectedCapabilityRoot; -use codex_protocol::protocol::EnvironmentConfigState; use std::collections::HashSet; pub(super) struct McpDesiredState { @@ -71,17 +69,15 @@ impl Session { pub(super) async fn latest_mcp_desired_state( &self, auth: Option, + environments: TurnEnvironmentSnapshot, ) -> McpDesiredState { - let (session_configuration, disabled_plugin_ids, environments) = { + let (session_configuration, disabled_plugin_ids) = { let state = self.state.lock().await; - // MCP tools must use current environments, not the selection saved for the next turn. ( state.session_configuration.clone(), state.active_disabled_plugin_ids.clone(), - self.services.turn_environments.snapshot(), ) }; - let environments = environments.await; let cwd = environments .primary() .and_then(|environment| environment.cwd().to_abs_path().ok()) @@ -100,8 +96,8 @@ impl Session { config: Arc::new(config), auth, submit_id: self.next_internal_sub_id(), - originator: session_configuration.originator.clone(), - session_source: session_configuration.session_source.clone(), + originator: session_configuration.originator, + session_source: session_configuration.session_source, environments, local_process_cwd, disabled_plugin_ids, @@ -272,35 +268,11 @@ impl Session { } if let Some(catalog) = catalog { - let selections = self.services.turn_environments.selections(); + let selections = environments.all_selections(); + let environment_scope = McpEnvironmentScope::Selected(&selections); projection.config.mcp_server_catalog = catalog.build_with_environment_authority(|environment_id| { - let Some(selection) = selections - .iter() - .find(|selection| selection.environment_id == environment_id) - else { - return if environment_id - == codex_config::DEFAULT_MCP_SERVER_ENVIRONMENT_ID - { - McpEnvironmentAuthority::Unrestricted - } else { - McpEnvironmentAuthority::SelectedPluginsOnly - }; - }; - match &selection.config { - EnvironmentConfigState::FromThread => { - McpEnvironmentAuthority::Unrestricted - } - EnvironmentConfigState::Pending | EnvironmentConfigState::Failed(_) => { - McpEnvironmentAuthority::Unavailable - } - EnvironmentConfigState::Ready(config) => config - .mcp_policy - .as_ref() - .map_or(McpEnvironmentAuthority::Unrestricted, |policy| { - McpEnvironmentAuthority::Restricted(policy) - }), - } + environment_scope.authority_for(environment_id) }); } projection @@ -379,16 +351,8 @@ impl Session { desired.local_process_cwd.clone(), ) .with_selected_environments( - desired - .environments - .turn_environments() - .map(|environment| { - ( - environment.selection.environment_id.clone(), - Arc::clone(&environment.environment), - ) - }) - .collect(), + desired.environments.all_selections().into(), + desired.environments.ready_environment_handles(), ); McpRuntimeInput { startup_policy: if matches!(desired.session_source, SessionSource::SubAgent(_)) { diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 97042ee7f6..60e61dc0ec 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -3862,12 +3862,13 @@ impl Session { } let (mcp, prepared_recommendations) = async { tokio::join!( - self.mcp_runtime_for_step( + // MCP refresh can be large; keep it off the sampling request's stack. + Box::pin(self.mcp_runtime_for_step( turn_context.as_ref(), &selected_capability_roots, required_servers, required_plugins, - ), + )), turn::prepare_tool_recommendations(self.as_ref(), turn_context.as_ref()), ) } diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index 8965f3b4ba..293cf99404 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -1130,7 +1130,7 @@ impl Session { session_source: &mcp_session_source, originator: &mcp_originator, disabled_plugin_ids: &mcp_disabled_plugin_ids, - environments: McpEnvironmentScope::Initial(environment_selections), + environments: McpEnvironmentScope::Selected(environment_selections), }, /*ready_selected_capability_roots*/ &[], /*executor_capability_discovery*/ None, @@ -1800,8 +1800,10 @@ impl Session { mcp_auth_changes.mark_unchanged(); } let latest_auth = sess.services.auth_manager.auth().await; + let resolved_environment_selections = resolved_environments.all_selections(); let mcp_projection = if startup_auth_changed || mcp_auth_changes.has_changed().unwrap_or(false) + || resolved_environment_selections.as_slice() != environment_selections { sess.services .mcp_manager @@ -1813,9 +1815,7 @@ impl Session { session_source: &session_configuration.session_source, originator: &session_configuration.originator, disabled_plugin_ids: &session_configuration.disabled_plugin_ids, - environments: McpEnvironmentScope::Live( - &sess.services.turn_environments, - ), + environments: McpEnvironmentScope::Selected(&resolved_environment_selections), }, /*ready_selected_capability_roots*/ &[], /*executor_capability_discovery*/ None, diff --git a/codex-rs/core/src/thread_manager_tests.rs b/codex-rs/core/src/thread_manager_tests.rs index d34e1439b7..9b892726ed 100644 --- a/codex-rs/core/src/thread_manager_tests.rs +++ b/codex-rs/core/src/thread_manager_tests.rs @@ -1696,7 +1696,9 @@ async fn start_thread_seeds_extension_data_for_mcp_and_lifecycle_contributors() session_source: &SessionSource::Exec, originator: &first_originator, disabled_plugin_ids: &[], - environments: McpEnvironmentScope::Live(&first_session.services.turn_environments), + environments: McpEnvironmentScope::Selected( + &first_session.services.turn_environments.selections(), + ), }, /*ready_selected_capability_roots*/ &[], /*executor_capability_discovery*/ None, @@ -1715,7 +1717,9 @@ async fn start_thread_seeds_extension_data_for_mcp_and_lifecycle_contributors() session_source: &second_session_source, originator: &second_originator, disabled_plugin_ids: &[], - environments: McpEnvironmentScope::Live(&second_session.services.turn_environments), + environments: McpEnvironmentScope::Selected( + &second_session.services.turn_environments.selections(), + ), }, /*ready_selected_capability_roots*/ &[], /*executor_capability_discovery*/ None, @@ -1787,8 +1791,8 @@ async fn start_thread_seeds_extension_data_for_mcp_and_lifecycle_contributors() session_source: &SessionSource::Exec, originator: &first_originator, disabled_plugin_ids: &disabled_plugin_ids, - environments: McpEnvironmentScope::Live( - &first_session.services.turn_environments, + environments: McpEnvironmentScope::Selected( + &first_session.services.turn_environments.selections(), ), }, /*ready_selected_capability_roots*/ &[], diff --git a/codex-rs/core/tests/suite/rmcp_client.rs b/codex-rs/core/tests/suite/rmcp_client.rs index cf9bfd27d7..2bb9eee369 100644 --- a/codex-rs/core/tests/suite/rmcp_client.rs +++ b/codex-rs/core/tests/suite/rmcp_client.rs @@ -83,6 +83,8 @@ use codex_protocol::protocol::TurnEnvironmentSelection; use codex_protocol::protocol::TurnEnvironmentSelections; use codex_protocol::protocol::TurnSettingsUpdate; use codex_protocol::protocol::TurnSettingsUpdateOutcome; +use codex_protocol::request_user_input::RequestUserInputAnswer; +use codex_protocol::request_user_input::RequestUserInputResponse; use codex_protocol::turn_input::TurnInput; use codex_protocol::user_input::UserInput; use codex_utils_cargo_bin::cargo_bin; @@ -885,6 +887,151 @@ async fn environment_mcp_policy_filters_runtime_config_and_model_tools( Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn future_environment_mcp_policy_applies_on_the_next_turn() -> anyhow::Result<()> { + skip_if_wine_exec!( + Ok(()), + "requires a Windows test_stdio_server in the Wine-exec environment" + ); + skip_if_no_network!(Ok(())); + + let server = responses::start_mock_server().await; + let response = responses::mount_sse_sequence( + &server, + vec![ + responses::sse(vec![ + responses::ev_function_call( + "pause", + "request_user_input", + &json!({ + "questions": [{ + "id": "continue", + "header": "Continue", + "question": "Continue?", + "options": [ + {"label": "Yes (Recommended)", "description": "Continue."}, + {"label": "No", "description": "Stop."} + ] + }] + }) + .to_string(), + ), + responses::ev_completed("paused"), + ]), + responses::sse(vec![responses::ev_completed("active-done")]), + responses::sse(vec![responses::ev_completed("next-done")]), + ], + ) + .await; + let command = remote_aware_stdio_server_bin()?; + let allowed_command = command.clone(); + let fixture = test_codex() + .with_model_info_override("gpt-5.4", |model| model.supports_search_tool = false) + .with_config(move |config| { + config + .features + .enable(Feature::DefaultModeRequestUserInput) + .unwrap(); + insert_mcp_server( + config, + "environment_policy", + stdio_transport(command, /*env*/ None, Vec::new()), + TestMcpServerOptions { + environment_id: remote_aware_environment_id(), + ..Default::default() + }, + ); + }) + .build_with_auto_env(&server) + .await?; + wait_for_mcp_server(&fixture.codex, "environment_policy").await?; + let selection = fixture.codex.environment_selections().await.remove(0); + let mut config = EnvironmentConfig { + allow_login_shell: fixture.config.permissions.allow_login_shell, + workspace_roots: selection.workspace_roots.clone(), + permission_profile: PermissionProfileSnapshot::legacy( + fixture.config.permissions.permission_profile().clone(), + ), + shell_environment_policy: Default::default(), + windows_sandbox_level: WindowsSandboxLevel::from_config(&fixture.config), + windows_sandbox_type: fixture.config.permissions.windows_sandbox_type, + windows_sandbox_private_desktop: fixture.config.permissions.windows_sandbox_private_desktop, + use_legacy_landlock: fixture.config.features.use_legacy_landlock(), + exec_policy: None, + mcp_policy: Some(EnvironmentMcpPolicy { + servers: Some(BTreeMap::from([( + "environment_policy".to_string(), + McpServerRequirement::Identity { + identity: McpServerIdentity::Command { + command: allowed_command, + }, + }, + )])), + plugins: None, + }), + network_policy: None, + selected_capability_roots: Vec::new(), + }; + let settings = |config| ThreadSettingsOverrides { + environments: Some(TurnEnvironmentSelections::new( + fixture.config.cwd.clone(), + vec![TurnEnvironmentSelection { + config: EnvironmentConfigState::Ready(config), + ..selection.clone() + }], + )), + ..Default::default() + }; + submit_thread_settings(&fixture.codex, settings(config.clone())).await?; + fixture + .codex + .start_or_steer_turn(TurnInputRequest::user_input(vec![UserInput::Text { + text: "pause before continuing".into(), + text_elements: Vec::new(), + }])) + .await?; + let EventMsg::RequestUserInput(request) = wait_for_event(&fixture.codex, |event| { + matches!(event, EventMsg::RequestUserInput(_)) + }) + .await + else { + unreachable!("wait_for_event should return the request-user-input event") + }; + + config.mcp_policy.as_mut().unwrap().servers = Some(BTreeMap::new()); + submit_thread_settings(&fixture.codex, settings(config)).await?; + fixture + .codex + .submit(Op::UserInputAnswer { + id: request.turn_id, + response: RequestUserInputResponse { + answers: HashMap::from([( + "continue".to_string(), + RequestUserInputAnswer { + answers: vec!["Yes (Recommended)".to_string()], + }, + )]), + }, + }) + .await?; + wait_for_event(&fixture.codex, |event| { + matches!(event, EventMsg::TurnComplete(_)) + }) + .await; + fixture.submit_text_turn("start the next turn").await?; + + let tool_is_visible = response + .requests() + .iter() + .map(|request| { + responses::namespace_child_tool(&request.body_json(), "mcp__environment_policy", "echo") + .is_some() + }) + .collect::>(); + assert_eq!(tool_is_visible, vec![true, true, false]); + Ok(()) +} + #[test_case("rmcp", "mcp__rmcp"; "simple name")] #[test_case("npm:@scope/package.name", "mcp__npm__scope_package_name"; "npm name")] #[tokio::test(flavor = "multi_thread", worker_threads = 1)]