diff --git a/codex-rs/core/src/session/turn.rs b/codex-rs/core/src/session/turn.rs index 9b3588fc1e..2004e6142f 100644 --- a/codex-rs/core/src/session/turn.rs +++ b/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 { diff --git a/codex-rs/core/src/turn_diff_tracker.rs b/codex-rs/core/src/turn_diff_tracker.rs index ed8e9684ad..94dfeeacdf 100644 --- a/codex-rs/core/src/turn_diff_tracker.rs +++ b/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, + ) { + 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; diff --git a/codex-rs/core/src/turn_diff_tracker_tests.rs b/codex-rs/core/src/turn_diff_tracker_tests.rs index 5f984a2c4f..5a18aa98c0 100644 --- a/codex-rs/core/src/turn_diff_tracker_tests.rs +++ b/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");