From c775dd3c332de1b69b25a4580f6c5bc44b94e284 Mon Sep 17 00:00:00 2001 From: sayan-oai Date: Thu, 17 Sep 2026 21:35:15 +0000 Subject: [PATCH] Defer environment selection changes until the next turn (#46310) ## Why Updating environment selections while a turn is running must not redirect its tools or prevent its pending environment setup from completing. ## What changed - Store future environment selections separately from active environments, and activate them when preparing new work with no task running. - Keep turn contexts and MCP workspace roots tied to the active environment snapshot. Report saved selections through thread settings and `environment_selections()`. - Route environment configuration and failure callbacks to matching active or future selections so each can finish setup independently. - Stop the running task before manual compaction adopts the saved environments. ## Testing Add regression tests that preserve a running task's working directory and workspace roots, then adopt the saved selection for the next task. Add remote environment coverage showing that future setup can complete while the active turn waits, with each turn receiving its own capability roots. GitOrigin-RevId: 09a6b90131d2ab2441ce2fe9b0368a3668889f31 --- codex-rs/core/src/agent/control/spawn.rs | 31 ++-- codex-rs/core/src/agent/control_tests.rs | 19 ++- codex-rs/core/src/codex_thread.rs | 3 +- codex-rs/core/src/environment_selection.rs | 87 +++++----- codex-rs/core/src/session/environment.rs | 125 ++++++++++++--- codex-rs/core/src/session/handlers.rs | 2 + codex-rs/core/src/session/mcp_runtime.rs | 22 ++- codex-rs/core/src/session/mod.rs | 33 ++-- codex-rs/core/src/session/session.rs | 4 + codex-rs/core/src/session/tests.rs | 90 +++++++++-- codex-rs/core/src/session/turn_context.rs | 43 +++-- codex-rs/core/src/session/turn_input_tests.rs | 2 +- codex-rs/core/tests/suite/hooks_executor.rs | 33 +++- codex-rs/core/tests/suite/remote_env.rs | 149 +++++++++++++++++- codex-rs/core/tests/suite/rmcp_client.rs | 14 +- 15 files changed, 522 insertions(+), 135 deletions(-) diff --git a/codex-rs/core/src/agent/control/spawn.rs b/codex-rs/core/src/agent/control/spawn.rs index 67081ab737..12ee357407 100644 --- a/codex-rs/core/src/agent/control/spawn.rs +++ b/codex-rs/core/src/agent/control/spawn.rs @@ -322,18 +322,13 @@ impl AgentControl { parent: Option>, ) -> CodexResult<()> { let state = self.upgrade()?; - let parent = if let Some(parent) = parent { + let owner_thread_id = parent.as_ref().map(|parent| parent.session.thread_id); + if let Some(parent) = &parent { let parent_thread_id = parent.session.thread_id; - let turn = parent.session.new_default_turn().await; - config = build_agent_resume_config(&turn).map_err(|_| { - CodexErr::InvalidRequest(format!( - "cannot resume multi-agent v2 child {thread_id} with the current parent settings" - )) - })?; let registered_parent = state.get_thread(parent_thread_id).await.ok(); if !registered_parent .as_ref() - .is_some_and(|registered| Arc::ptr_eq(registered, &parent)) + .is_some_and(|registered| Arc::ptr_eq(registered, parent)) || !parent.is_running() || parent.multi_agent_version() != Some(MultiAgentVersion::V2) || !Arc::ptr_eq(&self.state, &parent.session.services.agent_control.state) @@ -342,11 +337,7 @@ impl AgentControl { "cannot resume multi-agent v2 child {thread_id}: parent ownership is unavailable; resume the parent first" ))); } - Some((parent, turn.environments.clone())) - } else { - None - }; - let owner_thread_id = parent.as_ref().map(|(parent, _)| parent.session.thread_id); + } if owner_thread_id.is_none() && state.get_thread(thread_id).await.is_ok() { self.touch_loaded_v2_residency(&state, thread_id).await; return Ok(()); @@ -400,6 +391,20 @@ impl AgentControl { return Ok(()); } } + let parent = if let Some(parent) = parent { + let turn = parent + .session + .new_turn_with_default_settings(Uuid::now_v7().to_string(), Default::default()) + .await; + config = build_agent_resume_config(&turn).map_err(|_| { + CodexErr::InvalidRequest(format!( + "cannot resume multi-agent v2 child {thread_id} with the current parent settings" + )) + })?; + Some((parent, turn.environments.clone())) + } else { + None + }; config.model_reasoning_effort = stored_reasoning_effort; if let Some(role_name) = session_source.get_agent_role() { let runtime_approval_policy = config.permissions.approval_policy.value(); diff --git a/codex-rs/core/src/agent/control_tests.rs b/codex-rs/core/src/agent/control_tests.rs index 97d116b69e..f4bcc975a3 100644 --- a/codex-rs/core/src/agent/control_tests.rs +++ b/codex-rs/core/src/agent/control_tests.rs @@ -14,6 +14,7 @@ use crate::context::ManagedDeveloperInstructions; use crate::context::MultiAgentRoleInstructions; use crate::context::SubagentNotification; use crate::init_state_db; +use crate::session::SessionSettingsUpdate; use crate::thread_manager::StartThreadOptions; use crate::tools::handlers::multi_agents_common::thread_spawn_source; use assert_matches::assert_matches; @@ -66,6 +67,7 @@ use codex_protocol::protocol::TokenUsageRecord; use codex_protocol::protocol::TurnAbortReason; use codex_protocol::protocol::TurnAbortedEvent; use codex_protocol::protocol::TurnCompleteEvent; +use codex_protocol::protocol::TurnEnvironmentSelections; use codex_protocol::protocol::TurnStartedEvent; use codex_thread_store::ArchiveThreadParams; use codex_thread_store::InMemoryThreadStore; @@ -897,17 +899,21 @@ async fn check_v2_agent_reload(route: V2ReloadRoute) { .environments .primary() .expect("parent environment"); - let thread_config = environment.config().clone(); - let mut owner_config = thread_config.clone(); + let mut owner_config = environment.config().clone(); owner_config.allow_login_shell = false; let mut selection = environment.selection(); selection.config = EnvironmentConfigState::Ready(owner_config); parent_thread .session - .services - .turn_environments - .update_selections(std::slice::from_ref(&selection), &thread_config); - parent_turn = parent_thread.session.new_default_turn().await; + .update_settings(SessionSettingsUpdate { + environments: Some(TurnEnvironmentSelections::new( + parent_turn.config.cwd.clone(), + vec![selection], + )), + ..Default::default() + }) + .await + .expect("save parent environments"); parent_thread.session.mark_interrupted(); // The fixture has no task runner to finish the turn or consume child results. *parent_thread.session.active_turn.lock().await = None; @@ -921,6 +927,7 @@ async fn check_v2_agent_reload(route: V2ReloadRoute) { .ensure_multi_agent_v2_child_loaded(spawned_agent.thread_id) .await .expect("known child should reload through its parent"); + parent_turn = parent_thread.session.new_default_turn().await; assert!(harness.manager.get_thread(parent_thread_id).await.is_err()); } } diff --git a/codex-rs/core/src/codex_thread.rs b/codex-rs/core/src/codex_thread.rs index a7cd6cf38f..9bbbfe79db 100644 --- a/codex-rs/core/src/codex_thread.rs +++ b/codex-rs/core/src/codex_thread.rs @@ -908,8 +908,9 @@ impl CodexThread { self.session.refresh_codex_apps_tools().await } + /// Returns the environments configured for future turns. pub async fn environment_selections(&self) -> Vec { - self.session.services.turn_environments.selections() + self.session.configured_environment_selections().await } /// Installs resolved environment configuration and capability roots on this thread. diff --git a/codex-rs/core/src/environment_selection.rs b/codex-rs/core/src/environment_selection.rs index 3dec829d5b..1416e359c1 100644 --- a/codex-rs/core/src/environment_selection.rs +++ b/codex-rs/core/src/environment_selection.rs @@ -1,6 +1,7 @@ use std::collections::HashMap; use std::collections::HashSet; use std::fmt; +use std::future::Future; use std::sync::Arc; use std::sync::OnceLock; @@ -490,15 +491,6 @@ impl ThreadEnvironments { .collect() } - pub(crate) fn primary_workspace_roots(&self) -> Vec { - self.environments - .load() - .first() - .map_or_else(Vec::new, |environment| { - Self::primary_workspace_roots_for(std::slice::from_ref(&environment.selection)) - }) - } - /// Returns installed owner configuration without treating pending attachments as ready. pub(crate) fn primary_config_for( selections: &[TurnEnvironmentSelection], @@ -806,45 +798,53 @@ 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 = self.environments.load().len(), - non_blocking = self.non_blocking_snapshots, + environment_count = selected.len(), + non_blocking = non_blocking_snapshots, ) )] - pub(crate) async fn snapshot(&self) -> TurnEnvironmentSnapshot { + pub(crate) fn snapshot( + &self, + ) -> impl Future + Send + 'static { let selected = self.environments.load_full(); - let mut environments = Vec::with_capacity(selected.len()); - for environment in selected.iter() { - if let EnvironmentConfigState::Failed(error) = &environment.selection.config { - environments.push(TurnEnvironmentState::Failed { + let non_blocking_snapshots = self.non_blocking_snapshots; + async move { + let mut environments = Vec::with_capacity(selected.len()); + for environment in selected.iter() { + if let EnvironmentConfigState::Failed(error) = &environment.selection.config { + environments.push(TurnEnvironmentState::Failed { + selection: environment.selection.clone(), + error: error.clone(), + }); + continue; + } + let pending = matches!( + environment.selection.config, + EnvironmentConfigState::Pending + ); + let starting = StartingTurnEnvironment { selection: environment.selection.clone(), - error: error.clone(), - }); - continue; + config_origin: environment.config_origin, + resolution: environment.resolution.clone(), + }; + 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)); } - let pending = matches!( - environment.selection.config, - EnvironmentConfigState::Pending - ); - let starting = StartingTurnEnvironment { - selection: environment.selection.clone(), - config_origin: environment.config_origin, - resolution: environment.resolution.clone(), - }; - let resolved = if self.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 } } - TurnEnvironmentSnapshot { environments } } pub(crate) fn environment_manager(&self) -> Arc { @@ -1000,6 +1000,17 @@ impl TurnEnvironmentSnapshot { self.turn_environments().next() } + /// Returns the first selected environment's host folders, even if setup is not ready yet. + pub(crate) fn primary_workspace_roots(&self) -> Vec { + let selection = match self.environments.first() { + Some(TurnEnvironmentState::Ready(environment)) => &environment.selection, + Some(TurnEnvironmentState::Starting(environment)) => &environment.selection, + Some(TurnEnvironmentState::Failed { selection, .. }) => selection, + None => return Vec::new(), + }; + ThreadEnvironments::primary_workspace_roots_for(std::slice::from_ref(selection)) + } + /// Returns the primary environment's resolved permissions, or the provided fallback. pub(crate) fn permission_profile_or_else( &self, diff --git a/codex-rs/core/src/session/environment.rs b/codex-rs/core/src/session/environment.rs index 1be2fda061..2093251643 100644 --- a/codex-rs/core/src/session/environment.rs +++ b/codex-rs/core/src/session/environment.rs @@ -14,6 +14,7 @@ use codex_protocol::protocol::TurnEnvironmentSelection; use crate::config::ConstraintError; use crate::config::ConstraintResult; use crate::config::NetworkProxySpec; +use crate::environment_selection::TurnEnvironmentSnapshot; use crate::session::session::Session; use crate::session::session::SessionConfiguration; use crate::session::session::SessionSettingsUpdate; @@ -108,7 +109,7 @@ impl Session { current: &SessionConfiguration, updates: &SessionSettingsUpdate, ) -> ConstraintResult { - let current_environments = self.services.turn_environments.selections(); + let current_environments = ¤t.environments; if let Some(environments) = &updates.environments && let Some(environment) = environments.environments.iter().find(|environment| { environment.config == EnvironmentConfigState::FromThread @@ -126,7 +127,47 @@ impl Session { }); } - current.apply(updates, ¤t_environments) + current.apply(updates, current_environments) + } + + /// Activates the environments accepted for a new task. Configuration may have arrived for + /// those same environments while its settings were being prepared. + #[expect( + clippy::await_holding_invalid_type, + reason = "do not change environments while another task is running" + )] + pub(super) async fn activate_turn_environments( + &self, + configuration: &SessionConfiguration, + ) -> TurnEnvironmentSnapshot { + let snapshot = { + let active = self.active_turn.lock().await; + let state = self.state.lock().await; + // Another task may have started while this one was preparing its settings. + if active.as_ref().is_none_or(|turn| turn.task.is_none()) { + let mut environments = configuration.environments.clone(); + let latest_environments = &state.session_configuration.environments; + for environment in &mut environments { + if let Some(latest) = latest_environments.iter().find(|latest| { + latest.environment_id == environment.environment_id + && latest.cwd == environment.cwd + && latest.workspace_roots == environment.workspace_roots + }) { + environment.config = latest.config.clone(); + } + } + if environments != self.services.turn_environments.selections() { + self.mark_mcp_runtime_dirty(); + self.services.turn_environments.update_selections( + &environments, + &configuration.inferred_environment_config(), + ); + self.schedule_mcp_prewarm(); + } + } + self.services.turn_environments.snapshot() + }; + snapshot.await } pub(crate) async fn environment_ready( @@ -148,39 +189,79 @@ impl Session { .await } + /// Records configuration, or a failure to obtain it, for this environment and workspace. + /// The running turn and future turns may use different environments. Update the matching + /// selection or selections so work already waiting can still finish, and check successful + /// configuration before saving either list. async fn update_environment_configuration( &self, selection: &TurnEnvironmentSelection, config: EnvironmentConfigState, ) -> CodexResult<()> { // Serialize owner callbacks with ordinary thread settings updates. - let state = self.state.lock().await; - let mut environments = self.services.turn_environments.selections(); - let Some(environment) = environments.iter_mut().find(|environment| { + let mut state = self.state.lock().await; + let mut current = self.services.turn_environments.selections(); + let mut future = state.session_configuration.environments.clone(); + let matches = |environment: &TurnEnvironmentSelection| { environment.environment_id == selection.environment_id && environment.cwd == selection.cwd && environment.workspace_roots == selection.workspace_roots - }) else { - return Err(CodexErr::InvalidRequest(format!( - "environment `{}` is not selected on this thread with the requested workspace", - selection.environment_id - ))); }; + let (current_environment, future_environment) = match ( + current.iter_mut().find(|environment| matches(environment)), + future.iter_mut().find(|environment| matches(environment)), + ) { + (None, None) => { + return Err(CodexErr::InvalidRequest(format!( + "environment `{}` is not selected on this thread with the requested workspace", + selection.environment_id + ))); + } + (Some(current), Some(future)) if current.config != future.config => { + // Update the selection still waiting for configuration. If neither is waiting, + // keep the existing behavior of updating the current one. + if matches!(future.config, EnvironmentConfigState::Pending) { + (None, Some(future)) + } else { + (Some(current), None) + } + } + (Some(current), future) => (Some(current), future), + (None, Some(future)) => (None, Some(future)), + }; + let update_current = current_environment.is_some(); + let update_future = future_environment.is_some(); - environment.config = config; - if matches!(environment.config, EnvironmentConfigState::Ready(_)) { - state - .session_configuration - .validate(&environments) - .map_err(|error| CodexErr::InvalidRequest(error.to_string()))?; + if let Some(environment) = current_environment { + environment.config = config.clone(); + } + if let Some(environment) = future_environment { + environment.config = config.clone(); + } + if matches!(config, EnvironmentConfigState::Ready(_)) { + let validate = |environments: &[TurnEnvironmentSelection]| { + state + .session_configuration + .validate(environments) + .map_err(|error| CodexErr::InvalidRequest(error.to_string())) + }; + if update_current { + validate(¤t)?; + } + if update_future && (!update_current || current != future) { + validate(&future)?; + } } - // Invalidate MCP before installed configuration can wake a waiting turn. - self.mark_mcp_runtime_dirty(); - self.services.turn_environments.update_selections( - &environments, - &state.session_configuration.inferred_environment_config(), - ); + state.session_configuration.environments = future; + if update_current { + // Invalidate MCP before installed configuration can wake a waiting turn. + self.mark_mcp_runtime_dirty(); + self.services.turn_environments.update_selections( + ¤t, + &state.session_configuration.inferred_environment_config(), + ); + } Ok(()) } diff --git a/codex-rs/core/src/session/handlers.rs b/codex-rs/core/src/session/handlers.rs index dc739a59bf..dcea9a99d2 100644 --- a/codex-rs/core/src/session/handlers.rs +++ b/codex-rs/core/src/session/handlers.rs @@ -241,6 +241,8 @@ pub async fn reload_user_config(sess: &Arc) { } pub async fn compact(sess: &Arc, sub_id: String) { + // Stop the old turn before the compact task picks up the next turn's environments. + sess.abort_all_tasks(TurnAbortReason::Replaced).await; let turn_context = sess .new_turn_with_default_settings(sub_id, Default::default()) .await; diff --git a/codex-rs/core/src/session/mcp_runtime.rs b/codex-rs/core/src/session/mcp_runtime.rs index 7d6c5e2cfb..fe05899606 100644 --- a/codex-rs/core/src/session/mcp_runtime.rs +++ b/codex-rs/core/src/session/mcp_runtime.rs @@ -36,7 +36,6 @@ impl Session { &self, current: &SessionConfiguration, next: &SessionConfiguration, - updates: &SessionSettingsUpdate, ) -> bool { current.cwd() != next.cwd() || current.step_settings.approval_policy.value() @@ -44,9 +43,6 @@ impl Session { || current.step_settings.approvals_reviewer != next.step_settings.approvals_reviewer || current.permission_profile() != next.permission_profile() || current.windows_sandbox_level != next.windows_sandbox_level - || updates.environments.as_ref().is_some_and(|environments| { - environments.environments != self.services.turn_environments.selections() - }) } /// Waits on this session's refreshed server before tool execution is admitted. @@ -76,19 +72,25 @@ impl Session { &self, auth: Option, ) -> McpDesiredState { - let (session_configuration, disabled_plugin_ids) = { + let (session_configuration, disabled_plugin_ids, environments) = { let state = self.state.lock().await; + // MCP tools must use current environments, not the selection saved for the next turn. ( state.session_configuration.clone(), state.active_disabled_plugin_ids.clone(), + self.services.turn_environments.snapshot(), ) }; - let environments = self.services.turn_environments.snapshot().await; + let environments = environments.await; let cwd = environments .primary() .and_then(|environment| environment.cwd().to_abs_path().ok()) .unwrap_or_else(|| session_configuration.cwd().clone()); - let config = self.build_per_turn_config(&session_configuration, cwd); + let config = self.build_per_turn_config( + &session_configuration, + cwd, + environments.primary_workspace_roots(), + ); let local_process_cwd = environments .local_environment_cwd() .unwrap_or_else(|| session_configuration.cwd().clone()) @@ -116,7 +118,11 @@ impl Session { ) -> anyhow::Result<()> { let cwd = AbsolutePathBuf::from_absolute_path(mcp_runtime_cwd) .unwrap_or_else(|_| session_configuration.cwd().clone()); - let config = self.build_per_turn_config(session_configuration, cwd); + let config = self.build_per_turn_config( + session_configuration, + cwd, + resolved_environments.primary_workspace_roots(), + ); let local_process_cwd = resolved_environments .local_environment_cwd() .unwrap_or_else(|| session_configuration.cwd().clone()) diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 7a19a9078d..4b76ca52eb 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -835,6 +835,7 @@ impl Session { approval_policy: config.permissions.approval_policy.clone(), approvals_reviewer: config.approvals_reviewer, }), + environments: environment_selections.clone(), model_info_overrides: config.to_models_manager_config().into(), developer_instructions: config.developer_instructions.clone(), base_instructions, @@ -1882,22 +1883,15 @@ impl Session { let updated_permission_profile = updated.permission_profile(); let permission_profile_changed = previous_permission_profile != updated_permission_profile; - let mcp_inputs_changed = - self.mcp_inputs_differ(&state.session_configuration, &updated, &updates); let root_service_tier_changed = updated.parent_thread_id.is_none() && state.session_configuration.step_settings.service_tier != updated.step_settings.service_tier; + let environment_config = updated.inferred_environment_config(); + let mcp_inputs_changed = self.mcp_inputs_differ(&state.session_configuration, &updated); if mcp_inputs_changed { self.mark_mcp_runtime_dirty(); } - let environment_config = updated.inferred_environment_config(); - if let Some(environments) = &updates.environments { - self.services - .turn_environments - .update_selections(&environments.environments, &environment_config); - } else if state.session_configuration.inferred_environment_config() - != environment_config - { + if state.session_configuration.inferred_environment_config() != environment_config { self.services .turn_environments .update_thread_config(&environment_config); @@ -1918,7 +1912,7 @@ impl Session { configuration: state.session_configuration.clone(), snapshot: state .session_configuration - .thread_settings_snapshot(&self.services.turn_environments.selections()), + .thread_settings_snapshot(&state.session_configuration.environments), }; ( commit, @@ -1945,32 +1939,33 @@ impl Session { ) -> ConstraintResult { let state = self.state.lock().await; let configuration = self.apply_session_settings(&state.session_configuration, updates)?; - let environments = updates.environments.as_ref().map_or_else( - || self.services.turn_environments.selections(), - |environments| environments.environments.clone(), - ); - Ok(configuration.thread_config_snapshot(environments)) + Ok(configuration.thread_config_snapshot(configuration.environments.clone())) } pub(crate) async fn thread_config_snapshot(&self) -> ThreadConfigSnapshot { let state = self.state.lock().await; state .session_configuration - .thread_config_snapshot(self.services.turn_environments.selections()) + .thread_config_snapshot(state.session_configuration.environments.clone()) + } + + pub(crate) async fn configured_environment_selections(&self) -> Vec { + let state = self.state.lock().await; + state.session_configuration.environments.clone() } pub(crate) async fn thread_settings_snapshot(&self) -> ThreadSettingsSnapshot { let state = self.state.lock().await; state .session_configuration - .thread_settings_snapshot(&self.services.turn_environments.selections()) + .thread_settings_snapshot(&state.session_configuration.environments) } pub(crate) async fn restorable_thread_settings(&self) -> CodexThreadSettingsOverrides { let state = self.state.lock().await; state .session_configuration - .restorable_thread_settings(self.services.turn_environments.selections()) + .restorable_thread_settings(state.session_configuration.environments.clone()) } pub(crate) async fn set_app_server_client_info( diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index fc0f81acf6..8965f3b4ba 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -104,6 +104,8 @@ pub(crate) struct SessionConfiguration { /// Desired configured inputs inherited by future turns. pub(super) step_settings: Arc, + /// Ordered environments inherited by future turns. + pub(super) environments: Vec, /// Explicit startup overrides used when resolving effective model metadata. pub(super) model_info_overrides: ModelInfoOverrides, @@ -419,6 +421,7 @@ impl SessionConfiguration { let current_cwd = self.cwd().clone(); if let Some(environments) = &updates.environments { next_configuration.legacy_fallback_cwd = environments.legacy_fallback_cwd.clone(); + next_configuration.environments = environments.environments.clone(); } let cwd_changed = next_configuration.legacy_fallback_cwd != current_cwd; if let Some(runtime_workspace_roots) = &updates.runtime_workspace_roots { @@ -1402,6 +1405,7 @@ impl Session { environment_selections, &session_configuration.inferred_environment_config(), ); + session_configuration.environments = turn_environments.selections(); let resolved_environments = turn_environments.snapshot().await; let agents_md_manager = Arc::new(AgentsMdManager::new(instructions)); let plugin_skill_warmup = warm_plugins_and_skills_for_session_init( diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index d242d6d311..dad95c7001 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -4014,6 +4014,7 @@ async fn set_rate_limits_retains_previous_credits() { }; let session_configuration = SessionConfiguration { provider: create_model_provider(config.model_provider.clone(), /*auth_manager*/ None), + environments: Vec::new(), step_settings: Arc::new(StepSettings { collaboration_mode, reasoning_summary: config.model_reasoning_summary, @@ -4136,6 +4137,7 @@ async fn set_rate_limits_updates_plan_type_when_present() { }; let session_configuration = SessionConfiguration { provider: create_model_provider(config.model_provider.clone(), /*auth_manager*/ None), + environments: Vec::new(), step_settings: Arc::new(StepSettings { collaboration_mode, reasoning_summary: config.model_reasoning_summary, @@ -4755,6 +4757,7 @@ pub(crate) async fn make_session_configuration_for_tests() -> SessionConfigurati SessionConfiguration { provider: create_model_provider(config.model_provider.clone(), /*auth_manager*/ None), + environments: Vec::new(), step_settings: Arc::new(StepSettings { collaboration_mode, reasoning_summary: config.model_reasoning_summary, @@ -5742,7 +5745,7 @@ async fn relative_cwd_update_without_environments_resolves_under_session_cwd() { let state = session.state.lock().await; assert_eq!(state.session_configuration.cwd(), &updated_cwd); - assert!(session.services.turn_environments.selections().is_empty()); + assert!(state.session_configuration.environments.is_empty()); } #[tokio::test] @@ -5772,7 +5775,7 @@ async fn environment_settings_preserve_explicit_primary_cwd() { let state = session.state.lock().await; assert_eq!(state.session_configuration.cwd(), &updated_cwd); assert_eq!( - session.services.turn_environments.selections()[0].cwd, + state.session_configuration.environments[0].cwd, PathUri::from_abs_path(&environment_cwd) ); } @@ -5841,6 +5844,7 @@ async fn session_new_fails_when_zsh_fork_enabled_without_packaged_zsh() { config.model_provider.clone(), Some(Arc::clone(&auth_manager)), ), + environments: Vec::new(), step_settings: Arc::new(StepSettings { collaboration_mode, reasoning_summary: config.model_reasoning_summary, @@ -6018,6 +6022,7 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) { }; let default_environments = vec![local(config.cwd.clone())]; let session_configuration = SessionConfiguration { + environments: default_environments.clone(), provider: create_model_provider( config.model_provider.clone(), Some(Arc::clone(&auth_manager)), @@ -6234,8 +6239,13 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) { forked_from_ordinal_exclusive: None, next_internal_sub_id: AtomicU64::new(0), }; - let per_turn_config = - session.build_per_turn_config(&session_configuration, session_configuration.cwd().clone()); + let per_turn_config = session.build_per_turn_config( + &session_configuration, + session_configuration.cwd().clone(), + crate::environment_selection::ThreadEnvironments::primary_workspace_roots_for( + &session_configuration.environments, + ), + ); let plugins_input = per_turn_config.plugins_config_input(); let plugin_outcome = session .services @@ -6328,6 +6338,7 @@ async fn make_session_with_config_and_rx( }; let default_environments = vec![local(config.cwd.clone())]; let session_configuration = SessionConfiguration { + environments: default_environments.clone(), provider: create_model_provider( config.model_provider.clone(), Some(Arc::clone(&auth_manager)), @@ -6460,6 +6471,7 @@ async fn make_session_with_history_source_and_agent_control_and_rx( }; let default_environments = vec![local(config.cwd.clone())]; let session_configuration = SessionConfiguration { + environments: default_environments.clone(), provider: create_model_provider( config.model_provider.clone(), Some(Arc::clone(&auth_manager)), @@ -7536,7 +7548,55 @@ async fn turn_environments_set_primary_environment() { } #[tokio::test] -async fn default_turn_does_not_overlay_legacy_fallback_cwd_onto_stored_thread_environments() { +async fn new_task_does_not_change_an_already_running_tasks_environments() { + let (session, running_turn, events) = make_session_and_context_with_rx().await; + let current_cwd = running_turn.config.cwd.clone(); + let next_workspace = tempfile::tempdir().expect("next workspace"); + let next_cwd = next_workspace.path().abs(); + let finish = Arc::new(Notify::new()); + + session + .update_settings(SessionSettingsUpdate { + environments: Some(TurnEnvironmentSelections::new( + next_cwd.clone(), + vec![local(next_cwd.clone())], + )), + ..Default::default() + }) + .await + .expect("save next task's settings"); + session + .spawn_task( + running_turn, + Vec::new(), + HeldStepTask { + kind: TaskKind::Regular, + finish: Arc::clone(&finish), + }, + ) + .await; + + let candidate = session + .new_turn_with_default_settings("candidate".into(), Default::default()) + .await; + assert_eq!(candidate.config.cwd, current_cwd); + assert_eq!(candidate.config.workspace_roots, vec![current_cwd.clone()]); + assert_eq!( + session.services.turn_environments.selections(), + vec![local(current_cwd)] + ); + + finish.notify_one(); + recv_terminal_event(&events, TerminalEventKind::TurnComplete).await; + let next = session + .new_turn_with_default_settings("next".into(), Default::default()) + .await; + assert_eq!(next.config.cwd, next_cwd); + assert_eq!(next.config.workspace_roots, vec![next_cwd]); +} + +#[tokio::test] +async fn task_turn_does_not_overlay_legacy_fallback_cwd_onto_stored_thread_environments() { let (session, _initial_turn, _rx) = make_session_and_context_with_rx().await; let session_cwd = session.get_config().await.cwd.clone(); let selected_cwd = @@ -7551,7 +7611,9 @@ async fn default_turn_does_not_overlay_legacy_fallback_cwd_onto_stored_thread_en }) .await .expect("environment selection update should succeed"); - let turn_context = session.new_default_turn().await; + let turn_context = session + .new_turn_with_default_settings("task".into(), Default::default()) + .await; let turn_environments = &turn_context.environments; assert_eq!(turn_environments.turn_environments().count(), 1); @@ -7573,7 +7635,7 @@ async fn default_turn_does_not_overlay_legacy_fallback_cwd_onto_stored_thread_en } #[tokio::test] -async fn default_turn_honors_empty_stored_thread_environments() { +async fn task_turn_honors_empty_stored_thread_environments() { let (session, _initial_turn, _rx) = make_session_and_context_with_rx().await; let session_cwd = session.get_config().await.cwd.clone(); @@ -7587,7 +7649,9 @@ async fn default_turn_honors_empty_stored_thread_environments() { }) .await .expect("environment selection update should succeed"); - let turn_context = session.new_default_turn().await; + let turn_context = session + .new_turn_with_default_settings("task".into(), Default::default()) + .await; assert!(turn_context.environments.primary().is_none()); assert!( @@ -8202,6 +8266,7 @@ where }; let default_environments = vec![local(config.cwd.clone())]; let session_configuration = SessionConfiguration { + environments: default_environments.clone(), provider: create_model_provider( config.model_provider.clone(), Some(Arc::clone(&auth_manager)), @@ -8421,8 +8486,13 @@ where forked_from_ordinal_exclusive: None, next_internal_sub_id: AtomicU64::new(0), }); - let per_turn_config = - session.build_per_turn_config(&session_configuration, session_configuration.cwd().clone()); + let per_turn_config = session.build_per_turn_config( + &session_configuration, + session_configuration.cwd().clone(), + crate::environment_selection::ThreadEnvironments::primary_workspace_roots_for( + &session_configuration.environments, + ), + ); let plugins_input = per_turn_config.plugins_config_input(); let plugin_outcome = session .services diff --git a/codex-rs/core/src/session/turn_context.rs b/codex-rs/core/src/session/turn_context.rs index ca07a5e1c1..d092a6bcbe 100644 --- a/codex-rs/core/src/session/turn_context.rs +++ b/codex-rs/core/src/session/turn_context.rs @@ -769,6 +769,7 @@ impl Session { &self, session_configuration: &SessionConfiguration, cwd: AbsolutePathBuf, + workspace_roots: Vec, ) -> Config { // todo(aibrahim): store this state somewhere else so we don't need to mut config let config = session_configuration.original_config_do_not_use.clone(); @@ -776,7 +777,6 @@ impl Session { per_turn_config.cwd = cwd; per_turn_config.permissions.approval_policy = session_configuration.step_settings.approval_policy.clone(); - let workspace_roots = self.services.turn_environments.primary_workspace_roots(); per_turn_config.workspace_roots = workspace_roots.clone(); per_turn_config .permissions @@ -818,8 +818,15 @@ impl Session { &self, session_configuration: &SessionConfiguration, ) -> Config { - let mut config = - self.build_per_turn_config(session_configuration, session_configuration.cwd().clone()); + let workspace_roots = + crate::environment_selection::ThreadEnvironments::primary_workspace_roots_for( + &session_configuration.environments, + ); + let mut config = self.build_per_turn_config( + session_configuration, + session_configuration.cwd().clone(), + workspace_roots, + ); config.model = Some( session_configuration .step_settings @@ -1008,22 +1015,25 @@ impl Session { if let Some(service_tier) = service_tier_for_turn { Arc::make_mut(&mut configuration.step_settings).service_tier = Some(service_tier); } + let turn_environments = self.activate_turn_environments(&configuration).await; let turn_context = self - .new_turn_from_configuration(sub_id, configuration, options) + .new_turn_from_configuration(sub_id, configuration, turn_environments, options) .await; Ok(Some((turn_context, commit.snapshot))) } - /// Constructs a turn from the exact committed settings without starting a task. + /// Builds a context from the caller's chosen settings and environments without starting work. async fn new_turn_from_configuration( &self, sub_id: String, session_configuration: SessionConfiguration, + turn_environments: TurnEnvironmentSnapshot, options: NewTurnContextOptions, ) -> Arc { self.new_turn_context_from_configuration( sub_id, session_configuration, + turn_environments, options, TurnMultiAgentRuntime::ResolveAndStore, self.git_enrichment_policy, @@ -1036,9 +1046,11 @@ impl Session { sub_id: String, session_configuration: SessionConfiguration, ) -> Arc { + let turn_environments = self.services.turn_environments.snapshot().await; self.new_turn_context_from_configuration( sub_id, session_configuration, + turn_environments, NewTurnContextOptions::default(), TurnMultiAgentRuntime::Preview, GitEnrichmentPolicy::Skip, @@ -1051,11 +1063,11 @@ impl Session { &self, sub_id: String, session_configuration: SessionConfiguration, + turn_environments: TurnEnvironmentSnapshot, options: NewTurnContextOptions, multi_agent_runtime: TurnMultiAgentRuntime, git_enrichment_policy: GitEnrichmentPolicy, ) -> Arc { - let turn_environments = self.services.turn_environments.snapshot().await; let primary_turn_environment = turn_environments.primary(); // TODO(anp): Migrate per-turn config and legacy TurnContext cwd consumers to PathUri so // a foreign primary environment does not fall back to the session's host cwd. @@ -1063,7 +1075,11 @@ impl Session { .as_ref() .and_then(|turn_environment| turn_environment.cwd().to_abs_path().ok()) .unwrap_or_else(|| session_configuration.cwd().clone()); - let per_turn_config = self.build_per_turn_config(&session_configuration, cwd.clone()); + let per_turn_config = self.build_per_turn_config( + &session_configuration, + cwd.clone(), + turn_environments.primary_workspace_roots(), + ); let network_permission_profile = primary_turn_environment .map(TurnEnvironment::permission_profile) .cloned() @@ -1215,21 +1231,30 @@ impl Session { } } + /// Builds a context without starting work or changing the current environments. pub(crate) async fn new_default_turn(&self) -> Arc { - self.new_turn_with_default_settings( + let session_configuration = self.default_turn_configuration().await; + let turn_environments = self.services.turn_environments.snapshot().await; + self.new_turn_from_configuration( self.next_internal_sub_id(), + session_configuration, + turn_environments, NewTurnContextOptions::default(), ) .await } + /// Prepares new work on the saved environments; context-only callers keep the current ones. pub(crate) async fn new_turn_with_default_settings( &self, sub_id: String, options: NewTurnContextOptions, ) -> Arc { let session_configuration = self.default_turn_configuration().await; - self.new_turn_from_configuration(sub_id, session_configuration, options) + let turn_environments = self + .activate_turn_environments(&session_configuration) + .await; + self.new_turn_from_configuration(sub_id, session_configuration, turn_environments, options) .await } diff --git a/codex-rs/core/src/session/turn_input_tests.rs b/codex-rs/core/src/session/turn_input_tests.rs index 77dc3f4282..18f1368378 100644 --- a/codex-rs/core/src/session/turn_input_tests.rs +++ b/codex-rs/core/src/session/turn_input_tests.rs @@ -645,7 +645,7 @@ async fn automatic_admission_rechecks_plan_mode_without_committing_sparse_settin .await .expect("explicit settings update accepts the same patch"); assert_eq!( - session.services.turn_environments.selections(), + session.configured_environment_selections().await, proposed_environments.environments ); assert!(session.mcp_refresh.is_pending()); diff --git a/codex-rs/core/tests/suite/hooks_executor.rs b/codex-rs/core/tests/suite/hooks_executor.rs index 6af7942581..bfb265ba53 100644 --- a/codex-rs/core/tests/suite/hooks_executor.rs +++ b/codex-rs/core/tests/suite/hooks_executor.rs @@ -288,7 +288,12 @@ async fn executor_stop_hook_stops_after_disconnection() -> Result<()> { async fn executor_stop_hook_rejects_mismatched_environment() -> Result<()> { skip_if_no_network!(Ok(())); - let fixture = executor_stop_hook_fixture().await?; + let fixture = executor_hook_fixture( + ["first-turn", "switch-environments", "second-turn"] + .map(completed_turn_response) + .to_vec(), + ) + .await?; let selection = fixture.attach().await?; fixture .test @@ -372,6 +377,11 @@ async fn executor_stop_hook_rejects_mismatched_environment() -> Result<()> { }, ) .await?; + fixture + .test + .submit_text_turn("apply the new executor environment") + .await?; + fixture.wait_for_hook_call().await?; let mut mismatched_config = fixture.test.config.clone(); let mut node_repl = mismatched_config @@ -391,7 +401,16 @@ async fn executor_stop_hook_rejects_mismatched_environment() -> Result<()> { .codex .refresh_mcp_config(mismatched_config) .await; - wait_for_mcp_server(&fixture.test.codex, "node_repl").await?; + fixture + .test + .codex + .call_mcp_tool( + "node_repl", + "js", + Some(json!({"code": "1 + 1"})), + /*meta*/ None, + ) + .await?; assert_eq!( fixture .test @@ -409,7 +428,15 @@ async fn executor_stop_hook_rejects_mismatched_environment() -> Result<()> { executor.abort(); } - assert_eq!(fixture.calls().await?.len(), 1); + assert_eq!( + fixture + .calls() + .await? + .iter() + .filter(|call| call["params"]["name"] == "turn_ended") + .count(), + 2 + ); Ok(()) } diff --git a/codex-rs/core/tests/suite/remote_env.rs b/codex-rs/core/tests/suite/remote_env.rs index f4ce3b160c..6aa59c671d 100644 --- a/codex-rs/core/tests/suite/remote_env.rs +++ b/codex-rs/core/tests/suite/remote_env.rs @@ -9,6 +9,7 @@ use codex_core::EnvironmentConfig; use codex_core::EnvironmentNetworkPolicy; use codex_core::StartThreadOptions; use codex_core::TurnInputRequest; +use codex_core::TurnInputSubmission; use codex_core::WaitForEnvironmentToolConfig; use codex_core::compact::SUMMARIZATION_PROMPT; use codex_core::config::Config; @@ -1557,10 +1558,6 @@ async fn shared_executor_keeps_ready_capability_roots_scoped_to_each_attachment( ) .await?; - assert_eq!( - test.codex.inspect_selected_capability_roots().ready_roots, - vec![root("first-root")] - ); assert_eq!( second .thread @@ -1629,6 +1626,10 @@ async fn shared_executor_keeps_ready_capability_roots_scoped_to_each_attachment( wait_for_event(thread, |event| matches!(event, EventMsg::TurnComplete(_))).await; } + assert_eq!( + test.codex.inspect_selected_capability_roots().ready_roots, + vec![root("first-updated-root")] + ); let requests = response_mock.requests(); let root_fragments = requests .iter() @@ -1981,6 +1982,146 @@ async fn pending_attachment_installs_configuration_before_waiting_turn_resumes() Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn future_pending_environment_can_finish_without_retargeting_the_active_turn() -> Result<()> { + let server = start_mock_server().await; + let mut extensions = ExtensionRegistryBuilder::new(); + extensions.thread_lifecycle_contributor(Arc::new(WaitForEnvironmentTestExtension)); + extensions.prompt_contributor(Arc::new(ReadyCapabilityRootsTestExtension::default())); + let mut builder = test_codex() + .with_extensions(Arc::new(extensions.build())) + .with_config(|config| { + assert!(config.features.enable(Feature::DeferredExecutor).is_ok()); + }); + let test = builder.build_with_auto_env(&server).await?; + let selection = test.codex.environment_selections().await.remove(0); + let active = TurnEnvironmentSelection { + config: EnvironmentConfigState::Pending, + ..selection + }; + let future_cwd = active.cwd.join("future-environment")?; + test.fs() + .create_directory( + &future_cwd, + CreateDirectoryOptions { + recursive: false, + follow_symlinks: true, + }, + /*sandbox*/ None, + ) + .await?; + let future = TurnEnvironmentSelection { + cwd: future_cwd.clone(), + workspace_roots: vec![future_cwd], + ..active.clone() + }; + let owner_config = |selection: &TurnEnvironmentSelection, id: &str| EnvironmentConfig { + allow_login_shell: test.config.permissions.allow_login_shell, + workspace_roots: selection.workspace_roots.clone(), + permission_profile: PermissionProfileSnapshot::legacy(PermissionProfile::read_only()), + shell_environment_policy: Default::default(), + windows_sandbox_level: WindowsSandboxLevel::from_config(&test.config), + windows_sandbox_private_desktop: test.config.permissions.windows_sandbox_private_desktop, + use_legacy_landlock: test.config.features.use_legacy_landlock(), + exec_policy: None, + mcp_policy: None, + network_policy: None, + selected_capability_roots: vec![SelectedCapabilityRoot { + id: id.to_string(), + location: CapabilityRootLocation::Environment { + environment_id: selection.environment_id.clone(), + path: selection.cwd.clone(), + }, + }], + }; + let thread = test + .thread_manager + .start_thread(StartThreadOptions { + environments: Some(vec![active.clone()]), + ..StartThreadOptions::new(test.config.clone()) + }) + .await? + .thread; + let response_mock = mount_sse_sequence( + &server, + vec![ + sse(vec![ + ev_function_call( + "wait-for-active-environment", + "wait_for_environment", + &json!({ "environment_id": active.environment_id }).to_string(), + ), + ev_completed("waiting"), + ]), + sse(vec![ev_completed("active-done")]), + sse(vec![ev_completed("future-done")]), + ], + ) + .await; + let request = |text: &str| { + TurnInputRequest::user_input(vec![UserInput::Text { + text: text.to_string(), + text_elements: Vec::new(), + }]) + }; + thread.start_or_steer_turn(request("wait")).await?; + wait_for_response_request_count(&response_mock, /*expected_count*/ 1).await; + assert!(matches!( + thread + .start_or_steer_turn( + request("use the other environment next turn").with_thread_settings( + ThreadSettingsOverrides { + environments: Some(TurnEnvironmentSelections::new( + test.config.cwd.clone(), + vec![future.clone()], + )), + ..Default::default() + }, + ) + ) + .await?, + TurnInputSubmission::Steered { .. } + )); + thread + .environment_ready(&future, owner_config(&future, "future-root")) + .await?; + assert!(matches!( + thread.environment_selections().await[0].config, + EnvironmentConfigState::Ready(_) + )); + thread + .environment_ready(&active, owner_config(&active, "active-root")) + .await?; + wait_for_event(&thread, |event| matches!(event, EventMsg::TurnComplete(_))).await; + thread.start_or_steer_turn(request("next turn")).await?; + wait_for_event(&thread, |event| matches!(event, EventMsg::TurnComplete(_))).await; + + let requests = response_mock.requests(); + let (output, _) = requests[1] + .function_call_output_content_and_success("wait-for-active-environment") + .context("active wait output should be model visible")?; + assert_eq!( + serde_json::from_str::(&output.context("active wait output should contain JSON")?)?, + json!({ "environment_id": active.environment_id, "status": "ready" }) + ); + let ready_roots = requests[1..] + .iter() + .map(|request| { + request + .message_input_texts("user") + .into_iter() + .rfind(|text| text.contains("")) + }) + .collect::>(); + assert_eq!( + ready_roots, + ["active-root", "future-root"].map(|root| Some(format!( + "{root}" + ))) + ); + 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/rmcp_client.rs b/codex-rs/core/tests/suite/rmcp_client.rs index ccce14877a..cf9bfd27d7 100644 --- a/codex-rs/core/tests/suite/rmcp_client.rs +++ b/codex-rs/core/tests/suite/rmcp_client.rs @@ -704,7 +704,7 @@ async fn environment_mcp_policy_filters_runtime_config_and_model_tools( skip_if_no_network!(Ok(())); let server = responses::start_mock_server().await; - let response = mount_sse_once( + mount_sse_once( &server, responses::sse(vec![ responses::ev_response_created("resp-1"), @@ -786,6 +786,9 @@ async fn environment_mcp_policy_filters_runtime_config_and_model_tools( }, ) .await?; + fixture + .submit_text_turn("start with pending environment configuration") + .await?; let (pending_config, _) = fixture.codex.current_mcp_config_and_runtime_context().await; let pending_servers = pending_config.mcp_server_catalog.configured_servers(); assert!(!pending_servers["allowed"].enabled); @@ -856,6 +859,15 @@ async fn environment_mcp_policy_filters_runtime_config_and_model_tools( ) .await?; + let response = mount_sse_once( + &server, + responses::sse(vec![ + responses::ev_response_created("resp-2"), + responses::ev_assistant_message("msg-2", "done"), + responses::ev_completed("resp-2"), + ]), + ) + .await; fixture .submit_text_turn("show the available MCP tools") .await?;