Keep active-turn environments stable across settings updates (#33658)

## Why

With deferred execution enabled, rebuilding a step context from the latest
thread environment snapshot could apply a settings update to a turn that was
already in progress.

## What changed

- Preserve each turn's captured environment selections when creating later
  step contexts.
- Refresh only the readiness of environments that were already starting, while
  retaining their original selection order.
- Apply updated environment selections to the next turn.

## Testing

Add coverage for changing the selected workspace while a turn is paused,
verifying that the active turn keeps its original working directory and the
next turn uses the new one.

GitOrigin-RevId: c6f4388c843384ef045b7abdbc303dcfca11d230

codex-rs/core/src/environment_selection.rs

@@ -95,19 +95,22 @@ impl ThreadEnvironments {
     ) -> Self {
         // Reuse only attached environments from the supplied snapshot; drop starting entries.
         let environments = current
-            .turn_environments
+            .environments
             .into_iter()
-            .map(|environment| {
+            .filter_map(|environment| {
+                let TurnEnvironmentState::Ready(environment) = environment else {
+                    return None;
+                };
                 let selection = environment.selection();
                 let selected_environment = Arc::clone(&environment.environment);
                 let resolution: TurnEnvironmentResolution =
                     futures::future::ready(Ok(environment)).boxed().shared();
-                SelectedTurnEnvironment {
+                Some(SelectedTurnEnvironment {
                     selection,
                     environment: selected_environment,
                     connection_events_task: None,
                     resolution,
-                }
+                })
             })
             .collect();
         Self {
@@ -311,31 +314,25 @@ impl ThreadEnvironments {
 
     #[tracing::instrument(name = "environments.snapshot", skip_all)]
     pub(crate) async fn snapshot(&self) -> TurnEnvironmentSnapshot {
-        let current = self.environments.load_full();
-        let mut turn_environments = Vec::with_capacity(current.len());
-        let mut starting = Vec::new();
-        for environment in current.iter() {
+        let selected = self.environments.load_full();
+        let mut environments = Vec::with_capacity(selected.len());
+        for environment in selected.iter() {
             let resolved = if self.non_blocking_snapshots {
                 environment.resolution.clone().now_or_never()
             } else {
                 Some(environment.resolution.clone().await)
             };
-            match resolved {
-                Some(Ok(turn_environment)) => turn_environments.push(turn_environment),
-                Some(Err(err)) => tracing::debug!(
-                    environment_id = %environment.selection.environment_id,
-                    "skipping failed turn environment: {err}"
-                ),
-                None => starting.push(StartingTurnEnvironment {
+            if let Some(environment) = TurnEnvironmentState::from_resolution(
+                StartingTurnEnvironment {
                     selection: environment.selection.clone(),
                     resolution: environment.resolution.clone(),
-                }),
+                },
+                resolved,
+            ) {
+                environments.push(environment);
             }
         }
-        TurnEnvironmentSnapshot {
-            turn_environments,
-            starting,
-        }
+        TurnEnvironmentSnapshot { environments }
     }
 
     pub(crate) fn environment_manager(&self) -> Arc<EnvironmentManager> {
@@ -343,17 +340,79 @@ impl ThreadEnvironments {
     }
 }
 
+#[derive(Clone, Debug)]
+pub(crate) enum TurnEnvironmentState {
+    Ready(TurnEnvironment),
+    Starting(StartingTurnEnvironment),
+}
+
+impl TurnEnvironmentState {
+    fn from_resolution(
+        starting: StartingTurnEnvironment,
+        resolved: Option<TurnEnvironmentResult>,
+    ) -> Option<Self> {
+        match resolved {
+            Some(Ok(environment)) => Some(Self::Ready(environment)),
+            Some(Err(err)) => {
+                tracing::debug!(
+                    environment_id = %starting.selection.environment_id,
+                    "skipping failed turn environment: {err}"
+                );
+                None
+            }
+            None => Some(Self::Starting(starting)),
+        }
+    }
+}
+
 #[derive(Clone, Debug, Default)]
 pub(crate) struct TurnEnvironmentSnapshot {
-    pub(crate) turn_environments: Vec<TurnEnvironment>,
-    pub(crate) starting: Vec<StartingTurnEnvironment>,
+    // Keep ready and starting environments in their original selection order.
+    pub(crate) environments: Vec<TurnEnvironmentState>,
 }
 
 impl TurnEnvironmentSnapshot {
+    /// Promotes completed startup work without adopting newer thread selections.
+    pub(crate) fn refresh_readiness(&self) -> Self {
+        let environments = self
+            .environments
+            .iter()
+            .filter_map(|environment| match environment {
+                TurnEnvironmentState::Ready(environment) => {
+                    Some(TurnEnvironmentState::Ready(environment.clone()))
+                }
+                TurnEnvironmentState::Starting(environment) => {
+                    TurnEnvironmentState::from_resolution(
+                        environment.clone(),
+                        environment.resolution.clone().now_or_never(),
+                    )
+                }
+            })
+            .collect();
+        Self { environments }
+    }
+
+    pub(crate) fn turn_environments(&self) -> impl Iterator<Item = &TurnEnvironment> {
+        self.environments.iter().filter_map(|environment| {
+            let TurnEnvironmentState::Ready(environment) = environment else {
+                return None;
+            };
+            Some(environment)
+        })
+    }
+
+    pub(crate) fn starting(&self) -> impl Iterator<Item = &StartingTurnEnvironment> {
+        self.environments.iter().filter_map(|environment| {
+            let TurnEnvironmentState::Starting(environment) = environment else {
+                return None;
+            };
+            Some(environment)
+        })
+    }
+
     /// Maps each captured environment to its exact ready handle, or `None` when it was starting.
     pub(crate) fn captured_environments(&self) -> HashMap<String, Option<Arc<Environment>>> {
-        self.turn_environments
-            .iter()
+        self.turn_environments()
             .map(|environment| {
                 (
                     environment.environment_id.clone(),
@@ -361,20 +420,18 @@ impl TurnEnvironmentSnapshot {
                 )
             })
             .chain(
-                self.starting
-                    .iter()
+                self.starting()
                     .map(|environment| (environment.selection.environment_id.clone(), None)),
             )
             .collect()
     }
 
     pub(crate) fn primary(&self) -> Option<&TurnEnvironment> {
-        self.turn_environments.first()
+        self.turn_environments().next()
     }
 
     pub(crate) fn local(&self) -> Option<&TurnEnvironment> {
-        self.turn_environments
-            .iter()
+        self.turn_environments()
             .find(|environment| !environment.environment.is_remote())
     }
 
@@ -385,8 +442,7 @@ impl TurnEnvironmentSnapshot {
     }
 
     pub(crate) fn to_selections(&self) -> Vec<TurnEnvironmentSelection> {
-        self.turn_environments
-            .iter()
+        self.turn_environments()
             .map(TurnEnvironment::selection)
             .collect()
     }
@@ -397,12 +453,14 @@ impl TurnEnvironmentSnapshot {
     }
 
     pub(crate) fn single_local_environment(&self) -> Option<&TurnEnvironment> {
-        if !self.starting.is_empty() {
+        if self.starting().next().is_some() {
             return None;
         }
-        let [environment] = self.turn_environments.as_slice() else {
+        let mut environments = self.turn_environments();
+        let environment = environments.next()?;
+        if environments.next().is_some() {
             return None;
-        };
+        }
 
         (!environment.environment.is_remote()).then_some(environment)
     }
@@ -751,13 +809,13 @@ url = "ws://127.0.0.1:8765"
             .expect("snapshot should finish after the environment starts")
             .expect("snapshot task");
 
-        assert!(snapshot.starting.is_empty());
+        assert!(snapshot.starting().next().is_none());
         assert_eq!(snapshot.to_selections(), vec![selection]);
         server.await.expect("server task");
     }
 
     #[tokio::test]
-    async fn snapshot_keeps_starting_environment_until_it_can_be_attached() {
+    async fn snapshot_refreshes_readiness_in_selection_order() {
         let listener = TcpListener::bind("127.0.0.1:0")
             .await
             .expect("bind websocket listener");
@@ -790,36 +848,51 @@ url = "ws://127.0.0.1:8765"
             TurnEnvironmentSnapshot::default(),
             /*non_blocking_snapshots*/ true,
         );
+        turn_environments.update_selections(std::slice::from_ref(&local));
+        turn_environments.environments.load()[0]
+            .resolution
+            .clone()
+            .await
+            .expect("local environment should resolve");
         turn_environments.update_selections(&[remote.clone(), local.clone()]);
 
         let starting = turn_environments.snapshot().await;
-        assert!(starting.turn_environments.is_empty());
         assert_eq!(
             starting
-                .starting
-                .iter()
+                .turn_environments()
+                .map(TurnEnvironment::selection)
+                .collect::<Vec<_>>(),
+            vec![local.clone()]
+        );
+        assert_eq!(
+            starting
+                .starting()
                 .map(|environment| environment.selection.clone())
                 .collect::<Vec<_>>(),
-            vec![remote.clone(), local.clone()]
+            vec![remote.clone()]
         );
-        assert!(starting.to_selections().is_empty());
+        assert_eq!(starting.to_selections(), vec![local.clone()]);
         assert!(starting.single_local_environment().is_none());
 
         let server = tokio::spawn(serve_environment_info(listener));
         timeout(
             std::time::Duration::from_secs(5),
-            starting.starting[0].resolution.clone(),
+            starting
+                .starting()
+                .next()
+                .expect("starting environment")
+                .resolution
+                .clone(),
         )
         .await
         .expect("environment resolution should finish")
         .expect("environment resolution should succeed");
-        let attached = turn_environments.snapshot().await;
+        let attached = starting.refresh_readiness();
 
-        assert!(attached.starting.is_empty());
+        assert!(attached.starting().next().is_none());
         assert_eq!(
             attached
-                .turn_environments
-                .iter()
+                .turn_environments()
                 .map(TurnEnvironment::selection)
                 .collect::<Vec<_>>(),
             vec![remote.clone(), local.clone()]
@@ -866,9 +939,10 @@ url = "ws://127.0.0.1:8765"
         environments.update_selections(std::slice::from_ref(&selection));
 
         let replacement = environments.snapshot().await;
-        let [replacement] = replacement.starting.as_slice() else {
-            panic!("expected the replacement environment to be starting");
-        };
+        let replacement = replacement
+            .starting()
+            .next()
+            .expect("expected the replacement environment to be starting");
         assert_eq!(replacement.selection, selection);
         assert!(!failed_resolution.ptr_eq(&replacement.resolution));
     }
@@ -930,16 +1004,16 @@ url = "ws://127.0.0.1:8765"
         let changed_snapshot = environments.snapshot().await;
 
         let initial = initial_snapshot
-            .starting
-            .first()
+            .starting()
+            .next()
             .expect("initial environment");
         let reused = reused_snapshot
-            .starting
-            .first()
+            .starting()
+            .next()
             .expect("reused environment");
         let changed = changed_snapshot
-            .starting
-            .first()
+            .starting()
+            .next()
             .expect("changed environment");
         assert!(initial.resolution.ptr_eq(&reused.resolution));
         assert!(!reused.resolution.ptr_eq(&changed.resolution));
@@ -1001,8 +1075,7 @@ url = "ws://127.0.0.1:8765"
             crate::shell::default_user_shell(),
             ShellSnapshot::disabled(),
             TurnEnvironmentSnapshot {
-                turn_environments: vec![inherited],
-                starting: Vec::new(),
+                environments: vec![TurnEnvironmentState::Ready(inherited)],
             },
             /*non_blocking_snapshots*/ false,
         );
@@ -1039,27 +1112,25 @@ url = "ws://127.0.0.1:8765"
                 .expect("remote environment"),
         );
         let remote = TurnEnvironmentSnapshot {
-            turn_environments: vec![TurnEnvironment::new(
+            environments: vec![TurnEnvironmentState::Ready(TurnEnvironment::new(
                 REMOTE_ENVIRONMENT_ID.to_string(),
                 remote_environment.clone(),
                 cwd_uri.clone(),
                 Vec::new(),
                 /*shell*/ None,
-            )],
-            starting: Vec::new(),
+            ))],
         };
         let multiple = TurnEnvironmentSnapshot {
-            turn_environments: vec![
-                local.primary().expect("local environment").clone(),
-                TurnEnvironment::new(
+            environments: vec![
+                TurnEnvironmentState::Ready(local.primary().expect("local environment").clone()),
+                TurnEnvironmentState::Ready(TurnEnvironment::new(
                     REMOTE_ENVIRONMENT_ID.to_string(),
                     remote_environment,
                     cwd_uri,
                     Vec::new(),
                     /*shell*/ None,
-                ),
+                )),
             ],
-            starting: Vec::new(),
         };
 
         assert_eq!(local.single_local_environment_cwd(), Some(cwd));