diff --git a/codex-rs/core/src/session/mcp.rs b/codex-rs/core/src/session/mcp.rs index 1f8c752161..4a2e02f426 100644 --- a/codex-rs/core/src/session/mcp.rs +++ b/codex-rs/core/src/session/mcp.rs @@ -194,31 +194,45 @@ impl Session { &self, environments: &TurnEnvironmentSnapshot, ) -> Vec { - let mut seen_root_ids = self + let thread_root_count = self.services.selected_capability_roots.len(); + let mut root_locations_by_id = HashMap::new(); + let mut selected_capability_roots = Vec::new(); + let mut ready_environment_root_count = 0; + for (index, root) in self .services .selected_capability_roots .iter() - .map(|root| root.id.as_str()) - .collect::>(); - let mut selected_capability_roots = self.services.selected_capability_roots.clone(); - let mut ready_environment_root_count = 0; - for root in environments - .turn_environments - .iter() - .flat_map(|environment| environment.environment.selected_capability_roots()) + .chain( + environments + .turn_environments + .iter() + .flat_map(|environment| environment.environment.selected_capability_roots()), + ) + .enumerate() { - if !seen_root_ids.insert(root.id.as_str()) { + if let Some(kept_location) = root_locations_by_id.get(&root.id) { + if kept_location != &root.location { + tracing::warn!( + root_id = root.id, + ?kept_location, + ignored_location = ?root.location, + "ignoring selected capability root with conflicting location" + ); + } continue; } - if ready_environment_root_count == MAX_SELECTED_CAPABILITY_ROOTS { - tracing::warn!( - max_root_count = MAX_SELECTED_CAPABILITY_ROOTS, - "ignoring excess selected capability roots from ready environments" - ); - break; + if index >= thread_root_count { + if ready_environment_root_count == MAX_SELECTED_CAPABILITY_ROOTS { + tracing::warn!( + max_root_count = MAX_SELECTED_CAPABILITY_ROOTS, + "ignoring excess selected capability roots from ready environments" + ); + break; + } + ready_environment_root_count += 1; } + root_locations_by_id.insert(root.id.clone(), root.location.clone()); selected_capability_roots.push(root.clone()); - ready_environment_root_count += 1; } self.services .turn_environments diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index fef993626d..a28232229b 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -7688,27 +7688,24 @@ async fn refresh_mcp_servers_keeps_the_previous_runtime_alive() { ); } +struct PendingNoiseConnectProvider; + +impl codex_exec_server::NoiseRendezvousConnectProvider for PendingNoiseConnectProvider { + fn connect_bundle( + &self, + _: codex_exec_server::NoiseChannelPublicKey, + ) -> futures::future::BoxFuture< + '_, + Result, + > { + Box::pin(futures::future::pending()) + } +} + #[tokio::test] async fn deferred_environment_roots_refresh_plugin_availability() { struct ReadyPluginContributor; - struct PendingNoiseConnectProvider; - - impl codex_exec_server::NoiseRendezvousConnectProvider for PendingNoiseConnectProvider { - fn connect_bundle( - &self, - _: codex_exec_server::NoiseChannelPublicKey, - ) -> futures::future::BoxFuture< - '_, - Result< - codex_exec_server::NoiseRendezvousConnectBundle, - codex_exec_server::ExecServerError, - >, - > { - Box::pin(futures::future::pending()) - } - } - impl codex_extension_api::McpServerContributor for ReadyPluginContributor { fn id(&self) -> &'static str { "ready_plugin_test" @@ -7821,6 +7818,83 @@ async fn deferred_environment_roots_refresh_plugin_availability() { assert!(Arc::ptr_eq(&old_manager, &new_runtime.manager_arc())); } +#[tokio::test] +#[tracing_test::traced_test] +async fn conflicting_ready_environment_root_ids_keep_first_location() { + let (session, turn_context) = make_session_and_context().await; + let environment_manager = session.services.turn_environments.environment_manager(); + let selected_root = + |environment_id: &str, path: &str| codex_protocol::capabilities::SelectedCapabilityRoot { + id: "shared-root".to_string(), + location: codex_protocol::capabilities::CapabilityRootLocation::Environment { + environment_id: environment_id.to_string(), + path: PathUri::parse(path).expect("root URI"), + }, + }; + let selected_roots = [ + selected_root("executor-a", "file:///plugins/a"), + selected_root("executor-b", "file:///plugins/b"), + ]; + let local_environment = turn_context + .environments + .primary() + .expect("ready local environment"); + let mut turn_environments = Vec::new(); + for selected_root in &selected_roots { + let codex_protocol::capabilities::CapabilityRootLocation::Environment { + environment_id, + .. + } = &selected_root.location; + let registration = environment_manager + .register_deferred_noise_environment( + environment_id.clone(), + Arc::new(PendingNoiseConnectProvider), + ) + .expect("register deferred environment"); + let environment = environment_manager + .get_environment(environment_id) + .expect("deferred environment"); + registration + .complete(Ok(codex_exec_server::EnvironmentReadyInfo { + selected_capability_roots: vec![selected_root.clone()], + })) + .expect("complete deferred environment"); + turn_environments.push(TurnEnvironment::new( + environment_id.clone(), + environment, + local_environment.cwd().clone(), + local_environment.workspace_roots().to_vec(), + local_environment.shell.clone(), + )); + } + let environments = TurnEnvironmentSnapshot { + turn_environments, + starting: Vec::new(), + }; + + let resolved_roots = session + .resolve_selected_capability_roots_for_step(&environments) + .await; + + assert_eq!( + resolved_roots + .iter() + .map(|root| root.selected_root().clone()) + .collect::>(), + vec![selected_roots[0].clone()] + ); + logs_assert(|lines: &[&str]| { + lines + .iter() + .find(|line| { + line.contains("ignoring selected capability root with conflicting location") + && line.contains("root_id=\"shared-root\"") + }) + .map(|_| Ok(())) + .unwrap_or_else(|| Err("expected conflicting root location warning".to_string())) + }); +} + #[tokio::test] async fn built_tools_uses_the_step_mcp_runtime() -> anyhow::Result<()> { let (session, turn_context) = make_session_and_context().await;