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?;