openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-10-15 16:48:52 +00:00
Code Issues Packages Projects Releases Wiki Activity

Serialize environment updates and cancel removed pending attachments (#47077)

Browse Source 
## 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
This commit is contained in:
sayan-oai
2026-09-21 17:53:03 +00:00
committed by copyberry
parent e51aacad60
commit 62ea6d41e8
3 changed files with 248 additions and 83 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/core/src/environment_selection.rs
+88 -82
View File
@@ -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<Vec<SelectedTurnEnvironment>>,
environments: Mutex<Vec<SelectedTurnEnvironment>>,
connection_event_tx: OnceLock<Sender<Event>>,
}
@@ -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::<Vec<_>>();
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<TurnEnvironmentSelection> {
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<Event>) {
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<Output = TurnEnvironmentSnapshot> + 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::<Vec<_>>();
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);
codex-rs/core/src/environment_selection_snapshot_tests.rs
+155
View File
@@ -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"
);
}
codex-rs/core/tests/suite/remote_env.rs
+5 -1
View File
@@ -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::<Value>(&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()