diff --git a/codex-rs/core/src/environment_selection.rs b/codex-rs/core/src/environment_selection.rs index 29005bc6bf..f85efc1e79 100644 --- a/codex-rs/core/src/environment_selection.rs +++ b/codex-rs/core/src/environment_selection.rs @@ -84,17 +84,13 @@ impl EnvironmentConfigOrigin { /// 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. +/// A matching root from the current attachment refreshes its persisted location. Root IDs from +/// other environments and persisted roots absent from the attachment keep their existing location. pub(crate) fn combine_selected_capability_roots( thread_roots: &[SelectedCapabilityRoot], - sources: impl IntoIterator)>, + 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 attachment_roots = sources.into_iter().flatten().collect::>(); let mut combined_roots = thread_roots .iter() .map(|thread_root| { @@ -104,23 +100,21 @@ pub(crate) fn combine_selected_capability_roots( } = &thread_root.location; attachment_roots .iter() - .find_map(|(origin, attachment_root)| { - if *origin != EnvironmentConfigOrigin::Thread - || attachment_root.id != thread_root.id - { - return None; + .find(|attachment_root| { + if attachment_root.id != thread_root.id { + return false; } let CapabilityRootLocation::Environment { environment_id: attachment_environment_id, .. } = &attachment_root.location; - (attachment_environment_id == thread_environment_id).then_some(attachment_root) + attachment_environment_id == thread_environment_id }) .unwrap_or(thread_root) .clone() }) .collect::>(); - combined_roots.extend(attachment_roots.into_iter().map(|(_, root)| root)); + combined_roots.extend(attachment_roots); combined_roots } @@ -715,12 +709,11 @@ impl ThreadEnvironments { let EnvironmentConfigState::Ready(config) = &environment.selection.config else { return None; }; - Some(( - environment.config_origin, + Some( environment .config_origin .selected_capability_roots(&environment.environment, config), - )) + ) }), ); let mut seen_root_ids = HashSet::with_capacity(selected_capability_roots.len()); diff --git a/codex-rs/core/src/session/mcp.rs b/codex-rs/core/src/session/mcp.rs index 67e99d1c65..d2599a4d5c 100644 --- a/codex-rs/core/src/session/mcp.rs +++ b/codex-rs/core/src/session/mcp.rs @@ -512,12 +512,9 @@ impl Session { let combined_roots = combine_selected_capability_roots( &thread_roots, environments.turn_environments().map(|environment| { - ( - environment.config_origin, - environment - .config_origin - .selected_capability_roots(&environment.environment, environment.config()), - ) + environment + .config_origin + .selected_capability_roots(&environment.environment, environment.config()) }), ); for (index, root) in combined_roots.into_iter().enumerate() { diff --git a/codex-rs/core/tests/suite/remote_env.rs b/codex-rs/core/tests/suite/remote_env.rs index c658fd8faa..23acf8c18d 100644 --- a/codex-rs/core/tests/suite/remote_env.rs +++ b/codex-rs/core/tests/suite/remote_env.rs @@ -2622,11 +2622,14 @@ async fn active_environment_update_wakes_the_old_wait_with_the_new_selection() - Ok(()) } -#[test_case(true; "uses refreshed executor root")] -#[test_case(false; "preserves persisted root when executor reports none")] +#[test_case(true, false; "uses refreshed executor root")] +#[test_case(false, false; "preserves persisted root when executor reports none")] +#[test_case(true, true; "uses refreshed owner root")] +#[test_case(false, true; "preserves persisted root when owner reports none")] #[tokio::test(flavor = "multi_thread", worker_threads = 4)] async fn ready_before_selection_resolves_resumed_thread_capability_root_after_wait( - executor_reports_refreshed_root: bool, + reports_refreshed_root: bool, + owner_provided: bool, ) -> Result<()> { const WAIT_CALL_ID: &str = "wait-ready-before-selection"; @@ -2738,11 +2741,17 @@ async fn ready_before_selection_resolves_resumed_thread_capability_root_after_wa path: PathUri::parse("file:///stale-ready-first-root")?, }, }; - let expected_root = if executor_reports_refreshed_root { + let expected_root = if reports_refreshed_root { refreshed_root.clone() } else { stale_root.clone() }; + let mut executor_root = refreshed_root.clone(); + if owner_provided { + // Owner configuration must not inherit the executor's shared root location. + let CapabilityRootLocation::Environment { path, .. } = &mut executor_root.location; + *path = PathUri::parse("file:///shared-root-not-selected-by-owner")?; + } let provider = Arc::new(OfflineThenReadyNoiseConnectProvider { websocket_url: format!("{rendezvous_url}/relay?role=harness"), executor_public_key, @@ -2754,8 +2763,8 @@ async fn ready_before_selection_resolves_resumed_thread_capability_root_after_wa .report_environment_provisioning_status( REMOTE_ENVIRONMENT_ID.to_string(), Ok(EnvironmentReadyInfo { - selected_capability_roots: if executor_reports_refreshed_root { - vec![refreshed_root.clone()] + selected_capability_roots: if reports_refreshed_root || owner_provided { + vec![executor_root] } else { Vec::new() }, @@ -2793,7 +2802,7 @@ async fn ready_before_selection_resolves_resumed_thread_capability_root_after_wa anyhow::Ok(()) }); - let selection = TurnEnvironmentSelection::new( + let mut selection = TurnEnvironmentSelection::new( TurnEnvironmentRequest { environment_id: REMOTE_ENVIRONMENT_ID.to_string(), cwd: PathUri::from_abs_path(&test.config.cwd), @@ -2802,6 +2811,13 @@ async fn ready_before_selection_resolves_resumed_thread_capability_root_after_wa }, std::slice::from_ref(&stale_root), ); + if owner_provided { + let mut config = environment_config_for_selection(&test.config, &selection); + if reports_refreshed_root { + config.selected_capability_roots = vec![refreshed_root]; + } + selection.config = EnvironmentConfigState::Ready(config); + } let mut thread_extension_init = ExtensionDataInit::new(); thread_extension_init.insert(vec![stale_root]); let resumed = test