Make turn diff tracker multi-env aware

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

@@ -1,5 +1,6 @@
 use std::collections::HashMap;
 use std::collections::HashSet;
+use std::path::PathBuf;
 use std::sync::Arc;
 use std::sync::atomic::Ordering;
 
@@ -73,7 +74,6 @@ use codex_core_skills::injection::InjectedHostSkillPrompts;
 use codex_extension_api::TurnInputContext;
 use codex_extension_api::TurnInputEnvironment;
 use codex_features::Feature;
-use codex_git_utils::get_git_repo_root;
 use codex_git_utils::get_git_repo_root_with_fs;
 use codex_protocol::config_types::AutoCompactTokenLimitScope;
 use codex_protocol::config_types::ModeKind;
@@ -196,21 +196,11 @@ 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.
-    #[allow(deprecated)]
-    let display_root = match turn_context.environments.primary() {
-        Some(turn_environment) => get_git_repo_root_with_fs(
-            turn_environment.environment.get_filesystem().as_ref(),
-            &turn_environment.cwd,
-        )
-        .await
-        .unwrap_or_else(|| turn_environment.cwd.clone())
-        .into_path_buf(),
-        None => get_git_repo_root(turn_context.cwd.as_path())
-            .unwrap_or_else(|| turn_context.cwd.clone().into_path_buf()),
-    };
-    let turn_diff_tracker = Arc::new(tokio::sync::Mutex::new(TurnDiffTracker::with_display_root(
-        display_root,
-    )));
+    let turn_diff_tracker = Arc::new(tokio::sync::Mutex::new(
+        TurnDiffTracker::with_environment_display_roots(
+            turn_diff_display_roots(turn_context.as_ref()).await,
+        ),
+    ));
 
     // `ModelClientSession` is turn-scoped and caches WebSocket + sticky routing state, so we reuse
     // one instance across retries within this turn.
@@ -421,6 +411,21 @@ pub(crate) async fn run_turn(
     last_agent_message
 }
 
+async fn turn_diff_display_roots(turn_context: &TurnContext) -> Vec<(String, PathBuf)> {
+    let mut display_roots = Vec::new();
+    for turn_environment in &turn_context.environments.turn_environments {
+        let root = get_git_repo_root_with_fs(
+            turn_environment.environment.get_filesystem().as_ref(),
+            &turn_environment.cwd,
+        )
+        .await
+        .unwrap_or_else(|| turn_environment.cwd.clone())
+        .into_path_buf();
+        display_roots.push((turn_environment.environment_id.clone(), root));
+    }
+    display_roots
+}
+
 async fn run_hooks_and_record_inputs(
     sess: &Arc<Session>,
     turn_context: &Arc<TurnContext>,

codex-rs/core/src/tools/events.rs

@@ -70,16 +70,25 @@ pub(crate) enum ToolEventFailure<'a> {
 }
 
 enum TurnDiffTrackerUpdate<'a> {
-    Track(&'a AppliedPatchDelta),
+    Track {
+        environment_id: Option<String>,
+        delta: &'a AppliedPatchDelta,
+    },
     Invalidate,
     None,
 }
 
-fn tracker_update_for_known_delta(delta: &AppliedPatchDelta) -> TurnDiffTrackerUpdate<'_> {
+fn tracker_update_for_known_delta<'a>(
+    environment_id: Option<&str>,
+    delta: &'a AppliedPatchDelta,
+) -> TurnDiffTrackerUpdate<'a> {
     if delta.is_exact() && delta.is_empty() {
         TurnDiffTrackerUpdate::None
     } else {
-        TurnDiffTrackerUpdate::Track(delta)
+        TurnDiffTrackerUpdate::Track {
+            environment_id: environment_id.map(str::to_string),
+            delta,
+        }
     }
 }
 
@@ -120,6 +129,7 @@ pub(crate) enum ToolEmitter {
     ApplyPatch {
         changes: HashMap<PathBuf, FileChange>,
         auto_approved: bool,
+        environment_id: Option<String>,
     },
     UnifiedExec {
         command: Vec<String>,
@@ -141,10 +151,15 @@ impl ToolEmitter {
         }
     }
 
-    pub fn apply_patch(changes: HashMap<PathBuf, FileChange>, auto_approved: bool) -> Self {
+    pub fn apply_patch_for_environment(
+        changes: HashMap<PathBuf, FileChange>,
+        auto_approved: bool,
+        environment_id: String,
+    ) -> Self {
         Self::ApplyPatch {
             changes,
             auto_approved,
+            environment_id: Some(environment_id),
         }
     }
 
@@ -210,7 +225,11 @@ impl ToolEmitter {
                     .await;
             }
             (
-                Self::ApplyPatch { changes, .. },
+                Self::ApplyPatch {
+                    changes,
+                    environment_id,
+                    ..
+                },
                 ToolEventStage::Success {
                     output,
                     applied_patch_delta,
@@ -222,7 +241,7 @@ impl ToolEmitter {
                     PatchApplyStatus::Failed
                 };
                 let tracker_update = applied_patch_delta
-                    .map(tracker_update_for_known_delta)
+                    .map(|delta| tracker_update_for_known_delta(environment_id.as_deref(), delta))
                     .unwrap_or(TurnDiffTrackerUpdate::Invalidate);
                 emit_patch_end(
                     ctx,
@@ -267,7 +286,11 @@ impl ToolEmitter {
                 .await;
             }
             (
-                Self::ApplyPatch { changes, .. },
+                Self::ApplyPatch {
+                    changes,
+                    environment_id,
+                    ..
+                },
                 ToolEventStage::Failure(ToolEventFailure::Rejected {
                     message,
                     applied_patch_delta,
@@ -280,7 +303,9 @@ impl ToolEmitter {
                     (*message).to_string(),
                     PatchApplyStatus::Declined,
                     applied_patch_delta
-                        .map(tracker_update_for_known_delta)
+                        .map(|delta| {
+                            tracker_update_for_known_delta(environment_id.as_deref(), delta)
+                        })
                         .unwrap_or(TurnDiffTrackerUpdate::None),
                 )
                 .await;
@@ -565,8 +590,11 @@ async fn emit_patch_end(
             let mut guard = tracker.lock().await;
             let previous_diff = guard.get_unified_diff();
             let tracker_changed = match tracker_update {
-                TurnDiffTrackerUpdate::Track(delta) => {
-                    guard.track_delta(delta);
+                TurnDiffTrackerUpdate::Track {
+                    environment_id,
+                    delta,
+                } => {
+                    guard.track_delta(environment_id.as_deref().unwrap_or_default(), delta);
                     true
                 }
                 TurnDiffTrackerUpdate::Invalidate => {
@@ -627,14 +655,18 @@ mod tests {
         .await
         .expect("apply patch");
 
-        ToolEmitter::apply_patch(HashMap::new(), /*auto_approved*/ false)
-            .finish(
-                ToolEventCtx::new(session.as_ref(), turn.as_ref(), "call-id", Some(&tracker)),
-                out,
-                Some(&delta),
-            )
-            .await
-            .expect_err("failed patch");
+        ToolEmitter::ApplyPatch {
+            changes: HashMap::new(),
+            auto_approved: false,
+            environment_id: None,
+        }
+        .finish(
+            ToolEventCtx::new(session.as_ref(), turn.as_ref(), "call-id", Some(&tracker)),
+            out,
+            Some(&delta),
+        )
+        .await
+        .expect_err("failed patch");
 
         let completed = rx_event.recv().await.expect("item completed event");
         assert!(matches!(

codex-rs/core/src/tools/handlers/apply_patch.rs

@@ -378,8 +378,11 @@ impl ToolExecutor<ToolInvocation> for ApplyPatchHandler {
                     }
                     InternalApplyPatchInvocation::DelegateToRuntime(apply) => {
                         let changes = convert_apply_patch_to_protocol(&apply.action);
-                        let emitter =
-                            ToolEmitter::apply_patch(changes.clone(), apply.auto_approved);
+                        let emitter = ToolEmitter::apply_patch_for_environment(
+                            changes.clone(),
+                            apply.auto_approved,
+                            turn_environment.environment_id.clone(),
+                        );
                         let event_ctx = ToolEventCtx::new(
                             session.as_ref(),
                             turn.as_ref(),
@@ -537,7 +540,11 @@ pub(crate) async fn intercept_apply_patch(
                 }
                 InternalApplyPatchInvocation::DelegateToRuntime(apply) => {
                     let changes = convert_apply_patch_to_protocol(&apply.action);
-                    let emitter = ToolEmitter::apply_patch(changes.clone(), apply.auto_approved);
+                    let emitter = ToolEmitter::apply_patch_for_environment(
+                        changes.clone(),
+                        apply.auto_approved,
+                        turn_environment.environment_id.clone(),
+                    );
                     let event_ctx = ToolEventCtx::new(
                         session.as_ref(),
                         turn.as_ref(),