diff --git a/codex-rs/core/src/environment_selection.rs b/codex-rs/core/src/environment_selection.rs index 39a598bff9..0e7718bade 100644 --- a/codex-rs/core/src/environment_selection.rs +++ b/codex-rs/core/src/environment_selection.rs @@ -54,6 +54,59 @@ impl EnvironmentConfigOrigin { } selection } + + pub(crate) fn selected_capability_roots( + self, + environment: &Environment, + config: &EnvironmentConfig, + ) -> Vec { + match self { + Self::Thread => environment.selected_capability_roots(), + Self::Owner => config.selected_capability_roots.clone(), + } + } +} + +/// Combines persisted thread roots with attachment roots in selection order. +/// +/// Thread-owned attachments still rely on roots reported by legacy executors. A matching live root +/// refreshes the persisted location while explicit owner configuration retains the existing +/// thread-first collision behavior. +pub(crate) fn combine_selected_capability_roots( + thread_roots: &[SelectedCapabilityRoot], + sources: impl IntoIterator)>, +) -> Vec { + let attachment_roots = sources + .into_iter() + .flat_map(|(config_origin, roots)| roots.into_iter().map(move |root| (config_origin, root))) + .collect::>(); + let mut combined_roots = thread_roots + .iter() + .map(|thread_root| { + let CapabilityRootLocation::Environment { + environment_id: thread_environment_id, + .. + } = &thread_root.location; + attachment_roots + .iter() + .find_map(|(origin, attachment_root)| { + if *origin != EnvironmentConfigOrigin::Thread + || attachment_root.id != thread_root.id + { + return None; + } + let CapabilityRootLocation::Environment { + environment_id: attachment_environment_id, + .. + } = &attachment_root.location; + (attachment_environment_id == thread_environment_id).then_some(attachment_root) + }) + .unwrap_or(thread_root) + .clone() + }) + .collect::>(); + combined_roots.extend(attachment_roots.into_iter().map(|(_, root)| root)); + combined_roots } pub(crate) fn default_thread_environment_selections( @@ -410,13 +463,20 @@ impl ThreadEnvironments { thread_selected_capability_roots: &[SelectedCapabilityRoot], ) -> SelectedCapabilityRootsStatus { let environments = self.environments.load(); - let mut selected_capability_roots = thread_selected_capability_roots.to_vec(); - for environment in environments.iter() { - let EnvironmentConfigState::Ready(config) = &environment.selection.config else { - continue; - }; - selected_capability_roots.extend(config.selected_capability_roots.iter().cloned()); - } + let mut selected_capability_roots = combine_selected_capability_roots( + thread_selected_capability_roots, + environments.iter().filter_map(|environment| { + let EnvironmentConfigState::Ready(config) = &environment.selection.config else { + return None; + }; + Some(( + environment.config_origin, + environment + .config_origin + .selected_capability_roots(&environment.environment, config), + )) + }), + ); let mut seen_root_ids = HashSet::with_capacity(selected_capability_roots.len()); selected_capability_roots.retain(|root| seen_root_ids.insert(root.id.clone())); diff --git a/codex-rs/core/src/session/mcp.rs b/codex-rs/core/src/session/mcp.rs index 6b25214762..a46655e0d3 100644 --- a/codex-rs/core/src/session/mcp.rs +++ b/codex-rs/core/src/session/mcp.rs @@ -1,5 +1,6 @@ use super::mcp_refresh::McpRefreshInvalidationGuard; use super::*; +use crate::environment_selection::combine_selected_capability_roots; use crate::tools::sandboxing::executor_windows_sandbox_level; use codex_exec_server::ExecutorCapabilityDiscoveryCache; use codex_exec_server::ExecutorCapabilityDiscoverySnapshot; @@ -438,18 +439,18 @@ impl Session { 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() - .cloned() - .chain( - environments - .turn_environments() - .flat_map(|environment| environment.config().selected_capability_roots.clone()), - ) - .enumerate() - { + let combined_roots = combine_selected_capability_roots( + &self.services.selected_capability_roots, + environments.turn_environments().map(|environment| { + ( + environment.config_origin, + environment + .config_origin + .selected_capability_roots(&environment.environment, environment.config()), + ) + }), + ); + for (index, root) in combined_roots.into_iter().enumerate() { if let Some(kept_location) = root_locations_by_id.get(&root.id) { if kept_location != &root.location { tracing::warn!( diff --git a/codex-rs/core/tests/suite/remote_env.rs b/codex-rs/core/tests/suite/remote_env.rs index 43819b1e00..582de2f9d4 100644 --- a/codex-rs/core/tests/suite/remote_env.rs +++ b/codex-rs/core/tests/suite/remote_env.rs @@ -118,6 +118,7 @@ use std::fs; use std::path::PathBuf; use std::process::Command; use std::sync::Arc; +use std::sync::Mutex; use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; use std::time::Duration; @@ -158,13 +159,22 @@ impl ThreadLifecycleContributor for WaitForEnvironmentTestExtension { } } -struct ReadyCapabilityRootsTestExtension; +#[derive(Default)] +struct ReadyCapabilityRootsTestExtension { + observed_roots: Option>>>>, +} impl ContextContributor for ReadyCapabilityRootsTestExtension { fn contribute_world_state<'a>( &'a self, input: WorldStateContributionInput<'a>, ) -> ExtensionFuture<'a, Vec> { + if let Some(observed_roots) = &self.observed_roots { + observed_roots + .lock() + .expect("observed capability roots should not be poisoned") + .push(input.ready_selected_capability_roots.to_vec()); + } let root_ids = input .ready_selected_capability_roots .iter() @@ -1218,7 +1228,7 @@ async fn wait_for_response_request_count(response_mock: &ResponseMock, expected_ async fn shared_executor_keeps_ready_capability_roots_scoped_to_each_attachment() -> Result<()> { let server = start_mock_server().await; let mut extensions = ExtensionRegistryBuilder::new(); - extensions.prompt_contributor(Arc::new(ReadyCapabilityRootsTestExtension)); + extensions.prompt_contributor(Arc::new(ReadyCapabilityRootsTestExtension::default())); let mut builder = test_codex() .with_extensions(Arc::new(extensions.build())) .with_config(|config| { @@ -1484,7 +1494,7 @@ async fn pending_attachment_installs_configuration_before_waiting_turn_resumes() let server = start_mock_server().await; let mut extensions = ExtensionRegistryBuilder::new(); extensions.thread_lifecycle_contributor(Arc::new(WaitForEnvironmentTestExtension)); - extensions.prompt_contributor(Arc::new(ReadyCapabilityRootsTestExtension)); + extensions.prompt_contributor(Arc::new(ReadyCapabilityRootsTestExtension::default())); let mut builder = test_codex() .with_extensions(Arc::new(extensions.build())) .with_config(|config| { @@ -1718,9 +1728,12 @@ async fn pending_attachment_installs_configuration_before_waiting_turn_resumes() Ok(()) } +#[test_case(true; "uses refreshed executor root")] +#[test_case(false; "preserves persisted root when executor reports none")] #[tokio::test(flavor = "multi_thread", worker_threads = 4)] -async fn ready_before_selection_exposes_remote_tools_and_capability_context_after_wait() --> Result<()> { +async fn ready_before_selection_resolves_resumed_thread_capability_root_after_wait( + executor_reports_refreshed_root: bool, +) -> Result<()> { const WAIT_CALL_ID: &str = "wait-ready-before-selection"; let rendezvous = TcpListener::bind("127.0.0.1:0").await?; @@ -1804,9 +1817,12 @@ async fn ready_before_selection_exposes_remote_tools_and_capability_context_afte ], ) .await; + let observed_roots = Arc::new(Mutex::new(Vec::new())); let mut extensions = ExtensionRegistryBuilder::new(); extensions.thread_lifecycle_contributor(Arc::new(WaitForEnvironmentTestExtension)); - extensions.prompt_contributor(Arc::new(ReadyCapabilityRootsTestExtension)); + extensions.prompt_contributor(Arc::new(ReadyCapabilityRootsTestExtension { + observed_roots: Some(Arc::clone(&observed_roots)), + })); let mut builder = test_codex() .with_extensions(Arc::new(extensions.build())) .with_config(|config| { @@ -1816,20 +1832,36 @@ async fn ready_before_selection_exposes_remote_tools_and_capability_context_afte assert!(config.features.enable(Feature::UnifiedExec).is_ok()); }); let test = builder.build(&server).await?; - let ready_root = SelectedCapabilityRoot { + let refreshed_root = SelectedCapabilityRoot { id: "ready-first-root".to_string(), location: CapabilityRootLocation::Environment { environment_id: REMOTE_ENVIRONMENT_ID.to_string(), path: PathUri::parse("file:///ready-first-root")?, }, }; + let stale_root = SelectedCapabilityRoot { + id: refreshed_root.id.clone(), + location: CapabilityRootLocation::Environment { + environment_id: REMOTE_ENVIRONMENT_ID.to_string(), + path: PathUri::parse("file:///stale-ready-first-root")?, + }, + }; + let expected_root = if executor_reports_refreshed_root { + refreshed_root.clone() + } else { + stale_root.clone() + }; let environment = test .thread_manager .environment_manager() .report_environment_provisioning_status( REMOTE_ENVIRONMENT_ID.to_string(), Ok(EnvironmentReadyInfo { - selected_capability_roots: Vec::new(), + selected_capability_roots: if executor_reports_refreshed_root { + vec![refreshed_root.clone()] + } else { + Vec::new() + }, }), Arc::new(ReadyNoiseConnectProvider { websocket_url: format!("{rendezvous_url}/relay?role=harness"), @@ -1867,26 +1899,45 @@ async fn ready_before_selection_exposes_remote_tools_and_capability_context_afte anyhow::Ok(()) }); - test.submit_turn_with_environments( - "use the ready environment", - Some(vec![TurnEnvironmentSelection { - environment_id: REMOTE_ENVIRONMENT_ID.to_string(), - cwd: PathUri::from_abs_path(&test.config.cwd), - workspace_roots: vec![PathUri::from_abs_path(&test.config.cwd)], - config: EnvironmentConfigState::Ready(EnvironmentConfig { - allow_login_shell: true, - permission_profile: PermissionProfileSnapshot::legacy( - test.config.permissions.permission_profile().clone(), - ), - shell_environment_policy: Default::default(), - exec_policy: None, - mcp_policy: None, - network_policy: None, - selected_capability_roots: vec![ready_root], - }), - }]), - ) - .await?; + let selection = TurnEnvironmentSelection { + environment_id: REMOTE_ENVIRONMENT_ID.to_string(), + cwd: PathUri::from_abs_path(&test.config.cwd), + workspace_roots: vec![PathUri::from_abs_path(&test.config.cwd)], + config: EnvironmentConfigState::FromThread, + }; + let mut thread_extension_init = ExtensionDataInit::new(); + thread_extension_init.insert(vec![stale_root]); + let resumed = test + .thread_manager + .start_thread(StartThreadOptions { + environments: Some(vec![selection.clone()]), + thread_extension_init, + ..StartThreadOptions::new(test.config.clone()) + }) + .await?; + resumed + .thread + .start_or_steer_turn(TurnInputRequest::user_input(vec![UserInput::Text { + text: "use the ready environment".to_string(), + text_elements: Vec::new(), + }])) + .await?; + wait_for_event(&resumed.thread, |event| { + matches!(event, EventMsg::TurnComplete(_)) + }) + .await; + + assert_eq!( + resumed.thread.environment_selections().await, + vec![selection] + ); + assert_eq!( + resumed + .thread + .inspect_selected_capability_roots() + .ready_roots, + vec![expected_root.clone()] + ); let requests = response_mock.requests(); assert_eq!(requests.len(), 2); @@ -1930,6 +1981,14 @@ async fn ready_before_selection_exposes_remote_tools_and_capability_context_afte .iter() .any(|text| text.contains("ready-first-root")) ); + let observed_selected_roots = observed_roots + .lock() + .expect("observed capability roots should not be poisoned") + .iter() + .rfind(|roots| roots.iter().any(|root| root.id == expected_root.id)) + .cloned() + .context("selected capability root should reach world state")?; + assert_eq!(observed_selected_roots, vec![expected_root]); relay.abort(); remote_environment.abort();