From 62ea6d41e85d0c75181f2794d0f2dc40e55fce3c Mon Sep 17 00:00:00 2001 From: sayan-oai Date: Mon, 21 Sep 2026 17:40:30 +0000 Subject: [PATCH] Serialize environment updates and cancel removed pending attachments (#47077) ## Why Concurrent credential refreshes could restore a removed environment by overwriting the selection list with stale state. Unpolled snapshot futures could also keep pending configuration alive after an environment was removed. ## What changed - Protect environment selections with a mutex and update them in place to serialize selection changes, configuration refreshes, and connection listener setup. - Capture only selection metadata and resolution futures in snapshots so they do not retain pending configuration senders. ## Testing Add regression tests for removal during credential refresh and cancellation while an unpolled snapshot is held. Strengthen the pending attachment integration test to check that `wait_for_environment` returns the expected environment ID and `ready` status. GitOrigin-RevId: c6f933e8d3031b05335ddcfed00ab400aa5bad2a --- codex-rs/core/src/environment_selection.rs | 170 +++++++++--------- .../environment_selection_snapshot_tests.rs | 155 ++++++++++++++++ codex-rs/core/tests/suite/remote_env.rs | 6 +- 3 files changed, 248 insertions(+), 83 deletions(-) create mode 100644 codex-rs/core/src/environment_selection_snapshot_tests.rs diff --git a/codex-rs/core/src/environment_selection.rs b/codex-rs/core/src/environment_selection.rs index 2242a86de4..3f9c98d9ac 100644 --- a/codex-rs/core/src/environment_selection.rs +++ b/codex-rs/core/src/environment_selection.rs @@ -3,9 +3,10 @@ use std::collections::HashSet; use std::fmt; use std::future::Future; use std::sync::Arc; +use std::sync::Mutex; use std::sync::OnceLock; +use std::sync::PoisonError; -use arc_swap::ArcSwap; use async_channel::Sender; use codex_exec_server::Environment; use codex_exec_server::EnvironmentConnectionState; @@ -238,7 +239,7 @@ pub(crate) struct ThreadEnvironments { local_shell: Shell, shell_snapshot: ShellSnapshot, non_blocking_snapshots: bool, - environments: ArcSwap>, + environments: Mutex>, connection_event_tx: OnceLock>, } @@ -304,7 +305,7 @@ impl ThreadEnvironments { local_shell, shell_snapshot, non_blocking_snapshots, - environments: ArcSwap::from_pointee(environments), + environments: Mutex::new(environments), connection_event_tx: OnceLock::new(), } } @@ -344,7 +345,10 @@ impl ThreadEnvironments { environments: &[TurnEnvironmentSelection], thread_environment_config: &EnvironmentConfig, ) { - let previous = self.environments.load(); + let mut current = self + .environments + .lock() + .unwrap_or_else(PoisonError::into_inner); let mut seen_environment_ids = HashSet::with_capacity(environments.len()); let mut next = Vec::with_capacity(environments.len()); for selected_environment in environments { @@ -353,7 +357,7 @@ impl ThreadEnvironments { } let (selected_environment, config_origin) = resolve_selection_config(selected_environment.clone(), thread_environment_config); - if let Some(environment) = previous.iter().find(|environment| { + if let Some(environment) = current.iter().find(|environment| { let previous = &environment.selection; previous.environment_id == selected_environment.environment_id && previous.cwd == selected_environment.cwd @@ -399,7 +403,7 @@ impl ThreadEnvironments { continue; }; // Connection state belongs to the environment instance, not its cwd or roots. - let connection_events_task = previous + let connection_events_task = current .iter() .find(|previous| { previous.selection.environment_id.as_str() == environment_id.as_str() @@ -444,7 +448,7 @@ impl ThreadEnvironments { }; next.push(selected); } - let removed_connection_tasks = previous + let removed_connection_tasks = current .iter() .filter_map(|previous| { let task = previous.connection_events_task.as_ref()?; @@ -456,11 +460,10 @@ impl ThreadEnvironments { .then(|| Arc::clone(task)) }) .collect::>(); - let next = Arc::new(next); - self.environments.store(Arc::clone(&next)); + let previous = std::mem::replace(&mut *current, next); // Publish owner configuration before waking turns waiting on this attachment. - for environment in next.iter() { + for environment in current.iter() { let Some(completion) = &environment.pending_completion else { continue; }; @@ -472,16 +475,17 @@ impl ThreadEnvironments { let _ = completion.try_send(result); } - // ArcSwap readers may retain removed selections, so abort at logical removal. for task in removed_connection_tasks { task.abort(); } + drop(previous); } /// Projects canonical selections back to caller input form without losing config ownership. pub(crate) fn selections(&self) -> Vec { self.environments - .load() + .lock() + .unwrap_or_else(PoisonError::into_inner) .iter() .map(|environment| { environment @@ -520,52 +524,46 @@ impl ThreadEnvironments { /// Refreshes attachments whose configuration is inferred from the thread. pub(crate) fn update_thread_config(&self, config: &EnvironmentConfig) { - let environments = self + let mut environments = self .environments - .load() - .iter() - .map(|environment| { - let mut environment = environment.clone(); - let shell_settings_changed = environment.config_origin - == EnvironmentConfigOrigin::Thread - && matches!( - &environment.selection.config, - EnvironmentConfigState::Ready(previous) - if previous.allow_login_shell != config.allow_login_shell - || previous.shell_environment_policy != config.shell_environment_policy - ); - environment.refresh_thread_config(config); - if shell_settings_changed - && !environment.environment.is_remote() - && self.shell_snapshot.should_rebuild_inherited() - && let Some(Ok(resolved)) = environment.resolution.clone().now_or_never() - { - self.restart_shell_snapshot(&mut environment, resolved); - } - environment - }) - .collect(); - self.environments.store(Arc::new(environments)); + .lock() + .unwrap_or_else(PoisonError::into_inner); + for environment in environments.iter_mut() { + let shell_settings_changed = environment.config_origin + == EnvironmentConfigOrigin::Thread + && matches!( + &environment.selection.config, + EnvironmentConfigState::Ready(previous) + if previous.allow_login_shell != config.allow_login_shell + || previous.shell_environment_policy != config.shell_environment_policy + ); + environment.refresh_thread_config(config); + if shell_settings_changed + && !environment.environment.is_remote() + && self.shell_snapshot.should_rebuild_inherited() + && let Some(Ok(resolved)) = environment.resolution.clone().now_or_never() + { + self.restart_shell_snapshot(environment, resolved); + } + } } pub(crate) fn set_snapshot_credential_broker(&self, state: SnapshotCredentialBrokerState) { + let mut environments = self + .environments + .lock() + .unwrap_or_else(PoisonError::into_inner); if !self.shell_snapshot.set_credential_broker(state) { return; } - let mut environments = Vec::clone(&self.environments.load()); - let mut changed = false; - for selected in &mut environments { + for selected in environments.iter_mut() { if !selected.environment.is_remote() && let Some(Ok(resolved)) = selected.resolution.clone().now_or_never() { self.restart_shell_snapshot(selected, resolved); - changed = true; } } - if changed { - self.environments.store(Arc::new(environments)); - } } fn restart_shell_snapshot( @@ -596,7 +594,10 @@ impl ThreadEnvironments { &self, thread_selected_capability_roots: &[SelectedCapabilityRoot], ) -> SelectedCapabilityRootsStatus { - let environments = self.environments.load(); + let environments = self + .environments + .lock() + .unwrap_or_else(PoisonError::into_inner); let mut selected_capability_roots = combine_selected_capability_roots( thread_selected_capability_roots, environments.iter().filter_map(|environment| { @@ -674,23 +675,20 @@ impl ThreadEnvironments { } pub(crate) fn start_connection_event_forwarding(&self, tx_event: Sender) { + let mut environments = self + .environments + .lock() + .unwrap_or_else(PoisonError::into_inner); let tx_event = self.connection_event_tx.get_or_init(|| tx_event); - let current = self.environments.load_full(); - let environments = current - .iter() - .map(|selected| { - let mut selected = selected.clone(); - if selected.connection_events_task.is_none() { - selected.connection_events_task = Self::spawn_connection_event_listener( - selected.environment.as_ref(), - selected.selection.environment_id.clone(), - tx_event.clone(), - ); - } - selected - }) - .collect(); - self.environments.store(Arc::new(environments)); + for selected in environments.iter_mut() { + if selected.connection_events_task.is_none() { + selected.connection_events_task = Self::spawn_connection_event_listener( + selected.environment.as_ref(), + selected.selection.environment_id.clone(), + tx_event.clone(), + ); + } + } } #[tracing::instrument( @@ -812,29 +810,32 @@ impl ThreadEnvironments { pub(crate) fn snapshot( &self, ) -> impl Future + Send + 'static { - let selected = self.environments.load_full(); + // Keep the startup results, but not the senders that would keep removed Pending work alive. + let selected = self + .environments + .lock() + .unwrap_or_else(PoisonError::into_inner) + .iter() + .map(|environment| StartingTurnEnvironment { + selection: environment.selection.clone(), + config_origin: environment.config_origin, + resolution: environment.resolution.clone(), + }) + .collect::>(); let non_blocking_snapshots = self.non_blocking_snapshots; async move { let mut environments = Vec::with_capacity(selected.len()); - for environment in selected.iter() { - if let EnvironmentConfigState::Failed(error) = &environment.selection.config { + for starting in selected { + if let EnvironmentConfigState::Failed(error) = &starting.selection.config { environments.push(TurnEnvironmentState::Failed { - selection: environment + selection: starting .config_origin - .into_input_selection(environment.selection.clone()), + .into_input_selection(starting.selection.clone()), error: error.clone(), }); continue; } - let pending = matches!( - environment.selection.config, - EnvironmentConfigState::Pending - ); - let starting = StartingTurnEnvironment { - selection: environment.selection.clone(), - config_origin: environment.config_origin, - resolution: environment.resolution.clone(), - }; + let pending = matches!(starting.selection.config, EnvironmentConfigState::Pending); let resolved = if non_blocking_snapshots || pending { starting.resolution.clone().now_or_never() } else { @@ -1121,6 +1122,10 @@ impl TurnEnvironmentSnapshot { } } +#[cfg(test)] +#[path = "environment_selection_snapshot_tests.rs"] +mod snapshot_tests; + #[cfg(test)] mod tests { use std::time::Duration; @@ -1156,7 +1161,7 @@ mod tests { use super::*; - fn test_environment_config() -> EnvironmentConfig { + pub(super) fn test_environment_config() -> EnvironmentConfig { EnvironmentConfig { allow_login_shell: true, workspace_roots: Vec::new(), @@ -1641,11 +1646,10 @@ url = "ws://127.0.0.1:8765" ); turn_environments .update_selections(std::slice::from_ref(&local), &test_environment_config()); - turn_environments.environments.load()[0] + let resolution = turn_environments.environments.lock().unwrap()[0] .resolution - .clone() - .await - .expect("local environment should resolve"); + .clone(); + resolution.await.expect("local environment should resolve"); turn_environments.update_selections(&[remote.clone(), local.clone()], &expected_config); let starting = turn_environments.snapshot().await; @@ -1783,7 +1787,9 @@ url = "ws://127.0.0.1:8765" environments .update_selections(std::slice::from_ref(&selection), &test_environment_config()); let starting = environments.snapshot().await; - let failed_resolution = environments.environments.load()[0].resolution.clone(); + let failed_resolution = environments.environments.lock().unwrap()[0] + .resolution + .clone(); let error = failed_resolution .clone() .await @@ -1829,7 +1835,7 @@ url = "ws://127.0.0.1:8765" }; environments.update_thread_config(&next_config); assert!( - failed_resolution.ptr_eq(&environments.environments.load()[0].resolution), + failed_resolution.ptr_eq(&environments.environments.lock().unwrap()[0].resolution), "updating environment config must not retry a failed environment" ); environments.update_selections(std::slice::from_ref(&selection), &next_config); diff --git a/codex-rs/core/src/environment_selection_snapshot_tests.rs b/codex-rs/core/src/environment_selection_snapshot_tests.rs new file mode 100644 index 0000000000..b9bbf9d492 --- /dev/null +++ b/codex-rs/core/src/environment_selection_snapshot_tests.rs @@ -0,0 +1,155 @@ +//! Regression coverage for environment snapshots and pending configuration. + +use super::*; +use codex_exec_server::LOCAL_ENVIRONMENT_ID; +use codex_otel::SessionTelemetry; +use codex_protocol::ThreadId; +use codex_protocol::protocol::SessionSource; +use pretty_assertions::assert_eq; +use std::time::Duration; +use tokio::sync::oneshot; +use tokio::sync::watch; +use tokio::time::timeout; + +#[tokio::test] +async fn unpolled_snapshot_does_not_delay_canceling_a_removed_environment() { + let cwd = PathUri::from_abs_path(&AbsolutePathBuf::current_dir().expect("cwd")); + let selection = TurnEnvironmentSelection { + environment_id: LOCAL_ENVIRONMENT_ID.to_string(), + cwd: cwd.clone(), + workspace_roots: vec![cwd], + config: EnvironmentConfigState::Pending, + }; + let config = tests::test_environment_config(); + let environments = ThreadEnvironments::new( + Arc::new(EnvironmentManager::default_for_tests()), + crate::shell::default_user_shell(), + config.clone(), + ShellSnapshot::disabled(), + TurnEnvironmentSnapshot::default(), + /*non_blocking_snapshots*/ false, + ); + environments.update_selections(&[selection], &config); + let starting = environments + .snapshot() + .await + .starting() + .next() + .unwrap() + .clone(); + let held_snapshot = environments.snapshot(); + + environments.update_selections(&[], &config); + + let error = timeout(Duration::from_secs(5), starting.wait_until_ready()) + .await + .expect("removing the environment should cancel the wait before the old snapshot is polled") + .unwrap_err(); + assert!( + error + .to_string() + .contains("environment configuration was canceled") + ); + assert!(matches!( + held_snapshot.await.environments.as_slice(), + [TurnEnvironmentState::Failed { .. }] + )); +} + +#[tokio::test] +async fn credential_refresh_does_not_restore_a_removed_environment() { + let home = tempfile::tempdir().expect("home"); + let cwd = AbsolutePathBuf::from_absolute_path(home.path()).expect("cwd"); + let cwd_uri = PathUri::from_abs_path(&cwd); + let session_id = ThreadId::new(); + let telemetry = SessionTelemetry::new( + session_id, + "test", + "test", + /*account_id*/ None, + /*account_email*/ None, + /*auth_mode*/ None, + "test".to_string(), + /*log_user_prompts*/ false, + "test".to_string(), + SessionSource::Cli, + ); + let (broker, _) = watch::channel(SnapshotCredentialBrokerState::Starting); + let shell_snapshot = ShellSnapshot::new( + cwd, + session_id, + telemetry, + /*state_db*/ None, + Some(broker), + /*prefer_executor_snapshots*/ false, + ); + let config = tests::test_environment_config(); + let environments = Arc::new(ThreadEnvironments::new( + Arc::new(EnvironmentManager::default_for_tests()), + crate::shell::default_user_shell(), + config.clone(), + shell_snapshot, + TurnEnvironmentSnapshot::default(), + /*non_blocking_snapshots*/ false, + )); + environments.update_selections( + &[TurnEnvironmentSelection { + environment_id: LOCAL_ENVIRONMENT_ID.to_string(), + cwd: cwd_uri.clone(), + workspace_roots: vec![cwd_uri], + config: EnvironmentConfigState::FromThread, + }], + &config, + ); + let resolution = environments.environments.lock().unwrap()[0] + .resolution + .clone(); + let resolved = resolution.await.expect("local is ready"); + let (refresh_paused_tx, refresh_paused_rx) = oneshot::channel(); + let (resume_refresh_tx, resume_refresh_rx) = std::sync::mpsc::channel(); + // Pause the credential refresh while it is using the current list. + environments.environments.lock().unwrap()[0].resolution = async move { + refresh_paused_tx.send(()).expect("signal refresh paused"); + resume_refresh_rx.recv().expect("resume refresh"); + Ok(resolved) + } + .boxed() + .shared(); + + let refreshing = Arc::clone(&environments); + let refresh = tokio::task::spawn_blocking(move || { + refreshing.set_snapshot_credential_broker(SnapshotCredentialBrokerState::Unavailable); + }); + timeout(Duration::from_secs(/*secs*/ 5), refresh_paused_rx) + .await + .expect("credential refresh should reach the old list") + .expect("refresh paused"); + let (removal_started_tx, removal_started_rx) = oneshot::channel(); + let removing = Arc::clone(&environments); + let mut removal = tokio::task::spawn_blocking(move || { + removal_started_tx.send(()).expect("signal removal started"); + removing.update_selections(&[], &config); + }); + timeout(Duration::from_secs(/*secs*/ 5), removal_started_rx) + .await + .expect("removal should start") + .expect("removal started"); + let early_removal = timeout(Duration::from_secs(/*secs*/ 1), &mut removal) + .await + .ok(); + let removal_finished_early = early_removal.is_some(); + + resume_refresh_tx + .send(()) + .expect("release credential refresh"); + refresh.await.expect("credential refresh finished"); + match early_removal { + Some(result) => result.expect("removal finished"), + None => removal.await.expect("removal finished"), + } + assert_eq!(environments.selections(), Vec::new()); + assert!( + !removal_finished_early, + "removal should wait until the credential refresh finishes" + ); +} diff --git a/codex-rs/core/tests/suite/remote_env.rs b/codex-rs/core/tests/suite/remote_env.rs index 1a81fa4b04..4b3a952dda 100644 --- a/codex-rs/core/tests/suite/remote_env.rs +++ b/codex-rs/core/tests/suite/remote_env.rs @@ -1933,10 +1933,14 @@ async fn pending_attachment_installs_configuration_before_waiting_turn_resumes() let ready_request = response_mock .last_request() .context("waiting turn should resume")?; - let (_, wait_succeeded) = ready_request + let (wait_output, wait_succeeded) = ready_request .function_call_output_content_and_success(WAIT_CALL_ID) .context("wait_for_environment output should be model visible")?; assert_ne!(wait_succeeded, Some(false)); + assert_eq!( + serde_json::from_str::(&wait_output.context("wait output should contain JSON")?)?, + json!({ "environment_id": selection.environment_id, "status": "ready" }) + ); let body = ready_request.body_json(); let exec_command = body["tools"] .as_array()