Refresh resumed thread capability roots from executors (#39746)

## Why

Resumed threads can carry a persisted capability-root location that no longer
matches the location reported by the attached executor.

## What changed

- Refresh matching thread-owned capability roots from the live environment while
  preserving persisted roots when the executor reports none.
- Use the same merge behavior for capability-root inspection and per-step MCP
  resolution, without changing owner-configured root precedence.

## Testing

Extend the remote-environment resume test to cover both an executor-provided
replacement root and the persisted-root fallback.

GitOrigin-RevId: 0c66195b2406f431a2e156e00d5225a8d22d18b2

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<SelectedCapabilityRoot> {
+        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<Item = (EnvironmentConfigOrigin, Vec<SelectedCapabilityRoot>)>,
+) -> Vec<SelectedCapabilityRoot> {
+    let attachment_roots = sources
+        .into_iter()
+        .flat_map(|(config_origin, roots)| roots.into_iter().map(move |root| (config_origin, root)))
+        .collect::<Vec<_>>();
+    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::<Vec<_>>();
+    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();
+        let mut selected_capability_roots = combine_selected_capability_roots(
-        for environment in environments.iter() {
+            thread_selected_capability_roots,
-            let EnvironmentConfigState::Ready(config) = &environment.selection.config else {
+            environments.iter().filter_map(|environment| {
-                continue;
+                let EnvironmentConfigState::Ready(config) = &environment.selection.config else {
-            };
+                    return None;
-            selected_capability_roots.extend(config.selected_capability_roots.iter().cloned());
+                };
-        }
+                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()));
 

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
+        let combined_roots = combine_selected_capability_roots(
-            .services
+            &self.services.selected_capability_roots,
-            .selected_capability_roots
+            environments.turn_environments().map(|environment| {
-            .iter()
+                (
-            .cloned()
+                    environment.config_origin,
-            .chain(
+                    environment
-                environments
+                        .config_origin
-                    .turn_environments()
+                        .selected_capability_roots(&environment.environment, environment.config()),
-                    .flat_map(|environment| environment.config().selected_capability_roots.clone()),
+                )
-            )
+            }),
-            .enumerate()
+        );
-        {
+        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!(

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<Config> for WaitForEnvironmentTestExtension {
     }
 }
 
-struct ReadyCapabilityRootsTestExtension;
+#[derive(Default)]
+struct ReadyCapabilityRootsTestExtension {
+    observed_roots: Option<Arc<Mutex<Vec<Vec<SelectedCapabilityRoot>>>>>,
+}
 
 impl ContextContributor for ReadyCapabilityRootsTestExtension {
     fn contribute_world_state<'a>(
         &'a self,
         input: WorldStateContributionInput<'a>,
     ) -> ExtensionFuture<'a, Vec<WorldStateSectionContribution>> {
+        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()
+async fn ready_before_selection_resolves_resumed_thread_capability_root_after_wait(
--> Result<()> {
+    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(
+    let selection = TurnEnvironmentSelection {
-        "use the ready environment",
+        environment_id: REMOTE_ENVIRONMENT_ID.to_string(),
-        Some(vec![TurnEnvironmentSelection {
+        cwd: PathUri::from_abs_path(&test.config.cwd),
-            environment_id: REMOTE_ENVIRONMENT_ID.to_string(),
+        workspace_roots: vec![PathUri::from_abs_path(&test.config.cwd)],
-            cwd: PathUri::from_abs_path(&test.config.cwd),
+        config: EnvironmentConfigState::FromThread,
-            workspace_roots: vec![PathUri::from_abs_path(&test.config.cwd)],
+    };
-            config: EnvironmentConfigState::Ready(EnvironmentConfig {
+    let mut thread_extension_init = ExtensionDataInit::new();
-                allow_login_shell: true,
+    thread_extension_init.insert(vec![stale_root]);
-                permission_profile: PermissionProfileSnapshot::legacy(
+    let resumed = test
-                    test.config.permissions.permission_profile().clone(),
+        .thread_manager
-                ),
+        .start_thread(StartThreadOptions {
-                shell_environment_policy: Default::default(),
+            environments: Some(vec![selection.clone()]),
-                exec_policy: None,
+            thread_extension_init,
-                mcp_policy: None,
+            ..StartThreadOptions::new(test.config.clone())
-                network_policy: None,
+        })
-                selected_capability_roots: vec![ready_root],
+        .await?;
-            }),
+    resumed
-        }]),
+        .thread
-    )
+        .start_or_steer_turn(TurnInputRequest::user_input(vec![UserInput::Text {
-    .await?;
+            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_capability_roots>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();