Tie executor attachments to environment selection lifetimes (#59043)

## Why

Removing an executor must detach its selection even when older turn snapshots still reference it. Connection events queued for a removed attachment must not leak into a later selection of the same executor.

## What changed

- Introduce `EnvironmentAttachment` with a detach signal shared by turn snapshots and a guard owned by the selection.
- Preserve attachments across workspace and configuration changes, while giving reselections and child threads independent attachment lifetimes.
- Stop forwarding connection events on detach and discard buffered events from detached attachments.
- Activate initial attachments after `SessionConfigured`, using configured bindings for startup discovery.

## Testing

Add coverage for discarding buffered connection events on reselection and extend snapshot tests to verify immediate detachment, fresh attachments on reselection, and independent parent and child lifetimes.

GitOrigin-RevId: 6cf882e0d17a296030c13d80c6e2866c80e04638
This commit is contained in:
Adam Perry @ OpenAI
2026-11-03 02:37:20 +00:00
committed by copyberry
parent 16ddd538bd
commit 7e8143a897
38 changed files with 557 additions and 318 deletions
+1 -1
View File
@@ -458,7 +458,7 @@ impl LocalAgentControl {
if child_config.permission_profile == owner_config.permission_profile {
continue;
}
if owner_environment.environment.is_remote() {
if owner_environment.environment().is_remote() {
return Err(invalid_environment(
"permissions changed on a remote executor",
));
+1 -1
View File
@@ -71,7 +71,7 @@ pub(crate) async fn load_project_instructions(
break;
}
let filesystem = turn_environment.environment.get_filesystem();
let filesystem = turn_environment.environment().get_filesystem();
let sandbox = (!turn_environment
.permission_profile()
.file_system_sandbox_policy()
@@ -3,8 +3,11 @@
use std::sync::Arc;
use std::time::Duration;
use assert_matches::assert_matches;
use codex_exec_server::Environment;
use codex_exec_server::EnvironmentConnectionState;
use codex_protocol::protocol::EnvironmentConnectionEvent;
use codex_protocol::protocol::Event;
use codex_protocol::protocol::EventMsg;
use futures::SinkExt;
use futures::StreamExt;
@@ -17,9 +20,11 @@ use tokio::sync::oneshot;
use tokio_tungstenite::WebSocketStream;
use tokio_tungstenite::accept_async;
use tokio_tungstenite::tungstenite::Message;
use tokio_util::sync::CancellationToken;
use tokio_util::task::AbortOnDropHandle;
use super::ThreadEnvironments;
use super::ConnectionEventStreams;
use super::EnvironmentAttachment;
use super::tests::read_websocket_json;
// Keep paused time from advancing during socket I/O; only the test advances it.
@@ -77,13 +82,9 @@ async fn lifecycle_waits_for_recovery(retries: usize, recovery: Recovery) {
Environment::create_for_tests(Some(format!("ws://{}", listener.local_addr().unwrap())))
.unwrap(),
);
let (events, received) = async_channel::unbounded();
let _forwarding = ThreadEnvironments::spawn_connection_event_listener(
&environment,
"remote".to_owned(),
events,
)
.unwrap();
let (attachment, _detach_on_drop, events) =
EnvironmentAttachment::new("remote", Arc::clone(&environment));
let received = events.expect("remote connection events");
let mut raw = environment.subscribe_connection_state().unwrap();
let (disconnect, disconnected) = oneshot::channel();
let (resume, resumed) = oneshot::channel();
@@ -175,6 +176,48 @@ async fn lifecycle_waits_for_recovery(retries: usize, recovery: Recovery) {
}
}
assert!(received.try_recv().is_err());
drop(attachment);
assert!(
received.recv().await.is_err(),
"dropping the attachment should stop its listener"
);
finish.send(()).unwrap();
server.await.unwrap();
}
/// Reselecting an executor must discard events still queued for its previous attachment.
#[tokio::test(flavor = "current_thread")]
async fn reselection_discards_buffered_connection_events() {
let (tx_event, received) = async_channel::unbounded();
let streams = ConnectionEventStreams::new(tx_event);
let remote = EnvironmentConnectionEvent {
environment_id: "remote".to_owned(),
};
let connected = Event {
id: String::new(),
msg: EventMsg::EnvironmentConnected(remote.clone()),
};
// Queue both attachments' events before yielding to the forwarder.
let detached = CancellationToken::new();
let (old_tx, old_events) = async_channel::unbounded();
streams.add(Some(old_events), &detached);
old_tx.try_send(connected.clone()).unwrap();
old_tx
.try_send(Event {
id: String::new(),
msg: EventMsg::EnvironmentDisconnected(remote.clone()),
})
.unwrap();
detached.cancel();
drop(old_tx);
let (new_tx, new_events) = async_channel::unbounded();
streams.add(Some(new_events), &CancellationToken::new());
new_tx.try_send(connected).unwrap();
drop(new_tx);
streams.streams.close();
let messages = received.map(|event| event.msg).collect::<Vec<_>>().await;
assert_matches!(messages.as_slice(), [EventMsg::EnvironmentConnected(event)] if event == &remote);
}
+231 -198
View File
@@ -4,12 +4,11 @@ use std::fmt;
use std::future::Future;
use std::sync::Arc;
use std::sync::Mutex;
use std::sync::OnceLock;
use std::sync::PoisonError;
use async_channel::Receiver;
use async_channel::Sender;
use codex_exec_server::Environment;
use codex_exec_server::EnvironmentConnectionState;
use codex_exec_server::EnvironmentInfo;
use codex_exec_server::EnvironmentManager;
use codex_exec_server::ExecServerError;
@@ -23,18 +22,20 @@ use codex_protocol::models::PermissionProfile;
use codex_protocol::protocol::AskForApproval;
use codex_protocol::protocol::EnvironmentConfig;
use codex_protocol::protocol::EnvironmentConfigState;
use codex_protocol::protocol::EnvironmentConnectionEvent;
use codex_protocol::protocol::Event;
use codex_protocol::protocol::EventMsg;
use codex_protocol::protocol::TurnEnvironmentRequest;
use codex_protocol::protocol::TurnEnvironmentSelection;
use codex_skills_extension::EnvironmentSkillRequirements;
use codex_utils_absolute_path::AbsolutePathBuf;
use codex_utils_path_uri::PathUri;
use futures::FutureExt;
use futures::StreamExt;
use futures::future::BoxFuture;
use futures::future::Shared;
use futures::stream::BoxStream;
use tokio::sync::watch;
use tokio_util::sync::CancellationToken;
use tokio_util::sync::DropGuard;
use tokio_util::task::AbortOnDropHandle;
use tracing::Instrument;
use tracing::instrument::WithSubscriber;
@@ -47,6 +48,9 @@ use crate::shell::Shell;
use crate::shell_snapshot::ShellSnapshot;
use crate::shell_snapshot::SnapshotCredentialBrokerState;
mod attachment;
pub(crate) use attachment::EnvironmentAttachment;
// Reject pathological selected cwd values at the environment-selection boundary.
const MAX_TURN_ENVIRONMENT_CWD_BYTES: usize = 8 * 1024;
@@ -185,7 +189,6 @@ fn pending_configuration(
// without restarting the connection or shell resolution.
#[derive(Clone)]
struct ResolvedEnvironment {
environment: Arc<Environment>,
shell: Option<Shell>,
user_home_dir: Option<PathUri>,
executor_platform_os: Option<String>,
@@ -201,9 +204,9 @@ struct ResolvedEnvironment {
struct SelectedTurnEnvironment {
selection: TurnEnvironmentSelection,
config_origin: EnvironmentConfigOrigin,
environment: Arc<Environment>,
// Selection clones share one listener; the final handle drop aborts it.
connection_events_task: Option<Arc<AbortOnDropHandle<()>>>,
attachment: Arc<EnvironmentAttachment>,
// Only selections share this guard. A captured turn must not postpone detach.
detach_on_drop: Arc<DropGuard>,
resolution: TurnEnvironmentResolution,
// Executor retries keep this session's wait for its first accepted configuration.
pending_completion: Option<watch::Sender<Option<PendingConfigurationResult>>>,
@@ -216,7 +219,7 @@ pub(crate) struct StartingTurnEnvironment {
pub(crate) selection: TurnEnvironmentSelection,
config_origin: EnvironmentConfigOrigin,
resolution: TurnEnvironmentResolution,
environment: Arc<Environment>,
attachment: Arc<EnvironmentAttachment>,
owner_config_result: Option<PendingConfiguration>,
}
@@ -240,19 +243,30 @@ fn resolve_selection_config(
impl SelectedTurnEnvironment {
// Pending configuration belongs to each attachment set; completed startup can be shared.
fn for_turn(&self, local_shell: &Shell, shell_snapshot: &ShellSnapshot) -> Self {
fn for_turn(
&self,
attachment: Arc<EnvironmentAttachment>,
detach_on_drop: Arc<DropGuard>,
local_shell: &Shell,
shell_snapshot: &ShellSnapshot,
) -> Self {
if matches!(self.selection.config, EnvironmentConfigState::Pending) {
return ThreadEnvironments::start_environment(
self.selection.clone(),
self.config_origin,
Arc::clone(&self.environment),
attachment,
detach_on_drop,
self.owner_config_result.clone(),
/*pending_completion*/ None,
local_shell,
shell_snapshot,
);
}
self.clone()
Self {
attachment,
detach_on_drop,
..self.clone()
}
}
fn refresh_thread_config(&mut self, defaults: &ThreadEnvironmentDefaults) {
@@ -300,12 +314,47 @@ pub(crate) struct ThreadEnvironments {
state: Mutex<ThreadEnvironmentsState>,
// Retain configured bindings independently of active-turn membership and cancellation.
configured: Mutex<ThreadEnvironmentsState>,
connection_event_tx: OnceLock<Sender<Event>>,
}
struct ThreadEnvironmentsState {
thread_defaults: ThreadEnvironmentDefaults,
environments: Vec<SelectedTurnEnvironment>,
connection_events: Option<ConnectionEventStreams>,
}
/// Forwards the active attachments' streams to their thread's event channel.
struct ConnectionEventStreams {
streams: Sender<BoxStream<'static, Event>>,
_task: AbortOnDropHandle<()>,
}
impl ConnectionEventStreams {
fn new(tx_event: Sender<Event>) -> Self {
let (streams, incoming) = async_channel::unbounded::<BoxStream<'static, Event>>();
let task = tokio::spawn(async move {
let events = incoming
.flatten_unordered(/*limit*/ None)
.take_until(tx_event.closed());
tokio::pin!(events);
while let Some(event) = events.next().await {
if tx_event.send(event).await.is_err() {
break;
}
}
});
Self {
streams,
_task: AbortOnDropHandle::new(task),
}
}
fn add(&self, events: Option<Receiver<Event>>, detached: &CancellationToken) {
if let Some(events) = events {
// A closed receiver can still yield events queued before detach.
let events = events.take_until(detached.clone().cancelled_owned());
let _ = self.streams.try_send(events.boxed());
}
}
}
#[derive(Clone, Copy)]
@@ -372,10 +421,15 @@ impl ThreadEnvironments {
.into_input_selection(inherited.selection.clone()),
&thread_defaults,
);
let (attachment, detach_on_drop, _) = EnvironmentAttachment::new(
&selection.environment_id,
Arc::clone(&inherited.attachment.environment),
);
return Some(Self::start_environment(
selection,
config_origin,
Arc::clone(&inherited.environment),
attachment,
Arc::new(detach_on_drop),
inherited.owner_configuration(),
/*pending_completion*/ None,
&local_shell,
@@ -386,7 +440,7 @@ impl ThreadEnvironments {
};
let selection = environment.selection;
let config_origin = environment.config_origin;
let selected_environment = Arc::clone(&environment.environment);
let selected_environment = &environment.attachment.environment;
let inherited_snapshot = if !selected_environment.is_remote()
&& shell_snapshot.should_rebuild_inherited()
{
@@ -396,7 +450,6 @@ impl ThreadEnvironments {
};
let resolution: TurnEnvironmentResolution =
futures::future::ready(Ok(ResolvedEnvironment {
environment: environment.environment,
shell: environment.shell,
user_home_dir: environment.user_home_dir,
executor_platform_os: environment.executor_platform_os,
@@ -411,11 +464,15 @@ impl ThreadEnvironments {
}))
.boxed()
.shared();
let (attachment, detach_on_drop, _) = EnvironmentAttachment::new(
&selection.environment_id,
Arc::clone(selected_environment),
);
let mut inherited_environment = SelectedTurnEnvironment {
selection,
config_origin,
environment: selected_environment,
connection_events_task: None,
attachment,
detach_on_drop: Arc::new(detach_on_drop),
resolution,
pending_completion: None,
owner_config_result: None,
@@ -425,10 +482,6 @@ impl ThreadEnvironments {
Some(inherited_environment)
})
.collect();
let active = environments
.iter()
.map(|environment| environment.for_turn(&local_shell, &shell_snapshot))
.collect();
Self {
environment_manager,
local_shell,
@@ -436,16 +489,46 @@ impl ThreadEnvironments {
non_blocking_snapshots,
state: Mutex::new(ThreadEnvironmentsState {
thread_defaults: thread_defaults.clone(),
environments: active,
environments: Vec::new(),
connection_events: None,
}),
configured: Mutex::new(ThreadEnvironmentsState {
thread_defaults,
environments,
connection_events: None,
}),
connection_event_tx: OnceLock::new(),
}
}
/// Attach the initial active selection after `SessionConfigured` so clients see only later
/// connection changes. Startup discovery uses the configured bindings until this point.
pub(crate) fn activate_configured(&self, tx_event: Sender<Event>) {
let mut state = self.state.lock().unwrap_or_else(PoisonError::into_inner);
let configured = self
.configured
.lock()
.unwrap_or_else(PoisonError::into_inner);
let connection_events = ConnectionEventStreams::new(tx_event);
state.environments = configured
.environments
.iter()
.map(|environment| {
let (attachment, detach_on_drop, events) = EnvironmentAttachment::new(
&environment.selection.environment_id,
Arc::clone(&environment.attachment.environment),
);
connection_events.add(events, &attachment.detached);
environment.for_turn(
attachment,
Arc::new(detach_on_drop),
&self.local_shell,
&self.shell_snapshot,
)
})
.collect();
state.connection_events = Some(connection_events);
}
fn start_shell_snapshot_task(
shell_snapshot: ShellSnapshot,
environment: Arc<Environment>,
@@ -503,73 +586,82 @@ impl ThreadEnvironments {
}
fn apply_selections(&self, environments: &[TurnEnvironmentSelection], retry: StartupRetry<'_>) {
let bindings = self
.configured
.lock()
.unwrap_or_else(PoisonError::into_inner)
.environments
.clone();
self.update_selections_with_bindings(
&self.state,
environments,
&bindings,
self.connection_event_tx.get(),
retry,
);
{
let mut state = self.state.lock().unwrap_or_else(PoisonError::into_inner);
let configured = self
.configured
.lock()
.unwrap_or_else(PoisonError::into_inner);
self.update_selections_with_bindings(
&mut state,
environments,
&configured.environments,
retry,
);
}
// Retry configured bindings with their own settings, keeping turn-only changes temporary.
self.update_selections_with_bindings(
&self.configured,
&self.configured_snapshot_now().all_selections(),
&[],
/*tx_event*/ None,
retry,
);
self.apply_configured_selections(&self.configured_snapshot_now().all_selections(), retry);
}
/// Retains configured bindings without changing the active turn's attachments.
pub(crate) fn update_configured_selections(&self, environments: &[TurnEnvironmentSelection]) {
self.update_selections_with_bindings(
&self.configured,
environments,
&[],
/*tx_event*/ None,
StartupRetry::AfterSelectionUpdate,
);
self.apply_configured_selections(environments, StartupRetry::AfterSelectionUpdate);
}
fn apply_configured_selections(
&self,
environments: &[TurnEnvironmentSelection],
retry: StartupRetry<'_>,
) {
let mut configured = self
.configured
.lock()
.unwrap_or_else(PoisonError::into_inner);
self.update_selections_with_bindings(&mut configured, environments, &[], retry);
}
fn update_selections_with_bindings(
&self,
state: &Mutex<ThreadEnvironmentsState>,
state: &mut ThreadEnvironmentsState,
environments: &[TurnEnvironmentSelection],
bindings: &[SelectedTurnEnvironment],
tx_event: Option<&Sender<Event>>,
retry: StartupRetry<'_>,
) {
let mut state = state.lock().unwrap_or_else(PoisonError::into_inner);
let ThreadEnvironmentsState {
thread_defaults,
environments: current,
} = &mut *state;
connection_events,
} = state;
let connection_events = connection_events.as_ref();
let mut seen_environment_ids = HashSet::with_capacity(environments.len());
let mut next = Vec::with_capacity(environments.len());
let mut pending_completions = Vec::new();
// Connection state belongs to the environment instance, not its cwd or roots.
let connection_events_task = |environment_id: &str, environment: &Arc<Environment>| {
current
.iter()
.find(|previous| {
previous.selection.environment_id == environment_id
&& Arc::ptr_eq(&previous.environment, environment)
})
.and_then(|previous| previous.connection_events_task.clone())
.or_else(|| {
Self::spawn_connection_event_listener(
environment.as_ref(),
environment_id.to_owned(),
tx_event?.clone(),
)
})
};
// Workspace and configuration changes preserve the attachment to the same executor.
let attachment_for =
|environment_id: &str,
environment: &Arc<Environment>,
connection_events: Option<&ConnectionEventStreams>| {
current
.iter()
.find(|previous| {
previous.selection.environment_id == environment_id
&& Arc::ptr_eq(&previous.attachment.environment, environment)
})
.map(|previous| {
(
Arc::clone(&previous.attachment),
Arc::clone(&previous.detach_on_drop),
)
})
.unwrap_or_else(|| {
let (attachment, guard, events) =
EnvironmentAttachment::new(environment_id, Arc::clone(environment));
if let Some(connection_events) = connection_events {
connection_events.add(events, &attachment.detached);
}
(attachment, Arc::new(guard))
})
};
for selected_environment in environments {
if !seen_environment_ids.insert(selected_environment.environment_id.as_str()) {
continue;
@@ -590,10 +682,27 @@ impl ThreadEnvironments {
resolve_selection_config(selected_environment.clone(), thread_defaults);
if let Some(environment) = previous {
let failed = retry.should_restart(environment);
let restarting_as_pending =
matches!(selected_environment.config, EnvironmentConfigState::Pending)
&& !matches!(
environment.selection.config,
EnvironmentConfigState::Pending
);
let reuse_startup = !failed && !restarting_as_pending;
let mut environment = if current_environment.is_some() {
environment.clone()
} else {
environment.for_turn(&self.local_shell, &self.shell_snapshot)
let (attachment, detach_on_drop) = attachment_for(
&selected_environment.environment_id,
&environment.attachment.environment,
connection_events.filter(|_| reuse_startup),
);
environment.for_turn(
attachment,
detach_on_drop,
&self.local_shell,
&self.shell_snapshot,
)
};
if let Some(completion) = environment.pending_completion.take() {
match &selected_environment.config {
@@ -608,14 +717,7 @@ impl ThreadEnvironments {
}
}
}
let restarting_as_pending =
matches!(selected_environment.config, EnvironmentConfigState::Pending)
&& !matches!(
environment.selection.config,
EnvironmentConfigState::Pending
);
if !failed && !restarting_as_pending {
if reuse_startup {
let shell_settings_changed = matches!(
(&environment.selection.config, &selected_environment.config),
(EnvironmentConfigState::Ready(previous), EnvironmentConfigState::Ready(current))
@@ -625,18 +727,12 @@ impl ThreadEnvironments {
environment.selection = selected_environment;
environment.config_origin = config_origin;
if shell_settings_changed
&& !environment.environment.is_remote()
&& !environment.attachment.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);
}
if environment.connection_events_task.is_none() {
environment.connection_events_task = connection_events_task(
&environment.selection.environment_id,
&environment.environment,
);
}
next.push(environment);
continue;
}
@@ -647,36 +743,26 @@ impl ThreadEnvironments {
tracing::warn!("skipping unknown turn environment `{environment_id}`");
continue;
};
let connection_events_task = connection_events_task(environment_id, &environment);
let (attachment, detach_on_drop) =
attachment_for(environment_id, &environment, connection_events);
let pending = previous.filter(|previous| {
matches!(selected_environment.config, EnvironmentConfigState::Pending)
&& previous.pending_completion.is_some()
});
let mut selected = Self::start_environment(
let selected = Self::start_environment(
selected_environment,
config_origin,
environment,
attachment,
detach_on_drop,
pending.and_then(|previous| previous.owner_config_result.clone()),
current_environment.and_then(|previous| previous.pending_completion.clone()),
&self.local_shell,
&self.shell_snapshot,
);
selected.connection_events_task = connection_events_task;
next.push(selected);
}
let removed_connection_tasks = current
.iter()
.filter_map(|previous| {
let task = previous.connection_events_task.as_ref()?;
(!next.iter().any(|next| {
next.connection_events_task
.as_ref()
.is_some_and(|next_task| Arc::ptr_eq(task, next_task))
}))
.then(|| Arc::clone(task))
})
.collect::<Vec<_>>();
let previous = std::mem::replace(current, next);
// Install the new selection before dropping the guards for removed attachments.
drop(std::mem::replace(current, next));
// Publish after installing the new list, including pending bindings replaced by retries.
for (completion, result) in pending_completions {
@@ -688,18 +774,14 @@ impl ThreadEnvironments {
true
});
}
for task in removed_connection_tasks {
task.abort();
}
drop(previous);
}
#[allow(clippy::too_many_arguments)]
fn start_environment(
selection: TurnEnvironmentSelection,
config_origin: EnvironmentConfigOrigin,
environment: Arc<Environment>,
attachment: Arc<EnvironmentAttachment>,
detach_on_drop: Arc<DropGuard>,
owner_config_result: Option<PendingConfiguration>,
pending_completion: Option<watch::Sender<Option<PendingConfigurationResult>>>,
local_shell: &Shell,
@@ -717,7 +799,7 @@ impl ThreadEnvironments {
// The resolver waits for this Session; descendants can follow the original owner.
let (task, resolution) = Self::resolve_environment(
selection.clone(),
Arc::clone(&environment),
Arc::clone(&attachment.environment),
local_shell.clone(),
shell_snapshot.clone(),
configuration_ready,
@@ -729,8 +811,8 @@ impl ThreadEnvironments {
SelectedTurnEnvironment {
selection,
config_origin,
environment,
connection_events_task: None,
attachment,
detach_on_drop,
resolution: resolution.boxed().shared(),
pending_completion,
owner_config_result,
@@ -798,6 +880,7 @@ impl ThreadEnvironments {
let ThreadEnvironmentsState {
thread_defaults,
environments,
..
} = &mut *state;
if thread_defaults.common == defaults {
return;
@@ -814,7 +897,7 @@ impl ThreadEnvironments {
);
environment.refresh_thread_config(thread_defaults);
if shell_settings_changed
&& !environment.environment.is_remote()
&& !environment.attachment.environment.is_remote()
&& self.shell_snapshot.should_rebuild_inherited()
&& let Some(Ok(resolved)) = environment.resolution.clone().now_or_never()
{
@@ -838,7 +921,7 @@ impl ThreadEnvironments {
.iter_mut()
.chain(&mut configured.environments)
{
if !selected.environment.is_remote()
if !selected.attachment.environment.is_remote()
&& let Some(Ok(resolved)) = selected.resolution.clone().now_or_never()
{
self.restart_shell_snapshot(selected, resolved);
@@ -853,7 +936,7 @@ impl ThreadEnvironments {
) {
let shell_snapshot = Self::start_shell_snapshot_task(
self.shell_snapshot.clone(),
Arc::clone(&resolved.environment),
Arc::clone(&selected.attachment.environment),
selected.selection.cwd.clone(),
resolved.shell.clone(),
&selected.selection.config,
@@ -886,7 +969,7 @@ impl ThreadEnvironments {
Some(
environment
.config_origin
.selected_capability_roots(&environment.environment, config),
.selected_capability_roots(&environment.attachment.environment, config),
)
}),
);
@@ -908,65 +991,6 @@ impl ThreadEnvironments {
status
}
fn spawn_connection_event_listener(
environment: &Environment,
environment_id: String,
tx_event: Sender<Event>,
) -> Option<Arc<AbortOnDropHandle<()>>> {
let mut connection_state = environment.subscribe_connection_state()?;
let task = tokio::spawn(async move {
loop {
let state = tokio::select! {
_ = tx_event.closed() => return,
changed = connection_state.changed() => {
if changed.is_err() {
return;
}
*connection_state.borrow_and_update()
}
};
let msg = match state {
EnvironmentConnectionState::Recovering => continue,
EnvironmentConnectionState::Connected => {
EventMsg::EnvironmentConnected(EnvironmentConnectionEvent {
environment_id: environment_id.clone(),
})
}
EnvironmentConnectionState::Disconnected => {
EventMsg::EnvironmentDisconnected(EnvironmentConnectionEvent {
environment_id: environment_id.clone(),
})
}
};
if tx_event
.send(Event {
id: String::new(),
msg,
})
.await
.is_err()
{
return;
}
}
});
Some(Arc::new(AbortOnDropHandle::new(task)))
}
pub(crate) fn start_connection_event_forwarding(&self, tx_event: Sender<Event>) {
let mut state = self.state.lock().unwrap_or_else(PoisonError::into_inner);
let tx_event = self.connection_event_tx.get_or_init(|| tx_event);
for selected in &mut state.environments {
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(
name = "environments.resolve",
skip_all,
@@ -1061,7 +1085,6 @@ impl ThreadEnvironments {
let shell_snapshot_v2_supported = snapshot_v2
&& (environment.is_remote() || !shell_snapshot_builder.should_rebuild_inherited());
Ok(ResolvedEnvironment {
environment,
shell,
user_home_dir,
executor_platform_os,
@@ -1106,7 +1129,7 @@ impl ThreadEnvironments {
selection: environment.selection.clone(),
config_origin: environment.config_origin,
resolution: environment.resolution.clone(),
environment: Arc::clone(&environment.environment),
attachment: Arc::clone(&environment.attachment),
owner_config_result: environment.owner_config_result.clone().or_else(|| {
environment
.pending_completion
@@ -1128,7 +1151,7 @@ impl ThreadEnvironments {
non_blocking = non_blocking_snapshots,
)
)]
fn resolve_snapshot(
pub(crate) fn resolve_snapshot(
&self,
snapshot: TurnEnvironmentSnapshot,
) -> impl Future<Output = TurnEnvironmentSnapshot> + Send + 'static {
@@ -1169,7 +1192,7 @@ impl TurnEnvironmentState {
starting: StartingTurnEnvironment,
resolved: Option<TurnEnvironmentResult>,
) -> Self {
let required_skills = starting.environment.required_skills().to_vec();
let required_skills = starting.attachment.environment.required_skills().to_vec();
if let EnvironmentConfigState::Failed(error) = &starting.selection.config {
let error = error.clone();
return Self::Failed {
@@ -1193,10 +1216,10 @@ impl TurnEnvironmentState {
};
selection.config = EnvironmentConfigState::Ready(config);
}
let mut turn_environment = TurnEnvironment::new(
let mut turn_environment = TurnEnvironment::from_attachment(
selection,
starting.config_origin,
environment.environment,
starting.attachment,
environment.shell,
);
turn_environment.executor_platform_os = environment.executor_platform_os;
@@ -1313,7 +1336,7 @@ impl TurnEnvironmentSnapshot {
.map(|environment| {
(
environment.selection.environment_id.clone(),
Arc::clone(&environment.environment),
Arc::clone(environment.environment()),
)
})
.collect()
@@ -1325,7 +1348,7 @@ impl TurnEnvironmentSnapshot {
.map(|environment| {
(
environment.selection.environment_id.clone(),
Some(Arc::clone(&environment.environment)),
Some(Arc::clone(environment.environment())),
)
})
.chain(
@@ -1362,7 +1385,7 @@ impl TurnEnvironmentSnapshot {
pub(crate) fn local(&self) -> Option<&TurnEnvironment> {
self.turn_environments()
.find(|environment| !environment.environment.is_remote())
.find(|environment| !environment.environment().is_remote())
}
pub(crate) fn local_environment_cwd(&self) -> Option<AbsolutePathBuf> {
@@ -1376,7 +1399,7 @@ impl TurnEnvironmentSnapshot {
.iter()
.find_map(|environment| match environment {
TurnEnvironmentState::Ready(environment)
if !environment.environment.is_remote() =>
if !environment.environment().is_remote() =>
{
Some(&environment.selection)
}
@@ -1394,7 +1417,7 @@ impl TurnEnvironmentSnapshot {
#[cfg(test)]
pub(crate) fn primary_environment(&self) -> Option<Arc<codex_exec_server::Environment>> {
self.primary()
.map(|environment| Arc::clone(&environment.environment))
.map(|environment| Arc::clone(environment.environment()))
}
pub(crate) fn to_selections(&self) -> Vec<TurnEnvironmentSelection> {
@@ -1409,11 +1432,11 @@ impl TurnEnvironmentSnapshot {
let (selection, required) = match environment {
TurnEnvironmentState::Ready(environment) => (
&environment.selection,
environment.environment.required_skills(),
environment.environment().required_skills(),
),
TurnEnvironmentState::Starting(environment) => (
&environment.selection,
environment.environment.required_skills(),
environment.attachment.environment.required_skills(),
),
TurnEnvironmentState::Failed {
selection,
@@ -1473,7 +1496,7 @@ impl TurnEnvironmentSnapshot {
return None;
}
(!environment.environment.is_remote()).then_some(environment)
(!environment.environment().is_remote()).then_some(environment)
}
pub(crate) fn single_local_environment_cwd(&self) -> Option<AbsolutePathBuf> {
@@ -1507,6 +1530,8 @@ mod tests {
use codex_protocol::models::ActivePermissionProfile;
use codex_protocol::models::PermissionProfile;
use codex_protocol::permissions::FileSystemSandboxPolicyContext;
use codex_protocol::protocol::EnvironmentConnectionEvent;
use codex_protocol::protocol::EventMsg;
use codex_protocol::protocol::TurnEnvironmentRequest;
use codex_protocol::protocol::TurnEnvironmentSelection;
use codex_protocol::sandbox::SandboxType;
@@ -2124,6 +2149,7 @@ url = "ws://127.0.0.1:8765"
attached.clone(),
/*non_blocking_snapshots*/ true,
);
child.activate_configured(async_channel::unbounded().0);
assert_eq!(
child
.snapshot()
@@ -2320,7 +2346,7 @@ url = "ws://127.0.0.1:8765"
TurnEnvironmentSnapshot::default(),
/*non_blocking_snapshots*/ true,
));
environments.start_connection_event_forwarding(tx_event);
environments.activate_configured(tx_event);
environments.update_selections(std::slice::from_ref(&selection));
let initial_snapshot = environments.snapshot().await;
let second_listener = TcpListener::bind("127.0.0.1:0")
@@ -2361,6 +2387,9 @@ url = "ws://127.0.0.1:8765"
.expect("changed environment");
assert!(initial.resolution.ptr_eq(&reused.resolution));
assert!(!reused.resolution.ptr_eq(&changed.resolution));
assert!(Arc::ptr_eq(&initial.attachment, &reused.attachment));
assert!(initial.attachment.detached.is_cancelled());
assert!(!changed.attachment.detached.is_cancelled());
serve_environment_info(first_listener).await;
assert!(
@@ -2456,7 +2485,7 @@ url = "ws://127.0.0.1:8765"
let snapshot = environments.snapshot().await;
let inherited = snapshot.primary().expect("inherited environment");
assert!(Arc::ptr_eq(&inherited.environment, &inherited_environment));
assert!(Arc::ptr_eq(inherited.environment(), &inherited_environment));
assert_eq!(inherited.config(), &child_config);
assert_eq!(
inherited
@@ -2546,6 +2575,7 @@ url = "ws://127.0.0.1:8765"
parent.snapshot().await,
/*non_blocking_snapshots*/ false,
);
child.activate_configured(async_channel::unbounded().0);
let child_snapshot = child.snapshot().await;
let child_environment = child_snapshot.primary().expect("child environment");
@@ -2641,6 +2671,12 @@ url = "ws://127.0.0.1:8765"
let cwd = AbsolutePathBuf::current_dir()
.expect("cwd")
.join("starting-local");
let (attachment, _guard, _) = EnvironmentAttachment::new(
LOCAL_ENVIRONMENT_ID,
Arc::new(
Environment::create_for_tests(/*exec_server_url*/ None).expect("local environment"),
),
);
let snapshot = TurnEnvironmentSnapshot {
environments: vec![TurnEnvironmentState::Starting(StartingTurnEnvironment {
selection: TurnEnvironmentSelection {
@@ -2652,10 +2688,7 @@ url = "ws://127.0.0.1:8765"
},
config_origin: EnvironmentConfigOrigin::Thread,
resolution: futures::future::pending().boxed().shared(),
environment: Arc::new(
Environment::create_for_tests(/*exec_server_url*/ None)
.expect("local environment"),
),
attachment,
owner_config_result: None,
})],
};
@@ -0,0 +1,102 @@
//! The lifetime of one continuous selection of an executor.
use std::sync::Arc;
use async_channel::Receiver;
use async_channel::Sender;
use codex_exec_server::Environment;
use codex_exec_server::EnvironmentConnectionState;
use codex_protocol::protocol::EnvironmentConnectionEvent;
use codex_protocol::protocol::Event;
use codex_protocol::protocol::EventMsg;
use tokio::sync::watch;
use tokio_util::sync::CancellationToken;
use tokio_util::sync::DropGuard;
use tokio_util::task::AbortOnDropHandle;
/// Tracks one executor's membership in a thread's active or configured environment list.
/// Turn snapshots share the attachment to access the executor and observe when it leaves that list.
///
/// Removing the executor drops the selection's guard and cancels `detached`, even while
/// an older snapshot holds the attachment. Selecting the same executor again creates a new attachment.
pub(crate) struct EnvironmentAttachment {
pub(crate) environment: Arc<Environment>,
pub(crate) detached: CancellationToken,
_connection_events_task: Option<AbortOnDropHandle<()>>,
}
impl EnvironmentAttachment {
/// Returns the attachment, its detach guard, and a receiver for remote connection events.
/// Snapshots share the attachment while the selection owns the guard, so dropping the guard
/// signals detach immediately. The receiver starts at construction and can be dropped by
/// callers that do not publish events. Local executors have no connection stream.
pub(crate) fn new(
environment_id: &str,
environment: Arc<Environment>,
) -> (Arc<Self>, DropGuard, Option<Receiver<Event>>) {
let detached = CancellationToken::new();
let (connection_events_task, events) = environment
.subscribe_connection_state()
.map(|connection_state| {
let (tx_event, events) = async_channel::unbounded();
let task = tokio::spawn(forward_connection_events(
connection_state,
environment_id.to_owned(),
tx_event,
detached.clone(),
));
(AbortOnDropHandle::new(task), events)
})
.unzip();
let attachment = Arc::new(Self {
environment,
detached,
_connection_events_task: connection_events_task,
});
let detach_on_drop = attachment.detached.clone().drop_guard();
(attachment, detach_on_drop, events)
}
}
async fn forward_connection_events(
mut connection_state: watch::Receiver<EnvironmentConnectionState>,
environment_id: String,
tx_event: Sender<Event>,
detached: CancellationToken,
) {
loop {
let state = tokio::select! {
biased;
_ = detached.cancelled() => return,
_ = tx_event.closed() => return,
changed = connection_state.changed() => {
if changed.is_err() {
return;
}
*connection_state.borrow_and_update()
}
};
let msg = match state {
EnvironmentConnectionState::Recovering => continue,
EnvironmentConnectionState::Connected => {
EventMsg::EnvironmentConnected(EnvironmentConnectionEvent {
environment_id: environment_id.clone(),
})
}
EnvironmentConnectionState::Disconnected => {
EventMsg::EnvironmentDisconnected(EnvironmentConnectionEvent {
environment_id: environment_id.clone(),
})
}
};
if tx_event
.try_send(Event {
id: String::new(),
msg,
})
.is_err()
{
return;
}
}
}
@@ -69,6 +69,7 @@ async fn capability_roots_follow_attachments_through_configuration_changes() {
&snapshot.environments[1],
TurnEnvironmentState::Failed { .. }
));
let starting = Arc::clone(&snapshot.starting().next().unwrap().attachment);
let child = ThreadEnvironments::new(
Arc::clone(&manager),
@@ -78,6 +79,7 @@ async fn capability_roots_follow_attachments_through_configuration_changes() {
snapshot,
/*non_blocking_snapshots*/ false,
);
child.activate_configured(async_channel::unbounded().0);
assert_eq!(
child.snapshot().await.selected_capability_roots(),
vec![roots[1].clone()]
@@ -109,21 +111,20 @@ async fn capability_roots_follow_attachments_through_configuration_changes() {
independent_roots
);
assert_eq!(parent.snapshot().await.selected_capability_roots(), roots);
let child_snapshot = child.snapshot().await;
let child_local = child_snapshot.primary().expect("child executor");
let parent_snapshot = parent.snapshot().await;
let parent_local = parent_snapshot.primary().expect("parent executor");
assert!(Arc::ptr_eq(
&child
.snapshot()
.await
.primary()
.expect("child executor")
.environment,
&parent
.snapshot()
.await
.turn_environments()
.find(|environment| environment.selection.environment_id == LOCAL_ENVIRONMENT_ID)
.expect("parent executor")
.environment,
child_local.environment(),
parent_local.environment()
));
assert!(Arc::ptr_eq(&starting, &parent_local.attachment));
parent.update_selections(std::slice::from_ref(&local));
assert!(!parent_local.attachment.detached.is_cancelled());
drop(parent);
assert!(parent_local.attachment.detached.is_cancelled());
assert!(!child_local.attachment.detached.is_cancelled());
}
#[test_case::test_case(false; "standalone startup")]
@@ -149,7 +150,7 @@ async fn unpolled_snapshot_does_not_delay_canceling_a_removed_environment(shared
if shared_startup {
environments.update_configured_selections(std::slice::from_ref(&selection));
}
environments.update_selections(&[selection]);
environments.update_selections(std::slice::from_ref(&selection));
let starting = environments
.snapshot()
.await
@@ -158,8 +159,10 @@ async fn unpolled_snapshot_does_not_delay_canceling_a_removed_environment(shared
.unwrap()
.clone();
let held_snapshot = environments.snapshot();
let original = Arc::downgrade(&starting.attachment);
environments.update_selections(&[]);
assert!(starting.attachment.detached.is_cancelled());
let error = timeout(Duration::from_secs(5), starting.wait_until_ready())
.await
@@ -174,6 +177,19 @@ async fn unpolled_snapshot_does_not_delay_canceling_a_removed_environment(shared
held_snapshot.await.environments.as_slice(),
[TurnEnvironmentState::Failed { .. }]
));
environments.update_selections(&[selection]);
let snapshot = environments.snapshot_now();
let reattached = &snapshot.starting().next().unwrap().attachment;
assert!(Arc::ptr_eq(
&reattached.environment,
&starting.attachment.environment
));
assert!(!Arc::ptr_eq(reattached, &starting.attachment));
environments.update_configured_selections(&[]);
assert!(!reattached.detached.is_cancelled());
drop(starting);
assert!(original.upgrade().is_none());
}
#[tokio::test]
+1 -1
View File
@@ -78,7 +78,7 @@ pub(super) fn for_environment(
None => context.environments().primary(),
};
let native_cwd = environment
.filter(|environment| !environment.environment.is_remote())
.filter(|environment| !environment.environment().is_remote())
.and_then(|environment| environment.cwd().to_abs_path().ok());
let permission_profile = environment
.map(TurnEnvironment::permission_profile_with_workspace_roots)
@@ -146,7 +146,7 @@ pub(crate) async fn review_tool_call(
ExecCommandLifetime::OneShot
};
let prepared = parsed.prepare(lifetime).map_err(tool_error)?;
let fs = prepared.turn_environment.environment.get_filesystem();
let fs = prepared.turn_environment.environment().get_filesystem();
if let Some(patch) = prepare_intercepted_patch(
&prepared.resolved.command,
&prepared.cwd,
+2 -2
View File
@@ -68,8 +68,8 @@ async fn plugin_attribution_for_guardian_request(
.find(|environment| {
environment.selection.environment_id.as_str() == environment_id
})?;
if turn_environment.environment.is_remote() {
let file_system = turn_environment.environment.get_filesystem();
if turn_environment.environment().is_remote() {
let file_system = turn_environment.environment().get_filesystem();
turn.plugin_attribution_for_executor_command(command, cwd, file_system.as_ref())
.await
} else {
+8 -3
View File
@@ -8,6 +8,7 @@ use crate::config::NetworkProxySpec;
use crate::config::PermissionProfileSnapshot;
use crate::config::test_config;
use crate::context::ContextualUserFragment;
use crate::environment_selection::EnvironmentAttachment;
use crate::environment_selection::TurnEnvironmentState;
use crate::guardian::approval_request::guardian_request_target_item_id;
use crate::guardian::review::guardian_review_session_config;
@@ -679,10 +680,14 @@ async fn approval_permissions_use_the_owning_environment() -> anyhow::Result<()>
));
let mut windows = secondary.clone();
let windows_cwd = PathUri::parse("file:///C:/guardian-secondary")?;
windows.environment = Arc::new(codex_exec_server::Environment::create_for_tests(Some(
"ws://127.0.0.1:1".into(),
))?);
windows.selection.environment_id = "windows".into();
let (attachment, _guard, _) = EnvironmentAttachment::new(
&windows.selection.environment_id,
Arc::new(codex_exec_server::Environment::create_for_tests(Some(
"ws://127.0.0.1:1".into(),
))?),
);
windows.attachment = attachment;
windows.selection.cwd = windows_cwd.clone();
windows.config_mut().workspace_roots = vec![windows_cwd.clone()];
windows.config_mut().permission_profile =
@@ -208,7 +208,7 @@ impl HookEnvironmentExecutor for EnvironmentCommandExecutor {
ToolError::Rejected(message) => anyhow::anyhow!("{message}"),
ToolError::Codex(error) => anyhow::Error::new(error),
})?;
let backend = environment.environment.get_exec_backend();
let backend = environment.environment().get_exec_backend();
let process = tokio::time::timeout_at(deadline, backend.start(params))
.await
.context("environment command hook timed out during launch")??
+2 -2
View File
@@ -178,7 +178,7 @@ async fn build_uploaded_argument_value(
requires_sandbox.then(|| turn_environment.sandbox_context(additional_permissions));
if sandbox.is_some() {
let environment_info = turn_environment
.environment
.environment()
.info()
.await
.map_err(|error| contextualize_error(error.to_string()))?;
@@ -188,7 +188,7 @@ async fn build_uploaded_argument_value(
));
}
}
let fs = turn_environment.environment.get_filesystem();
let fs = turn_environment.environment().get_filesystem();
let metadata = fs
.get_metadata(&path_uri, GetMetadataOptions::default(), sandbox.as_ref())
.await
+1 -1
View File
@@ -1373,7 +1373,7 @@ async fn mcp_sandbox_cwd_uses_matching_server_environment_uri() -> anyhow::Resul
.initial_environments
.primary()
.expect("primary environment")
.environment
.environment()
.clone();
turn_context
.initial_environments
+5 -5
View File
@@ -558,17 +558,17 @@ impl Session {
.get_or_init(ExecutorCapabilityDiscoveryV2Cache::default);
let remote_environment_count = environments
.turn_environments()
.filter(|selected| selected.environment.is_remote())
.filter(|selected| selected.environment().is_remote())
.count();
cache.ensure_capacity(remote_environment_count);
let mut snapshots = Vec::new();
let mut skipped_incompatible_executor = false;
for selected in environments.turn_environments() {
if !selected.environment.is_remote() {
if !selected.environment().is_remote() {
continue;
}
let environment_id = &selected.selection.environment_id;
let info = match selected.environment.info().await {
let info = match selected.environment().info().await {
Ok(info) => info,
Err(_) => {
tracing::warn!(
@@ -587,7 +587,7 @@ impl Session {
sandbox: Some(selected.sandbox_context(/*additional_permissions*/ None)),
};
let result = cache
.discover(environment_id, &selected.environment, request)
.discover(environment_id, selected.environment(), request)
.await;
snapshots.push(ExecutorCapabilityDiscoveryV2Snapshot {
environment_id: environment_id.clone(),
@@ -632,7 +632,7 @@ impl Session {
environments.turn_environments().map(|environment| {
environment
.config_origin
.selected_capability_roots(&environment.environment, environment.config())
.selected_capability_roots(environment.environment(), environment.config())
}),
);
for (index, root) in combined_roots.into_iter().enumerate() {
+1 -1
View File
@@ -154,7 +154,7 @@ impl Session {
let mut catalog = None;
let mut registered = HashSet::new();
for selected in environments.turn_environments() {
let environment = &selected.environment;
let environment = selected.environment();
if !environment.is_remote() {
continue;
}
+7 -5
View File
@@ -1535,9 +1535,11 @@ impl Session {
config.features.enabled(Feature::DeferredExecutor),
));
turn_environments.update_configured_selections(environment_selections);
turn_environments.update_selections(environment_selections);
session_configuration.environments = turn_environments.selections();
let resolved_environments = turn_environments.snapshot().await;
let configured_environments = turn_environments.configured_snapshot_now();
session_configuration.environments = configured_environments.all_selections();
let resolved_environments = turn_environments
.resolve_snapshot(configured_environments)
.await;
let agents_md_manager = Arc::new(AgentsMdManager::new(instructions));
// Keep discovery's future out of the large embedded startup/resume frame.
let plugin_skill_warmup = Box::pin(warm_plugins_and_skills_for_session_init(
@@ -1976,7 +1978,7 @@ impl Session {
let thread_config = sess
.effective_thread_config_snapshot(
&session_configuration,
turn_environments.selections(),
session_configuration.environments.clone(),
)
.await;
let events = std::iter::once(Event {
@@ -2009,7 +2011,7 @@ impl Session {
for event in events {
sess.send_event_raw(event).await;
}
turn_environments.start_connection_event_forwarding(tx_event.clone());
turn_environments.activate_configured(tx_event.clone());
let startup_auth_changed = mcp_auth_changes.has_changed().unwrap_or(false);
if startup_auth_changed {
+25 -22
View File
@@ -21,6 +21,7 @@ use crate::config::RuntimeConfigRefresh;
use crate::context::ContextualUserFragment;
use crate::context::DeveloperInstructions;
use crate::context::TurnAborted;
use crate::environment_selection::EnvironmentAttachment;
use crate::environment_selection::EnvironmentConfigOrigin;
use crate::environment_selection::ThreadEnvironments;
use crate::environment_selection::TurnEnvironmentState;
@@ -7026,11 +7027,12 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) {
resolved_environments,
/*non_blocking_snapshots*/ false,
));
turn_environments.activate_configured(tx_event.clone());
let environment = Arc::clone(
&resolved_turn_environments
resolved_turn_environments
.primary()
.expect("primary environment")
.environment,
.environment(),
);
let plugins_manager = Arc::new(plugins_manager_for_config(
&config,
@@ -7958,7 +7960,7 @@ async fn request_permissions_tool_resolves_legacy_paths_against_selected_environ
config: EnvironmentConfigState::Ready(current_environment_config),
},
current_environment.config_origin,
current_environment.environment,
Arc::clone(current_environment.environment()),
current_environment.shell,
)
};
@@ -8466,11 +8468,11 @@ async fn turn_environments_set_primary_environment() {
.primary()
.expect("primary environment should be set");
assert!(std::sync::Arc::ptr_eq(
&turn_environment.environment,
&turn_environments
turn_environment.environment(),
turn_environments
.primary()
.expect("primary environment")
.environment
.environment()
));
assert!(
turn_context
@@ -8495,17 +8497,17 @@ async fn turn_environments_set_primary_environment() {
};
assert!(Arc::ptr_eq(
&stored_environment,
&turn_environment.environment
turn_environment.environment()
));
let default_turn = session.new_default_turn().await;
assert!(Arc::ptr_eq(
&stored_environment,
&default_turn
default_turn
.initial_environments
.primary()
.expect("default turn primary environment")
.environment
.environment()
));
}
@@ -8584,11 +8586,11 @@ async fn task_turn_does_not_overlay_legacy_fallback_cwd_onto_stored_thread_envir
.primary()
.expect("primary environment should be set");
assert!(std::sync::Arc::ptr_eq(
&turn_environment.environment,
&turn_environments
turn_environment.environment(),
turn_environments
.primary()
.expect("primary environment")
.environment
.environment()
));
#[allow(deprecated)]
let turn_cwd = turn_context.cwd.clone();
@@ -8660,7 +8662,7 @@ async fn primary_environment_uses_first_turn_environment() {
config: EnvironmentConfigState::Ready(first_environment_config),
},
first_environment.config_origin,
Arc::clone(&first_environment.environment),
Arc::clone(first_environment.environment()),
/*shell*/ None,
)));
@@ -9355,11 +9357,12 @@ where
resolved_turn_environments.clone(),
/*non_blocking_snapshots*/ false,
));
turn_environments.activate_configured(tx_event.clone());
let environment = Arc::clone(
&resolved_turn_environments
resolved_turn_environments
.primary()
.expect("primary environment")
.environment,
.environment(),
);
let plugins_manager = Arc::new(plugins_manager_for_config(
&config,
@@ -10630,7 +10633,7 @@ async fn record_context_updates_emits_environment_item_for_cwd_changes() {
.clone();
let environment_config = environment.config().clone();
current_context.initial_environments.environments[0] =
TurnEnvironmentState::Ready(TurnEnvironment::new(
TurnEnvironmentState::Ready(TurnEnvironment::from_attachment(
TurnEnvironmentSelection {
selected_capability_roots: Default::default(),
environment_id: environment.selection.environment_id,
@@ -10639,7 +10642,7 @@ async fn record_context_updates_emits_environment_item_for_cwd_changes() {
config: EnvironmentConfigState::Ready(environment_config),
},
environment.config_origin,
environment.environment,
environment.attachment,
environment.shell,
));
@@ -10697,7 +10700,7 @@ async fn record_context_updates_use_environment_permission_profile_and_workspace
environment_config.permission_profile =
PermissionProfileSnapshot::legacy(PermissionProfile::workspace_write());
current_context.initial_environments.environments[0] =
TurnEnvironmentState::Ready(TurnEnvironment::new(
TurnEnvironmentState::Ready(TurnEnvironment::from_attachment(
TurnEnvironmentSelection {
selected_capability_roots: Default::default(),
environment_id: environment.selection.environment_id,
@@ -10706,7 +10709,7 @@ async fn record_context_updates_use_environment_permission_profile_and_workspace
config: EnvironmentConfigState::Ready(environment_config),
},
environment.config_origin,
environment.environment,
environment.attachment,
environment.shell,
));
@@ -10777,7 +10780,7 @@ async fn record_context_updates_omits_environment_item_when_disabled() {
.clone();
let environment_config = environment.config().clone();
current_context.initial_environments.environments[0] =
TurnEnvironmentState::Ready(TurnEnvironment::new(
TurnEnvironmentState::Ready(TurnEnvironment::from_attachment(
TurnEnvironmentSelection {
selected_capability_roots: Default::default(),
environment_id: environment.selection.environment_id,
@@ -10786,7 +10789,7 @@ async fn record_context_updates_omits_environment_item_when_disabled() {
config: EnvironmentConfigState::Ready(environment_config),
},
environment.config_origin,
environment.environment,
environment.attachment,
environment.shell,
));
@@ -11353,7 +11356,7 @@ async fn turn_context_item_stores_local_cwd() {
config: EnvironmentConfigState::Ready(environment_config),
},
environment.config_origin,
environment.environment,
Arc::clone(environment.environment()),
environment.shell,
));
+1 -1
View File
@@ -947,7 +947,7 @@ async fn turn_diff_display_roots(step_context: &StepContext) -> Vec<(String, Pat
// A turn cwd is expected to be a directory. If it is a file, the failed `<cwd>/.git` probe
// is ignored and ancestor search continues from its parent.
let root = find_nearest_ancestor_with_markers(
turn_environment.environment.get_filesystem().as_ref(),
turn_environment.environment().get_filesystem().as_ref(),
cwd,
vec![".git".to_string()],
FindUpErrorPolicy::Ignore,
+25 -6
View File
@@ -4,6 +4,7 @@ use super::token_budget::resolve_token_budget;
use super::*;
use crate::config::TokenBudgetConfig;
use crate::cyber_access_program;
use crate::environment_selection::EnvironmentAttachment;
use crate::environment_selection::EnvironmentConfigOrigin;
use crate::environment_selection::TurnEnvironmentSnapshot;
use crate::exec_policy::AllowPrefixRules;
@@ -69,7 +70,7 @@ pub(crate) struct ShellSnapshotCacheKey {
pub(crate) struct TurnEnvironment {
pub(crate) selection: TurnEnvironmentSelection,
pub(crate) config_origin: EnvironmentConfigOrigin,
pub(crate) environment: Arc<Environment>,
pub(crate) attachment: Arc<EnvironmentAttachment>,
/// Cached from the selected executor; `None` means it did not report one.
pub(crate) user_home_dir: Option<PathUri>,
/// Cached from the selected executor; `None` means it did not report them.
@@ -84,17 +85,31 @@ pub(crate) struct TurnEnvironment {
}
impl TurnEnvironment {
#[cfg(test)]
pub(crate) fn new(
selection: TurnEnvironmentSelection,
config_origin: EnvironmentConfigOrigin,
environment: Arc<Environment>,
shell: Option<shell::Shell>,
) -> Self {
let (attachment, guard, _) =
EnvironmentAttachment::new(&selection.environment_id, environment);
// Standalone test snapshots have no thread selection to signal detach.
guard.disarm();
Self::from_attachment(selection, config_origin, attachment, shell)
}
pub(crate) fn from_attachment(
selection: TurnEnvironmentSelection,
config_origin: EnvironmentConfigOrigin,
attachment: Arc<EnvironmentAttachment>,
shell: Option<shell::Shell>,
) -> Self {
debug_assert!(matches!(selection.config, EnvironmentConfigState::Ready(_)));
Self {
selection,
config_origin,
environment,
attachment,
user_home_dir: None,
temporary_directories: None,
shell,
@@ -106,6 +121,10 @@ impl TurnEnvironment {
}
}
pub(crate) fn environment(&self) -> &Arc<Environment> {
&self.attachment.environment
}
pub(crate) fn config(&self) -> &EnvironmentConfig {
let EnvironmentConfigState::Ready(config) = &self.selection.config else {
unreachable!("ready turn environments always carry resolved configuration")
@@ -182,7 +201,7 @@ impl TurnEnvironment {
.as_ref()
.clone()
.build(
Arc::clone(&self.environment),
Arc::clone(self.environment()),
PathUri::from_abs_path(cwd),
Some(shell.clone()),
use_login_shell,
@@ -226,7 +245,7 @@ impl TurnEnvironment {
additional_permissions: Option<AdditionalPermissionProfile>,
) -> impl EnvironmentAccess + '_ {
FileSystemEnvironmentAccessor::new(
self.environment.filesystem_ref(),
self.environment().filesystem_ref(),
self.sandbox_context(additional_permissions),
)
}
@@ -264,7 +283,7 @@ impl TurnEnvironment {
&self,
) -> codex_file_system::WindowsSandboxSelection {
let config = self.config();
if self.environment.is_remote() {
if self.environment().is_remote() {
executor_windows_sandbox_selection(
config.windows_sandbox_type,
config.windows_sandbox_level,
@@ -299,7 +318,7 @@ impl std::fmt::Debug for TurnEnvironment {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("TurnEnvironment")
.field("environment_id", &self.selection.environment_id)
.field("environment", &self.environment)
.field("environment", self.environment())
.field("cwd", &self.selection.cwd)
.field("workspace_roots", &self.config().workspace_roots)
.field("user_home_dir", &self.user_home_dir)
+1 -1
View File
@@ -238,7 +238,7 @@ impl Session {
if turn_context.config.include_permissions_instructions {
let environment = step_context.environments.primary();
let sandbox = environment
.filter(|environment| environment.environment.is_remote())
.filter(|environment| environment.environment().is_remote())
.map(|environment| {
environment.sandbox_context(/*additional_permissions*/ None)
});
@@ -449,6 +449,7 @@ async fn inactive_profiles_keep_snapshots_but_active_brokers_require_sandbox() -
updated.clone(),
/*non_blocking_snapshots*/ false,
);
child.activate_configured(async_channel::unbounded().0);
let child = child.snapshot().await;
assert!(!Arc::ptr_eq(
&updated.primary().expect("parent turn").shell_snapshot_cache,
+1 -1
View File
@@ -56,7 +56,7 @@ pub(crate) fn skills_load_input_for_environments(
.cwd()
.to_abs_path()
.ok()
.map(|cwd| (cwd, environment.environment.get_filesystem()))
.map(|cwd| (cwd, environment.environment().get_filesystem()))
});
let (cwd, fs) = match workspace {
Some((cwd, fs)) => (cwd, Some(fs)),
@@ -488,7 +488,7 @@ pub(crate) async fn prepare_patch_request(
)));
}
};
let fs = turn_environment.environment.get_filesystem();
let fs = turn_environment.environment().get_filesystem();
let sandbox = turn_environment.sandbox_context(/*additional_permissions*/ None);
match codex_apply_patch::verify_apply_patch_args(
args,
@@ -528,7 +528,7 @@ fn prepare_verified_patch(
let cwd = action.cwd.clone();
let sandbox_context = turn_environment.sandbox_context(/*additional_permissions*/ None);
let policy_context = file_system_sandbox_policy_context_for_cwd(&sandbox_context, &cwd);
let sandbox_route = if turn_environment.environment.is_remote() {
let sandbox_route = if turn_environment.environment().is_remote() {
PatchSandboxRoute::ExecutorManaged
} else {
PatchSandboxRoute::Platform(windows_sandbox_level_for_legacy_checks(
@@ -8,6 +8,7 @@ use crate::agent::child_config::build_agent_spawn_config;
use crate::config::AgentRoleConfig;
use crate::config::DEFAULT_AGENT_MAX_DEPTH;
use crate::config::PermissionProfileSnapshot;
use crate::environment_selection::EnvironmentAttachment;
use crate::environment_selection::EnvironmentConfigOrigin;
use crate::environment_selection::TurnEnvironmentState;
use crate::function_tool::FunctionCallError;
@@ -2337,9 +2338,13 @@ async fn spawn_agent_reapplies_runtime_sandbox_after_role_config() {
else {
panic!("parent environment should be ready");
};
environment.environment = environment_manager
.default_environment()
.expect("sandbox-capable test environment should exist");
let (attachment, _guard, _) = EnvironmentAttachment::new(
&environment.selection.environment_id,
environment_manager
.default_environment()
.expect("sandbox-capable test environment should exist"),
);
environment.attachment = attachment;
environment.config_mut().permission_profile =
PermissionProfileSnapshot::legacy(expected_permission_profile.clone());
environment.config_origin = EnvironmentConfigOrigin::Owner;
@@ -176,7 +176,7 @@ impl ExecCommandHandler {
.await;
let prepared = parsed.prepare(self.lifetime)?;
let max_output_tokens = prepared.max_output_tokens;
let fs = prepared.turn_environment.environment.get_filesystem();
let fs = prepared.turn_environment.environment().get_filesystem();
let intercepted_patch = intercept_apply_patch(
&prepared.resolved.command,
&prepared.cwd,
@@ -223,7 +223,7 @@ impl ExecCommandHandler {
let process_id = manager.allocate_process_id().await;
let hook_command = prepared.hook_command.clone();
let tool_cwd =
(!prepared.turn_environment.environment.is_remote()).then(|| prepared.cwd.clone());
(!prepared.turn_environment.environment().is_remote()).then(|| prepared.cwd.clone());
let request = ExecCommandRequest {
prepared,
process_id,
@@ -361,7 +361,7 @@ impl CoreToolRuntime for ExecCommandHandler {
let cwd = args.resolve_cwd(environment)?;
// Hook scripts execute on the host. Remote command paths cannot replace
// the host cwd in hook input, even when both machines use the same OS.
Ok((!environment.environment.is_remote()).then_some(cwd))
Ok((!environment.environment().is_remote()).then_some(cwd))
}
fn with_updated_hook_input(
@@ -142,7 +142,7 @@ impl<'a> ParsedExecCommand<'a> {
) -> Result<Self, FunctionCallError> {
let workdir: ExecCommandWorkdirArgs = parse_arguments(arguments)?;
let cwd = workdir.resolve_cwd(turn_environment)?;
let environment = &turn_environment.environment;
let environment = turn_environment.environment();
// Remote executors enforce URI-native sandbox policy themselves. Only a host-local
// sandbox needs a native cwd for resolving paths nested in the permissions config.
@@ -222,7 +222,7 @@ impl<'a> ParsedExecCommand<'a> {
.unwrap_or_else(|| session.user_shell());
// TODO(anp): Resolve requested shells in remote environments instead of restricting
// commands to the reported default shell.
if turn_environment.environment.is_remote()
if turn_environment.environment().is_remote()
&& let Some(requested_shell) = args.shell.take()
{
let Some(remote_shell) = turn_environment.shell.as_ref() else {
@@ -375,6 +375,7 @@ async fn exec_command_rejects_login_when_selected_environment_disallows_it() {
#[tokio::test]
#[cfg(not(windows))]
async fn exec_command_reuses_foreign_windows_grant() {
use crate::environment_selection::EnvironmentAttachment;
use crate::session::tests::make_session_and_context_with_auth_and_config_and_rx;
use codex_features::Feature;
use codex_protocol::models::AdditionalPermissionProfile;
@@ -404,6 +405,13 @@ async fn exec_command_reuses_foreign_windows_grant() {
)),
..Default::default()
};
let (attachment, _guard, _) = EnvironmentAttachment::new(
codex_exec_server::REMOTE_ENVIRONMENT_ID,
Arc::new(
Environment::create_for_tests(Some("ws://127.0.0.1:1/remote-exec-server".to_string()))
.expect("remote environment"),
),
);
{
let turn = Arc::get_mut(&mut turn).expect("turn should be uniquely owned");
let TurnEnvironmentState::Ready(environment) = turn
@@ -418,10 +426,7 @@ async fn exec_command_reuses_foreign_windows_grant() {
environment.selection.cwd = cwd.clone();
environment.selection.workspace_roots = vec![cwd.clone()];
environment.config_mut().workspace_roots = vec![cwd];
environment.environment = Arc::new(
Environment::create_for_tests(Some("ws://127.0.0.1:1/remote-exec-server".to_string()))
.expect("remote environment"),
);
environment.attachment = attachment;
}
let step_context = StepContext::for_test(Arc::clone(&turn));
@@ -149,7 +149,7 @@ impl ViewImageHandler {
})?;
let model_visible_path = path_uri.inferred_native_path_string();
let sandbox = turn_environment.sandbox_context(/*additional_permissions*/ None);
let fs = turn_environment.environment.get_filesystem();
let fs = turn_environment.environment().get_filesystem();
let metadata = fs
.get_metadata(&path_uri, GetMetadataOptions::default(), Some(&sandbox))
@@ -306,10 +306,10 @@ mod tests {
selection.cwd = PathUri::from_abs_path(&cwd);
selection.workspace_roots.clear();
turn.initial_environments.environments[0] =
TurnEnvironmentState::Ready(TurnEnvironment::new(
TurnEnvironmentState::Ready(TurnEnvironment::from_attachment(
selection,
current.config_origin,
current.environment,
current.attachment,
current.shell,
));
}
+1 -1
View File
@@ -1366,7 +1366,7 @@ pub(crate) async fn remote_network_proxy_launch(
turn: &TurnContext,
) -> Result<(RemoteNetworkProxyLaunchConfig, RemoteCredentialMasking), ToolError> {
let network_proxy = environment
.environment
.environment()
.info()
.await
.map_err(|err| {
+3 -3
View File
@@ -145,7 +145,7 @@ impl ToolOrchestrator {
#[cfg(target_os = "macos")]
let sandbox_manager = sandbox_manager.with_allowed_symlinked_codex_home(
environment
.environment
.environment()
.local_runtime_paths()
.and_then(|paths| paths.allowed_symlinked_codex_home.clone()),
);
@@ -269,9 +269,9 @@ impl ToolOrchestrator {
(None, false)
};
if (escalated_profile.is_some() || writable_root_metadata)
&& environment.environment.is_remote()
&& environment.environment().is_remote()
{
let info = environment.environment.info().await.map_err(|err| {
let info = environment.environment().info().await.map_err(|err| {
ToolError::Rejected(format!(
"could not verify exec-server support for escalation with denied reads: {err}"
))
@@ -162,7 +162,7 @@ impl ToolRuntime<ApplyPatchRequest, ApplyPatchRuntimeOutput> for ApplyPatchRunti
}
fn uses_executor_managed_process_sandbox(&self, req: &ApplyPatchRequest) -> bool {
req.turn_environment.environment.is_remote()
req.turn_environment.environment().is_remote()
}
fn sandbox_cwd<'a>(&self, req: &'a ApplyPatchRequest) -> Option<&'a PathUri> {
@@ -176,7 +176,7 @@ impl ToolRuntime<ApplyPatchRequest, ApplyPatchRuntimeOutput> for ApplyPatchRunti
ctx: &ToolCtx,
) -> Result<ApplyPatchRuntimeOutput, ToolError> {
let started_at = Instant::now();
let fs = req.turn_environment.environment.get_filesystem();
let fs = req.turn_environment.environment().get_filesystem();
let sandbox = Self::file_system_sandbox_context_for_attempt(req, attempt);
let mut stdout = Vec::new();
let mut stderr = Vec::new();
@@ -187,7 +187,7 @@ impl<'a> ToolRuntime<UnifiedExecRequest<'_>, UnifiedExecAttempt> for UnifiedExec
}
fn uses_executor_managed_process_sandbox(&self, req: &UnifiedExecRequest<'_>) -> bool {
req.prepared.turn_environment.environment.is_remote() || req.shell_snapshot.is_some()
req.prepared.turn_environment.environment().is_remote() || req.shell_snapshot.is_some()
}
fn sandbox_cwd<'b>(&self, req: &'b UnifiedExecRequest<'_>) -> Option<&'b PathUri> {
@@ -263,7 +263,7 @@ impl<'a> ToolRuntime<UnifiedExecRequest<'_>, UnifiedExecAttempt> for UnifiedExec
req.network.as_ref(),
launch_sandbox_permissions,
));
let environment_is_remote = prepared.turn_environment.environment.is_remote();
let environment_is_remote = prepared.turn_environment.environment().is_remote();
let credential_broker_available = !environment_is_remote
&& req.network.is_some()
&& ctx
@@ -446,7 +446,7 @@ impl<'a> ToolRuntime<UnifiedExecRequest<'_>, UnifiedExecAttempt> for UnifiedExec
ctx,
&prepared.resolved.command,
&prepared.cwd,
prepared.turn_environment.environment.as_ref(),
prepared.turn_environment.environment().as_ref(),
)
.await;
if let Some(sidecar) = metrics_sidecar.as_ref() {
@@ -473,7 +473,7 @@ impl<'a> ToolRuntime<UnifiedExecRequest<'_>, UnifiedExecAttempt> for UnifiedExec
if ctx.session.enabled(Feature::LoginShellPackagePath)
&& prepared.resolved.shell.is_posix_login()
&& !explicit_env_overrides.contains_key("PATH")
&& let Ok(info) = prepared.turn_environment.environment.info().await
&& let Ok(info) = prepared.turn_environment.environment().info().await
&& let Some(command_with_path_prepends) =
prepared.resolved.shell.derive_exec_args_with_path_prepends(
&prepared.hook_command,
@@ -642,7 +642,7 @@ impl<'a> ToolRuntime<UnifiedExecRequest<'_>, UnifiedExecAttempt> for UnifiedExec
req.shell_snapshot.clone(),
windows_sandbox_proxy_settings_mode,
prepared.tty,
prepared.turn_environment.environment.as_ref(),
prepared.turn_environment.environment().as_ref(),
)
.await;
let mut process = with_launch_failure_events(process, prepared, ctx).await?;
@@ -22,10 +22,10 @@ pub(super) async fn with_launch_failure_events(
let Err(ToolError::Rejected(message)) = &result else {
return result;
};
let plugin_attribution = if req.turn_environment.environment.is_remote() {
let plugin_attribution = if req.turn_environment.environment().is_remote() {
let file_system = req
.turn_environment
.environment
.environment()
.get_filesystem_without_reconnect();
ctx.step_context
.turn
@@ -27,7 +27,7 @@ impl SnapshotMetrics {
// remote legacy, and non-login invocations are outside this denominator.
if !cfg!(unix)
|| req.shell_snapshot.is_some()
|| prepared.turn_environment.environment.is_remote()
|| prepared.turn_environment.environment().is_remote()
|| !config.features.enabled(Feature::ShellSnapshot)
|| !prepared.resolved.shell.is_posix_login()
|| (prepared.cwd != prepared.turn_environment.selection.cwd && !brokered)
+11 -6
View File
@@ -41,6 +41,7 @@ use serde_json::json;
use crate::WaitForEnvironmentToolConfig;
use crate::config::CurrentTimeReminderConfig;
use crate::environment_selection::EnvironmentAttachment;
use crate::environment_selection::TurnEnvironmentState;
use crate::responses_metadata::TurnToolFunctionInfo;
use crate::responses_metadata::TurnToolNamespacesInfo;
@@ -1017,6 +1018,15 @@ async fn exec_command_guidance_follows_executor_platform_and_fallbacks() {
#[tokio::test]
async fn disabling_shell_tools_disables_command_tools_for_all_environments() {
let (attachment, _guard, _) = EnvironmentAttachment::new(
"remote",
Arc::new(
codex_exec_server::Environment::create_for_tests(Some(
"ws://127.0.0.1:1/remote-exec-server".to_string(),
))
.expect("remote test environment"),
),
);
let remote_environment = probe(|turn| {
set_feature(turn, Feature::ShellTool, /*enabled*/ false);
update_turn_settings_for_test(turn, |settings| {
@@ -1032,12 +1042,7 @@ async fn disabling_shell_tools_disables_command_tools_for_all_environments() {
panic!("primary environment should be ready");
};
environment.selection.environment_id = "remote".to_string();
environment.environment = Arc::new(
codex_exec_server::Environment::create_for_tests(Some(
"ws://127.0.0.1:1/remote-exec-server".to_string(),
))
.expect("remote test environment"),
);
environment.attachment = attachment;
})
.await;
remote_environment.assert_visible_lacks(&["exec_command", "write_stdin"]);
+1 -1
View File
@@ -122,7 +122,7 @@ async fn exec_command_with_tty(
turn.initial_environments
.primary()
.expect("turn environment")
.environment
.environment()
.as_ref(),
)
.await?,
@@ -579,8 +579,8 @@ impl UnifiedExecProcessManager {
/*turn_diff_tracker*/ None,
);
event_ctx.model_context = Some(&model_context);
let plugin_attribution = if prepared.turn_environment.environment.is_remote() {
let file_system = prepared.turn_environment.environment.get_filesystem();
let plugin_attribution = if prepared.turn_environment.environment().is_remote() {
let file_system = prepared.turn_environment.environment().get_filesystem();
context
.step_context
.turn
@@ -868,7 +868,7 @@ impl UnifiedExecProcessManager {
exit_code,
original_token_count: Some(original_token_count),
hook_command: Some(prepared.hook_command.clone()),
tool_cwd: (!prepared.turn_environment.environment.is_remote())
tool_cwd: (!prepared.turn_environment.environment().is_remote())
.then(|| prepared.cwd.clone()),
};
@@ -65,7 +65,7 @@ impl Session {
for environment in environments.turn_environments() {
// Remote starts cannot cancel an in-flight capture. Leave them, deferred
// environments, and managed networking on the normal lazy capture path.
if environment.environment.is_remote()
if environment.environment().is_remote()
|| !environment.shell_snapshot_v2_supported
|| !environment.config().allow_login_shell
|| environment.config().network_policy.is_some()
@@ -154,7 +154,7 @@ impl Session {
};
let result = environment
.environment
.environment()
.get_exec_backend()
.prewarm_shell_snapshot(params)
.await;