diff --git a/codex-rs/core/src/agents_md.rs b/codex-rs/core/src/agents_md.rs index e0b80b2f61a4..37e0c0dd2da4 100644 --- a/codex-rs/core/src/agents_md.rs +++ b/codex-rs/core/src/agents_md.rs @@ -56,7 +56,7 @@ pub(crate) async fn load_project_instructions( environments: &TurnEnvironmentSnapshot, ) -> Option { let mut loaded = LoadedAgentsMd::from_user_instructions(user_instructions); - for turn_environment in &environments.turn_environments { + for turn_environment in environments.turn_environments() { let filesystem = turn_environment.environment.get_filesystem(); match read_agents_md( config, diff --git a/codex-rs/core/src/agents_md_tests.rs b/codex-rs/core/src/agents_md_tests.rs index 0d98e946965a..8bb628d9da7f 100644 --- a/codex-rs/core/src/agents_md_tests.rs +++ b/codex-rs/core/src/agents_md_tests.rs @@ -2,6 +2,7 @@ use super::*; use crate::config::ConfigBuilder; use crate::context::ContextualUserFragment; use crate::environment_selection::TurnEnvironmentSnapshot; +use crate::environment_selection::TurnEnvironmentState; use crate::session::turn_context::TurnEnvironment; use codex_config::ConfigLayerEntry; use codex_config::ConfigLayerStack; @@ -324,10 +325,10 @@ fn resolved_local_environments( environments: [(&str, AbsolutePathBuf); N], ) -> TurnEnvironmentSnapshot { TurnEnvironmentSnapshot { - turn_environments: environments + environments: environments .into_iter() .map(|(environment_id, cwd)| { - TurnEnvironment::new( + TurnEnvironmentState::Ready(TurnEnvironment::new( environment_id.to_string(), Arc::new( Environment::create_for_tests(/*exec_server_url*/ None) @@ -336,10 +337,9 @@ fn resolved_local_environments( PathUri::from_abs_path(&cwd), Vec::new(), /*shell*/ None, - ) + )) }) .collect(), - starting: Vec::new(), } } diff --git a/codex-rs/core/src/codex_delegate_tests.rs b/codex-rs/core/src/codex_delegate_tests.rs index 5cfbf252d014..6f67902aa1f6 100644 --- a/codex-rs/core/src/codex_delegate_tests.rs +++ b/codex-rs/core/src/codex_delegate_tests.rs @@ -1,4 +1,5 @@ use super::*; +use crate::environment_selection::TurnEnvironmentState; use crate::mcp_tool_call::MCP_TOOL_APPROVAL_DECLINE_SYNTHETIC; use crate::mcp_tool_call::MCP_TOOL_APPROVAL_QUESTION_ID_PREFIX; use async_channel::bounded; @@ -211,7 +212,11 @@ async fn handle_request_permissions_uses_tool_call_id_for_round_trip() { crate::session::tests::make_session_and_context_with_rx().await; *parent_session.active_turn.lock().await = Some(crate::state::ActiveTurn::default()); let parent_ctx_mut = Arc::get_mut(&mut parent_ctx).expect("single turn context ref"); - parent_ctx_mut.environments.turn_environments[0].environment_id = "remote".to_string(); + let TurnEnvironmentState::Ready(environment) = &mut parent_ctx_mut.environments.environments[0] + else { + panic!("expected ready primary environment"); + }; + environment.environment_id = "remote".to_string(); let (tx_sub, rx_sub) = bounded(SUBMISSION_CHANNEL_CAPACITY); let (_tx_events, rx_events_child) = bounded(SUBMISSION_CHANNEL_CAPACITY); diff --git a/codex-rs/core/src/context/world_state/environment.rs b/codex-rs/core/src/context/world_state/environment.rs index 0ed9f2d91c1f..d5a8abf78d32 100644 --- a/codex-rs/core/src/context/world_state/environment.rs +++ b/codex-rs/core/src/context/world_state/environment.rs @@ -323,8 +323,7 @@ enum EnvironmentStatus { fn environment_states(snapshot: &TurnEnvironmentSnapshot) -> BTreeMap { let mut environments = snapshot - .turn_environments - .iter() + .turn_environments() .map(|environment| { ( environment.environment_id.clone(), @@ -339,7 +338,7 @@ fn environment_states(snapshot: &TurnEnvironmentSnapshot) -> BTreeMap>(); - for environment in &snapshot.starting { + for environment in snapshot.starting() { environments .entry(environment.selection.environment_id.clone()) .or_insert_with(|| EnvironmentState { diff --git a/codex-rs/core/src/environment_selection.rs b/codex-rs/core/src/environment_selection.rs index eaf305d7f641..97c9c0d3bdb4 100644 --- a/codex-rs/core/src/environment_selection.rs +++ b/codex-rs/core/src/environment_selection.rs @@ -95,19 +95,22 @@ impl ThreadEnvironments { ) -> Self { // Reuse only attached environments from the supplied snapshot; drop starting entries. let environments = current - .turn_environments + .environments .into_iter() - .map(|environment| { + .filter_map(|environment| { + let TurnEnvironmentState::Ready(environment) = environment else { + return None; + }; let selection = environment.selection(); let selected_environment = Arc::clone(&environment.environment); let resolution: TurnEnvironmentResolution = futures::future::ready(Ok(environment)).boxed().shared(); - SelectedTurnEnvironment { + Some(SelectedTurnEnvironment { selection, environment: selected_environment, connection_events_task: None, resolution, - } + }) }) .collect(); Self { @@ -311,31 +314,25 @@ impl ThreadEnvironments { #[tracing::instrument(name = "environments.snapshot", skip_all)] pub(crate) async fn snapshot(&self) -> TurnEnvironmentSnapshot { - let current = self.environments.load_full(); - let mut turn_environments = Vec::with_capacity(current.len()); - let mut starting = Vec::new(); - for environment in current.iter() { + let selected = self.environments.load_full(); + let mut environments = Vec::with_capacity(selected.len()); + for environment in selected.iter() { let resolved = if self.non_blocking_snapshots { environment.resolution.clone().now_or_never() } else { Some(environment.resolution.clone().await) }; - match resolved { - Some(Ok(turn_environment)) => turn_environments.push(turn_environment), - Some(Err(err)) => tracing::debug!( - environment_id = %environment.selection.environment_id, - "skipping failed turn environment: {err}" - ), - None => starting.push(StartingTurnEnvironment { + if let Some(environment) = TurnEnvironmentState::from_resolution( + StartingTurnEnvironment { selection: environment.selection.clone(), resolution: environment.resolution.clone(), - }), + }, + resolved, + ) { + environments.push(environment); } } - TurnEnvironmentSnapshot { - turn_environments, - starting, - } + TurnEnvironmentSnapshot { environments } } pub(crate) fn environment_manager(&self) -> Arc { @@ -343,17 +340,79 @@ impl ThreadEnvironments { } } +#[derive(Clone, Debug)] +pub(crate) enum TurnEnvironmentState { + Ready(TurnEnvironment), + Starting(StartingTurnEnvironment), +} + +impl TurnEnvironmentState { + fn from_resolution( + starting: StartingTurnEnvironment, + resolved: Option, + ) -> Option { + match resolved { + Some(Ok(environment)) => Some(Self::Ready(environment)), + Some(Err(err)) => { + tracing::debug!( + environment_id = %starting.selection.environment_id, + "skipping failed turn environment: {err}" + ); + None + } + None => Some(Self::Starting(starting)), + } + } +} + #[derive(Clone, Debug, Default)] pub(crate) struct TurnEnvironmentSnapshot { - pub(crate) turn_environments: Vec, - pub(crate) starting: Vec, + // Keep ready and starting environments in their original selection order. + pub(crate) environments: Vec, } impl TurnEnvironmentSnapshot { + /// Promotes completed startup work without adopting newer thread selections. + pub(crate) fn refresh_readiness(&self) -> Self { + let environments = self + .environments + .iter() + .filter_map(|environment| match environment { + TurnEnvironmentState::Ready(environment) => { + Some(TurnEnvironmentState::Ready(environment.clone())) + } + TurnEnvironmentState::Starting(environment) => { + TurnEnvironmentState::from_resolution( + environment.clone(), + environment.resolution.clone().now_or_never(), + ) + } + }) + .collect(); + Self { environments } + } + + pub(crate) fn turn_environments(&self) -> impl Iterator { + self.environments.iter().filter_map(|environment| { + let TurnEnvironmentState::Ready(environment) = environment else { + return None; + }; + Some(environment) + }) + } + + pub(crate) fn starting(&self) -> impl Iterator { + self.environments.iter().filter_map(|environment| { + let TurnEnvironmentState::Starting(environment) = environment else { + return None; + }; + Some(environment) + }) + } + /// Maps each captured environment to its exact ready handle, or `None` when it was starting. pub(crate) fn captured_environments(&self) -> HashMap>> { - self.turn_environments - .iter() + self.turn_environments() .map(|environment| { ( environment.environment_id.clone(), @@ -361,20 +420,18 @@ impl TurnEnvironmentSnapshot { ) }) .chain( - self.starting - .iter() + self.starting() .map(|environment| (environment.selection.environment_id.clone(), None)), ) .collect() } pub(crate) fn primary(&self) -> Option<&TurnEnvironment> { - self.turn_environments.first() + self.turn_environments().next() } pub(crate) fn local(&self) -> Option<&TurnEnvironment> { - self.turn_environments - .iter() + self.turn_environments() .find(|environment| !environment.environment.is_remote()) } @@ -385,8 +442,7 @@ impl TurnEnvironmentSnapshot { } pub(crate) fn to_selections(&self) -> Vec { - self.turn_environments - .iter() + self.turn_environments() .map(TurnEnvironment::selection) .collect() } @@ -397,12 +453,14 @@ impl TurnEnvironmentSnapshot { } pub(crate) fn single_local_environment(&self) -> Option<&TurnEnvironment> { - if !self.starting.is_empty() { + if self.starting().next().is_some() { return None; } - let [environment] = self.turn_environments.as_slice() else { + let mut environments = self.turn_environments(); + let environment = environments.next()?; + if environments.next().is_some() { return None; - }; + } (!environment.environment.is_remote()).then_some(environment) } @@ -751,13 +809,13 @@ url = "ws://127.0.0.1:8765" .expect("snapshot should finish after the environment starts") .expect("snapshot task"); - assert!(snapshot.starting.is_empty()); + assert!(snapshot.starting().next().is_none()); assert_eq!(snapshot.to_selections(), vec![selection]); server.await.expect("server task"); } #[tokio::test] - async fn snapshot_keeps_starting_environment_until_it_can_be_attached() { + async fn snapshot_refreshes_readiness_in_selection_order() { let listener = TcpListener::bind("127.0.0.1:0") .await .expect("bind websocket listener"); @@ -790,36 +848,51 @@ url = "ws://127.0.0.1:8765" TurnEnvironmentSnapshot::default(), /*non_blocking_snapshots*/ true, ); + turn_environments.update_selections(std::slice::from_ref(&local)); + turn_environments.environments.load()[0] + .resolution + .clone() + .await + .expect("local environment should resolve"); turn_environments.update_selections(&[remote.clone(), local.clone()]); let starting = turn_environments.snapshot().await; - assert!(starting.turn_environments.is_empty()); assert_eq!( starting - .starting - .iter() + .turn_environments() + .map(TurnEnvironment::selection) + .collect::>(), + vec![local.clone()] + ); + assert_eq!( + starting + .starting() .map(|environment| environment.selection.clone()) .collect::>(), - vec![remote.clone(), local.clone()] + vec![remote.clone()] ); - assert!(starting.to_selections().is_empty()); + assert_eq!(starting.to_selections(), vec![local.clone()]); assert!(starting.single_local_environment().is_none()); let server = tokio::spawn(serve_environment_info(listener)); timeout( std::time::Duration::from_secs(5), - starting.starting[0].resolution.clone(), + starting + .starting() + .next() + .expect("starting environment") + .resolution + .clone(), ) .await .expect("environment resolution should finish") .expect("environment resolution should succeed"); - let attached = turn_environments.snapshot().await; + let attached = starting.refresh_readiness(); - assert!(attached.starting.is_empty()); + assert!(attached.starting().next().is_none()); assert_eq!( attached - .turn_environments - .iter() + .turn_environments() .map(TurnEnvironment::selection) .collect::>(), vec![remote.clone(), local.clone()] @@ -866,9 +939,10 @@ url = "ws://127.0.0.1:8765" environments.update_selections(std::slice::from_ref(&selection)); let replacement = environments.snapshot().await; - let [replacement] = replacement.starting.as_slice() else { - panic!("expected the replacement environment to be starting"); - }; + let replacement = replacement + .starting() + .next() + .expect("expected the replacement environment to be starting"); assert_eq!(replacement.selection, selection); assert!(!failed_resolution.ptr_eq(&replacement.resolution)); } @@ -930,16 +1004,16 @@ url = "ws://127.0.0.1:8765" let changed_snapshot = environments.snapshot().await; let initial = initial_snapshot - .starting - .first() + .starting() + .next() .expect("initial environment"); let reused = reused_snapshot - .starting - .first() + .starting() + .next() .expect("reused environment"); let changed = changed_snapshot - .starting - .first() + .starting() + .next() .expect("changed environment"); assert!(initial.resolution.ptr_eq(&reused.resolution)); assert!(!reused.resolution.ptr_eq(&changed.resolution)); @@ -1001,8 +1075,7 @@ url = "ws://127.0.0.1:8765" crate::shell::default_user_shell(), ShellSnapshot::disabled(), TurnEnvironmentSnapshot { - turn_environments: vec![inherited], - starting: Vec::new(), + environments: vec![TurnEnvironmentState::Ready(inherited)], }, /*non_blocking_snapshots*/ false, ); @@ -1039,27 +1112,25 @@ url = "ws://127.0.0.1:8765" .expect("remote environment"), ); let remote = TurnEnvironmentSnapshot { - turn_environments: vec![TurnEnvironment::new( + environments: vec![TurnEnvironmentState::Ready(TurnEnvironment::new( REMOTE_ENVIRONMENT_ID.to_string(), remote_environment.clone(), cwd_uri.clone(), Vec::new(), /*shell*/ None, - )], - starting: Vec::new(), + ))], }; let multiple = TurnEnvironmentSnapshot { - turn_environments: vec![ - local.primary().expect("local environment").clone(), - TurnEnvironment::new( + environments: vec![ + TurnEnvironmentState::Ready(local.primary().expect("local environment").clone()), + TurnEnvironmentState::Ready(TurnEnvironment::new( REMOTE_ENVIRONMENT_ID.to_string(), remote_environment, cwd_uri, Vec::new(), /*shell*/ None, - ), + )), ], - starting: Vec::new(), }; assert_eq!(local.single_local_environment_cwd(), Some(cwd)); diff --git a/codex-rs/core/src/mcp_openai_file.rs b/codex-rs/core/src/mcp_openai_file.rs index 02c1587e04ea..59ea2c751dad 100644 --- a/codex-rs/core/src/mcp_openai_file.rs +++ b/codex-rs/core/src/mcp_openai_file.rs @@ -208,6 +208,7 @@ async fn build_uploaded_argument_value( #[cfg(test)] mod tests { use super::*; + use crate::environment_selection::TurnEnvironmentState; use crate::session::tests::make_session_and_context; use crate::session::turn_context::TurnEnvironment; use codex_utils_absolute_path::AbsolutePathBuf; @@ -220,11 +221,10 @@ mod tests { fn set_primary_environment_cwd(turn_context: &mut TurnContext, cwd: &Path) { let cwd = AbsolutePathBuf::try_from(cwd).expect("absolute path"); turn_context.permission_profile = codex_protocol::models::PermissionProfile::Disabled; - let primary = turn_context - .environments - .turn_environments - .first_mut() - .expect("primary environment"); + let TurnEnvironmentState::Ready(primary) = &mut turn_context.environments.environments[0] + else { + panic!("expected ready primary environment"); + }; *primary = TurnEnvironment::new( primary.environment_id.clone(), Arc::clone(&primary.environment), diff --git a/codex-rs/core/src/mcp_tool_call.rs b/codex-rs/core/src/mcp_tool_call.rs index 266d67b4853d..5f9d56d6cd1a 100644 --- a/codex-rs/core/src/mcp_tool_call.rs +++ b/codex-rs/core/src/mcp_tool_call.rs @@ -777,8 +777,7 @@ async fn augment_mcp_tool_request_meta_with_sandbox_state( fn sandbox_cwd_for_mcp_server(step_context: &StepContext, environment_id: &str) -> Option { if let Some(environment) = step_context .environments - .turn_environments - .iter() + .turn_environments() .find(|environment| environment.environment_id == environment_id) { return Some(environment.cwd().clone()); diff --git a/codex-rs/core/src/mcp_tool_call_tests.rs b/codex-rs/core/src/mcp_tool_call_tests.rs index f91901cad3cd..3f74188ca133 100644 --- a/codex-rs/core/src/mcp_tool_call_tests.rs +++ b/codex-rs/core/src/mcp_tool_call_tests.rs @@ -1,6 +1,7 @@ use super::*; use crate::config::ConfigBuilder; use crate::config::ManagedFeatures; +use crate::environment_selection::TurnEnvironmentState; use crate::session::step_context::StepContext; use crate::session::tests::make_session_and_context; use crate::session::tests::make_session_and_context_with_rx; @@ -1158,19 +1159,22 @@ async fn mcp_tool_call_request_meta_includes_turn_started_at_unix_ms() { async fn mcp_sandbox_cwd_uses_matching_server_environment_uri() -> anyhow::Result<()> { let (_, mut turn_context) = make_session_and_context().await; let secondary_cwd = PathUri::parse("file:///C:/remote/project")?; - let environment = turn_context.environments.turn_environments[0] + let environment = turn_context + .environments + .primary() + .expect("primary environment") .environment .clone(); turn_context .environments - .turn_environments - .push(TurnEnvironment::new( + .environments + .push(TurnEnvironmentState::Ready(TurnEnvironment::new( "remote".to_string(), environment, secondary_cwd.clone(), Vec::new(), /*shell*/ None, - )); + ))); let step_context = StepContext::for_test(Arc::new(turn_context)); let sandbox_cwd = sandbox_cwd_for_mcp_server(&step_context, "remote"); diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index eb9cb33684cf..9a1016453429 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -2430,8 +2430,7 @@ impl Session { let turn_environment = match args.environment_id.as_deref() { Some(environment_id) => turn_context .environments - .turn_environments - .iter() + .turn_environments() .find(|environment| environment.environment_id == environment_id), None => turn_context.environments.primary(), }; @@ -2854,9 +2853,9 @@ impl Session { .config .features .enabled(Feature::DeferredExecutor); - // Keep the old turn-frozen environment view unless deferred executors are enabled. let environments = if deferred_executor_enabled { - self.services.turn_environments.snapshot().await + // Keep selections fixed for the turn while allowing their startup work to finish. + turn_context.environments.refresh_readiness() } else { turn_context.environments.clone() }; diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index 02bcad277ac2..889b1d59f4cb 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -8,6 +8,7 @@ use crate::config::test_config; use crate::context::ContextualUserFragment; use crate::context::TurnAborted; use crate::environment_selection::ThreadEnvironments; +use crate::environment_selection::TurnEnvironmentState; use crate::function_tool::FunctionCallError; use crate::session::step_context::StepContext; use crate::shell::default_user_shell; @@ -5072,7 +5073,7 @@ async fn absolute_cwd_update_with_turn_environment_is_allowed() { let turn_cwd = turn_context.cwd.clone(); assert_eq!(turn_cwd, absolute_cwd); assert_eq!(turn_context.config.cwd, absolute_cwd); - assert_eq!(turn_context.environments.turn_environments.len(), 1); + assert_eq!(turn_context.environments.turn_environments().count(), 1); } #[tokio::test] @@ -6055,14 +6056,19 @@ async fn request_permissions_tool_resolves_relative_paths_against_selected_envir mcp_elicitations: true, })) .expect("test setup should allow updating approval policy"); - let current_environment = turn_context_mut.environments.turn_environments[0].clone(); - turn_context_mut.environments.turn_environments[0] = TurnEnvironment::new( - "remote".to_string(), - current_environment.environment, - PathUri::from_abs_path(&environment_cwd), - Vec::new(), - current_environment.shell, - ); + let current_environment = turn_context_mut + .environments + .primary() + .expect("primary environment") + .clone(); + turn_context_mut.environments.environments[0] = + TurnEnvironmentState::Ready(TurnEnvironment::new( + "remote".to_string(), + current_environment.environment, + PathUri::from_abs_path(&environment_cwd), + Vec::new(), + current_environment.shell, + )); let call_id = "call-1".to_string(); let handler = RequestPermissionsHandler; @@ -6581,16 +6587,25 @@ async fn turn_environments_set_primary_environment() { .expect("turn should start"); let turn_environments = &turn_context.environments; - assert_eq!(turn_environments.turn_environments.len(), 1); + assert_eq!(turn_environments.turn_environments().count(), 1); let turn_environment = turn_context .environments .primary() .expect("primary environment should be set"); assert!(std::sync::Arc::ptr_eq( &turn_environment.environment, - &turn_environments.turn_environments[0].environment + &turn_environments + .primary() + .expect("primary environment") + .environment )); - assert!(!turn_context.environments.turn_environments.is_empty()); + assert!( + turn_context + .environments + .turn_environments() + .next() + .is_some() + ); #[allow(deprecated)] let turn_cwd = turn_context.cwd.clone(); assert_eq!(turn_cwd.as_path(), selected_cwd.as_path()); @@ -6639,14 +6654,17 @@ async fn default_turn_does_not_overlay_legacy_fallback_cwd_onto_stored_thread_en let turn_context = session.new_default_turn().await; let turn_environments = &turn_context.environments; - assert_eq!(turn_environments.turn_environments.len(), 1); + assert_eq!(turn_environments.turn_environments().count(), 1); let turn_environment = turn_context .environments .primary() .expect("primary environment should be set"); assert!(std::sync::Arc::ptr_eq( &turn_environment.environment, - &turn_environments.turn_environments[0].environment + &turn_environments + .primary() + .expect("primary environment") + .environment )); #[allow(deprecated)] let turn_cwd = turn_context.cwd.clone(); @@ -6668,31 +6686,41 @@ async fn default_turn_honors_empty_stored_thread_environments() { let turn_context = session.new_default_turn().await; assert!(turn_context.environments.primary().is_none()); - assert!(turn_context.environments.turn_environments.is_empty()); + assert!( + turn_context + .environments + .turn_environments() + .next() + .is_none() + ); #[allow(deprecated)] let turn_cwd = turn_context.cwd.clone(); assert_eq!(turn_cwd, session_cwd); assert_eq!(turn_context.config.cwd, session_cwd); - assert_eq!(turn_context.environments.turn_environments.len(), 0); + assert_eq!(turn_context.environments.turn_environments().count(), 0); } #[tokio::test] async fn primary_environment_uses_first_turn_environment() { let (_session, mut turn_context) = make_session_and_context().await; - let first_environment = turn_context.environments.turn_environments[0].clone(); + let first_environment = turn_context + .environments + .primary() + .expect("primary environment") + .clone(); #[allow(deprecated)] let second_cwd = turn_context.cwd.join("second"); let second_cwd_uri = codex_utils_path_uri::PathUri::from_abs_path(&second_cwd); turn_context .environments - .turn_environments - .push(TurnEnvironment::new( + .environments + .push(TurnEnvironmentState::Ready(TurnEnvironment::new( "second".to_string(), Arc::clone(&first_environment.environment), second_cwd_uri.clone(), Vec::new(), /*shell*/ None, - )); + ))); assert_eq!( turn_context @@ -6705,16 +6733,20 @@ async fn primary_environment_uses_first_turn_environment() { assert_eq!( turn_context .environments - .turn_environments - .iter() + .turn_environments() .find(|environment| environment.environment_id == "second") .expect("second environment") .cwd(), &second_cwd_uri ); - assert_eq!(turn_context.environments.turn_environments.len(), 2); + assert_eq!(turn_context.environments.turn_environments().count(), 2); assert_eq!( - turn_context.environments.turn_environments[1].cwd(), + turn_context + .environments + .turn_environments() + .nth(1) + .expect("second environment") + .cwd(), &second_cwd_uri ); } @@ -6738,7 +6770,13 @@ async fn empty_turn_environments_clear_primary_environment() { .expect("turn should start"); assert!(turn_context.environments.primary().is_none()); - assert!(turn_context.environments.turn_environments.is_empty()); + assert!( + turn_context + .environments + .turn_environments() + .next() + .is_none() + ); #[allow(deprecated)] let turn_cwd = turn_context.cwd.clone(); assert_eq!(turn_cwd, session.get_config().await.cwd); @@ -8018,14 +8056,19 @@ async fn record_context_updates_emits_environment_item_for_cwd_changes() { ) .await; let cwd = test_path_buf("/new-repo").abs(); - let environment = current_context.environments.turn_environments[0].clone(); - current_context.environments.turn_environments[0] = TurnEnvironment::new( - environment.environment_id, - environment.environment, - PathUri::from_abs_path(&cwd), - Vec::new(), - environment.shell, - ); + let environment = current_context + .environments + .primary() + .expect("primary environment") + .clone(); + current_context.environments.environments[0] = + TurnEnvironmentState::Ready(TurnEnvironment::new( + environment.environment_id, + environment.environment, + PathUri::from_abs_path(&cwd), + Vec::new(), + environment.shell, + )); let update_items = record_context_update_items(&session, previous_context, current_context).await; @@ -8078,14 +8121,19 @@ async fn record_context_updates_omits_environment_item_when_disabled() { let mut config = (*current_context.config).clone(); config.include_environment_context = false; current_context.config = Arc::new(config); - let environment = current_context.environments.turn_environments[0].clone(); - current_context.environments.turn_environments[0] = TurnEnvironment::new( - environment.environment_id, - environment.environment, - PathUri::from_abs_path(&test_path_buf("/new-repo").abs()), - Vec::new(), - environment.shell, - ); + let environment = current_context + .environments + .primary() + .expect("primary environment") + .clone(); + current_context.environments.environments[0] = + TurnEnvironmentState::Ready(TurnEnvironment::new( + environment.environment_id, + environment.environment, + PathUri::from_abs_path(&test_path_buf("/new-repo").abs()), + Vec::new(), + environment.shell, + )); let update_items = record_context_update_items(&session, previous_context, current_context).await; @@ -8806,15 +8854,19 @@ fn file_system_policy_with_unreadable_glob(turn_context: &TurnContext) -> FileSy #[tokio::test] async fn turn_context_item_stores_local_cwd() { let (_session, mut turn_context) = make_session_and_context().await; - let environment = turn_context.environments.turn_environments[0].clone(); + let environment = turn_context + .environments + .primary() + .expect("primary environment") + .clone(); let cwd = PathUri::parse("file:///C:/windows").expect("Windows cwd URI"); - turn_context.environments.turn_environments[0] = TurnEnvironment::new( + turn_context.environments.environments[0] = TurnEnvironmentState::Ready(TurnEnvironment::new( "remote".to_string(), environment.environment, cwd, Vec::new(), environment.shell, - ); + )); #[allow(deprecated)] let local_cwd = turn_context.cwd.clone(); diff --git a/codex-rs/core/src/session/turn.rs b/codex-rs/core/src/session/turn.rs index 36930606ab32..72c7cba9e0e7 100644 --- a/codex-rs/core/src/session/turn.rs +++ b/codex-rs/core/src/session/turn.rs @@ -480,7 +480,7 @@ pub(crate) async fn run_turn( #[instrument(level = "trace", skip_all)] async fn turn_diff_display_roots(turn_context: &TurnContext) -> Vec<(String, PathBuf)> { let mut display_roots = Vec::new(); - for turn_environment in &turn_context.environments.turn_environments { + for turn_environment in turn_context.environments.turn_environments() { // TODO(anp): Migrate git-root discovery and diff display roots to PathUri so foreign // environment roots can participate without host-native conversion. let Ok(cwd) = turn_environment.cwd().to_abs_path() else { @@ -711,8 +711,7 @@ async fn build_extension_turn_input_items( let environments = turn_context .environments - .turn_environments - .iter() + .turn_environments() .enumerate() .filter_map(|(index, environment)| { // TODO(anp): Migrate extension turn-input environments to PathUri so foreign cwd diff --git a/codex-rs/core/src/thread_manager_tests.rs b/codex-rs/core/src/thread_manager_tests.rs index df2f00e839eb..3980471611f7 100644 --- a/codex-rs/core/src/thread_manager_tests.rs +++ b/codex-rs/core/src/thread_manager_tests.rs @@ -894,13 +894,21 @@ async fn resume_and_fork_do_not_restore_thread_environments_from_rollout() { .new_turn_with_sub_id("resume-turn".to_string(), SessionSettingsUpdate::default()) .await .expect("build resumed turn context"); - assert_eq!(resumed_turn.environments.turn_environments.len(), 1); + assert_eq!(resumed_turn.environments.turn_environments().count(), 1); assert_eq!( - resumed_turn.environments.turn_environments[0].cwd(), + resumed_turn + .environments + .primary() + .expect("primary environment") + .cwd(), &PathUri::from_abs_path(&default_cwd) ); assert_ne!( - resumed_turn.environments.turn_environments[0].cwd(), + resumed_turn + .environments + .primary() + .expect("primary environment") + .cwd(), &PathUri::from_abs_path(&selected_cwd) ); @@ -920,13 +928,21 @@ async fn resume_and_fork_do_not_restore_thread_environments_from_rollout() { .new_turn_with_sub_id("fork-turn".to_string(), SessionSettingsUpdate::default()) .await .expect("build forked turn context"); - assert_eq!(forked_turn.environments.turn_environments.len(), 1); + assert_eq!(forked_turn.environments.turn_environments().count(), 1); assert_eq!( - forked_turn.environments.turn_environments[0].cwd(), + forked_turn + .environments + .primary() + .expect("primary environment") + .cwd(), &PathUri::from_abs_path(&default_cwd) ); assert_ne!( - forked_turn.environments.turn_environments[0].cwd(), + forked_turn + .environments + .primary() + .expect("primary environment") + .cwd(), &PathUri::from_abs_path(&selected_cwd) ); } diff --git a/codex-rs/core/src/tools/handlers/agent_jobs/spawn_agents_on_csv.rs b/codex-rs/core/src/tools/handlers/agent_jobs/spawn_agents_on_csv.rs index a519e5db1b08..2ad0e22fb8a2 100644 --- a/codex-rs/core/src/tools/handlers/agent_jobs/spawn_agents_on_csv.rs +++ b/codex-rs/core/src/tools/handlers/agent_jobs/spawn_agents_on_csv.rs @@ -300,7 +300,13 @@ pub async fn handle( } fn single_local_environment_cwd(turn: &TurnContext) -> Result { - let [turn_environment] = turn.environments.turn_environments.as_slice() else { + let mut environments = turn.environments.turn_environments(); + let Some(turn_environment) = environments.next() else { + return Err(FunctionCallError::RespondToModel( + "spawn_agents_on_csv requires exactly one local environment".to_string(), + )); + }; + if environments.next().is_some() { return Err(FunctionCallError::RespondToModel( "spawn_agents_on_csv requires exactly one local environment".to_string(), )); diff --git a/codex-rs/core/src/tools/handlers/extension_tools.rs b/codex-rs/core/src/tools/handlers/extension_tools.rs index 13a75bf4c2e6..6abacef2c9d6 100644 --- a/codex-rs/core/src/tools/handlers/extension_tools.rs +++ b/codex-rs/core/src/tools/handlers/extension_tools.rs @@ -124,9 +124,8 @@ async fn to_extension_call(invocation: &ToolInvocation) -> ExtensionToolCall { reasoning_effort: invocation.turn.effective_reasoning_effort(), }) .and_then(|metadata| to_ascii_json_string(&metadata).ok()); - let mut environments = - Vec::with_capacity(invocation.step_context.environments.turn_environments.len()); - for environment in &invocation.step_context.environments.turn_environments { + let mut environments = Vec::new(); + for environment in invocation.step_context.environments.turn_environments() { // TODO(anp): Migrate extension ToolEnvironment and granted-permission lookup to PathUri // so extensions can receive foreign environment cwd values. let Ok(native_cwd) = environment.cwd().to_abs_path() else { @@ -337,8 +336,7 @@ mod tests { let truncation_policy = turn.model_info.truncation_policy.into(); let expected_sandbox_cwds = turn .environments - .turn_environments - .iter() + .turn_environments() .map(|environment| Some(environment.cwd().clone())) .collect::>(); let history_item = ResponseItem::Message { diff --git a/codex-rs/core/src/tools/handlers/mod.rs b/codex-rs/core/src/tools/handlers/mod.rs index 29ad1729793f..2b5559fb27b4 100644 --- a/codex-rs/core/src/tools/handlers/mod.rs +++ b/codex-rs/core/src/tools/handlers/mod.rs @@ -162,8 +162,7 @@ fn resolve_tool_environment<'a>( || Ok(environments.primary()), |environment_id| { environments - .turn_environments - .iter() + .turn_environments() .find(|environment| environment.environment_id == environment_id) .map(Some) .ok_or_else(|| { diff --git a/codex-rs/core/src/tools/handlers/view_image.rs b/codex-rs/core/src/tools/handlers/view_image.rs index 7ee6f53ec5ca..20efe1b5498d 100644 --- a/codex-rs/core/src/tools/handlers/view_image.rs +++ b/codex-rs/core/src/tools/handlers/view_image.rs @@ -244,6 +244,7 @@ impl ToolOutput for ViewImageOutput { #[cfg(test)] mod tests { use super::*; + use crate::environment_selection::TurnEnvironmentState; use crate::session::step_context::StepContext; use crate::session::tests::make_session_and_context; use crate::session::turn_context::TurnEnvironment; @@ -262,17 +263,17 @@ mod tests { fn replace_primary_environment_cwd(turn: &mut crate::TurnContext, cwd: AbsolutePathBuf) { let current = turn .environments - .turn_environments - .first() + .turn_environments() + .next() .cloned() .expect("default local turn environment"); - turn.environments.turn_environments[0] = TurnEnvironment::new( + turn.environments.environments[0] = TurnEnvironmentState::Ready(TurnEnvironment::new( current.environment_id, current.environment, PathUri::from_abs_path(&cwd), Vec::new(), current.shell, - ); + )); } #[test] diff --git a/codex-rs/core/src/tools/handlers/wait_for_environment.rs b/codex-rs/core/src/tools/handlers/wait_for_environment.rs index 8cdaf34aa3af..6bcb6f3c62aa 100644 --- a/codex-rs/core/src/tools/handlers/wait_for_environment.rs +++ b/codex-rs/core/src/tools/handlers/wait_for_environment.rs @@ -72,14 +72,12 @@ impl ToolExecutor for WaitForEnvironmentHandler { let environment_id = args.environment_id; let already_ready = step_context .environments - .turn_environments - .iter() + .turn_environments() .any(|environment| environment.environment_id == environment_id); if !already_ready { let Some(environment) = step_context .environments - .starting - .iter() + .starting() .find(|environment| environment.selection.environment_id == environment_id) .cloned() else { diff --git a/codex-rs/core/src/tools/network_approval.rs b/codex-rs/core/src/tools/network_approval.rs index 9a2edfb3c59a..73f3d3b15fbc 100644 --- a/codex-rs/core/src/tools/network_approval.rs +++ b/codex-rs/core/src/tools/network_approval.rs @@ -666,8 +666,7 @@ impl NetworkApprovalService { } else { turn_context .environments - .turn_environments - .iter() + .turn_environments() .find(|environment| environment.environment_id == environment_id) .and_then(|environment| environment.cwd().to_abs_path().ok()) .unwrap_or_else(|| { diff --git a/codex-rs/core/src/tools/spec_plan.rs b/codex-rs/core/src/tools/spec_plan.rs index 347b70935f81..bd7db25737f6 100644 --- a/codex-rs/core/src/tools/spec_plan.rs +++ b/codex-rs/core/src/tools/spec_plan.rs @@ -629,7 +629,7 @@ fn standalone_web_search_enabled(turn_context: &TurnContext) -> bool { } fn tool_environment_mode(step_context: &StepContext) -> ToolEnvironmentMode { - ToolEnvironmentMode::from_count(step_context.environments.turn_environments.len()) + ToolEnvironmentMode::from_count(step_context.environments.turn_environments().count()) } #[instrument(level = "trace", skip_all)] @@ -685,8 +685,7 @@ fn unified_exec_should_include_shell_parameter( UnifiedExecShellMode::ZshFork(_) ) || step_context .environments - .turn_environments - .iter() + .turn_environments() .any(|environment| environment.environment.is_remote()) } diff --git a/codex-rs/core/src/tools/spec_plan_tests.rs b/codex-rs/core/src/tools/spec_plan_tests.rs index 3f9fa3bc1b57..acc08cb3a491 100644 --- a/codex-rs/core/src/tools/spec_plan_tests.rs +++ b/codex-rs/core/src/tools/spec_plan_tests.rs @@ -31,6 +31,7 @@ use pretty_assertions::assert_eq; use serde_json::json; use crate::config::CurrentTimeReminderConfig; +use crate::environment_selection::TurnEnvironmentState; use crate::session::step_context::StepContext; use crate::session::tests::make_session_and_context; use crate::session::turn_context::TurnContext; @@ -341,9 +342,15 @@ impl ToolExecutor for DeferredExtensionTool { } fn duplicate_primary_environment(turn: &mut TurnContext) { - let mut second_environment = turn.environments.turn_environments[0].clone(); + let mut second_environment = turn + .environments + .primary() + .expect("primary environment") + .clone(); second_environment.environment_id = "secondary".to_string(); - turn.environments.turn_environments.push(second_environment); + turn.environments + .environments + .push(TurnEnvironmentState::Ready(second_environment)); } fn mcp_tool(server: &str, namespace: &str, name: &str) -> ToolInfo { @@ -587,20 +594,22 @@ async fn zsh_fork_unified_exec_keeps_shell_parameter_when_remote_environment_ava .expect("primary environment") .cwd() .clone(); - turn.environments.turn_environments.push( - crate::session::turn_context::TurnEnvironment::new( - "remote".to_string(), - Arc::new( - codex_exec_server::Environment::create_for_tests(Some( - "ws://127.0.0.1:1/remote-exec-server".to_string(), - )) - .expect("remote test environment"), + turn.environments + .environments + .push(TurnEnvironmentState::Ready( + crate::session::turn_context::TurnEnvironment::new( + "remote".to_string(), + Arc::new( + codex_exec_server::Environment::create_for_tests(Some( + "ws://127.0.0.1:1/remote-exec-server".to_string(), + )) + .expect("remote test environment"), + ), + remote_cwd, + Vec::new(), + /*shell*/ None, ), - remote_cwd, - Vec::new(), - /*shell*/ None, - ), - ); + )); }) .await; @@ -615,7 +624,7 @@ async fn zsh_fork_unified_exec_keeps_shell_parameter_when_remote_environment_ava #[tokio::test] async fn environment_count_controls_environment_backed_tools() { let no_environment = probe(|turn| { - turn.environments.turn_environments.clear(); + turn.environments.environments.clear(); set_feature(turn, Feature::ShellTool, /*enabled*/ true); set_feature(turn, Feature::RequestPermissionsTool, /*enabled*/ true); turn.model_info.apply_patch_tool_type = Some(ApplyPatchToolType::Freeform); @@ -670,7 +679,7 @@ async fn environment_tools_follow_the_step_context() { turn.model_info.apply_patch_tool_type = Some(ApplyPatchToolType::Freeform); let environments = turn.environments.clone(); - turn.environments.turn_environments.clear(); + turn.environments.environments.clear(); let turn = Arc::new(turn); let step_context = Arc::new(StepContext::new( Arc::clone(&turn), diff --git a/codex-rs/core/src/unified_exec/mod_tests.rs b/codex-rs/core/src/unified_exec/mod_tests.rs index 7e3cf8083aa8..389164a43913 100644 --- a/codex-rs/core/src/unified_exec/mod_tests.rs +++ b/codex-rs/core/src/unified_exec/mod_tests.rs @@ -1,6 +1,7 @@ use super::head_tail_buffer::HeadTailBuffer; use super::*; use crate::codex_thread::BackgroundTerminalInfo; +use crate::environment_selection::TurnEnvironmentState; use crate::exec::ExecCapturePolicy; use crate::exec::ExecExpiration; use crate::sandboxing::ExecRequest; @@ -891,8 +892,10 @@ async fn remote_exec_server_rejects_inherited_fd_launches() -> anyhow::Result<() let remote_test_env = remote_test_env().await?; let (_, mut turn) = make_session_and_context().await; - turn.environments.turn_environments[0].environment = - Arc::new(remote_test_env.environment().clone()); + let TurnEnvironmentState::Ready(environment) = &mut turn.environments.environments[0] else { + panic!("expected ready primary environment"); + }; + environment.environment = Arc::new(remote_test_env.environment().clone()); #[allow(deprecated)] let cwd = turn.cwd.clone(); diff --git a/codex-rs/core/tests/suite/remote_env.rs b/codex-rs/core/tests/suite/remote_env.rs index cdeaef1832cd..6af56301bab3 100644 --- a/codex-rs/core/tests/suite/remote_env.rs +++ b/codex-rs/core/tests/suite/remote_env.rs @@ -64,6 +64,7 @@ use core_test_support::responses::start_mock_server; use core_test_support::skip_if_no_network; use core_test_support::skip_if_no_remote_env; use core_test_support::skip_if_target_windows; +use core_test_support::submit_thread_settings; use core_test_support::test_codex::TestCodex; use core_test_support::test_codex::local; use core_test_support::test_codex::test_codex; @@ -367,6 +368,130 @@ async fn deferred_executor_does_not_duplicate_initial_environment_context() -> R Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn settings_update_does_not_retarget_active_turn_environment() -> Result<()> { + let server = start_mock_server().await; + let response_mock = mount_sse_sequence( + &server, + vec![ + sse(vec![ + ev_response_created("resp-1"), + ev_function_call( + "pause-turn", + "request_user_input", + &json!({ + "questions": [{ + "id": "continue", + "header": "Continue", + "question": "Continue after settings update?", + "options": [{ + "label": "Yes (Recommended)", + "description": "Continue the test." + }, { + "label": "No", + "description": "Stop the test." + }] + }] + }) + .to_string(), + ), + ev_completed("resp-1"), + ]), + sse(vec![ + ev_response_created("resp-2"), + ev_assistant_message("msg-2", "first turn done"), + ev_completed("resp-2"), + ]), + sse(vec![ + ev_response_created("resp-3"), + ev_assistant_message("msg-3", "second turn done"), + ev_completed("resp-3"), + ]), + ], + ) + .await; + let mut builder = test_codex().with_config(|config| { + assert!(config.features.enable(Feature::DeferredExecutor).is_ok()); + assert!( + config + .features + .enable(Feature::DefaultModeRequestUserInput) + .is_ok() + ); + }); + let test = builder.build(&server).await?; + let initial_cwd = test.config.cwd.clone(); + let next_workspace = TempDir::new()?; + let next_cwd = next_workspace.path().abs(); + + test.codex + .submit(Op::UserInput { + items: vec![UserInput::Text { + text: "pause before continuing".into(), + text_elements: Vec::new(), + }], + final_output_json_schema: None, + responsesapi_client_metadata: None, + additional_context: Default::default(), + thread_settings: Default::default(), + }) + .await?; + let request = wait_for_event_match(&test.codex, |event| match event { + EventMsg::RequestUserInput(request) => Some(request.clone()), + _ => None, + }) + .await; + + submit_thread_settings( + &test.codex, + ThreadSettingsOverrides { + environments: Some(TurnEnvironmentSelections::new( + next_cwd.clone(), + vec![local(next_cwd.clone())], + )), + ..Default::default() + }, + ) + .await?; + test.codex + .submit(Op::UserInputAnswer { + id: request.turn_id, + response: RequestUserInputResponse { + answers: HashMap::from([( + "continue".to_string(), + RequestUserInputAnswer { + answers: vec!["Yes (Recommended)".to_string()], + }, + )]), + }, + }) + .await?; + wait_for_event(&test.codex, |event| { + matches!(event, EventMsg::TurnComplete(_)) + }) + .await; + test.submit_turn("start the next turn").await?; + + let request_texts = response_mock + .requests() + .iter() + .map(|request| request.message_input_texts("user").join("\n")) + .collect::>(); + let initial_cwd = format!("{}", initial_cwd.as_path().display()); + let next_cwd = format!("{}", next_cwd.as_path().display()); + assert_eq!( + request_texts + .iter() + .map(|text| text.contains(&next_cwd)) + .collect::>(), + vec![false, false, true] + ); + assert!(request_texts[0].contains(&initial_cwd)); + assert!(request_texts[1].contains(&initial_cwd)); + + Ok(()) +} + async fn read_exec_server_json(websocket: &mut WebSocketStream) -> Value { loop { match timeout(Duration::from_secs(5), websocket.next())