core: refresh turn diff roots from step context

codex-rs/core/src/session/turn.rs

@@ -204,7 +204,8 @@ pub(crate) async fn run_turn(
     let mut stop_hook_active = false;
     // Although from the perspective of codex.rs, TurnDiffTracker has the lifecycle of a Task which contains
     // many turns, from the perspective of the user, it is a single turn.
-    let display_roots = turn_diff_display_roots(turn_context.as_ref()).await;
+    let mut diff_environment_selections = first_step_context.environments.to_selections();
+    let display_roots = turn_diff_display_roots(first_step_context.as_ref()).await;
     let turn_diff_tracker = Arc::new(tokio::sync::Mutex::new(
         TurnDiffTracker::with_environment_display_roots(display_roots),
     ));
@@ -244,6 +245,16 @@ pub(crate) async fn run_turn(
             Some(step_context) => step_context,
             None => sess.capture_step_context(Arc::clone(&turn_context)).await,
         };
+        let current_environment_selections = step_context.environments.to_selections();
+        if current_environment_selections != diff_environment_selections {
+            // Refresh turn-owned derived state before tools use this request's environments.
+            let display_roots = turn_diff_display_roots(step_context.as_ref()).await;
+            turn_diff_tracker
+                .lock()
+                .await
+                .set_environment_display_roots(display_roots);
+            diff_environment_selections = current_environment_selections;
+        }
         let sampling_request_result: CodexResult<_> = async {
             super::time_reminder::maybe_record_current_time_reminder(
                 sess.as_ref(),
@@ -454,9 +465,9 @@ pub(crate) async fn run_turn(
 }
 
 #[instrument(level = "trace", skip_all)]
-async fn turn_diff_display_roots(turn_context: &TurnContext) -> Vec<(String, PathBuf)> {
+async fn turn_diff_display_roots(step_context: &StepContext) -> Vec<(String, PathBuf)> {
     let mut display_roots = Vec::new();
-    for turn_environment in &turn_context.environments.turn_environments {
+    for turn_environment in &step_context.environments.turn_environments {
         // TODO(anp): Migrate git-root discovery and diff display roots to PathUri so foreign
         // environment roots can participate without host-native conversion.
         let Ok(cwd) = turn_environment.cwd().to_abs_path() else {

codex-rs/core/src/turn_diff_tracker.rs

@@ -90,6 +90,18 @@ impl TurnDiffTracker {
         tracker
     }
 
+    pub(crate) fn set_environment_display_roots(
+        &mut self,
+        display_roots: impl IntoIterator<Item = (String, PathBuf)>,
+    ) {
+        self.display_roots_by_environment = display_roots.into_iter().collect();
+        if self.valid {
+            // Rendered diffs include display paths, so rebuild them for the new roots.
+            self.rendered_diffs.clear();
+            self.refresh_unified_diff();
+        }
+    }
+
     pub fn track_delta(&mut self, environment_id: &str, delta: &AppliedPatchDelta) {
         if !self.valid {
             return;

codex-rs/core/src/turn_diff_tracker_tests.rs

@@ -84,6 +84,46 @@ index {ZERO_OID}..{right_oid}
     assert_eq!(tracker.get_unified_diff(), Some(expected));
 }
 
+#[tokio::test]
+async fn updating_display_roots_rerenders_existing_diff() {
+    let dir = tempdir().expect("tempdir");
+    let workspace = dir.path().join("workspace");
+    fs::create_dir(&workspace).expect("workspace directory");
+    let add = apply_verified_patch(
+        &workspace,
+        "*** Begin Patch\n*** Add File: a.txt\n+foo\n*** End Patch",
+    )
+    .await;
+    let mut tracker = tracker_with_root(dir.path());
+    tracker.track_delta("", &add);
+    let right_oid = git_blob_sha1_hex("foo\n");
+    let before = format!(
+        r#"diff --git a/workspace/a.txt b/workspace/a.txt
+new file mode {REGULAR_FILE_MODE}
+index {ZERO_OID}..{right_oid}
+--- {DEV_NULL}
++++ b/workspace/a.txt
+@@ -0,0 +1 @@
++foo
+"#,
+    );
+    assert_eq!(tracker.get_unified_diff(), Some(before));
+
+    tracker.set_environment_display_roots([("".to_string(), workspace)]);
+
+    let after = format!(
+        r#"diff --git a/a.txt b/a.txt
+new file mode {REGULAR_FILE_MODE}
+index {ZERO_OID}..{right_oid}
+--- {DEV_NULL}
++++ b/a.txt
+@@ -0,0 +1 @@
++foo
+"#,
+    );
+    assert_eq!(tracker.get_unified_diff(), Some(after));
+}
+
 #[tokio::test]
 async fn invalidated_tracker_suppresses_existing_diff() {
     let dir = tempdir().expect("tempdir");