From aab4f9614877c45e651229446b52358e1de6368c Mon Sep 17 00:00:00 2001 From: teddywyly-oai Date: Fri, 9 Oct 2026 18:55:59 +0000 Subject: [PATCH] Refresh persisted capability roots from owner-provided configuration (#52676) ## Why Resumed threads can retain stale capability root locations when the current environment supplies an explicit owner configuration. Matching roots from that configuration previously could not refresh persisted thread roots. ## What changed Refresh persisted roots from the current attachment when both the root ID and environment ID match, including roots selected by owner-provided configuration. Preserve persisted locations when no matching root is supplied, and continue using the owner's selected roots rather than the executor's shared roots. Apply the same merge behavior to capability root inspection and per-step resolution. ## Testing Extend the remote environment regression test to cover owner-provided configurations with and without refreshed roots, including an executor root location that differs from the owner's selection. GitOrigin-RevId: ad3a639924b348d6d5337ccc64dc8c08234276e3 --- codex-rs/core/src/environment_selection.rs | 29 ++++++++------------- codex-rs/core/src/session/mcp.rs | 9 +++---- codex-rs/core/tests/suite/remote_env.rs | 30 +++++++++++++++++----- 3 files changed, 37 insertions(+), 31 deletions(-) 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