From 370cd7cba87aeef8cf39bb194e51829b210c26a6 Mon Sep 17 00:00:00 2001 From: sayan-oai Date: Wed, 23 Sep 2026 03:00:37 +0000 Subject: [PATCH] Allow environment selection updates during an active turn (#47435) ## What changed - Add `environments` to `TurnSettingsUpdate` to replace the selection for subsequent steps without changing future turns. - Capture model settings and environments together so combined updates apply or reject atomically, while existing steps retain their captured selection. - Validate environment IDs, working directories, configuration ownership, and permission requirements against the running turn's selections. - Preserve executor startup behavior when switching environments and use the step's local working directory for tool hooks. ## Testing Add coverage for workspace switching and restoration on the next turn, atomic model and environment updates, permission validation, executor connection waits, and replacement of pending environment selections. GitOrigin-RevId: 02b5c57a78fd594b2aed84fc09d9e99eb97be28c --- codex-rs/app-server/src/request_processors.rs | 4 +- .../src/request_processors/turn_processor.rs | 1 + codex-rs/core/src/codex_thread.rs | 27 +- codex-rs/core/src/environment_selection.rs | 119 ++++--- codex-rs/core/src/hook_runtime.rs | 19 +- codex-rs/core/src/lib.rs | 1 + codex-rs/core/src/session/daemon_recovery.rs | 8 +- codex-rs/core/src/session/environment.rs | 40 ++- codex-rs/core/src/session/mcp.rs | 2 +- codex-rs/core/src/session/mod.rs | 46 ++- .../core/src/session/notification_tests.rs | 5 +- codex-rs/core/src/session/review.rs | 6 +- codex-rs/core/src/session/session.rs | 4 +- codex-rs/core/src/session/step_activation.rs | 179 +++++++---- .../core/src/session/step_activation_tests.rs | 78 ++++- codex-rs/core/src/session/step_context.rs | 8 - codex-rs/core/src/session/step_settings.rs | 3 +- codex-rs/core/src/session/tests.rs | 55 ++-- .../core/src/session/tests/guardian_tests.rs | 12 +- codex-rs/core/src/session/turn_context.rs | 19 +- codex-rs/core/src/thread_manager.rs | 33 -- codex-rs/core/src/tools/network_approval.rs | 13 +- codex-rs/core/src/tools/orchestrator.rs | 6 +- .../runtimes/zsh_fork/unix_escalation.rs | 4 +- .../core/src/unified_exec/process_manager.rs | 6 +- codex-rs/core/tests/common/test_codex.rs | 34 ++ codex-rs/core/tests/suite/remote_env.rs | 293 ++++++++++++++++-- codex-rs/core/tests/suite/scenarios.rs | 128 ++++++++ ...s__astra_active_environment_selection.snap | 129 ++++++++ .../step_settings/environment_selection.rs | 247 ++++++++++++++- codex-rs/protocol/src/protocol.rs | 6 +- 31 files changed, 1203 insertions(+), 332 deletions(-) create mode 100644 codex-rs/core/tests/suite/snapshots/all__suite__scenarios__astra_active_environment_selection.snap diff --git a/codex-rs/app-server/src/request_processors.rs b/codex-rs/app-server/src/request_processors.rs index 3338d98af3..c6b3693383 100644 --- a/codex-rs/app-server/src/request_processors.rs +++ b/codex-rs/app-server/src/request_processors.rs @@ -366,6 +366,7 @@ use codex_core::read_head_for_summary; use codex_core::sandboxing::SandboxPermissions; use codex_core::truncate_rollout_after_turn_id; use codex_core::truncate_rollout_before_turn_id; +use codex_core::validate_environment_ids_and_cwds; use codex_core::windows_sandbox::WindowsSandboxLevelExt; use codex_core::windows_sandbox::WindowsSandboxSetupMode as CoreWindowsSandboxSetupMode; use codex_core::windows_sandbox::WindowsSandboxSetupRequest; @@ -675,8 +676,7 @@ fn resolve_turn_environment_selections( config: EnvironmentConfigState::FromThread, }); } - thread_manager - .validate_environment_selections(&selections) + validate_environment_ids_and_cwds(&thread_manager.environment_manager(), &selections) .map_err(environment_selection_error)?; Ok(Some(selections)) } diff --git a/codex-rs/app-server/src/request_processors/turn_processor.rs b/codex-rs/app-server/src/request_processors/turn_processor.rs index 243f4b122c..e3b0d0b5bc 100644 --- a/codex-rs/app-server/src/request_processors/turn_processor.rs +++ b/codex-rs/app-server/src/request_processors/turn_processor.rs @@ -227,6 +227,7 @@ impl TurnRequestProcessor { approvals_reviewer: params .approvals_reviewer .map(codex_app_server_protocol::ApprovalsReviewer::to_core), + environments: None, model: params.model, // Match thread/settings/update: public null does not clear effort. effort: params.effort.map(Some), diff --git a/codex-rs/core/src/codex_thread.rs b/codex-rs/core/src/codex_thread.rs index 29c37de904..24d5a8902a 100644 --- a/codex-rs/core/src/codex_thread.rs +++ b/codex-rs/core/src/codex_thread.rs @@ -3,7 +3,6 @@ use crate::config::ConstraintResult; use crate::context::ContextualUserFragment; use crate::context::GuardianReviewEvidence; use crate::elicitation::ElicitationRegistration; -use crate::environment_selection::TurnEnvironmentState; use crate::session::SessionIo; use crate::session::SessionSettingsUpdate; use crate::session::Submission; @@ -531,28 +530,6 @@ impl CodexThread { self.session.inject_if_running(items).await } - /// Environment selections captured by the active turn, before later settings updates. - /// Includes environments that are still starting or have failed. Hosts use this snapshot - /// to authorize steering against every executor that the active turn selected. - pub async fn active_turn_environment_selections( - &self, - ) -> Option> { - let active = self.session.active_turn.lock().await; - let task = active.as_ref()?.task.as_ref()?; - Some( - task.turn_context - .initial_environments - .environments - .iter() - .map(|environment| match environment { - TurnEnvironmentState::Ready(environment) => environment.selection(), - TurnEnvironmentState::Starting(environment) => environment.selection.clone(), - TurnEnvironmentState::Failed { selection, .. } => selection.clone(), - }) - .collect(), - ) - } - /// Captures a regular turn only after its input is recorded. The caller must flush the rollout. pub async fn interrupted_turn( &self, @@ -837,13 +814,13 @@ impl CodexThread { /// Returns the active turn's reviewer, including live updates, or the thread default. pub async fn approvals_reviewer_for_turn(&self, turn_id: &str) -> ApprovalsReviewer { - if let Some((turn, inputs, _)) = self + if let Some((turn, settings, _, _)) = self .session .active_turn_context_and_strict_auto_review() .await && turn.sub_id == turn_id { - inputs.settings.approvals_reviewer() + settings.approvals_reviewer() } else { self.config_snapshot().await.approvals_reviewer } diff --git a/codex-rs/core/src/environment_selection.rs b/codex-rs/core/src/environment_selection.rs index f9ad1b782e..ea3d31931b 100644 --- a/codex-rs/core/src/environment_selection.rs +++ b/codex-rs/core/src/environment_selection.rs @@ -18,6 +18,8 @@ use codex_exec_server::SelectedCapabilityRootsStatus; use codex_protocol::capabilities::CapabilityRootLocation; use codex_protocol::capabilities::SelectedCapabilityRoot; use codex_protocol::config_types::ShellEnvironmentPolicy; +use codex_protocol::error::CodexErr; +use codex_protocol::error::Result as CodexResult; use codex_protocol::models::PermissionProfile; use codex_protocol::protocol::AskForApproval; use codex_protocol::protocol::EnvironmentConfig; @@ -44,6 +46,9 @@ use crate::shell::Shell; use crate::shell_snapshot::ShellSnapshot; use crate::shell_snapshot::SnapshotCredentialBrokerState; +// Reject pathological selected cwd values at the environment-selection boundary. +const MAX_TURN_ENVIRONMENT_CWD_BYTES: usize = 8 * 1024; + /// Records whether a normalized config should follow later thread setting updates. #[derive(Clone, Copy, Debug, Eq, PartialEq)] pub(crate) enum EnvironmentConfigOrigin { @@ -134,6 +139,36 @@ pub(crate) fn default_thread_environment_selections( .collect() } +/// Checks that environment IDs are registered and unique and that working directories are not too long. +pub fn validate_environment_ids_and_cwds( + environment_manager: &EnvironmentManager, + environments: &[TurnEnvironmentSelection], +) -> CodexResult<()> { + let mut environment_ids = HashSet::with_capacity(environments.len()); + for environment in environments { + if environment.cwd.inferred_native_path_string().len() > MAX_TURN_ENVIRONMENT_CWD_BYTES { + return Err(CodexErr::InvalidRequest( + "turn environment working directory exceeds the maximum size".to_string(), + )); + } + if !environment_ids.insert(environment.environment_id.as_str()) { + return Err(CodexErr::InvalidRequest(format!( + "duplicate turn environment id `{}`", + environment.environment_id + ))); + } + environment_manager + .get_environment(&environment.environment_id) + .ok_or_else(|| { + CodexErr::InvalidRequest(format!( + "unknown turn environment id `{}`", + environment.environment_id + )) + })?; + } + Ok(()) +} + type TurnEnvironmentResult = Result>; type TurnEnvironmentResolution = Shared>; type PendingConfigurationResult = Result; @@ -334,6 +369,7 @@ impl ThreadEnvironments { shell_snapshot } + /// Updates the selected list before waking work that was waiting on the previous selection. pub(crate) fn update_selections(&self, environments: &[TurnEnvironmentSelection]) { let mut state = self.state.lock().unwrap_or_else(PoisonError::into_inner); let ThreadEnvironmentsState { @@ -471,6 +507,7 @@ impl ThreadEnvironments { for task in removed_connection_tasks { task.abort(); } + drop(previous); } @@ -791,55 +828,58 @@ impl ThreadEnvironments { /// Captures the selected list immediately, then returns a future that may wait for setup. /// This lets callers release their locks before waiting without accidentally using a newer /// selection. The trace still covers the setup wait itself. - #[tracing::instrument( - name = "environments.snapshot", - skip_all, - fields( - environment_count = selected.len(), - non_blocking = non_blocking_snapshots, - ) - )] pub(crate) fn snapshot( &self, ) -> impl Future + Send + 'static { + self.resolve_snapshot(self.snapshot_now()) + } + + /// Copies the current selection without waiting for any executor to connect. + pub(crate) fn snapshot_now(&self) -> TurnEnvironmentSnapshot { // Keep the startup results, but not the senders that would keep removed Pending work alive. - let selected = self + let environments = self .state .lock() .unwrap_or_else(PoisonError::into_inner) .environments .iter() - .map(|environment| StartingTurnEnvironment { - selection: environment.selection.clone(), - config_origin: environment.config_origin, - resolution: environment.resolution.clone(), + .map(|environment| { + let starting = StartingTurnEnvironment { + selection: environment.selection.clone(), + config_origin: environment.config_origin, + resolution: environment.resolution.clone(), + }; + let resolved = starting.resolution.clone().now_or_never(); + TurnEnvironmentState::from_resolution(starting, resolved) }) - .collect::>(); + .collect(); + TurnEnvironmentSnapshot { environments } + } + + /// Applies the same startup policy to an already captured selection without adopting newer ones. + #[tracing::instrument( + name = "environments.snapshot", + skip_all, + fields( + environment_count = snapshot.environments.len(), + non_blocking = non_blocking_snapshots, + ) + )] + fn resolve_snapshot( + &self, + snapshot: TurnEnvironmentSnapshot, + ) -> impl Future + Send + 'static { let non_blocking_snapshots = self.non_blocking_snapshots; async move { - let mut environments = Vec::with_capacity(selected.len()); - for starting in selected { - if let EnvironmentConfigState::Failed(error) = &starting.selection.config { - environments.push(TurnEnvironmentState::Failed { - selection: starting - .config_origin - .into_input_selection(starting.selection.clone()), - error: error.clone(), - }); - continue; + if !non_blocking_snapshots { + for starting in snapshot.starting() { + if !matches!(starting.selection.config, EnvironmentConfigState::Pending) { + let _ = starting.wait_until_ready().await; + } } - let pending = matches!(starting.selection.config, EnvironmentConfigState::Pending); - let resolved = if non_blocking_snapshots || pending { - starting.resolution.clone().now_or_never() - } else { - Some(match starting.wait_until_ready().await { - Ok(()) => starting.resolution.clone().await, - Err(error) => Err(error), - }) - }; - environments.push(TurnEnvironmentState::from_resolution(starting, resolved)); } - TurnEnvironmentSnapshot { environments } + // An earlier environment may have become ready while we waited for a later one. + snapshot.refresh_readiness() } } @@ -865,6 +905,15 @@ impl TurnEnvironmentState { starting: StartingTurnEnvironment, resolved: Option, ) -> Self { + if let EnvironmentConfigState::Failed(error) = &starting.selection.config { + let error = error.clone(); + return Self::Failed { + selection: starting + .config_origin + .into_input_selection(starting.selection), + error, + }; + } match resolved { Some(Ok(environment)) => { let mut selection = starting.selection; diff --git a/codex-rs/core/src/hook_runtime.rs b/codex-rs/core/src/hook_runtime.rs index d4d568d975..7808d695b0 100644 --- a/codex-rs/core/src/hook_runtime.rs +++ b/codex-rs/core/src/hook_runtime.rs @@ -53,6 +53,7 @@ use codex_protocol::protocol::WarningEvent; use codex_rollout::state_db; use codex_thread_store::PersistContext; use codex_thread_store::ReadThreadParams; +use codex_utils_absolute_path::AbsolutePathBuf; use serde_json::Map; use serde_json::Value; use tokio::sync::Mutex; @@ -60,6 +61,7 @@ use tracing::instrument; use crate::context::ContextualUserFragment; use crate::context::HookAdditionalContext; +use crate::environment_selection::TurnEnvironmentSnapshot; use crate::event_mapping::parse_turn_item; use crate::guardian::GuardianReviewContext; use crate::session::TurnInput; @@ -195,8 +197,7 @@ pub(crate) async fn run_pre_tool_use_hooks( session_id: sess.session_id().into(), turn_id: turn_context.sub_id.clone(), subagent: thread_spawn_subagent_hook_context(sess, turn_context), - #[allow(deprecated)] - cwd: turn_context.cwd.clone(), + cwd: tool_hook_cwd(&step_context.environments, turn_context), transcript_path: sess.hook_transcript_path().await, model: step_context.settings.model_info.slug.clone(), permission_mode: hook_permission_mode(step_context.settings.approval_policy()), @@ -243,6 +244,14 @@ pub(crate) async fn run_pre_tool_use_hooks( } } +#[allow(deprecated)] +fn tool_hook_cwd(environments: &TurnEnvironmentSnapshot, turn: &TurnContext) -> AbsolutePathBuf { + // Hooks run on the host, so a remote workspace cannot replace the local fallback. + environments + .local_environment_cwd() + .unwrap_or_else(|| turn.cwd.clone()) +} + // PermissionRequest hooks share the same preview/start/completed event flow as // other hook types, but they return an optional decision instead of mutating // tool input or post-run state. @@ -257,8 +266,7 @@ pub(crate) async fn run_permission_request_hooks( session_id: sess.session_id().into(), turn_id: turn_context.sub_id.clone(), subagent: thread_spawn_subagent_hook_context(sess, turn_context), - #[allow(deprecated)] - cwd: turn_context.cwd.to_path_buf(), + cwd: tool_hook_cwd(review_context.environments(), turn_context).to_path_buf(), transcript_path: sess.hook_transcript_path().await, model: review_context.model_info.slug.clone(), permission_mode: hook_permission_mode(review_context.approval_policy), @@ -300,8 +308,7 @@ pub(crate) async fn run_post_tool_use_hooks( session_id: sess.session_id().into(), turn_id: turn_context.sub_id.clone(), subagent: thread_spawn_subagent_hook_context(sess, turn_context), - #[allow(deprecated)] - cwd: turn_context.cwd.clone(), + cwd: tool_hook_cwd(&step_context.environments, turn_context), transcript_path: sess.hook_transcript_path().await, model: step_context.settings.model_info.slug.clone(), permission_mode: hook_permission_mode(step_context.settings.approval_policy()), diff --git a/codex-rs/core/src/lib.rs b/codex-rs/core/src/lib.rs index 3402a27dfb..e7a2a4a3c1 100644 --- a/codex-rs/core/src/lib.rs +++ b/codex-rs/core/src/lib.rs @@ -244,3 +244,4 @@ pub mod otel_init; // Captured environment bindings can be passed back to ThreadManager by internal reviewers. pub use environment_selection::TurnEnvironmentSnapshot; +pub use environment_selection::validate_environment_ids_and_cwds; diff --git a/codex-rs/core/src/session/daemon_recovery.rs b/codex-rs/core/src/session/daemon_recovery.rs index 84ceca138d..5fdefff870 100644 --- a/codex-rs/core/src/session/daemon_recovery.rs +++ b/codex-rs/core/src/session/daemon_recovery.rs @@ -20,12 +20,12 @@ impl Session { } let context = &task.turn_context; context.extension_data.get::()?; - let inputs = context.next_step_input.load(); + let settings = context.next_step_settings.load(); + let environments = self.services.turn_environments.snapshot_now(); // Remote identities/configuration are not persisted across daemon restarts. - if inputs.environments.environments.len() != 1 { + if environments.environments.len() != 1 { return None; } - let environments = inputs.environments.refresh_readiness(); let environment = environments.single_local_environment()?.selection(); if environment.config != EnvironmentConfigState::FromThread { return None; @@ -34,7 +34,7 @@ impl Session { context.sub_id.clone(), crate::TurnStartOptions { final_output_json_schema: context.final_output_json_schema.clone(), - service_tier: Some(inputs.settings.service_tier.clone().unwrap_or_else(|| { + service_tier: Some(settings.service_tier.clone().unwrap_or_else(|| { codex_protocol::config_types::SERVICE_TIER_DEFAULT_REQUEST_VALUE.to_string() })), cyber_access_program: context.cyber_access_program, diff --git a/codex-rs/core/src/session/environment.rs b/codex-rs/core/src/session/environment.rs index f0134ddc41..28f339c12e 100644 --- a/codex-rs/core/src/session/environment.rs +++ b/codex-rs/core/src/session/environment.rs @@ -45,7 +45,7 @@ impl ThreadEnvironmentDefaults { } } -pub(super) fn validate_environment_selections( +pub(super) fn validate_environment_configs( selections: &[TurnEnvironmentSelection], ) -> ConstraintResult<()> { for selection in selections { @@ -68,6 +68,27 @@ pub(super) fn validate_environment_selections( Ok(()) } +pub(super) fn ensure_configs_stay_owner_provided( + current: &[TurnEnvironmentSelection], + proposed: &[TurnEnvironmentSelection], +) -> ConstraintResult<()> { + if let Some(environment) = proposed.iter().find(|environment| { + environment.config == EnvironmentConfigState::FromThread + && current.iter().any(|current| { + current.environment_id == environment.environment_id + && current.config != EnvironmentConfigState::FromThread + }) + }) { + return Err(ConstraintError::InvalidValue { + field_name: "environments", + candidate: environment.environment_id.clone(), + allowed: "owner-provided environment configuration".to_string(), + requirement_source: codex_config::RequirementSource::Unknown, + }); + } + Ok(()) +} + fn validate_environment_config( selection: &TurnEnvironmentSelection, config: &EnvironmentConfig, @@ -137,21 +158,8 @@ impl Session { updates: &SessionSettingsUpdate, ) -> ConstraintResult { let current_environments = ¤t.environments; - if let Some(environments) = &updates.environments - && let Some(environment) = environments.environments.iter().find(|environment| { - environment.config == EnvironmentConfigState::FromThread - && current_environments.iter().any(|current| { - current.environment_id == environment.environment_id - && current.config != EnvironmentConfigState::FromThread - }) - }) - { - return Err(ConstraintError::InvalidValue { - field_name: "environments", - candidate: environment.environment_id.clone(), - allowed: "owner-provided environment configuration".to_string(), - requirement_source: codex_config::RequirementSource::Unknown, - }); + if let Some(environments) = &updates.environments { + ensure_configs_stay_owner_provided(current_environments, &environments.environments)?; } current.apply(updates, current_environments) diff --git a/codex-rs/core/src/session/mcp.rs b/codex-rs/core/src/session/mcp.rs index 036e90afc4..bcb699e6b2 100644 --- a/codex-rs/core/src/session/mcp.rs +++ b/codex-rs/core/src/session/mcp.rs @@ -744,7 +744,7 @@ async fn review_guardian_mcp_elicitation( let Some(mcp_config) = session.services.mcp_runtime.current_config() else { return Ok(None); }; - let step_settings = Arc::clone(&turn_context.next_step_input.load().settings); + let step_settings = turn_context.next_step_settings.load_full(); // User approval skips ordinary CUA checks, not separate sensitive requests. let user_cua_execution = step_settings.approvals_reviewer() == ApprovalsReviewer::User diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 54efc41335..b977bb16f6 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -3314,14 +3314,32 @@ impl Session { )] pub(crate) async fn active_turn_context_and_strict_auto_review( &self, - ) -> Option<(Arc, Arc, bool)> { + ) -> Option<( + Arc, + Arc, + TurnEnvironmentSnapshot, + bool, + )> { let active = self.active_turn.lock().await; let active = active.as_ref()?; let task = active.task.as_ref()?; let turn_context = Arc::clone(&task.turn_context); - let step_inputs = turn_context.next_step_input.load_full(); - let ts = active.turn_state.lock().await; - Some((turn_context, step_inputs, ts.strict_auto_review_enabled())) + let settings = turn_context.next_step_settings.load_full(); + let strict_auto_review = active.turn_state.lock().await.strict_auto_review_enabled(); + let environments = self.services.turn_environments.snapshot_now(); + Some((turn_context, settings, environments, strict_auto_review)) + } + + #[expect( + clippy::await_holding_invalid_type, + reason = "active turn reads must stay consistent with the matching turn state" + )] + pub(crate) async fn strict_auto_review_enabled(&self) -> bool { + let active = self.active_turn.lock().await; + let Some(active) = active.as_ref().filter(|active| active.task.is_some()) else { + return false; + }; + active.turn_state.lock().await.strict_auto_review_enabled() } pub(crate) async fn granted_session_permissions( @@ -3764,10 +3782,19 @@ impl Session { required_servers: &[String], required_plugins: &HashSet, ) -> CodexResult> { - // Capture settings and selection together before asynchronous planning. - // Existing steps retain this version even if the turn is updated. - let inputs = turn_context.next_step_input.load_full(); - let mut settings = Arc::clone(&inputs.settings); + // Read the step's model and record its environments together so an update cannot split them. + // Wait for executor startup below, after releasing the lock. + let (mut settings, environments) = { + let _active = self + .active_turn + .lock() + .or_cancel(cancellation_token) + .await?; + ( + turn_context.next_step_settings.load_full(), + self.services.turn_environments.snapshot(), + ) + }; if matches!( turn_context.session_source, SessionSource::SubAgent(SubAgentSource::ThreadSpawn { .. }) @@ -3792,8 +3819,7 @@ impl Session { settings.model_info.as_ref(), ); let session_telemetry = settings.telemetry(&turn_context.session_telemetry); - // Refresh only the captured step selection, without adopting newer inputs. - let environments = inputs.environments.refresh_readiness(); + let environments = environments.or_cancel(cancellation_token).await?; let (loaded_agents_md, warnings) = self .services .agents_md_manager diff --git a/codex-rs/core/src/session/notification_tests.rs b/codex-rs/core/src/session/notification_tests.rs index 42f24ace69..456970a7bb 100644 --- a/codex-rs/core/src/session/notification_tests.rs +++ b/codex-rs/core/src/session/notification_tests.rs @@ -45,10 +45,7 @@ async fn task_finish_preserves_notification_budget_and_queued_user_input() { let mut current = turn.initial_settings.as_ref().clone(); Arc::make_mut(&mut current.model_info).truncation_policy = codex_protocol::openai_models::TruncationPolicyConfig::tokens(/*limit*/ 400); - turn.next_step_input.store(Arc::new(StepInputs { - settings: Arc::new(current), - environments: turn.next_step_input.load().environments.clone(), - })); + turn.next_step_settings.store(Arc::new(current)); session .on_task_finished(Arc::clone(&turn), /*task_result*/ Ok(None)) diff --git a/codex-rs/core/src/session/review.rs b/codex-rs/core/src/session/review.rs index 4728c750d9..1762e6b05d 100644 --- a/codex-rs/core/src/session/review.rs +++ b/codex-rs/core/src/session/review.rs @@ -1,4 +1,3 @@ -use super::step_context::StepInputs; use super::step_settings::ResolvedStepSettings; use super::*; use arc_swap::ArcSwap; @@ -164,10 +163,7 @@ pub(super) async fn spawn_review_thread( initial_settings: Arc::clone(&step_settings), disabled_plugin_ids: parent_turn_context.disabled_plugin_ids.clone(), active_host_plugin_identities: None, - next_step_input: ArcSwap::from_pointee(StepInputs { - settings: step_settings, - environments: parent_turn_context.initial_environments.clone(), - }), + next_step_settings: ArcSwap::from(step_settings), session_telemetry: session_telemetry_for_context, provider: provider_for_context, session_source, diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index bccda5e63b..da91f5ed34 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -366,7 +366,7 @@ impl SessionConfiguration { ) -> ConstraintResult<()> { self.step_settings .validate(&self.step_settings_constraints(environments))?; - super::environment::validate_environment_selections(environments) + super::environment::validate_environment_configs(environments) } pub(super) fn step_settings_constraints( @@ -537,7 +537,7 @@ impl SessionConfiguration { .map_or(current_environments, |environments| { environments.environments.as_slice() }); - super::environment::validate_environment_selections(next_environments)?; + super::environment::validate_environment_configs(next_environments)?; // Apply step settings last: the proposed permissions and environment // selections must be complete before deriving their validation constraints. next_configuration.step_settings = Arc::new(self.step_settings.apply( diff --git a/codex-rs/core/src/session/step_activation.rs b/codex-rs/core/src/session/step_activation.rs index 1086d937b8..c50b4b5377 100644 --- a/codex-rs/core/src/session/step_activation.rs +++ b/codex-rs/core/src/session/step_activation.rs @@ -1,14 +1,16 @@ //! Restricted updates to a running turn's immutable settings snapshots. +use super::environment::ensure_configs_stay_owner_provided; +use super::environment::validate_environment_configs; use super::session::Session; use super::session::SessionConfiguration; -use super::step_context::StepInputs; use super::step_settings::ResolvedStepSettings; use super::step_settings::StepSettingsConstraints; use super::step_settings::StepSettingsUpdate; use super::turn_context::TurnContext; use crate::config::Config; use crate::config::ConstraintResult; +use crate::environment_selection::validate_environment_ids_and_cwds; use crate::exec_policy::AllowPrefixRules; use codex_features::Feature; use codex_prompts::ResolvedModelMessages; @@ -16,6 +18,8 @@ use codex_protocol::openai_models::GuardianV2ModelConfig; use codex_protocol::openai_models::GuardianV2TranscriptModelConfig; use codex_protocol::openai_models::MODEL_SPECIALTY_CYBER; use codex_protocol::openai_models::ModelInfo; +use codex_protocol::protocol::EnvironmentConfigState; +use codex_protocol::protocol::TurnEnvironmentSelection; use codex_protocol::protocol::TurnSettingsUpdate; use codex_protocol::protocol::TurnSettingsUpdateOutcome; use std::sync::Arc; @@ -227,20 +231,21 @@ impl Session { turn_id: &str, update: TurnSettingsUpdate, ) -> TurnSettingsUpdateOutcome { - let reviewer_only = update.approvals_reviewer.is_some() - && update.model.is_none() - && update.effort.is_none() - && update.summary.is_none() - && update.service_tier.is_none(); - if !reviewer_only && !self.features.enabled(Feature::StepModelSwitching) { + let updates_model_settings = update.model.is_some() + || update.effort.is_some() + || update.summary.is_some() + || update.service_tier.is_some(); + let requires_model_switching = updates_model_settings + || (update.approvals_reviewer.is_none() && update.environments.is_none()); + if requires_model_switching && !self.features.enabled(Feature::StepModelSwitching) { return TurnSettingsUpdateOutcome::Rejected { reason: "turn settings updates require the step_model_switching feature" .to_string(), }; } - // Capture the exact live task and its settings, then release the - // lock. A task that starts during preparation is never a new target. + // Capture the running task named by this update and its settings, then release the lock. + // A task that starts during preparation is never a new target. let target = { let active = self.active_turn.lock().await; active.as_ref().and_then(|active| { @@ -250,7 +255,7 @@ impl Session { ( Arc::clone(&task.turn_context), Arc::clone(&task.done), - task.turn_context.next_step_input.load_full(), + task.turn_context.next_step_settings.load_full(), ) }) }) @@ -261,11 +266,13 @@ impl Session { }; let TurnSettingsUpdate { approvals_reviewer, + environments, model, effort, summary, service_tier, } = update; + let updates_step_settings = updates_model_settings || approvals_reviewer.is_some(); let update = StepSettingsUpdate { approvals_reviewer, model, @@ -274,65 +281,95 @@ impl Session { service_tier, ..Default::default() }; - // Apply the sparse patch to the captured active base using the shared - // settings rules. The task can progress, finish, or be cancelled while - // preparation awaits; no publication locks are held here. - let prepared = self - .prepare_step_settings_activation(&turn_context, ¤t.settings, &update) - .await; + // Build the proposed model and reviewer settings. + // Apply only the fields the caller supplied to the settings captured above. The task can + // progress, finish, or be cancelled while this awaits; we don't hold the update locks here. + let prepared = if updates_step_settings { + let current_environments = self.services.turn_environments.selections(); + let proposed = environments.as_deref().unwrap_or(¤t_environments); + self.prepare_step_settings_activation(&turn_context, ¤t, &update, proposed) + .await + .map(Some) + } else { + Ok(None) + }; + + // Validate the environment configuration supplied in the update. + let environment_config_validation = + environments.as_deref().map(validate_environment_configs); + + // Confirm the task we originally targeted is still running. let active = self.active_turn.lock().await; let Some(task) = active.as_ref().and_then(|active| active.task.as_ref()) else { return TurnSettingsUpdateOutcome::TargetUnavailable; }; - // A later task may reuse the same context and turn ID. `done` is - // allocated per task, so matching only the ID/context is insufficient. - // A mismatch abandons the update without retrying or retargeting. + // A later task may reuse the same context and turn ID. `done` is a completion signal + // created for each task, so matching only the ID/context is insufficient. A mismatch + // abandons the update without retrying or applying it to the replacement. if !Arc::ptr_eq(&task.done, &task_done) || !Arc::ptr_eq(&task.turn_context, &turn_context) - || !Arc::ptr_eq(&task.turn_context.next_step_input.load_full(), ¤t) || task.cancellation_token.is_cancelled() { return TurnSettingsUpdateOutcome::TargetUnavailable; } - let destination = match prepared { - Ok(destination) => destination, + let updated_settings = match prepared { + Ok(settings) => settings, Err(reason) => return TurnSettingsUpdateOutcome::Rejected { reason }, }; - // Managed requirements can change during resolution. Keep the live - // authorization and safety checks atomic with publication under state - // and active_turn; no asynchronous preparation runs under these locks. + let candidate_settings = updated_settings.as_ref().unwrap_or(current.as_ref()); + + // Recheck the latest rules before applying the update. + // Managed requirements can change during model lookup. Keep the live authorization and + // safety checks together with applying the update under state and active_turn; no model + // lookup or other preparation runs under these locks. let state = self.state.lock().await; - if let Err(reason) = self - .validate_active_step_settings( + // Environment configuration can arrive before its executor connects. If this update + // omits environments, use what the running turn's manager knows now. + let current_environments = self.services.turn_environments.selections(); + let proposed = environments.as_deref().unwrap_or(¤t_environments); + let validation = (|| { + if let Some(configs) = environment_config_validation { + validate_environment_ids_and_cwds( + &self.services.turn_environments.environment_manager(), + proposed, + ) + .map_err(|error| error.to_string())?; + ensure_configs_stay_owner_provided(¤t_environments, proposed) + .map_err(|error| error.to_string())?; + configs.map_err(|error| error.to_string())?; + } + self.validate_active_step_settings( &turn_context, - &destination, + candidate_settings, &state.session_configuration, + proposed, ) - .map_err(|error| error.to_string()) - .and_then(|()| { - // A reviewer-only patch cannot change any model-owned authority. - // Managed reviewer restrictions were checked above. - if reviewer_only { - return Ok(()); - } + .map_err(|error| error.to_string())?; + // Neither a reviewer nor an environment change changes model-owned authority. + if requires_model_switching { check_legacy_turn_safety( &turn_context, - ¤t.settings, - &destination, + ¤t, + candidate_settings, &state.session_configuration.original_config_do_not_use, - ) - }) - { + )?; + } + Ok::<_, String>(()) + })(); + if let Err(reason) = validation { return TurnSettingsUpdateOutcome::Rejected { reason }; } - // Publish the immutable snapshot. Frozen initial settings, existing step - // captures, and future thread settings are not changed. - task.turn_context - .next_step_input - .store(Arc::new(StepInputs { - settings: Arc::new(destination), - environments: current.environments.clone(), - })); + + // Save changed settings first. The manager then installs any new environment list before + // waking waiting work. The next step cannot read either until we release this lock. + if let Some(settings) = updated_settings { + task.turn_context + .next_step_settings + .store(Arc::new(settings)); + } + if environments.is_some() { + self.services.turn_environments.update_selections(proposed); + } TurnSettingsUpdateOutcome::Applied } @@ -341,6 +378,7 @@ impl Session { turn_context: &TurnContext, current: &ResolvedStepSettings, update: &StepSettingsUpdate, + environments: &[TurnEnvironmentSelection], ) -> Result { let (requirements, overrides, trusted_guardian_reviewer) = { let state = self.state.lock().await; @@ -356,9 +394,10 @@ impl Session { requirements: &requirements, guardian_approval_enabled: self.features.enabled(Feature::GuardianApproval), trusted_guardian_reviewer, - has_full_disk_write_access: turn_context - .file_system_sandbox_policy() - .has_full_disk_write_access(), + has_full_disk_write_access: any_environment_has_full_disk_write( + turn_context, + environments, + ), }; current .apply_update( @@ -380,6 +419,7 @@ impl Session { turn_context: &TurnContext, settings: &ResolvedStepSettings, configuration: &SessionConfiguration, + environments: &[TurnEnvironmentSelection], ) -> ConstraintResult<()> { let requirements = configuration .original_config_do_not_use @@ -389,13 +429,44 @@ impl Session { requirements, guardian_approval_enabled: self.features.enabled(Feature::GuardianApproval), trusted_guardian_reviewer: configuration.trusted_guardian_reviewer, - has_full_disk_write_access: turn_context - .file_system_sandbox_policy() - .has_full_disk_write_access(), + has_full_disk_write_access: any_environment_has_full_disk_write( + turn_context, + environments, + ), }) } } +fn any_environment_has_full_disk_write( + turn: &TurnContext, + environments: &[TurnEnvironmentSelection], +) -> bool { + let fallback = turn.config.permissions.permission_profile(); + let mut known = environments.iter().filter_map(|environment| { + let (profile, roots) = match &environment.config { + EnvironmentConfigState::FromThread => (fallback, &environment.workspace_roots), + EnvironmentConfigState::Ready(config) => ( + config.permission_profile.permission_profile(), + &config.workspace_roots, + ), + EnvironmentConfigState::Pending | EnvironmentConfigState::Failed(_) => return None, + }; + Some( + profile + .clone() + .materialize_project_roots_with_path_uris(roots) + .file_system_sandbox_policy() + .has_full_disk_write_access_for_convention(environment.cwd.infer_path_convention()), + ) + }); + match known.next() { + Some(first) => first || known.any(|full_write| full_write), + None => fallback + .file_system_sandbox_policy() + .has_full_disk_write_access(), + } +} + #[cfg(test)] #[path = "step_activation_tests.rs"] mod tests; diff --git a/codex-rs/core/src/session/step_activation_tests.rs b/codex-rs/core/src/session/step_activation_tests.rs index cae3f507c7..8ab26dc96f 100644 --- a/codex-rs/core/src/session/step_activation_tests.rs +++ b/codex-rs/core/src/session/step_activation_tests.rs @@ -9,6 +9,7 @@ use crate::session::tests::HeldStepTask; use crate::session::tests::make_session_and_context; use crate::session::tests::update_selected_settings_for_test; use crate::session::tests::update_turn_settings_for_test; +use crate::session::turn_context::TurnEnvironment; use crate::state::TaskKind; use codex_config::AutoReviewRequirementsToml; use codex_config::ConfigLayerStack; @@ -161,6 +162,45 @@ async fn instruction_refresh_serializes_reads_and_releases_on_cancellation() { ); } +#[tokio::test] +async fn next_step_waits_for_environment_publication_before_capturing() { + let ActivationFixture { session, turn, .. } = activation_fixture(activation_models()).await; + let next_cwd = turn.config.cwd.join("next-step-environment"); + std::fs::create_dir_all(&next_cwd).expect("create next workspace"); + let next_cwd = codex_utils_path_uri::PathUri::from_abs_path(&next_cwd); + let mut selections = session.services.turn_environments.selections(); + let primary = selections.first_mut().expect("initial environment"); + assert_ne!(primary.cwd, next_cwd); + primary.cwd = next_cwd.clone(); + primary.workspace_roots = vec![next_cwd.clone()]; + let mut next_settings = turn.next_step_settings.load_full().as_ref().clone(); + Arc::make_mut(&mut next_settings.model_info).slug = "next-step-model".to_string(); + + let cancellation = CancellationToken::new(); + let capture = session.capture_step_context(Arc::clone(&turn), &cancellation); + tokio::pin!(capture); + { + let _active = session.active_turn.lock().await; + let mut context = std::task::Context::from_waker(futures::task::noop_waker_ref()); + assert!(std::future::Future::poll(capture.as_mut(), &mut context).is_pending()); + turn.next_step_settings.store(Arc::new(next_settings)); + session + .services + .turn_environments + .update_selections(&selections); + } + + let step = capture.await.expect("capture after publication"); + assert_eq!( + ( + step.settings.model_info.slug.as_str(), + step.environments.primary().map(TurnEnvironment::cwd), + ), + ("next-step-model", Some(&next_cwd)), + ); + session.abort_all_tasks(TurnAbortReason::Interrupted).await; +} + #[test_case(None; "thread only")] #[test_case(Some("global instructions"); "with global instructions")] #[tokio::test] @@ -740,7 +780,7 @@ async fn delayed_activation_does_not_retarget_a_task(change: TaskChangeDuringLoo lookup, } = activation_fixture(models).await; let desired = desired_step_settings(&session).await; - let original = turn.next_step_input.load_full(); + let original = turn.next_step_settings.load_full(); let update_session = Arc::clone(&session); let turn_id = turn.sub_id.clone(); let update = tokio::spawn(async move { @@ -788,7 +828,7 @@ async fn delayed_activation_does_not_retarget_a_task(change: TaskChangeDuringLoo TaskChangeDuringLookup::FinishedAndReusedContext => Arc::clone(&turn), TaskChangeDuringLookup::CancelledWithRejectedDestination => unreachable!(), }; - let inputs = replacement.next_step_input.load_full(); + let inputs = replacement.next_step_settings.load_full(); session .spawn_task( Arc::clone(&replacement), @@ -799,7 +839,7 @@ async fn delayed_activation_does_not_retarget_a_task(change: TaskChangeDuringLoo }, ) .await; - assert!(Arc::ptr_eq(&turn.next_step_input.load_full(), &original)); + assert!(Arc::ptr_eq(&turn.next_step_settings.load_full(), &original)); (replacement, inputs) } }; @@ -809,7 +849,7 @@ async fn delayed_activation_does_not_retarget_a_task(change: TaskChangeDuringLoo TurnSettingsUpdateOutcome::TargetUnavailable ); assert!(Arc::ptr_eq( - &expected_turn.next_step_input.load_full(), + &expected_turn.next_step_settings.load_full(), &expected_inputs, )); assert_eq!(desired_step_settings(&session).await, desired); @@ -827,7 +867,7 @@ enum ManagedAuthorizationChange { #[tokio::test] async fn reviewer_only_activation_enforces_managed_authority(required_review: bool) { let ActivationFixture { session, turn, .. } = activation_fixture(activation_models()).await; - let original = turn.next_step_input.load_full(); + let original = turn.next_step_settings.load_full(); let desired = desired_step_settings(&session).await; let source = RequirementSource::Unknown; let mut sourced = ConfigRequirementsWithSources::default(); @@ -871,7 +911,7 @@ async fn reviewer_only_activation_enforces_managed_authority(required_review: bo .await, TurnSettingsUpdateOutcome::Rejected { .. } )); - assert!(Arc::ptr_eq(&turn.next_step_input.load_full(), &original)); + assert!(Arc::ptr_eq(&turn.next_step_settings.load_full(), &original)); assert_eq!(desired_step_settings(&session).await, desired); session.abort_all_tasks(TurnAbortReason::Replaced).await; } @@ -888,7 +928,7 @@ async fn delayed_activation_rechecks_live_managed_authorization( lookup, .. } = activation_fixture(activation_models()).await; - let original = turn.next_step_input.load_full(); + let original = turn.next_step_settings.load_full(); let desired = desired_step_settings(&session).await; let update_session = Arc::clone(&session); let turn_id = turn.sub_id.clone(); @@ -914,7 +954,7 @@ async fn delayed_activation_rechecks_live_managed_authorization( }; let sourced = match change { ManagedAuthorizationChange::ApprovalPolicy => { - let allowed = if original.settings.approval_policy() == AskForApproval::Never { + let allowed = if original.approval_policy() == AskForApproval::Never { AskForApproval::OnRequest } else { AskForApproval::Never @@ -925,7 +965,7 @@ async fn delayed_activation_rechecks_live_managed_authorization( } } ManagedAuthorizationChange::ApprovalsReviewer => { - let allowed = if original.settings.approvals_reviewer() == ApprovalsReviewer::User { + let allowed = if original.approvals_reviewer() == ApprovalsReviewer::User { ApprovalsReviewer::AutoReview } else { ApprovalsReviewer::User @@ -955,10 +995,10 @@ async fn delayed_activation_rechecks_live_managed_authorization( match change { ManagedAuthorizationChange::ApprovalPolicy => requirements .approval_policy - .can_set(&original.settings.approval_policy()), + .can_set(&original.approval_policy()), ManagedAuthorizationChange::ApprovalsReviewer => requirements .approvals_reviewer - .can_set(&original.settings.approvals_reviewer()), + .can_set(&original.approvals_reviewer()), } .expect_err("the refreshed constraint rejects the admitted value") }; @@ -969,7 +1009,7 @@ async fn delayed_activation_rechecks_live_managed_authorization( reason: expected_error.to_string(), } ); - assert!(Arc::ptr_eq(&turn.next_step_input.load_full(), &original)); + assert!(Arc::ptr_eq(&turn.next_step_settings.load_full(), &original)); assert_eq!(desired_step_settings(&session).await, desired); session.abort_all_tasks(TurnAbortReason::Replaced).await; } @@ -1029,7 +1069,12 @@ async fn activation_must_match_the_retained_turn_authority(change: ManagedPolicy let mut live = session.state.lock().await.session_configuration.clone(); live.original_config_do_not_use = Arc::clone(&prepared.config); assert_eq!( - session.validate_active_step_settings(&prepared, &destination, &live,), + session.validate_active_step_settings( + &prepared, + &destination, + &live, + &session.services.turn_environments.selections(), + ), Ok(()) ); assert_eq!( @@ -1083,7 +1128,12 @@ async fn activation_must_match_the_retained_turn_authority(change: ManagedPolicy // The destination satisfies today's managed policy. Only the temporary // compatibility check rejects the mismatch with retained turn consumers. assert_eq!( - session.validate_active_step_settings(&prepared, &destination, &live,), + session.validate_active_step_settings( + &prepared, + &destination, + &live, + &session.services.turn_environments.selections(), + ), Ok(()) ); assert_eq!( diff --git a/codex-rs/core/src/session/step_context.rs b/codex-rs/core/src/session/step_context.rs index 080243f955..855d854fd0 100644 --- a/codex-rs/core/src/session/step_context.rs +++ b/codex-rs/core/src/session/step_context.rs @@ -15,14 +15,6 @@ use codex_otel::SessionTelemetry; use codex_protocol::items::ModelInvocationContext; use codex_protocol::protocol::TurnContextItem; -/// Inputs for the next step, published together so capture cannot mix versions. -#[derive(Debug)] -pub(crate) struct StepInputs { - pub(crate) settings: Arc, - /// Selection is fixed within this version; attachment startup may still finish. - pub(crate) environments: TurnEnvironmentSnapshot, -} - /// Request-scoped state that may change between model sampling requests. pub(crate) struct StepContext { pub(crate) turn: Arc, diff --git a/codex-rs/core/src/session/step_settings.rs b/codex-rs/core/src/session/step_settings.rs index c915037a17..cd00b6d216 100644 --- a/codex-rs/core/src/session/step_settings.rs +++ b/codex-rs/core/src/session/step_settings.rs @@ -232,8 +232,7 @@ pub(crate) struct StepSettingsUpdate { } /// Constraints used when applying and validating a candidate settings version. -/// Future settings use the proposed environment; active settings use the -/// environment already admitted for that execution. +/// Each caller supplies the permissions of the target environment selection. pub(crate) struct StepSettingsConstraints<'a> { pub(crate) requirements: &'a ConfigRequirements, pub(crate) guardian_approval_enabled: bool, diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index 0091a77a07..1a14903b62 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -2,8 +2,6 @@ mod notification_tests; use super::mcp_refresh::McpRefresh; -use super::step_context::StepInputs; - #[path = "turn_start_mcp_tests.rs"] mod turn_start_mcp_tests; use super::step_settings::ResolvedStepSettings; @@ -254,10 +252,7 @@ pub(crate) fn update_turn_settings_for_test( update(&mut settings); let settings = Arc::new(settings); turn.initial_settings = Arc::clone(&settings); - turn.next_step_input.store(Arc::new(StepInputs { - settings, - environments: turn.initial_environments.clone(), - })); + turn.next_step_settings.store(settings); } impl StepContext { @@ -4363,23 +4358,20 @@ async fn turn_context_with_model_updates_model_fields() { }); Arc::make_mut(&mut turn_context.config).service_tier = turn_context.initial_settings.service_tier.clone(); - let captured = turn_context.next_step_input.load_full(); - let mut current_selection = captured.settings.selected().clone(); + let captured = turn_context.next_step_settings.load_full(); + let mut current_selection = captured.selected().clone(); current_selection.reasoning_summary = Some(ReasoningSummaryConfig::None); current_selection.service_tier = None; current_selection .collaboration_mode .settings .reasoning_effort = Some(ReasoningEffortConfig::High); - let current = Arc::new(StepInputs { - settings: Arc::new(ResolvedStepSettings::new( - Arc::new(current_selection), - Arc::clone(turn_context.model_info()), - /*fast_mode_enabled*/ true, - )), - environments: captured.environments.clone(), - }); - turn_context.next_step_input.store(Arc::clone(¤t)); + let current = Arc::new(ResolvedStepSettings::new( + Arc::new(current_selection), + Arc::clone(turn_context.model_info()), + /*fast_mode_enabled*/ true, + )); + turn_context.next_step_settings.store(Arc::clone(¤t)); let updated = turn_context .with_model("gpt-5.5".to_string(), &session.services.models_manager) .await; @@ -4414,22 +4406,19 @@ async fn turn_context_with_model_updates_model_fields() { Some(ServiceTier::Fast.request_value()) ), ); - assert!(Arc::ptr_eq( - &captured.settings, - &turn_context.initial_settings - )); + assert!(Arc::ptr_eq(&captured, &turn_context.initial_settings)); assert!(Arc::ptr_eq( ¤t, - &turn_context.next_step_input.load_full() + &turn_context.next_step_settings.load_full() )); - assert!(!Arc::ptr_eq(&captured.settings, &updated.initial_settings)); + assert!(!Arc::ptr_eq(&captured, &updated.initial_settings)); assert!(Arc::ptr_eq( &updated.initial_settings, - &updated.next_step_input.load().settings + &updated.next_step_settings.load_full() )); assert!(!Arc::ptr_eq( - &updated.next_step_input.load_full(), - &turn_context.next_step_input.load_full() + &updated.next_step_settings.load_full(), + &turn_context.next_step_settings.load_full() )); assert_eq!(updated.config.model.as_deref(), Some("gpt-5.5")); assert_eq!(updated.collaboration_mode().model(), "gpt-5.5"); @@ -10786,14 +10775,13 @@ async fn build_initial_context_uses_retained_step_after_model_change() { .as_mut() .unwrap() .instructions_template = Some("B instructions".to_string()); - turn_context.next_step_input.store(Arc::new(StepInputs { - settings: Arc::new(ResolvedStepSettings::new( + turn_context + .next_step_settings + .store(Arc::new(ResolvedStepSettings::new( Arc::new(selected_b), Arc::new(model_b), /*fast_mode_enabled*/ false, - )), - environments: step_a.environments.clone(), - })); + ))); let step_b = session .capture_step_context(Arc::clone(&turn_context), &CancellationToken::new()) .await @@ -11862,10 +11850,7 @@ async fn task_finish_emits_turn_item_lifecycle_for_leftover_pending_user_input() let mut current = tc.initial_settings.as_ref().clone(); Arc::make_mut(&mut current.model_info).supports_image_detail_original = true; - tc.next_step_input.store(Arc::new(StepInputs { - settings: Arc::new(current), - environments: tc.next_step_input.load().environments.clone(), - })); + tc.next_step_settings.store(Arc::new(current)); sess.on_task_finished(Arc::clone(&tc), /*task_result*/ Ok(None)) .await; diff --git a/codex-rs/core/src/session/tests/guardian_tests.rs b/codex-rs/core/src/session/tests/guardian_tests.rs index a2888d782b..103f45fbca 100644 --- a/codex-rs/core/src/session/tests/guardian_tests.rs +++ b/codex-rs/core/src/session/tests/guardian_tests.rs @@ -116,7 +116,7 @@ async fn activate_turn_with_new_review_authority(session: &Arc) -> Arc< ) .await; - let (active_turn, _, _) = session + let (active_turn, _, _, _) = session .active_turn_context_and_strict_auto_review() .await .expect("next turn should have active review authority"); @@ -798,8 +798,7 @@ async fn network_approval_uses_published_task_authority_within_same_turn( .task .as_ref() .expect("active task"); - let current = task.turn_context.next_step_input.load_full(); - let mut settings = Arc::clone(¤t.settings); + let mut settings = task.turn_context.next_step_settings.load_full(); update_selected_settings_for_test(Arc::make_mut(&mut settings), |selected| { selected .approval_policy @@ -807,12 +806,7 @@ async fn network_approval_uses_published_task_authority_within_same_turn( .expect("update policy"); selected.approvals_reviewer = ApprovalsReviewer::User; }); - task.turn_context - .next_step_input - .store(Arc::new(StepInputs { - settings, - environments: current.environments.clone(), - })); + task.turn_context.next_step_settings.store(settings); } let decision = session .services diff --git a/codex-rs/core/src/session/turn_context.rs b/codex-rs/core/src/session/turn_context.rs index 165be62947..76befa42c6 100644 --- a/codex-rs/core/src/session/turn_context.rs +++ b/codex-rs/core/src/session/turn_context.rs @@ -1,4 +1,3 @@ -use super::step_context::StepInputs; use super::step_settings::ResolvedStepSettings; use super::token_budget::has_explicit_settings; use super::token_budget::resolve_token_budget; @@ -321,8 +320,8 @@ pub struct TurnContext { /// Thread-owned plugin selection captured when this turn was admitted. pub(crate) disabled_plugin_ids: Vec, pub(super) active_host_plugin_identities: Option>, - /// Inputs for the next step; request consumers use their captured StepContext. - pub(super) next_step_input: ArcSwap, + /// Settings for the next step; environments are owned by `ThreadEnvironments`. + pub(super) next_step_settings: ArcSwap, /// Turn-wide telemetry; model-attributed step work should use `StepContext::session_telemetry`. pub(crate) session_telemetry: SessionTelemetry, pub(crate) provider: SharedModelProvider, @@ -331,7 +330,7 @@ pub struct TurnContext { pub(crate) parent_thread_id: Option, pub(crate) originator: String, /// Initial selection retained for legacy turn consumers. Step work uses StepContext. - // TODO(sayan): Migrate all remaining consumers to next_step_input's environments. + // TODO(sayan): Migrate remaining step consumers to their StepContext's environments. pub(crate) initial_environments: TurnEnvironmentSnapshot, /// The session's absolute working directory. All relative paths provided /// by the model as well as sandbox policies are resolved against this path @@ -381,7 +380,7 @@ enum TurnContextBuildMode { impl TurnContext { /// Captures current model metadata without preparing a step. pub(crate) fn capture_current_model_info(&self) -> Arc { - Arc::clone(&self.next_step_input.load().settings.model_info) + Arc::clone(&self.next_step_settings.load().model_info) } /// Legacy: returns the frozen initial-turn model metadata. @@ -692,10 +691,7 @@ impl TurnContext { initial_settings: Arc::clone(&step_settings), disabled_plugin_ids: self.disabled_plugin_ids.clone(), active_host_plugin_identities: self.active_host_plugin_identities.clone(), - next_step_input: ArcSwap::from_pointee(StepInputs { - settings: step_settings, - environments: self.initial_environments.clone(), - }), + next_step_settings: ArcSwap::from(step_settings), session_telemetry, provider: self.provider.clone(), session_source: self.session_source.clone(), @@ -1011,10 +1007,7 @@ impl Session { initial_settings: Arc::clone(&step_settings), disabled_plugin_ids: session_configuration.disabled_plugin_ids.clone(), active_host_plugin_identities: None, - next_step_input: ArcSwap::from_pointee(StepInputs { - settings: step_settings, - environments: environments.clone(), - }), + next_step_settings: ArcSwap::from(step_settings), session_telemetry: session_telemetry_for_context, provider, session_source, diff --git a/codex-rs/core/src/thread_manager.rs b/codex-rs/core/src/thread_manager.rs index 521388dbbe..df153259e6 100644 --- a/codex-rs/core/src/thread_manager.rs +++ b/codex-rs/core/src/thread_manager.rs @@ -109,8 +109,6 @@ use tracing::instrument; use tracing::warn; const THREAD_CREATED_CHANNEL_CAPACITY: usize = 1024; -// Reject pathological selected cwd values at the environment-selection boundary. -const MAX_TURN_ENVIRONMENT_CWD_BYTES: usize = 8 * 1024; /// Test-only override for enabling thread-manager behaviors used by integration /// tests. @@ -848,37 +846,6 @@ impl ThreadManager { ) } - pub fn validate_environment_selections( - &self, - environments: &[TurnEnvironmentSelection], - ) -> CodexResult<()> { - let mut environment_ids = HashSet::with_capacity(environments.len()); - for environment in environments { - if environment.cwd.inferred_native_path_string().len() > MAX_TURN_ENVIRONMENT_CWD_BYTES - { - return Err(CodexErr::InvalidRequest( - "turn environment working directory exceeds the maximum size".to_string(), - )); - } - if !environment_ids.insert(environment.environment_id.as_str()) { - return Err(CodexErr::InvalidRequest(format!( - "duplicate turn environment id `{}`", - environment.environment_id - ))); - } - self.state - .environment_manager - .get_environment(&environment.environment_id) - .ok_or_else(|| { - CodexErr::InvalidRequest(format!( - "unknown turn environment id `{}`", - environment.environment_id - )) - })?; - } - Ok(()) - } - pub(crate) fn git_root_discovery(&self) -> Arc { Arc::clone(&self.state.git_root_discovery) } diff --git a/codex-rs/core/src/tools/network_approval.rs b/codex-rs/core/src/tools/network_approval.rs index de37326e0c..78924010a2 100644 --- a/codex-rs/core/src/tools/network_approval.rs +++ b/codex-rs/core/src/tools/network_approval.rs @@ -648,13 +648,7 @@ impl NetworkApprovalService { { return NetworkDecision::deny(REASON_NOT_ALLOWED); } - let active_turn = session - .active_turn_context_and_strict_auto_review() - .await - .map(|(turn, inputs, strict)| { - let environments = inputs.environments.refresh_readiness(); - (turn, inputs, environments, strict) - }); + let active_turn = session.active_turn_context_and_strict_auto_review().await; let Some(environment_id) = active_environment_id.or_else(|| { active_turn .as_ref() @@ -685,8 +679,7 @@ impl NetworkApprovalService { format!("Network access to \"{target}\" was blocked by policy."); let prompt_reason = format!("{} is not in the allowed_domains", request.host); - let Some((turn_context, step_inputs, active_environments, strict_auto_review)) = - active_turn + let Some((turn_context, settings, active_environments, strict_auto_review)) = active_turn else { if let Some(owner_call) = owner_call.as_ref() { self.record_call_outcome(&owner_call.registration_id, policy_denial_message); @@ -739,7 +732,7 @@ impl NetworkApprovalService { } let review_context = GuardianReviewContext::from_resolved_settings( Arc::clone(&turn_context), - &step_inputs.settings, + &settings, review_environments, ); if !allows_network_approval_flow(review_context.approval_policy) { diff --git a/codex-rs/core/src/tools/orchestrator.rs b/codex-rs/core/src/tools/orchestrator.rs index f17241a94c..87092d755f 100644 --- a/codex-rs/core/src/tools/orchestrator.rs +++ b/codex-rs/core/src/tools/orchestrator.rs @@ -133,11 +133,7 @@ impl ToolOrchestrator { let otel = turn_ctx.session_telemetry.clone(); let otel_tn = flat_tool_name(&tool_ctx.tool_name).into_owned(); let otel_ci = &tool_ctx.call_id; - let strict_auto_review = tool_ctx - .session - .active_turn_context_and_strict_auto_review() - .await - .is_some_and(|(_, _, strict_auto_review)| strict_auto_review); + let strict_auto_review = tool_ctx.session.strict_auto_review_enabled().await; // 1) Approval let mut already_approved = false; diff --git a/codex-rs/core/src/tools/runtimes/zsh_fork/unix_escalation.rs b/codex-rs/core/src/tools/runtimes/zsh_fork/unix_escalation.rs index 1849e8c2d2..f0aaf9d5f4 100644 --- a/codex-rs/core/src/tools/runtimes/zsh_fork/unix_escalation.rs +++ b/codex-rs/core/src/tools/runtimes/zsh_fork/unix_escalation.rs @@ -282,7 +282,7 @@ impl CoreShellActionProvider { }; match stopwatch .pause_for(async { - let (turn_context, step_inputs, strict_auto_review) = self + let (turn_context, settings, _, strict_auto_review) = self .session .active_turn_context_and_strict_auto_review() .await @@ -295,7 +295,7 @@ impl CoreShellActionProvider { let approval_ctx = ApprovalContext { review_context: GuardianReviewContext::from_resolved_settings( Arc::clone(&turn_context), - &step_inputs.settings, + &settings, self.review_context.environments(), ), // The running process can outlive its launching tool or code-mode cell. diff --git a/codex-rs/core/src/unified_exec/process_manager.rs b/codex-rs/core/src/unified_exec/process_manager.rs index 105783518d..a21dec6485 100644 --- a/codex-rs/core/src/unified_exec/process_manager.rs +++ b/codex-rs/core/src/unified_exec/process_manager.rs @@ -897,11 +897,7 @@ impl UnifiedExecProcessManager { }; let _interaction_guard = locked_process.interaction_lock().lock_owned().await; // A queued write must observe strict review enabled while it was waiting. - let strict_auto_review = context - .session - .active_turn_context_and_strict_auto_review() - .await - .is_some_and(|(_, _, strict)| strict); + let strict_auto_review = context.session.strict_auto_review_enabled().await; let approval = { let store = self.process_store.lock().await; let entry = store diff --git a/codex-rs/core/tests/common/test_codex.rs b/codex-rs/core/tests/common/test_codex.rs index 1564fe4252..b475bb35e0 100644 --- a/codex-rs/core/tests/common/test_codex.rs +++ b/codex-rs/core/tests/common/test_codex.rs @@ -26,6 +26,7 @@ use codex_core::resolve_installation_id; use codex_core::shell::Shell; use codex_core::shell::get_shell_by_model_provided_path; use codex_core::thread_store_from_config; +use codex_core::windows_sandbox::WindowsSandboxLevelExt; use codex_exec_server::CreateDirectoryOptions; use codex_exec_server::ExecutorFileSystem; use codex_exec_server::RemoveOptions; @@ -45,14 +46,17 @@ use codex_protocol::config_types::CollaborationMode; use codex_protocol::config_types::ModeKind; use codex_protocol::config_types::ReasoningSummary; use codex_protocol::config_types::Settings; +use codex_protocol::config_types::WindowsSandboxLevel; use codex_protocol::mcp::ClientMcpExtensions; use codex_protocol::mcp::OPENAI_FORM_EXTENSION_ID; use codex_protocol::models::PermissionProfile; +use codex_protocol::models::PermissionProfileSnapshot; use codex_protocol::openai_models::ModelInfo; use codex_protocol::openai_models::ModelsResponse; use codex_protocol::openai_models::TruncationPolicyConfig; use codex_protocol::openai_models::WebSearchToolType; use codex_protocol::protocol::AskForApproval; +use codex_protocol::protocol::EnvironmentConfig; use codex_protocol::protocol::EnvironmentConfigState; use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::RealtimeConversationVersion as RealtimeWsVersion; @@ -132,6 +136,36 @@ pub fn local(cwd: AbsolutePathBuf) -> TurnEnvironmentSelection { } } +/// Builds explicit environment configuration with the test thread's permissions and shell settings. +pub fn environment_config_for_selection( + config: &Config, + selection: &TurnEnvironmentSelection, +) -> EnvironmentConfig { + let permissions = &config.permissions; + let profile = permissions.permission_profile().clone(); + let permission_profile = match permissions.active_permission_profile() { + Some(active) => PermissionProfileSnapshot::active_with_profile_workspace_roots( + profile, + active, + permissions.profile_workspace_roots().to_vec(), + ), + None => PermissionProfileSnapshot::legacy(profile), + }; + EnvironmentConfig { + allow_login_shell: permissions.allow_login_shell, + workspace_roots: selection.workspace_roots.clone(), + permission_profile, + shell_environment_policy: permissions.shell_environment_policy.clone(), + windows_sandbox_level: WindowsSandboxLevel::from_config(config), + windows_sandbox_type: permissions.windows_sandbox_type, + use_legacy_landlock: config.features.use_legacy_landlock(), + exec_policy: None, + mcp_policy: None, + network_policy: None, + selected_capability_roots: Vec::new(), + } +} + /// Converts the host-shaped /C:/... cwd projection used by Wine tests back /// into the selected executor's Windows URI. pub fn executor_path_uri(path: impl AsRef) -> Result { diff --git a/codex-rs/core/tests/suite/remote_env.rs b/codex-rs/core/tests/suite/remote_env.rs index 4b3a952dda..a3ba5a2a11 100644 --- a/codex-rs/core/tests/suite/remote_env.rs +++ b/codex-rs/core/tests/suite/remote_env.rs @@ -70,6 +70,8 @@ use codex_protocol::protocol::SandboxPolicy; use codex_protocol::protocol::ThreadSettingsOverrides; use codex_protocol::protocol::TurnEnvironmentSelection; use codex_protocol::protocol::TurnEnvironmentSelections; +use codex_protocol::protocol::TurnSettingsUpdate; +use codex_protocol::protocol::TurnSettingsUpdateOutcome; use codex_protocol::request_permissions::PermissionGrantScope; use codex_protocol::request_permissions::RequestPermissionProfile; use codex_protocol::request_permissions::RequestPermissionsResponse; @@ -101,6 +103,7 @@ use core_test_support::startup::expect_startup; use core_test_support::submit_thread_settings; use core_test_support::test_codex::TestCodex; use core_test_support::test_codex::TestCodexBuilder; +use core_test_support::test_codex::environment_config_for_selection; use core_test_support::test_codex::local; use core_test_support::test_codex::test_codex; use core_test_support::test_codex::test_env; @@ -907,7 +910,6 @@ async fn settings_update_does_not_retarget_active_turn_environment() -> Result<( let test = builder.build(&server).await?; let initial_cwd = test.config.cwd.clone(); let initial_environments = test.codex.environment_selections().await; - assert_eq!(test.codex.active_turn_environment_selections().await, None); let next_workspace = TempDir::new()?; let next_cwd = next_workspace.path().abs(); let next_environments = @@ -925,11 +927,6 @@ async fn settings_update_does_not_retarget_active_turn_environment() -> Result<( }) .await; - assert_eq!( - test.codex.active_turn_environment_selections().await, - Some(initial_environments.clone()) - ); - let preview = test .codex .preview_thread_settings_overrides(CodexThreadSettingsOverrides { @@ -960,10 +957,6 @@ async fn settings_update_does_not_retarget_active_turn_environment() -> Result<( test.codex.environment_selections().await, next_environments.environments ); - assert_eq!( - test.codex.active_turn_environment_selections().await, - Some(initial_environments) - ); let snapshot = test.codex.config_snapshot().await; assert_eq!( snapshot.environment_selections(), @@ -988,7 +981,6 @@ async fn settings_update_does_not_retarget_active_turn_environment() -> Result<( matches!(event, EventMsg::TurnComplete(_)) }) .await; - assert_eq!(test.codex.active_turn_environment_selections().await, None); test.submit_turn("start the next turn").await?; let request_texts = response_mock @@ -1011,6 +1003,150 @@ async fn settings_update_does_not_retarget_active_turn_environment() -> Result<( Ok(()) } +#[test_case(false; "configuration supplied with the selection")] +#[test_case(true; "configuration arrives after the selection")] +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn active_environment_update_waits_for_a_configured_executor_to_connect( + configuration_arrives_later: bool, +) -> Result<()> { + let listener = TcpListener::bind("127.0.0.1:0").await?; + let exec_server_url = format!("ws://{}", listener.local_addr()?); + let (attach, connection) = tokio::sync::oneshot::channel(); + let (shutdown, stop) = tokio::sync::oneshot::channel(); + let executor = tokio::spawn(serve_environment_with_agents_md( + listener, "", connection, stop, + )); + let server = start_mock_server().await; + let responses = mount_sse_sequence( + &server, + vec![ + sse(vec![ + ev_function_call( + "pause-for-switch", + "request_user_input", + &json!({"questions": [{ + "id": "continue", "header": "Continue", "question": "Continue?", + "options": [ + {"label": "Yes", "description": "Continue the test."}, + {"label": "No", "description": "Stop the test."} + ] + }]}) + .to_string(), + ), + ev_completed("first"), + ]), + sse(vec![ev_completed("second")]), + ], + ) + .await; + let mut builder = test_codex().with_config(|config| { + config.project_doc_max_bytes = 0; + assert!(config.features.disable(Feature::DeferredExecutor).is_ok()); + assert!( + config + .features + .enable(Feature::DefaultModeRequestUserInput) + .is_ok() + ); + }); + let test = builder.build_with_auto_env(&server).await?; + let remote_environment_id = "connecting"; + test.thread_manager + .environment_manager() + .upsert_environment( + remote_environment_id.to_string(), + exec_server_url, + /*connect_timeout*/ None, + )?; + let remote_cwd = test + .executor_environment() + .selection() + .cwd + .join("remote-workspace")?; + let mut remote_selection = TurnEnvironmentSelection { + environment_id: remote_environment_id.to_string(), + cwd: remote_cwd.clone(), + workspace_roots: vec![remote_cwd.clone()], + config: EnvironmentConfigState::Pending, + }; + let owner_config = environment_config_for_selection(&test.config, &remote_selection); + if !configuration_arrives_later { + remote_selection.config = EnvironmentConfigState::Ready(owner_config.clone()); + } + test.codex + .start_or_steer_turn(TurnInputRequest::user_input(vec![UserInput::Text { + text: "pause before switching".to_string(), + text_elements: Vec::new(), + }])) + .await?; + let request = wait_for_event_match(&test.codex, |event| match event { + EventMsg::RequestUserInput(request) => Some(request.clone()), + _ => None, + }) + .await; + + let (reply, outcome) = tokio::sync::oneshot::channel(); + test.codex + .submit(Op::TurnSettings { + turn_id: request.turn_id.clone(), + update: TurnSettingsUpdate { + environments: Some(vec![remote_selection.clone()]), + ..Default::default() + }, + reply, + }) + .await?; + assert_eq!(outcome.await?, TurnSettingsUpdateOutcome::Applied); + if configuration_arrives_later { + test.codex + .environment_ready(&remote_selection, owner_config) + .await?; + } + test.codex + .submit(Op::UserInputAnswer { + id: request.turn_id, + response: RequestUserInputResponse { + answers: HashMap::from([( + "continue".to_string(), + RequestUserInputAnswer { + answers: vec!["Yes".to_string()], + }, + )]), + }, + }) + .await?; + + assert!( + timeout( + Duration::from_secs(/*secs*/ 1), + wait_for_response_request_count(&responses, /*expected_count*/ 2) + ) + .await + .is_err(), + "the follow-up must not run before the configured executor connects" + ); + attach.send(()).expect("release executor startup"); + wait_for_event(&test.codex, |event| { + matches!(event, EventMsg::TurnComplete(_)) + }) + .await; + let requests = responses.requests(); + let context = requests[1] + .message_input_texts("user") + .into_iter() + .rfind(|text| text.contains("")) + .context("remote environment should be model visible")?; + assert!(context.contains(&format!( + "{}", + remote_cwd.inferred_native_path_string() + ))); + assert!(!context.contains("starting")); + + let _ = shutdown.send(()); + executor.await?; + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn deferred_executor_promotes_primary_environment_when_startup_completes() -> Result<()> { let listener = TcpListener::bind("127.0.0.1:0").await?; @@ -1116,19 +1252,6 @@ async fn deferred_executor_promotes_primary_environment_when_startup_completes() }) .await; - let active_environments = test - .codex - .active_turn_environment_selections() - .await - .context("active turn environments")?; - assert_eq!( - active_environments - .iter() - .map(|selection| selection.environment_id.as_str()) - .collect::>(), - vec![REMOTE_ENVIRONMENT_ID, "local"] - ); - let requests = response_mock.requests(); let initial_context = requests[1] .message_input_texts("user") @@ -1188,10 +1311,6 @@ async fn deferred_executor_promotes_primary_environment_when_startup_completes() } }); core_test_support::wait_for_mcp_server(&test.codex, "deferred").await?; - assert_eq!( - test.codex.active_turn_environment_selections().await, - Some(active_environments) - ); test.codex .submit(Op::UserInputAnswer { id: request.turn_id, @@ -1210,8 +1329,6 @@ async fn deferred_executor_promotes_primary_environment_when_startup_completes() }) .await; - assert_eq!(test.codex.active_turn_environment_selections().await, None); - let requests = response_mock.requests(); assert!( requests[2] @@ -2134,6 +2251,122 @@ async fn future_pending_environment_can_finish_without_retargeting_the_active_tu Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn active_environment_update_wakes_the_old_wait_with_the_new_selection() -> Result<()> { + let server = start_mock_server().await; + let mut builder = test_codex_with_wait_for_environment().with_config(|config| { + assert!(config.features.enable(Feature::DeferredExecutor).is_ok()); + assert!(config.features.disable(Feature::StepModelSwitching).is_ok()); + }); + let test = builder.build_with_auto_env(&server).await?; + let mut original = test.codex.environment_selections().await.remove(0); + original.config = EnvironmentConfigState::Pending; + let switched_cwd = original.cwd.join("active-environment")?; + test.fs() + .create_directory( + &switched_cwd, + CreateDirectoryOptions { + recursive: false, + follow_symlinks: true, + }, + /*sandbox*/ None, + ) + .await?; + let switched = TurnEnvironmentSelection { + cwd: switched_cwd.clone(), + workspace_roots: vec![switched_cwd], + ..original.clone() + }; + let thread = test + .thread_manager + .start_thread(StartThreadOptions { + environments: Some(vec![original.clone()]), + ..StartThreadOptions::new(test.config.clone()) + }) + .await? + .thread; + let wait = |id: &str| { + sse(vec![ + ev_function_call( + id, + "wait_for_environment", + &json!({ "environment_id": original.environment_id }).to_string(), + ), + ev_completed(id), + ]) + }; + let responses = mount_sse_sequence( + &server, + vec![ + wait("wait-for-original"), + wait("wait-for-switched"), + sse(vec![ev_completed("active-done")]), + ], + ) + .await; + let TurnInputSubmission::Started { turn_id } = thread + .start_or_steer_turn(TurnInputRequest::user_input(vec![UserInput::Text { + text: "use the selected environment".to_string(), + text_elements: Vec::new(), + }])) + .await? + else { + anyhow::bail!("expected to start a turn"); + }; + wait_for_response_request_count(&responses, /*expected_count*/ 1).await; + + let (reply, outcome) = tokio::sync::oneshot::channel(); + thread + .submit(Op::TurnSettings { + turn_id, + update: TurnSettingsUpdate { + environments: Some(vec![switched.clone()]), + ..Default::default() + }, + reply, + }) + .await?; + assert_eq!(outcome.await?, TurnSettingsUpdateOutcome::Applied); + wait_for_response_request_count(&responses, /*expected_count*/ 2).await; + thread + .environment_ready( + &switched, + environment_config_for_selection(&test.config, &switched), + ) + .await?; + wait_for_event(&thread, |event| matches!(event, EventMsg::TurnComplete(_))).await; + + let requests = responses.requests(); + let wait_output = |request: &ResponsesRequest, id| -> Result { + let (output, _) = request + .function_call_output_content_and_success(id) + .context("wait output should be model visible")?; + output.context("wait output should contain text") + }; + assert!(wait_output(&requests[1], "wait-for-original")?.contains("configuration was canceled")); + assert_eq!( + serde_json::from_str::(&wait_output(&requests[2], "wait-for-switched")?)?, + json!({ "environment_id": original.environment_id, "status": "ready" }) + ); + for (request, selection, is_starting) in [ + (&requests[0], &original, true), + (&requests[1], &switched, true), + (&requests[2], &switched, false), + ] { + let context = request + .message_input_texts("user") + .into_iter() + .rfind(|text| text.contains("")) + .context("environment context should be model visible")?; + assert!(context.contains(&format!( + "{}", + selection.cwd.inferred_native_path_string() + ))); + assert_eq!(context.contains("starting"), is_starting); + } + Ok(()) +} + #[test_case(true; "uses refreshed executor root")] #[test_case(false; "preserves persisted root when executor reports none")] #[tokio::test(flavor = "multi_thread", worker_threads = 4)] diff --git a/codex-rs/core/tests/suite/scenarios.rs b/codex-rs/core/tests/suite/scenarios.rs index e40a21006e..92447fa9ea 100644 --- a/codex-rs/core/tests/suite/scenarios.rs +++ b/codex-rs/core/tests/suite/scenarios.rs @@ -19,6 +19,7 @@ use codex_context_fragments::ContextualUserFragment; use codex_core::StartThreadOptions; use codex_core::TurnInputRequest; use codex_core::config::Config; +use codex_exec_server::CreateDirectoryOptions; use codex_exec_server::EnvironmentManager; use codex_exec_server::LOCAL_ENVIRONMENT_ID; use codex_extension_api::ExtensionDataInit; @@ -35,8 +36,13 @@ use codex_protocol::models::ImageReference; use codex_protocol::openai_models::CodeModeToolMessages; use codex_protocol::openai_models::ToolMessage; use codex_protocol::openai_models::ToolMode; +use codex_protocol::protocol::EnvironmentConfigState; use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::Op; +use codex_protocol::protocol::TurnSettingsUpdate; +use codex_protocol::protocol::TurnSettingsUpdateOutcome; +use codex_protocol::request_user_input::RequestUserInputAnswer; +use codex_protocol::request_user_input::RequestUserInputResponse; use codex_protocol::user_input::UserInput; use codex_skills_extension::ExecutorSkillProvider; use codex_skills_extension::SkillProviders; @@ -60,6 +66,7 @@ use core_test_support::skip_if_wine_exec; use core_test_support::stdio_server_bin; use core_test_support::streaming_sse::StreamingSseChunk; use core_test_support::streaming_sse::start_streaming_sse_server; +use core_test_support::test_codex::environment_config_for_selection; use core_test_support::test_codex::executor_path_uri; use core_test_support::test_codex::test_codex; use core_test_support::wait_for_event; @@ -326,6 +333,127 @@ async fn astra_asks_an_async_question_and_receives_the_answer_while_working() -> Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn astra_switches_environments_for_the_rest_of_the_active_turn() -> Result<()> { + skip_if_no_network!(Ok(())); + + let server = start_mock_server().await; + let mock = mount_sse_sequence( + &server, + vec![ + sse(vec![ + ev_function_call_with_namespace( + "pause-for-switch", + "functions", + "request_user_input", + &json!({"questions": [{ + "id": "continue", "header": "Continue", "question": "Continue in the other workspace?", + "options": [ + {"label": "Yes", "description": "Continue working."}, + {"label": "No", "description": "Stop working."} + ] + }]}).to_string(), + ), + ev_completed("first-step"), + ]), + sse(vec![ + ev_assistant_message("active-turn", "I continued in the other workspace."), + ev_completed("second-step"), + ]), + sse(vec![ + ev_assistant_message("next-turn", "I am back in the original workspace."), + ev_completed("next-turn"), + ]), + ], + ) + .await; + let test = test_codex() + .with_model("gpt-6-astra") + .with_auth(CodexAuth::create_dummy_chatgpt_auth_for_testing()) + .with_config(|config| { + configure_scenario_catalog(config); + config.workspace_roots = vec![config.cwd.clone()]; + config + .features + .enable(Feature::DefaultModeRequestUserInput) + .expect("enable request_user_input"); + }) + .build_with_auto_env(&server) + .await?; + let original = test.executor_environment().selection().clone(); + let mut other = original.clone(); + other.cwd = original.cwd.join("other-workspace")?; + other.workspace_roots = vec![other.cwd.clone()]; + other.config = + EnvironmentConfigState::Ready(environment_config_for_selection(&test.config, &other)); + test.fs() + .create_directory( + &other.cwd, + CreateDirectoryOptions { + recursive: false, + follow_symlinks: true, + }, + /*sandbox*/ None, + ) + .await?; + + test.codex + .start_or_steer_turn(TurnInputRequest::user_input(vec![text( + "Ask before continuing in the other workspace.", + )])) + .await?; + let request = wait_for_event_match(&test.codex, |event| match event { + EventMsg::RequestUserInput(request) => Some(request.clone()), + _ => None, + }) + .await; + let (reply, outcome) = oneshot::channel(); + test.codex + .submit(Op::TurnSettings { + turn_id: request.turn_id.clone(), + update: TurnSettingsUpdate { + environments: Some(vec![other.clone()]), + ..Default::default() + }, + reply, + }) + .await?; + assert_eq!( + tokio::time::timeout(Duration::from_secs(/*secs*/ 10), outcome).await??, + TurnSettingsUpdateOutcome::Applied + ); + test.codex + .submit(Op::UserInputAnswer { + id: request.turn_id, + response: RequestUserInputResponse { + answers: HashMap::from([( + "continue".to_string(), + RequestUserInputAnswer { + answers: vec!["Yes".to_string()], + }, + )]), + }, + }) + .await?; + wait_for_event(&test.codex, |event| { + matches!(event, EventMsg::TurnComplete(_)) + }) + .await; + test.submit_text_turn("Which workspace are you using now?") + .await?; + + let requests = mock.requests(); + insta::assert_snapshot!( + "astra_active_environment_selection", + context_snapshot::format_request_history_snapshot( + "Astra starts in one workspace, switches for the next step of the active turn, and returns to the original workspace on the next turn.", + &requests, + &ContextSnapshotOptions::default(), + ) + ); + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn astra_kickoff_with_skills_plugins_and_remote_compaction() -> Result<()> { skip_if_no_network!(Ok(())); diff --git a/codex-rs/core/tests/suite/snapshots/all__suite__scenarios__astra_active_environment_selection.snap b/codex-rs/core/tests/suite/snapshots/all__suite__scenarios__astra_active_environment_selection.snap new file mode 100644 index 0000000000..a2677a2da0 --- /dev/null +++ b/codex-rs/core/tests/suite/snapshots/all__suite__scenarios__astra_active_environment_selection.snap @@ -0,0 +1,129 @@ +--- +source: core/tests/suite/scenarios.rs +expression: "context_snapshot::format_request_history_snapshot(\"Astra starts in one workspace, switches for the next step of the active turn, and returns to the original workspace on the next turn.\",\n&requests, &ContextSnapshotOptions::default(),)" +--- +Scenario: Astra starts in one workspace, switches for the next step of the active turn, and returns to the original workspace on the next turn. + +## Window 1 +-- request 1 (turn) -- +00:additional_tools/developer (3; hash=b1e101cdca5088a0): + - namespace/functions; hash=dead25325ac79c67 + - custom/exec + - function/wait + - function/request_user_input + - function/request_user_input_async + - namespace/clock: Tools for reading and waiting on time.; hash=f402e8c5e9b5e317 + - function/sleep + - namespace/collaboration: Tools for spawning and managing sub-agents.; hash=fd994b11525da73c + - function/followup_task + - function/interrupt_agent + - function/list_agents + - function/send_message + - function/spawn_agent + - function/wait_agent +01:message/developer: + You are Codex, an agent based on GPT-6. You and the user share one workspace, and your job is to collabo...them until their intended goal is completely handled. [hash=efa50f074efa5803] + + # When to ask the user for permission + + Use your best judgement given task context for when you really need user permission, like a competent co...ork without ending the turn to clarify with the user. [hash=ca7367e658b51388] + + User authorization and preferences persist across turns. Do not request permission again when the user h... any guidelines provided in skills or external files. [hash=911f91c92c72882c] + + + + - Skill naming: If a plugin contributes skills, those skill entries are prefixed with plugin_name: in the Skills list. + - MCP naming: Plugin-provided MCP tools keep standard MCP identifiers such as mcp__server__tool; use tool provenance to tell which plugin they come from. + - Trigger rules: If the user explicitly names a plugin, prefer capabilities associated with that plugin for that turn. + - Relationship to capabilities: Plugins are not invoked directly. Use their underlying skills, MCP tools, and app tools to help solve the task. + - Relevance: Determine what a plugin can help with from explicit user mention or from the plugin-associa..., MCP tools, and apps exposed elsewhere in this turn. [hash=620459a4e7b590f0] + - Missing/blocked: If the user requests a plugin that does not have relevant callable capabilities for t..., say so briefly and continue with the best fallback. [hash=746b0469a8f6cb3e] + +02:message/developer[2]: + [01] + Filesystem sandboxing defines which files can be read or written. `sandbox_mode` is `read-only`: The san... permits reading files. Network access is restricted. [hash=be7da4425f89053c] + # Escalation Requests + + Commands are run outside the sandbox if they are approved by the user, or match an existing rule that al...hell control operators, including but not limited to: [hash=d24ebf5fa4e4b9bb] + + - Pipes: | + - Logical operators: &&, || + + NEVER provide a prefix_rule if your command uses a heredoc or herestring. + + ### Examples + Good examples of prefixes: + - ["npm", "run", "dev"] + - ["gh", "pr", "check"] + - ["cargo", "test"] + + [02] # Collaboration Mode: Default + + You are now in Default mode. Any previous instructions for other modes (e.g. Plan mode) are no longer active. + + Your active mode changes only when new developer instructions with a different `...<...by themselves. Known mode names are Default and Plan. [hash=e75f6b12fc52d812] + + ## request_user_input availability + + Use the `request_user_input` tool only when it is listed in the available tools for this turn. + + Use the `request_user_input` tool only for optional questions where the answer would materially improve the quality of the work. + + If `request_user_input` returns no answers, continue with best judgment instead of asking again or treating the turn as blocked. + + Never use the `request_user_input` tool for permission requests or permission-related escalations. + +03:message/developer: + You are `/root`, the primary agent in a team of agents collaborating to fulfill the user's goals. + + At the start of your turn, you are the active agent. + You can spawn sub-agents to handle subtasks, and those sub-agents can spawn their own sub-agents. + All agents in the team, including the agents that you can assign tasks to, are equally intelligent and capable, and have access to the same set of tools. + + You can use `spawn_agent` to create a new agent, `followup_task` to give an existing agent a new task an...message to a running agent without triggering a turn. [hash=7d40ff8b58977202] + `send_message` calls may be read by a human, so ensure they are legible. Always put proper spaces between words and/or numbers. + + - All agents use the same current working directory. + - As a result, edits made by one agent are immediately visible to all other agents. + + When calling `wait_agent`, prefer longer waits (minutes) to avoid busy polling. + + There are 4 available concurrency slots, meaning that up to 4 agents can be active at once, including you. + + Full-history forks (`fork_turns` omitted or `"all"`) inherit the parent model and reasoning effort and d...ne"` or a positive integer string. [hash=c7f52e90688299ce] +04:message/developer: + Any earlier instruction enabling proactive multi-agent delegation no longer applies. D...elegation, or parallel agent work. [hash=2b6e037067b4390c] +05:message/user: + + + + + + [hash=e8c1de60307d0a3d] + +06:message/user: + Ask before continuing in the other workspace. +-- request 2 (turn) -- +07:function_call/functions.request_user_input:{"questions":[{"header":"Continue","id":"continue","options":[{"description":"Continue working.","label"...o"}],"question":"Continue in the other workspace?"}]} [hash=028e565351844d96] +08:function_call_output:{"answers":{"continue":{"answers":["Yes"]}}} +09:message/user: + + /other-workspace + + + + /other-workspace [hash=0be8191a0451377a] + +-- request 3 (turn) -- +10:message/assistant: + I continued in the other workspace. +11:message/user: + + + + + + [hash=e8c1de60307d0a3d] + +12:message/user: + Which workspace are you using now? 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 e2601485a1..f2fa41b82e 100644 --- a/codex-rs/core/tests/suite/step_settings/environment_selection.rs +++ b/codex-rs/core/tests/suite/step_settings/environment_selection.rs @@ -1,9 +1,13 @@ -//! Active model updates preserve the captured environment until the next turn. +//! Active settings keep environment selection and model changes scoped to the running turn. use super::*; +use codex_config::test_support::CloudConfigBundleFixture; use codex_exec_server::CreateDirectoryOptions; +use codex_protocol::models::PermissionProfileSnapshot; +use codex_protocol::protocol::EnvironmentConfigState; use codex_protocol::protocol::TurnEnvironmentSelections; use core_test_support::submit_thread_settings; +use core_test_support::test_codex::environment_config_for_selection; use pretty_assertions::assert_eq; #[tokio::test(flavor = "multi_thread", worker_threads = 2)] @@ -115,6 +119,235 @@ async fn model_update_preserves_active_environment_and_next_turn_uses_new_select Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn combined_active_model_and_environment_updates_apply_or_reject_together() -> Result<()> { + skip_if_no_network!(Ok(())); + + let server = start_mock_server().await; + let responses = mount_sse_sequence( + &server, + vec![ + paused_response("pause-before-invalid-update", "pause"), + paused_response("pause-before-valid-update", "pause-again"), + sse(vec![ + ev_response_created("write-in-new-environment"), + ev_apply_patch_custom_tool_call( + "write-in-new-environment", + "*** Begin Patch\n*** Add File: marker.txt\n+selected B\n*** End Patch\n", + ), + ev_completed("write-in-new-environment"), + ]), + sse_completed("after-valid-update"), + ], + ) + .await; + let test = direct_tool_settings_test() + .with_config(|config| { + config + .permissions + .set_permission_profile(PermissionProfile::Disabled) + .expect("test permissions"); + }) + .build_with_auto_env(&server) + .await?; + let original = test.executor_environment().selection().clone(); + let mut next_environment = original.clone(); + next_environment.cwd = original.cwd.join("selected-environment")?; + next_environment.workspace_roots = vec![next_environment.cwd.clone()]; + next_environment.config = EnvironmentConfigState::Ready(environment_config_for_selection( + &test.config, + &next_environment, + )); + test.fs() + .create_directory( + &next_environment.cwd, + CreateDirectoryOptions { + recursive: false, + follow_symlinks: true, + }, + /*sandbox*/ None, + ) + .await?; + let next_marker = next_environment.cwd.join("marker.txt")?; + let next_cwd = next_environment.cwd.clone(); + + let paused = start_paused_turn(&test.codex).await?; + let mut owned_environment = original.clone(); + owned_environment.config = + EnvironmentConfigState::Ready(environment_config_for_selection(&test.config, &original)); + apply_turn_settings( + &test.codex, + &paused.turn_id, + TurnSettingsUpdate { + environments: Some(vec![owned_environment]), + ..Default::default() + }, + ) + .await?; + let mut inherited_environment = next_environment.clone(); + inherited_environment.config = EnvironmentConfigState::FromThread; + let inherited = submit_turn_settings( + &test.codex, + &paused.turn_id, + TurnSettingsUpdate { + model: Some(MODEL_B.to_string()), + environments: Some(vec![inherited_environment]), + ..Default::default() + }, + ) + .await?; + let TurnSettingsUpdateOutcome::Rejected { reason } = inherited else { + anyhow::bail!( + "an owner-configured environment cannot go back to inheriting: {inherited:?}" + ); + }; + assert!( + reason.contains("owner-provided environment configuration"), + "{reason}" + ); + answer_paused_turn(&test.codex, &paused.turn_id).await?; + let paused = wait_for_event_match(&test.codex, |event| match event { + EventMsg::RequestUserInput(request) => Some(request.clone()), + _ => None, + }) + .await; + apply_turn_settings( + &test.codex, + &paused.turn_id, + TurnSettingsUpdate { + model: Some(MODEL_B.to_string()), + environments: Some(vec![next_environment]), + ..Default::default() + }, + ) + .await?; + answer_paused_turn(&test.codex, &paused.turn_id).await?; + wait_for_event(&test.codex, |event| { + matches!(event, EventMsg::TurnComplete(_)) + }) + .await; + + let requests = responses.requests(); + assert_eq!( + requests + .iter() + .map(|request| request.body_json()["model"].clone()) + .collect::>(), + [MODEL_A, MODEL_A, MODEL_B, MODEL_B].map(|model| json!(model)), + ); + let context = requests[1].message_input_texts("user"); + let original_cwd = format!("{}", original.cwd.inferred_native_path_string()); + let rejected_cwd = format!("{}", next_cwd.inferred_native_path_string()); + assert!(context.iter().any(|text| text.contains(&original_cwd))); + assert!(!context.iter().any(|text| text.contains(&rejected_cwd))); + assert_eq!( + test.fs() + .read_file_text(&next_marker, Default::default(), /*sandbox*/ None) + .await?, + "selected B\n", + ); + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn active_updates_check_known_permissions_from_the_running_turn() -> Result<()> { + skip_if_no_network!(Ok(())); + + let server = start_mock_server().await; + mount_sse_sequence( + &server, + vec![ + paused_response("pause-before-environment-policy", "pause"), + sse_completed("after-environment-policy"), + ], + ) + .await; + let test = direct_tool_settings_test() + .with_cloud_config_bundle( + CloudConfigBundleFixture::loader_with_enterprise_requirement(format!( + "[auto_review]\nrequired_on_models = [\"{MODEL_A}\"]\n" + )), + ) + .with_config(|config| { + config.approvals_reviewer = ApprovalsReviewer::AutoReview; + config + .permissions + .set_permission_profile(PermissionProfile::read_only()) + .expect("set restricted thread permissions"); + for feature in [Feature::GuardianApproval, Feature::DeferredExecutor] { + config + .features + .enable(feature) + .expect("enable test feature"); + } + }) + .build_with_auto_env(&server) + .await?; + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await?; + test.thread_manager + .environment_manager() + .upsert_environment( + "policy-secondary".to_string(), + format!("ws://{}", listener.local_addr()?), + /*connect_timeout*/ None, + )?; + let mut primary = test.executor_environment().selection().clone(); + primary.config = + EnvironmentConfigState::Ready(environment_config_for_selection(&test.config, &primary)); + let mut secondary = primary.clone(); + secondary.environment_id = "policy-secondary".to_string(); + let mut unrestricted = environment_config_for_selection(&test.config, &secondary); + unrestricted.permission_profile = + PermissionProfileSnapshot::legacy(PermissionProfile::Disabled); + let paused = start_paused_turn(&test.codex).await?; + let current_model = TurnSettingsUpdate { + model: Some(MODEL_A.to_string()), + ..Default::default() + }; + apply_turn_settings(&test.codex, &paused.turn_id, current_model.clone()).await?; + secondary.config = EnvironmentConfigState::Ready(unrestricted.clone()); + let supplied = submit_turn_settings( + &test.codex, + &paused.turn_id, + TurnSettingsUpdate { + environments: Some(vec![primary.clone(), secondary.clone()]), + ..current_model.clone() + }, + ) + .await?; + let TurnSettingsUpdateOutcome::Rejected { reason } = supplied else { + anyhow::bail!("unrestricted secondary must reject the reviewed model: {supplied:?}"); + }; + assert!(reason.contains("you need to use auto review"), "{reason}"); + + secondary.config = EnvironmentConfigState::Pending; + apply_turn_settings( + &test.codex, + &paused.turn_id, + TurnSettingsUpdate { + environments: Some(vec![primary, secondary.clone()]), + ..Default::default() + }, + ) + .await?; + test.codex + .environment_ready(&secondary, unrestricted) + .await?; + // Its permissions are known even though the executor hasn't connected yet. + let without_selection = + submit_turn_settings(&test.codex, &paused.turn_id, current_model).await?; + assert_eq!( + without_selection, + TurnSettingsUpdateOutcome::Rejected { reason } + ); + answer_paused_turn(&test.codex, &paused.turn_id).await?; + wait_for_event(&test.codex, |event| { + matches!(event, EventMsg::TurnComplete(_)) + }) + .await; + 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(())); @@ -148,6 +381,18 @@ async fn inherited_permission_update_applies_to_the_next_turn_not_the_next_step( }, ) .await?; + // Recreate the inherited selection after saving future permissions; it must use this turn's. + let mut inherited = test.executor_environment().selection().clone(); + inherited.workspace_roots.clear(); + apply_turn_settings( + &test.codex, + &paused.turn_id, + TurnSettingsUpdate { + environments: Some(vec![inherited]), + ..Default::default() + }, + ) + .await?; answer_paused_turn(&test.codex, &paused.turn_id).await?; wait_for_event(&test.codex, |event| { matches!(event, EventMsg::TurnComplete(_)) diff --git a/codex-rs/protocol/src/protocol.rs b/codex-rs/protocol/src/protocol.rs index 614c201c70..397e2d078b 100644 --- a/codex-rs/protocol/src/protocol.rs +++ b/codex-rs/protocol/src/protocol.rs @@ -476,10 +476,14 @@ pub struct ConversationSpeechParams { /// Supported sparse changes to one live task's current settings, regardless of /// task kind. Child sessions and consumers of frozen initial settings are unchanged. -#[derive(Debug, Clone, Default, PartialEq, Eq)] +#[derive(Debug, Clone, Default, PartialEq)] pub struct TurnSettingsUpdate { /// Changes the reviewer for subsequent approval requests, not pending reviews. pub approvals_reviewer: Option, + /// Replaces the selection for subsequent steps, without changing future turns. + /// Environments may inherit the running turn's defaults or provide their own configuration, + /// which can be pending. An already-selected environment with its own cannot switch back. + pub environments: Option>, pub model: Option, /// `None` preserves the selection; `Some(None)` clears it. pub effort: Option>,