diff --git a/codex-rs/core/src/environment_selection.rs b/codex-rs/core/src/environment_selection.rs index 5f1a29a43f..f9ad1b782e 100644 --- a/codex-rs/core/src/environment_selection.rs +++ b/codex-rs/core/src/environment_selection.rs @@ -36,6 +36,7 @@ use tokio_util::task::AbortOnDropHandle; use tracing::Instrument; use tracing::instrument::WithSubscriber; +use crate::session::ThreadEnvironmentDefaults; use crate::session::turn_context::ShellSnapshotCache; use crate::session::turn_context::ShellSnapshotTask; use crate::session::turn_context::TurnEnvironment; @@ -171,17 +172,14 @@ pub(crate) struct StartingTurnEnvironment { resolution: TurnEnvironmentResolution, } -/// Resolves config once at the attachment boundary and records who owns future updates. +/// Uses the current thread defaults when requested and records who owns later config updates. fn resolve_selection_config( mut selection: TurnEnvironmentSelection, - thread_config: &EnvironmentConfig, + defaults: &ThreadEnvironmentDefaults, ) -> (TurnEnvironmentSelection, EnvironmentConfigOrigin) { let (config, origin) = match selection.config { EnvironmentConfigState::FromThread => ( - EnvironmentConfigState::Ready(thread_config_for_selection( - &selection.workspace_roots, - thread_config, - )), + EnvironmentConfigState::Ready(defaults.for_selection(&selection)), EnvironmentConfigOrigin::Thread, ), config @ (EnvironmentConfigState::Ready(_) @@ -192,23 +190,11 @@ fn resolve_selection_config( (selection, origin) } -fn thread_config_for_selection( - workspace_roots: &[PathUri], - thread_config: &EnvironmentConfig, -) -> EnvironmentConfig { - EnvironmentConfig { - workspace_roots: workspace_roots.to_vec(), - ..thread_config.clone() - } -} - impl SelectedTurnEnvironment { - fn refresh_thread_config(&mut self, config: &EnvironmentConfig) { + fn refresh_thread_config(&mut self, defaults: &ThreadEnvironmentDefaults) { if self.config_origin == EnvironmentConfigOrigin::Thread { - self.selection.config = EnvironmentConfigState::Ready(thread_config_for_selection( - &self.selection.workspace_roots, - config, - )); + self.selection.config = + EnvironmentConfigState::Ready(defaults.for_selection(&self.selection)); } } } @@ -239,15 +225,20 @@ pub(crate) struct ThreadEnvironments { local_shell: Shell, shell_snapshot: ShellSnapshot, non_blocking_snapshots: bool, - environments: Mutex>, + state: Mutex, connection_event_tx: OnceLock>, } +struct ThreadEnvironmentsState { + thread_defaults: ThreadEnvironmentDefaults, + environments: Vec, +} + impl ThreadEnvironments { pub(crate) fn new( environment_manager: Arc, local_shell: Shell, - thread_environment_config: impl Fn(&Environment) -> EnvironmentConfig, + thread_defaults: ThreadEnvironmentDefaults, shell_snapshot: ShellSnapshot, current: TurnEnvironmentSnapshot, non_blocking_snapshots: bool, @@ -296,8 +287,7 @@ impl ThreadEnvironments { pending_completion: None, }; // Child threads re-infer thread-owned config while preserving owner config. - let config = thread_environment_config(&inherited_environment.environment); - inherited_environment.refresh_thread_config(&config); + inherited_environment.refresh_thread_config(&thread_defaults); Some(inherited_environment) }) .collect(); @@ -306,7 +296,10 @@ impl ThreadEnvironments { local_shell, shell_snapshot, non_blocking_snapshots, - environments: Mutex::new(environments), + state: Mutex::new(ThreadEnvironmentsState { + thread_defaults, + environments, + }), connection_event_tx: OnceLock::new(), } } @@ -341,15 +334,12 @@ impl ThreadEnvironments { shell_snapshot } - pub(crate) fn update_selections( - &self, - environments: &[TurnEnvironmentSelection], - thread_environment_config: impl Fn(&Environment) -> EnvironmentConfig, - ) { - let mut current = self - .environments - .lock() - .unwrap_or_else(PoisonError::into_inner); + pub(crate) fn update_selections(&self, environments: &[TurnEnvironmentSelection]) { + let mut state = self.state.lock().unwrap_or_else(PoisonError::into_inner); + let ThreadEnvironmentsState { + thread_defaults, + environments: current, + } = &mut *state; let mut seen_environment_ids = HashSet::with_capacity(environments.len()); let mut next = Vec::with_capacity(environments.len()); for selected_environment in environments { @@ -375,10 +365,8 @@ impl ThreadEnvironments { ); if !failed && !restarting_as_pending { - let (selected_environment, config_origin) = resolve_selection_config( - selected_environment.clone(), - &thread_environment_config(&environment.environment), - ); + let (selected_environment, config_origin) = + resolve_selection_config(selected_environment.clone(), thread_defaults); let shell_settings_changed = matches!( (&environment.selection.config, &selected_environment.config), (EnvironmentConfigState::Ready(previous), EnvironmentConfigState::Ready(current)) @@ -405,10 +393,8 @@ impl ThreadEnvironments { tracing::warn!("skipping unknown turn environment `{environment_id}`"); continue; }; - let (selected_environment, config_origin) = resolve_selection_config( - selected_environment.clone(), - &thread_environment_config(&environment), - ); + let (selected_environment, config_origin) = + resolve_selection_config(selected_environment.clone(), thread_defaults); // Connection state belongs to the environment instance, not its cwd or roots. let connection_events_task = current .iter() @@ -467,7 +453,7 @@ impl ThreadEnvironments { .then(|| Arc::clone(task)) }) .collect::>(); - let previous = std::mem::replace(&mut *current, next); + let previous = std::mem::replace(current, next); // Publish owner configuration before waking turns waiting on this attachment. for environment in current.iter() { @@ -490,9 +476,10 @@ impl ThreadEnvironments { /// Projects canonical selections back to caller input form without losing config ownership. pub(crate) fn selections(&self) -> Vec { - self.environments + self.state .lock() .unwrap_or_else(PoisonError::into_inner) + .environments .iter() .map(|environment| { environment @@ -529,26 +516,28 @@ impl ThreadEnvironments { .unwrap_or_default() } - /// Refreshes attachments whose configuration is inferred from the thread. - pub(crate) fn update_thread_config( - &self, - thread_environment_config: impl Fn(&Environment) -> EnvironmentConfig, - ) { - let mut environments = self - .environments - .lock() - .unwrap_or_else(PoisonError::into_inner); + /// Changes the defaults inherited by all selected and later-added environments. A step that + /// already captured its environments keeps the previous values. + pub(crate) fn set_active_thread_defaults(&self, defaults: EnvironmentConfig) { + let mut state = self.state.lock().unwrap_or_else(PoisonError::into_inner); + let ThreadEnvironmentsState { + thread_defaults, + environments, + } = &mut *state; + if thread_defaults.common == defaults { + return; + } + thread_defaults.common = defaults; for environment in environments.iter_mut() { - let config = thread_environment_config(&environment.environment); 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 + if previous.allow_login_shell != thread_defaults.common.allow_login_shell + || previous.shell_environment_policy != thread_defaults.common.shell_environment_policy ); - environment.refresh_thread_config(&config); + environment.refresh_thread_config(thread_defaults); if shell_settings_changed && !environment.environment.is_remote() && self.shell_snapshot.should_rebuild_inherited() @@ -559,16 +548,13 @@ impl ThreadEnvironments { } } - 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) { + pub(crate) fn set_snapshot_credential_broker(&self, broker: SnapshotCredentialBrokerState) { + let mut state = self.state.lock().unwrap_or_else(PoisonError::into_inner); + if !self.shell_snapshot.set_credential_broker(broker) { return; } - for selected in environments.iter_mut() { + for selected in &mut state.environments { if !selected.environment.is_remote() && let Some(Ok(resolved)) = selected.resolution.clone().now_or_never() { @@ -605,10 +591,8 @@ impl ThreadEnvironments { &self, thread_selected_capability_roots: &[SelectedCapabilityRoot], ) -> SelectedCapabilityRootsStatus { - let environments = self - .environments - .lock() - .unwrap_or_else(PoisonError::into_inner); + let state = self.state.lock().unwrap_or_else(PoisonError::into_inner); + let environments = &state.environments; let mut selected_capability_roots = combine_selected_capability_roots( thread_selected_capability_roots, environments.iter().filter_map(|environment| { @@ -686,12 +670,9 @@ 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 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 environments.iter_mut() { + 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(), @@ -823,9 +804,10 @@ impl ThreadEnvironments { ) -> impl Future + Send + 'static { // Keep the startup results, but not the senders that would keep removed Pending work alive. let selected = self - .environments + .state .lock() .unwrap_or_else(PoisonError::into_inner) + .environments .iter() .map(|environment| StartingTurnEnvironment { selection: environment.selection.clone(), @@ -1196,12 +1178,12 @@ mod tests { let turn_environments = Arc::new(ThreadEnvironments::new( environment_manager, crate::shell::default_user_shell(), - |_| test_environment_config(), + ThreadEnvironmentDefaults::new(test_environment_config(), SandboxType::None), ShellSnapshot::disabled(), TurnEnvironmentSnapshot::default(), /*non_blocking_snapshots*/ false, )); - turn_environments.update_selections(selections, |_| test_environment_config()); + turn_environments.update_selections(selections); turn_environments.snapshot().await; turn_environments } @@ -1375,20 +1357,20 @@ url = "ws://127.0.0.1:8765" let turn_environments = ThreadEnvironments::new( Arc::new(EnvironmentManager::default_for_tests()), local_shell.clone(), - |_| expected_config.clone(), + ThreadEnvironmentDefaults::new( + expected_config.clone(), + expected_config.windows_sandbox_type, + ), ShellSnapshot::disabled(), TurnEnvironmentSnapshot::default(), /*non_blocking_snapshots*/ false, ); - turn_environments.update_selections( - &[TurnEnvironmentSelection { - environment_id: LOCAL_ENVIRONMENT_ID.to_string(), - cwd: PathUri::from_abs_path(&cwd), - workspace_roots: Vec::new(), - config: EnvironmentConfigState::FromThread, - }], - |_| expected_config.clone(), - ); + turn_environments.update_selections(&[TurnEnvironmentSelection { + environment_id: LOCAL_ENVIRONMENT_ID.to_string(), + cwd: PathUri::from_abs_path(&cwd), + workspace_roots: Vec::new(), + config: EnvironmentConfigState::FromThread, + }]); let snapshot = turn_environments.snapshot().await; let environment = snapshot.primary().expect("local environment"); @@ -1550,14 +1532,12 @@ url = "ws://127.0.0.1:8765" let environments = Arc::new(ThreadEnvironments::new( manager, crate::shell::default_user_shell(), - |_| test_environment_config(), + ThreadEnvironmentDefaults::new(test_environment_config(), SandboxType::None), ShellSnapshot::disabled(), TurnEnvironmentSnapshot::default(), /*non_blocking_snapshots*/ false, )); - environments.update_selections(std::slice::from_ref(&selection), |_| { - test_environment_config() - }); + environments.update_selections(std::slice::from_ref(&selection)); // Exercise first-use callsite registration from a thread without a subscriber. std::thread::spawn({ let environments = Arc::clone(&environments); @@ -1648,15 +1628,6 @@ url = "ws://127.0.0.1:8765" windows_sandbox_type: local_sandbox, ..expected_config.clone() }; - let config_for = - |config: &EnvironmentConfig, environment: &Environment| EnvironmentConfig { - windows_sandbox_type: if environment.is_remote() { - config.windows_sandbox_type - } else { - local_sandbox - }, - ..config.clone() - }; let cwd = PathUri::from_abs_path(&cwd); let remote = TurnEnvironmentSelection { environment_id: REMOTE_ENVIRONMENT_ID.to_string(), @@ -1673,21 +1644,30 @@ url = "ws://127.0.0.1:8765" let turn_environments = ThreadEnvironments::new( Arc::clone(&manager), crate::shell::default_user_shell(), - |environment| config_for(&test_environment_config(), environment), + ThreadEnvironmentDefaults::new(test_environment_config(), local_sandbox), ShellSnapshot::disabled(), TurnEnvironmentSnapshot::default(), /*non_blocking_snapshots*/ true, ); - turn_environments.update_selections(std::slice::from_ref(&local), |environment| { - config_for(&test_environment_config(), environment) - }); - let resolution = turn_environments.environments.lock().unwrap()[0] + turn_environments.update_selections(std::slice::from_ref(&local)); + let resolution = turn_environments.state.lock().unwrap().environments[0] .resolution .clone(); resolution.await.expect("local environment should resolve"); - turn_environments.update_selections(&[remote.clone(), local.clone()], |environment| { - config_for(&expected_config, environment) - }); + turn_environments.update_selections(&[remote.clone(), local.clone()]); + let initial = turn_environments.snapshot().await; + assert_eq!( + initial.primary().expect("local environment").config(), + &EnvironmentConfig { + windows_sandbox_type: local_sandbox, + ..test_environment_config() + }, + ); + assert_eq!( + initial.starting().next().unwrap().selection.config, + EnvironmentConfigState::Ready(test_environment_config()), + ); + turn_environments.set_active_thread_defaults(expected_config.clone()); let starting = turn_environments.snapshot().await; let resolved_remote = TurnEnvironmentSelection { @@ -1726,7 +1706,7 @@ url = "ws://127.0.0.1:8765" permission_profile: PermissionProfileSnapshot::legacy(PermissionProfile::Disabled), ..test_environment_config() }; - turn_environments.update_thread_config(|environment| config_for(&next_config, environment)); + turn_environments.set_active_thread_defaults(next_config.clone()); let next_starting = turn_environments.snapshot().await; assert!( !next_starting.has_full_access(AskForApproval::Never, &PermissionProfile::Disabled) @@ -1753,7 +1733,7 @@ url = "ws://127.0.0.1:8765" let child = ThreadEnvironments::new( manager, crate::shell::default_user_shell(), - |environment| config_for(&next_config, environment), + ThreadEnvironmentDefaults::new(next_config.clone(), local_sandbox), ShellSnapshot::disabled(), attached.clone(), /*non_blocking_snapshots*/ true, @@ -1843,16 +1823,14 @@ url = "ws://127.0.0.1:8765" let environments = ThreadEnvironments::new( Arc::clone(&manager), crate::shell::default_user_shell(), - |_| test_environment_config(), + ThreadEnvironmentDefaults::new(test_environment_config(), SandboxType::None), ShellSnapshot::disabled(), TurnEnvironmentSnapshot::default(), /*non_blocking_snapshots*/ true, ); - environments.update_selections(std::slice::from_ref(&selection), |_| { - test_environment_config() - }); + environments.update_selections(std::slice::from_ref(&selection)); let starting = environments.snapshot().await; - let failed_resolution = environments.environments.lock().unwrap()[0] + let failed_resolution = environments.state.lock().unwrap().environments[0] .resolution .clone(); let error = failed_resolution @@ -1898,12 +1876,13 @@ url = "ws://127.0.0.1:8765" allow_login_shell: false, ..test_environment_config() }; - environments.update_thread_config(|_| next_config.clone()); + environments.set_active_thread_defaults(next_config.clone()); assert!( - failed_resolution.ptr_eq(&environments.environments.lock().unwrap()[0].resolution), + failed_resolution + .ptr_eq(&environments.state.lock().unwrap().environments[0].resolution), "updating environment config must not retry a failed environment" ); - environments.update_selections(std::slice::from_ref(&selection), |_| next_config.clone()); + environments.update_selections(std::slice::from_ref(&selection)); let replacement = environments.snapshot().await; let replacement = replacement @@ -1947,15 +1926,13 @@ url = "ws://127.0.0.1:8765" let environments = Arc::new(ThreadEnvironments::new( Arc::clone(&manager), crate::shell::default_user_shell(), - |_| test_environment_config(), + ThreadEnvironmentDefaults::new(test_environment_config(), SandboxType::None), ShellSnapshot::disabled(), TurnEnvironmentSnapshot::default(), /*non_blocking_snapshots*/ true, )); environments.start_connection_event_forwarding(tx_event); - environments.update_selections(std::slice::from_ref(&selection), |_| { - test_environment_config() - }); + environments.update_selections(std::slice::from_ref(&selection)); let initial_snapshot = environments.snapshot().await; let second_listener = TcpListener::bind("127.0.0.1:0") .await @@ -1973,17 +1950,12 @@ url = "ws://127.0.0.1:8765" ) .expect("replace environment"); - environments.update_selections(std::slice::from_ref(&selection), |_| { - test_environment_config() - }); + environments.update_selections(std::slice::from_ref(&selection)); let reused_snapshot = environments.snapshot().await; - environments.update_selections( - &[TurnEnvironmentSelection { - cwd: PathUri::from_abs_path(&cwd.join("changed")), - ..selection - }], - |_| test_environment_config(), - ); + environments.update_selections(&[TurnEnvironmentSelection { + cwd: PathUri::from_abs_path(&cwd.join("changed")), + ..selection + }]); let changed_snapshot = environments.snapshot().await; let initial = initial_snapshot @@ -2081,7 +2053,7 @@ url = "ws://127.0.0.1:8765" let environments = ThreadEnvironments::new( manager, crate::shell::default_user_shell(), - |_| child_config.clone(), + ThreadEnvironmentDefaults::new(child_config.clone(), child_config.windows_sandbox_type), ShellSnapshot::disabled(), TurnEnvironmentSnapshot { environments: vec![TurnEnvironmentState::Ready(inherited)], @@ -2089,7 +2061,7 @@ url = "ws://127.0.0.1:8765" /*non_blocking_snapshots*/ false, ); - environments.update_selections(std::slice::from_ref(&selection), |_| child_config.clone()); + environments.update_selections(std::slice::from_ref(&selection)); let snapshot = environments.snapshot().await; let inherited = snapshot.primary().expect("inherited environment"); @@ -2129,8 +2101,16 @@ url = "ws://127.0.0.1:8765" }, }; let manager = Arc::new(EnvironmentManager::default_for_tests()); - let parent = - resolve_turn_environments(Arc::clone(&manager), std::slice::from_ref(&selection)).await; + let parent = ThreadEnvironments::new( + Arc::clone(&manager), + crate::shell::default_user_shell(), + ThreadEnvironmentDefaults::new(test_environment_config(), SandboxType::WindowsMxc), + ShellSnapshot::disabled(), + TurnEnvironmentSnapshot::default(), + /*non_blocking_snapshots*/ false, + ); + parent.update_selections(std::slice::from_ref(&selection)); + parent.snapshot().await; let parent_owner_config = EnvironmentConfig { allow_login_shell: false, workspace_roots: selection.workspace_roots.clone(), @@ -2146,9 +2126,7 @@ url = "ws://127.0.0.1:8765" }; let mut owner_selection = selection.clone(); owner_selection.config = EnvironmentConfigState::Ready(parent_owner_config.clone()); - parent.update_selections(std::slice::from_ref(&owner_selection), |_| { - test_environment_config() - }); + parent.update_selections(std::slice::from_ref(&owner_selection)); let child_thread_config = EnvironmentConfig { windows_sandbox_type: SandboxType::WindowsMxc, @@ -2157,7 +2135,7 @@ url = "ws://127.0.0.1:8765" ), ..test_environment_config() }; - parent.update_thread_config(|_| child_thread_config.clone()); + parent.set_active_thread_defaults(child_thread_config.clone()); assert_eq!( parent .snapshot() @@ -2170,7 +2148,7 @@ url = "ws://127.0.0.1:8765" let child = ThreadEnvironments::new( manager, crate::shell::default_user_shell(), - |_| child_thread_config.clone(), + ThreadEnvironmentDefaults::new(child_thread_config.clone(), SandboxType::WindowsMxc), ShellSnapshot::disabled(), parent.snapshot().await, /*non_blocking_snapshots*/ false, @@ -2183,9 +2161,7 @@ url = "ws://127.0.0.1:8765" ..parent_owner_config.clone() }; owner_selection.config = EnvironmentConfigState::Ready(cleared_parent_owner_config.clone()); - parent.update_selections(std::slice::from_ref(&owner_selection), |_| { - child_thread_config.clone() - }); + parent.update_selections(std::slice::from_ref(&owner_selection)); let cleared_snapshot = parent.snapshot().await; assert_eq!( cleared_snapshot @@ -2200,13 +2176,10 @@ url = "ws://127.0.0.1:8765" cwd: PathUri::from_abs_path(&cwd.join("changed")), ..selection }; - let new_thread_config = test_environment_config(); - parent.update_selections(std::slice::from_ref(&changed_selection), |_| { - new_thread_config.clone() - }); + parent.update_selections(std::slice::from_ref(&changed_selection)); let changed_snapshot = parent.snapshot().await; let changed_environment = changed_snapshot.primary().expect("changed environment"); - assert_eq!(changed_environment.config(), &new_thread_config); + assert_eq!(changed_environment.config(), &child_thread_config); } #[tokio::test] diff --git a/codex-rs/core/src/environment_selection_snapshot_tests.rs b/codex-rs/core/src/environment_selection_snapshot_tests.rs index 1fe5204116..8302c52923 100644 --- a/codex-rs/core/src/environment_selection_snapshot_tests.rs +++ b/codex-rs/core/src/environment_selection_snapshot_tests.rs @@ -5,6 +5,7 @@ use codex_exec_server::LOCAL_ENVIRONMENT_ID; use codex_otel::SessionTelemetry; use codex_protocol::ThreadId; use codex_protocol::protocol::SessionSource; +use codex_protocol::sandbox::SandboxType; use pretty_assertions::assert_eq; use std::time::Duration; use tokio::sync::oneshot; @@ -20,16 +21,15 @@ async fn unpolled_snapshot_does_not_delay_canceling_a_removed_environment() { 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(), + ThreadEnvironmentDefaults::new(tests::test_environment_config(), SandboxType::None), ShellSnapshot::disabled(), TurnEnvironmentSnapshot::default(), /*non_blocking_snapshots*/ false, ); - environments.update_selections(&[selection], |_| config.clone()); + environments.update_selections(&[selection]); let starting = environments .snapshot() .await @@ -39,7 +39,7 @@ async fn unpolled_snapshot_does_not_delay_canceling_a_removed_environment() { .clone(); let held_snapshot = environments.snapshot(); - environments.update_selections(&[], |_| config.clone()); + environments.update_selections(&[]); let error = timeout(Duration::from_secs(5), starting.wait_until_ready()) .await @@ -83,32 +83,28 @@ async fn credential_refresh_does_not_restore_a_removed_environment() { 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(), + ThreadEnvironmentDefaults::new(tests::test_environment_config(), SandboxType::None), 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.clone(), - ); - let resolution = environments.environments.lock().unwrap()[0] + environments.update_selections(&[TurnEnvironmentSelection { + environment_id: LOCAL_ENVIRONMENT_ID.to_string(), + cwd: cwd_uri.clone(), + workspace_roots: vec![cwd_uri], + config: EnvironmentConfigState::FromThread, + }]); + let resolution = environments.state.lock().unwrap().environments[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 { + environments.state.lock().unwrap().environments[0].resolution = async move { refresh_paused_tx.send(()).expect("signal refresh paused"); resume_refresh_rx.recv().expect("resume refresh"); Ok(resolved) @@ -128,7 +124,7 @@ async fn credential_refresh_does_not_restore_a_removed_environment() { 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.clone()); + removing.update_selections(&[]); }); timeout(Duration::from_secs(/*secs*/ 5), removal_started_rx) .await diff --git a/codex-rs/core/src/mcp_openai_file.rs b/codex-rs/core/src/mcp_openai_file.rs index 04f4e9f9bf..110a585fce 100644 --- a/codex-rs/core/src/mcp_openai_file.rs +++ b/codex-rs/core/src/mcp_openai_file.rs @@ -385,17 +385,19 @@ mod tests { .expect("ready primary environment"); let selection = environment.selection(); let environment_config = environment.config().clone(); + let local_windows_sandbox_type = environment_config.windows_sandbox_type; let environments = crate::environment_selection::ThreadEnvironments::new( session.services.turn_environments.environment_manager(), crate::shell::default_user_shell(), - |_| environment_config.clone(), + crate::session::ThreadEnvironmentDefaults::new( + environment_config, + local_windows_sandbox_type, + ), crate::shell_snapshot::ShellSnapshot::disabled(), Default::default(), /*non_blocking_snapshots*/ true, ); - environments.update_selections(std::slice::from_ref(&selection), |_| { - environment_config.clone() - }); + environments.update_selections(std::slice::from_ref(&selection)); turn_context.initial_environments = environments.snapshot().await; turn_context .initial_environments diff --git a/codex-rs/core/src/session/environment.rs b/codex-rs/core/src/session/environment.rs index bd4217d116..f0134ddc41 100644 --- a/codex-rs/core/src/session/environment.rs +++ b/codex-rs/core/src/session/environment.rs @@ -10,6 +10,7 @@ use codex_protocol::error::Result as CodexResult; use codex_protocol::protocol::EnvironmentConfig; use codex_protocol::protocol::EnvironmentConfigState; use codex_protocol::protocol::TurnEnvironmentSelection; +use codex_protocol::sandbox::SandboxType; use crate::config::ConstraintError; use crate::config::ConstraintResult; @@ -19,6 +20,31 @@ use crate::session::session::Session; use crate::session::session::SessionConfiguration; use crate::session::session::SessionSettingsUpdate; +/// Defaults for environments that inherit their configuration from the running turn. +pub(crate) struct ThreadEnvironmentDefaults { + pub(crate) common: EnvironmentConfig, + // Chosen once when the session starts; it does not apply to remote environments. + local_windows_sandbox_type: SandboxType, +} + +impl ThreadEnvironmentDefaults { + pub(crate) fn new(common: EnvironmentConfig, local_windows_sandbox_type: SandboxType) -> Self { + Self { + common, + local_windows_sandbox_type, + } + } + + pub(crate) fn for_selection(&self, selection: &TurnEnvironmentSelection) -> EnvironmentConfig { + let mut config = self.common.clone(); + config.workspace_roots = selection.workspace_roots.clone(); + if selection.environment_id == LOCAL_ENVIRONMENT_ID { + config.windows_sandbox_type = self.local_windows_sandbox_type; + } + config + } +} + pub(super) fn validate_environment_selections( selections: &[TurnEnvironmentSelection], ) -> ConstraintResult<()> { @@ -157,12 +183,14 @@ impl Session { environment.config = latest.config.clone(); } } + // Apply the saved defaults even if this turn selects the same environments. + self.services + .turn_environments + .set_active_thread_defaults(configuration.inferred_environment_config()); if environments != self.services.turn_environments.selections() { self.services .turn_environments - .update_selections(&environments, |environment| { - configuration.inferred_environment_config_for(environment) - }); + .update_selections(&environments); } } self.services.turn_environments.snapshot() @@ -257,13 +285,7 @@ impl Session { if update_current { // Invalidate MCP before installed configuration can wake a waiting turn. self.mark_mcp_runtime_dirty(); - self.services - .turn_environments - .update_selections(¤t, |environment| { - state - .session_configuration - .inferred_environment_config_for(environment) - }); + self.services.turn_environments.update_selections(¤t); } Ok(()) } diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 961f36fe1d..54efc41335 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -264,6 +264,7 @@ mod turn_input; mod turn_suspension; mod world_state; use self::code_mode_warning::unsupported_code_mode_warning; +pub(crate) use self::environment::ThreadEnvironmentDefaults; #[cfg(test)] use self::handlers::submission_dispatch_span; use self::handlers::submission_loop; @@ -1901,18 +1902,11 @@ impl Session { let root_service_tier_changed = updated.parent_thread_id.is_none() && state.session_configuration.step_settings.service_tier != updated.step_settings.service_tier; - let environment_config = updated.inferred_environment_config(); let mcp_inputs_changed = self.mcp_inputs_differ(&state.session_configuration, &updated); if mcp_inputs_changed { self.mark_mcp_runtime_dirty(); } - if state.session_configuration.inferred_environment_config() != environment_config { - self.services - .turn_environments - .update_thread_config(|environment| { - updated.inferred_environment_config_for(environment) - }); - } + // Save new environment defaults for future turns. The running turn keeps its own. state.session_configuration = updated; if root_service_tier_changed { self.services.agent_control.set_root_service_tier( diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index 43a0af3c36..bccda5e63b 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -189,18 +189,6 @@ impl SessionConfiguration { } } - /// Apply the local rollout only when inferring config for a local environment. - pub(super) fn inferred_environment_config_for( - &self, - environment: &Environment, - ) -> EnvironmentConfig { - let mut config = self.inferred_environment_config(); - if !environment.is_remote() { - config.windows_sandbox_type = self.windows_sandbox_type; - } - config - } - pub(super) fn permission_profile(&self) -> PermissionProfile { self.permission_profile_state.permission_profile().clone() } @@ -1426,15 +1414,15 @@ impl Session { let turn_environments = Arc::new(ThreadEnvironments::new( environment_manager, default_shell.clone(), - |environment| session_configuration.inferred_environment_config_for(environment), + ThreadEnvironmentDefaults::new( + session_configuration.inferred_environment_config(), + session_configuration.windows_sandbox_type, + ), shell_snapshot, inherited_environments.unwrap_or_default(), config.features.enabled(Feature::DeferredExecutor), )); - turn_environments.update_selections( - environment_selections, - |environment| session_configuration.inferred_environment_config_for(environment), - ); + turn_environments.update_selections(environment_selections); session_configuration.environments = turn_environments.selections(); let resolved_environments = turn_environments.snapshot().await; let agents_md_manager = Arc::new(AgentsMdManager::new(instructions)); diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index 68bbc8ff29..0091a77a07 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -5002,14 +5002,15 @@ async fn resolved_environments_for_configuration( let turn_environments = ThreadEnvironments::new( Arc::clone(&environment_manager), default_user_shell(), - |environment| session_configuration.inferred_environment_config_for(environment), + ThreadEnvironmentDefaults::new( + session_configuration.inferred_environment_config(), + session_configuration.windows_sandbox_type, + ), ShellSnapshot::disabled(), TurnEnvironmentSnapshot::default(), /*non_blocking_snapshots*/ false, ); - turn_environments.update_selections(environment_selections, |environment| { - session_configuration.inferred_environment_config_for(environment) - }); + turn_environments.update_selections(environment_selections); (environment_manager, turn_environments.snapshot().await) } @@ -5959,7 +5960,23 @@ async fn permission_profile_updates_apply_to_next_turn_environment() { .update_settings(updates) .await .expect("permission profile update should succeed"); - session.new_default_turn().await + assert_eq!( + session + .services + .turn_environments + .snapshot() + .await + .primary() + .expect("current environment") + .config(), + &active_environment_config, + ); + session + .new_turn_with_default_settings( + "permission-profile-update".to_string(), + Default::default(), + ) + .await }; let next_environment = next_turn .initial_environments @@ -6400,7 +6417,10 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) { let turn_environments = Arc::new(ThreadEnvironments::new( environment_manager, default_user_shell(), - |environment| session_configuration.inferred_environment_config_for(environment), + ThreadEnvironmentDefaults::new( + session_configuration.inferred_environment_config(), + session_configuration.windows_sandbox_type, + ), ShellSnapshot::disabled(), resolved_environments, /*non_blocking_snapshots*/ false, @@ -8667,7 +8687,10 @@ where let turn_environments = Arc::new(ThreadEnvironments::new( environment_manager, default_user_shell(), - |environment| session_configuration.inferred_environment_config_for(environment), + ThreadEnvironmentDefaults::new( + session_configuration.inferred_environment_config(), + session_configuration.windows_sandbox_type, + ), ShellSnapshot::disabled(), resolved_turn_environments.clone(), /*non_blocking_snapshots*/ false, @@ -9100,15 +9123,12 @@ async fn mcp_elicitation_reviewer_uses_active_reviewer_and_latest_runtime_policy ) .await; assert_eq!(old_turn.config.approvals_reviewer, ApprovalsReviewer::User); + let task = NeverEndingTask { + kind: TaskKind::Regular, + listen_to_cancellation_token: true, + }; session - .spawn_task( - Arc::clone(&old_turn), - Vec::new(), - NeverEndingTask { - kind: TaskKind::Regular, - listen_to_cancellation_token: true, - }, - ) + .spawn_task(Arc::clone(&old_turn), Vec::new(), task) .await; session.mark_mcp_runtime_dirty(); @@ -9205,6 +9225,38 @@ async fn mcp_elicitation_reviewer_uses_active_reviewer_and_latest_runtime_policy .await .expect("permission profile should update"); session.refresh_mcp_if_dirty().await; + assert_eq!( + session + .mcp_elicitation_reviewer() + .review(request.clone()) + .await + .expect("elicitation review should succeed"), + Some(ElicitationResponse { + action: ElicitationAction::Decline, + content: None, + meta: Some(json!({ "approvals_reviewer": "auto_review" })), + }) + ); + + session.abort_all_tasks(TurnAbortReason::Interrupted).await; + let (next_turn, _) = session + .new_turn_with_sub_id( + "next-turn".to_string(), + SessionSettingsUpdate { + step_settings: StepSettingsUpdate { + approvals_reviewer: Some(ApprovalsReviewer::AutoReview), + ..Default::default() + }, + ..Default::default() + }, + Default::default(), + ) + .await + .expect("next turn should start with the same reviewer"); + session + .spawn_task(Arc::clone(&next_turn), Vec::new(), task) + .await; + session.refresh_mcp_if_dirty().await; assert_eq!( session .mcp_elicitation_reviewer() @@ -9214,7 +9266,7 @@ async fn mcp_elicitation_reviewer_uses_active_reviewer_and_latest_runtime_policy Some(ElicitationResponse { action: ElicitationAction::Accept, content: Some(json!({})), - meta: None, + meta: Some(json!({ "approvals_reviewer": "auto_review" })), }) ); @@ -9225,7 +9277,7 @@ async fn mcp_elicitation_reviewer_uses_active_reviewer_and_latest_runtime_policy .into_iter() .next() .expect("session should select its executor environment"); - let mut owner_config = old_turn + let mut owner_config = next_turn .initial_environments .primary() .expect("ready environment") @@ -9237,6 +9289,11 @@ async fn mcp_elicitation_reviewer_uses_active_reviewer_and_latest_runtime_policy .environment_ready(&selection, owner_config) .await .expect("attachment owner should install its restricted permissions"); + session.abort_all_tasks(TurnAbortReason::Interrupted).await; + let restricted_turn = session + .new_turn_with_default_settings("owner-restricted".to_string(), Default::default()) + .await; + session.spawn_task(restricted_turn, Vec::new(), task).await; session.refresh_mcp_if_dirty().await; assert_eq!( session diff --git a/codex-rs/core/src/shell_snapshot_tests.rs b/codex-rs/core/src/shell_snapshot_tests.rs index ead4186256..45b2b9cdae 100644 --- a/codex-rs/core/src/shell_snapshot_tests.rs +++ b/codex-rs/core/src/shell_snapshot_tests.rs @@ -2,6 +2,8 @@ use super::*; #[cfg(unix)] use crate::config::NetworkProxySpec; #[cfg(unix)] +use crate::session::ThreadEnvironmentDefaults; +#[cfg(unix)] use crate::tools::runtimes::RuntimePathPrepends; #[cfg(unix)] use crate::tools::runtimes::maybe_wrap_shell_lc_with_snapshot; @@ -313,20 +315,17 @@ async fn inactive_profiles_keep_snapshots_but_active_brokers_require_sandbox() - let environments = ThreadEnvironments::new( Arc::clone(&manager), shell.clone(), - |_| config.clone(), + ThreadEnvironmentDefaults::new(config.clone(), config.windows_sandbox_type), snapshot_builder.clone(), TurnEnvironmentSnapshot::default(), /*non_blocking_snapshots*/ false, ); - environments.update_selections( - &[TurnEnvironmentSelection { - environment_id: LOCAL_ENVIRONMENT_ID.to_string(), - cwd: PathUri::from_abs_path(&dir.path().abs()), - workspace_roots: Vec::new(), - config: EnvironmentConfigState::FromThread, - }], - |_| config.clone(), - ); + environments.update_selections(&[TurnEnvironmentSelection { + environment_id: LOCAL_ENVIRONMENT_ID.to_string(), + cwd: PathUri::from_abs_path(&dir.path().abs()), + workspace_roots: Vec::new(), + config: EnvironmentConfigState::FromThread, + }]); for (state, expect_snapshot) in [ (SnapshotCredentialBrokerState::Starting, false), (SnapshotCredentialBrokerState::Inactive, true), @@ -354,7 +353,7 @@ async fn inactive_profiles_keep_snapshots_but_active_brokers_require_sandbox() - .shell_environment_policy .r#set .insert("CORP_REGION".into(), "west".into()); - environments.update_thread_config(|_| config.clone()); + environments.set_active_thread_defaults(config.clone()); let updated = environments.snapshot().await; assert!(!Arc::ptr_eq( &environment.shell_snapshot_cache, @@ -366,7 +365,7 @@ async fn inactive_profiles_keep_snapshots_but_active_brokers_require_sandbox() - let child = ThreadEnvironments::new( Arc::clone(&manager), shell.clone(), - |_| config.clone(), + ThreadEnvironmentDefaults::new(config.clone(), config.windows_sandbox_type), snapshot_builder.clone(), updated.clone(), /*non_blocking_snapshots*/ false, @@ -393,7 +392,7 @@ async fn inactive_profiles_keep_snapshots_but_active_brokers_require_sandbox() - environments.set_snapshot_credential_broker(SnapshotCredentialBrokerState::Ready( started_proxy.proxy(), )); - environments.update_thread_config(|_| config.clone()); + environments.set_active_thread_defaults(config.clone()); let turn = environments.snapshot().await; let environment = turn.primary().expect("brokered environment"); let mut tool_config = crate::config::ConfigBuilder::without_managed_config_for_tests() diff --git a/codex-rs/core/tests/suite/step_settings/environment_selection.rs b/codex-rs/core/tests/suite/step_settings/environment_selection.rs index 6e6b4697f4..e2601485a1 100644 --- a/codex-rs/core/tests/suite/step_settings/environment_selection.rs +++ b/codex-rs/core/tests/suite/step_settings/environment_selection.rs @@ -114,3 +114,59 @@ async fn model_update_preserves_active_environment_and_next_turn_uses_new_select ); Ok(()) } + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn inherited_permission_update_applies_to_the_next_turn_not_the_next_step() -> Result<()> { + skip_if_no_network!(Ok(())); + + let server = start_mock_server().await; + let responses = mount_sse_sequence( + &server, + vec![ + paused_response("before-permission-update", "pause"), + sse_completed("current-turn-done"), + sse_completed("next-turn-done"), + ], + ) + .await; + let test = direct_tool_settings_test() + .with_config(|config| { + config + .permissions + .set_permission_profile(PermissionProfile::read_only()) + .expect("set restricted thread permissions"); + }) + .build_with_auto_env(&server) + .await?; + + let paused = start_paused_turn(&test.codex).await?; + submit_thread_settings( + &test.codex, + ThreadSettingsOverrides { + permission_profile: Some(PermissionProfile::Disabled), + ..Default::default() + }, + ) + .await?; + answer_paused_turn(&test.codex, &paused.turn_id).await?; + wait_for_event(&test.codex, |event| { + matches!(event, EventMsg::TurnComplete(_)) + }) + .await; + test.submit_text_turn("start the next turn").await?; + + let unrestricted = responses + .requests() + .iter() + .map(|request| { + let filesystem = request + .message_input_texts("user") + .into_iter() + .rfind(|text| text.contains("")) + .expect("model request includes filesystem permissions"); + filesystem.contains(">(); + assert_eq!(unrestricted, [false, false, true]); + Ok(()) +}