From b4e6758e756123a7c2ee52be888ab6e6a544dbff Mon Sep 17 00:00:00 2001 From: pakrym-oai Date: Fri, 12 Jun 2026 14:26:48 -0700 Subject: [PATCH 01/14] core: retain resolved turn environments --- codex-rs/core/src/agents_md.rs | 4 +- codex-rs/core/src/agents_md_tests.rs | 8 +- codex-rs/core/src/codex_delegate.rs | 3 +- codex-rs/core/src/environment_selection.rs | 195 ++++++++++++------ codex-rs/core/src/mcp_tool_call_tests.rs | 11 +- codex-rs/core/src/session/mcp.rs | 4 +- codex-rs/core/src/session/mod.rs | 20 +- codex-rs/core/src/session/session.rs | 52 ++--- codex-rs/core/src/session/tests.rs | 109 ++++++---- .../core/src/session/tests/guardian_tests.rs | 6 +- codex-rs/core/src/session/turn_context.rs | 103 +++++---- codex-rs/core/src/state/service.rs | 4 - codex-rs/core/src/state/session.rs | 9 +- codex-rs/core/src/state/session_tests.rs | 30 +-- codex-rs/core/src/thread_manager.rs | 10 +- 15 files changed, 328 insertions(+), 240 deletions(-) diff --git a/codex-rs/core/src/agents_md.rs b/codex-rs/core/src/agents_md.rs index d63342bdded4..107bc3145aa7 100644 --- a/codex-rs/core/src/agents_md.rs +++ b/codex-rs/core/src/agents_md.rs @@ -18,7 +18,7 @@ use crate::config::Config; use crate::context::ContextualUserFragment; use crate::context::UserInstructions as ContextUserInstructions; -use crate::environment_selection::ResolvedTurnEnvironments; +use crate::environment_selection::TurnEnvironments; use codex_app_server_protocol::ConfigLayerSource; use codex_config::ConfigLayerStackOrdering; use codex_config::default_project_root_markers; @@ -48,7 +48,7 @@ const AGENTS_MD_SEPARATOR: &str = "\n\n--- project-doc ---\n\n"; pub(crate) async fn load_project_instructions( config: &mut Config, user_instructions: Option, - environments: &ResolvedTurnEnvironments, + environments: &TurnEnvironments, ) -> Option { let mut loaded = LoadedAgentsMd::from_user_instructions(user_instructions); for turn_environment in &environments.turn_environments { diff --git a/codex-rs/core/src/agents_md_tests.rs b/codex-rs/core/src/agents_md_tests.rs index 8d0987ca4dcd..1a49af30d761 100644 --- a/codex-rs/core/src/agents_md_tests.rs +++ b/codex-rs/core/src/agents_md_tests.rs @@ -1,6 +1,6 @@ use super::*; use crate::config::ConfigBuilder; -use crate::environment_selection::ResolvedTurnEnvironments; +use crate::environment_selection::TurnEnvironments; use crate::session::turn_context::TurnEnvironment; use codex_config::ConfigLayerEntry; use codex_config::ConfigLayerStack; @@ -9,6 +9,7 @@ use codex_config::ConfigRequirementsToml; use codex_exec_server::CopyOptions; use codex_exec_server::CreateDirectoryOptions; use codex_exec_server::Environment; +use codex_exec_server::EnvironmentManager; use codex_exec_server::ExecutorFileSystemFuture; use codex_exec_server::FileMetadata; use codex_exec_server::FileSystemSandboxContext; @@ -254,8 +255,9 @@ async fn agents_md_paths(config: &TestConfig) -> std::io::Result( environments: [(&str, AbsolutePathBuf); N], -) -> ResolvedTurnEnvironments { - ResolvedTurnEnvironments { +) -> TurnEnvironments { + TurnEnvironments { + environment_manager: Arc::new(EnvironmentManager::default_for_tests()), turn_environments: environments .into_iter() .map(|(environment_id, cwd)| { diff --git a/codex-rs/core/src/codex_delegate.rs b/codex-rs/core/src/codex_delegate.rs index 135909bf0037..2e38c4e7d342 100644 --- a/codex-rs/core/src/codex_delegate.rs +++ b/codex-rs/core/src/codex_delegate.rs @@ -90,7 +90,6 @@ pub(crate) async fn run_codex_thread_interactive( installation_id: parent_session.installation_id.clone(), auth_manager, models_manager, - environment_manager: Arc::clone(&parent_session.services.environment_manager), skills_manager: Arc::clone(&parent_session.services.skills_manager), plugins_manager: Arc::clone(&parent_session.services.plugins_manager), mcp_manager: Arc::clone(&parent_session.services.mcp_manager), @@ -108,7 +107,7 @@ pub(crate) async fn run_codex_thread_interactive( inherited_exec_policy: Some(Arc::clone(&parent_session.services.exec_policy)), parent_rollout_thread_trace: codex_rollout_trace::ThreadTraceContext::disabled(), parent_trace: None, - environment_selections: parent_ctx.environments.clone(), + turn_environments: parent_ctx.environments.clone(), thread_extension_init: codex_extension_api::ExtensionDataInit::default(), analytics_events_client: Some(parent_session.services.analytics_events_client.clone()), thread_store: Arc::clone(&parent_session.services.thread_store), diff --git a/codex-rs/core/src/environment_selection.rs b/codex-rs/core/src/environment_selection.rs index 31c7494cae26..57c1fb90a8f5 100644 --- a/codex-rs/core/src/environment_selection.rs +++ b/codex-rs/core/src/environment_selection.rs @@ -26,58 +26,60 @@ pub(crate) fn default_thread_environment_selections( .collect() } -#[derive(Clone, Debug, Default)] -pub(crate) struct ResolvedTurnEnvironments { +#[derive(Clone, Debug)] +pub(crate) struct TurnEnvironments { + pub(crate) environment_manager: Arc, pub(crate) turn_environments: Vec, } -impl ResolvedTurnEnvironments { - pub(crate) fn to_selections(&self) -> Vec { - self.turn_environments - .iter() - .map(TurnEnvironment::selection) - .collect() - } - - pub(crate) fn primary(&self) -> Option<&TurnEnvironment> { - self.turn_environments.first() - } - - #[cfg(test)] - pub(crate) fn primary_environment(&self) -> Option> { - self.primary() - .map(|environment| Arc::clone(&environment.environment)) - } - - pub(crate) fn primary_filesystem(&self) -> Option> { - self.primary() - .map(|environment| environment.environment.get_filesystem()) +impl TurnEnvironments { + pub(crate) async fn resolve( + environment_manager: Arc, + environments: &[TurnEnvironmentSelection], + ) -> CodexResult { + Self { + environment_manager, + turn_environments: Vec::new(), + } + .with_selections(environments) + .await } - pub(crate) fn single_local_environment_cwd(&self) -> Option<&AbsolutePathBuf> { - let [environment] = self.turn_environments.as_slice() else { - return None; - }; - - (!environment.environment.is_remote()).then_some(environment.cwd()) + pub(crate) async fn with_selections( + &self, + environments: &[TurnEnvironmentSelection], + ) -> CodexResult { + let mut seen_environment_ids = HashSet::with_capacity(environments.len()); + let mut turn_environments = Vec::with_capacity(environments.len()); + for selected_environment in environments { + if !seen_environment_ids.insert(selected_environment.environment_id.as_str()) { + return Err(CodexErr::InvalidRequest(format!( + "duplicate turn environment id `{}`", + selected_environment.environment_id + ))); + } + let turn_environment = match self.turn_environments.iter().find(|environment| { + environment.environment_id == selected_environment.environment_id + && environment.cwd_uri() == &selected_environment.cwd + }) { + Some(environment) => environment.clone(), + None => self.resolve_selection(selected_environment).await?, + }; + turn_environments.push(turn_environment); + } + Ok(Self { + environment_manager: Arc::clone(&self.environment_manager), + turn_environments, + }) } -} -pub(crate) async fn resolve_environment_selections( - environment_manager: &EnvironmentManager, - environments: &[TurnEnvironmentSelection], -) -> CodexResult { - let mut seen_environment_ids = HashSet::with_capacity(environments.len()); - let mut turn_environments = Vec::with_capacity(environments.len()); - for selected_environment in environments { - if !seen_environment_ids.insert(selected_environment.environment_id.as_str()) { - return Err(CodexErr::InvalidRequest(format!( - "duplicate turn environment id `{}`", - selected_environment.environment_id - ))); - } + async fn resolve_selection( + &self, + selected_environment: &TurnEnvironmentSelection, + ) -> CodexResult { let environment_id = selected_environment.environment_id.clone(); - let environment = environment_manager + let environment = self + .environment_manager .get_environment(&environment_id) .ok_or_else(|| { CodexErr::InvalidRequest(format!("unknown turn environment id `{environment_id}`")) @@ -97,7 +99,7 @@ pub(crate) async fn resolve_environment_selections( None } }; - turn_environments.push(TurnEnvironment::new( + Ok(TurnEnvironment::new( environment_id, environment, selected_environment.cwd.to_abs_path().map_err(|err| { @@ -107,9 +109,38 @@ pub(crate) async fn resolve_environment_selections( )) })?, shell, - )); + )) + } + + pub(crate) fn to_selections(&self) -> Vec { + self.turn_environments + .iter() + .map(TurnEnvironment::selection) + .collect() + } + + pub(crate) fn primary(&self) -> Option<&TurnEnvironment> { + self.turn_environments.first() + } + + #[cfg(test)] + pub(crate) fn primary_environment(&self) -> Option> { + self.primary() + .map(|environment| Arc::clone(&environment.environment)) + } + + pub(crate) fn primary_filesystem(&self) -> Option> { + self.primary() + .map(|environment| environment.environment.get_filesystem()) + } + + pub(crate) fn single_local_environment_cwd(&self) -> Option<&AbsolutePathBuf> { + let [environment] = self.turn_environments.as_slice() else { + return None; + }; + + (!environment.environment.is_remote()).then_some(environment.cwd()) } - Ok(ResolvedTurnEnvironments { turn_environments }) } #[cfg(test)] @@ -201,10 +232,10 @@ url = "ws://127.0.0.1:8765" async fn resolve_environment_selections_rejects_duplicate_ids() { let cwd = AbsolutePathBuf::current_dir().expect("cwd"); let cwd_uri = PathUri::from_abs_path(&cwd); - let manager = EnvironmentManager::default_for_tests(); + let manager = Arc::new(EnvironmentManager::default_for_tests()); - let err = resolve_environment_selections( - &manager, + let err = TurnEnvironments::resolve( + manager, &[ TurnEnvironmentSelection { environment_id: "local".to_string(), @@ -227,10 +258,10 @@ url = "ws://127.0.0.1:8765" let cwd = AbsolutePathBuf::current_dir().expect("cwd"); let selected_cwd = cwd.join("selected"); let selected_cwd_uri = PathUri::from_abs_path(&selected_cwd); - let manager = EnvironmentManager::default_for_tests(); + let manager = Arc::new(EnvironmentManager::default_for_tests()); - let resolved = resolve_environment_selections( - &manager, + let resolved = TurnEnvironments::resolve( + Arc::clone(&manager), &[TurnEnvironmentSelection { environment_id: "local".to_string(), cwd: selected_cwd_uri, @@ -263,13 +294,59 @@ url = "ws://127.0.0.1:8765" ); } + #[tokio::test] + async fn matching_environment_id_and_cwd_reuse_resolved_environment() { + let cwd = AbsolutePathBuf::current_dir().expect("cwd"); + let manager = Arc::new( + EnvironmentManager::create_for_tests( + Some("ws://127.0.0.1:8765".to_string()), + Some(test_runtime_paths()), + ) + .await, + ); + let selection = TurnEnvironmentSelection { + environment_id: REMOTE_ENVIRONMENT_ID.to_string(), + cwd: PathUri::from_abs_path(&cwd), + }; + let initial = TurnEnvironments::resolve(Arc::clone(&manager), &[selection.clone()]) + .await + .expect("environment selection should resolve"); + manager + .upsert_environment( + REMOTE_ENVIRONMENT_ID.to_string(), + "ws://127.0.0.1:9876".to_string(), + ) + .expect("replace environment"); + + let reused = initial + .with_selections(std::slice::from_ref(&selection)) + .await + .expect("matching environment selection should resolve"); + let changed = reused + .with_selections(&[TurnEnvironmentSelection { + cwd: PathUri::from_abs_path(&cwd.join("changed")), + ..selection + }]) + .await + .expect("changed environment selection should resolve"); + + assert!(Arc::ptr_eq( + &initial.primary().expect("initial environment").environment, + &reused.primary().expect("reused environment").environment, + )); + assert!(!Arc::ptr_eq( + &reused.primary().expect("reused environment").environment, + &changed.primary().expect("changed environment").environment, + )); + } + #[tokio::test] async fn single_local_environment_cwd_requires_exactly_one_local_environment() { let cwd = AbsolutePathBuf::current_dir().expect("cwd"); let cwd_uri = PathUri::from_abs_path(&cwd); - let local_manager = EnvironmentManager::default_for_tests(); - let local = resolve_environment_selections( - &local_manager, + let local_manager = Arc::new(EnvironmentManager::default_for_tests()); + let local = TurnEnvironments::resolve( + Arc::clone(&local_manager), &[TurnEnvironmentSelection { environment_id: LOCAL_ENVIRONMENT_ID.to_string(), cwd: cwd_uri, @@ -281,7 +358,8 @@ url = "ws://127.0.0.1:8765" Environment::create_for_tests(Some("ws://127.0.0.1:8765".to_string())) .expect("remote environment"), ); - let remote = ResolvedTurnEnvironments { + let remote = TurnEnvironments { + environment_manager: Arc::clone(&local_manager), turn_environments: vec![TurnEnvironment::new( REMOTE_ENVIRONMENT_ID.to_string(), remote_environment.clone(), @@ -289,7 +367,8 @@ url = "ws://127.0.0.1:8765" /*shell*/ None, )], }; - let multiple = ResolvedTurnEnvironments { + let multiple = TurnEnvironments { + environment_manager: local_manager, turn_environments: vec![ local.primary().expect("local environment").clone(), TurnEnvironment::new( diff --git a/codex-rs/core/src/mcp_tool_call_tests.rs b/codex-rs/core/src/mcp_tool_call_tests.rs index b413704cd54b..bb8d714b5baf 100644 --- a/codex-rs/core/src/mcp_tool_call_tests.rs +++ b/codex-rs/core/src/mcp_tool_call_tests.rs @@ -1267,10 +1267,13 @@ async fn install_host_owned_codex_apps_manager(session: &Session, turn_context: session.get_tx_event(), CancellationToken::new(), turn_context.permission_profile(), - codex_mcp::McpRuntimeContext::new(Arc::clone(&session.services.environment_manager), { - #[allow(deprecated)] - turn_context.cwd.to_path_buf() - }), + codex_mcp::McpRuntimeContext::new( + Arc::clone(&turn_context.environments.environment_manager), + { + #[allow(deprecated)] + turn_context.cwd.to_path_buf() + }, + ), turn_context.config.codex_home.to_path_buf(), codex_mcp::codex_apps_tools_cache_key(auth.as_ref()), /*host_owned_codex_apps_enabled*/ true, diff --git a/codex-rs/core/src/session/mcp.rs b/codex-rs/core/src/session/mcp.rs index c033511a759b..0427900024c2 100644 --- a/codex-rs/core/src/session/mcp.rs +++ b/codex-rs/core/src/session/mcp.rs @@ -319,11 +319,11 @@ impl Session { .await; let mcp_runtime_context = match turn_context.environments.primary() { Some(turn_environment) => McpRuntimeContext::new( - Arc::clone(&self.services.environment_manager), + Arc::clone(&turn_context.environments.environment_manager), turn_environment.cwd().to_path_buf(), ), None => McpRuntimeContext::new( - Arc::clone(&self.services.environment_manager), + Arc::clone(&turn_context.environments.environment_manager), #[allow(deprecated)] turn_context.cwd.to_path_buf(), ), diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index ed9e58113ea9..641f0c52ce93 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -31,7 +31,7 @@ use crate::context::NetworkRuleSaved; use crate::context::PermissionsInstructions; use crate::context::PersonalitySpecInstructions; use crate::default_skill_metadata_budget; -use crate::environment_selection::ResolvedTurnEnvironments; +use crate::environment_selection::TurnEnvironments; use crate::exec_policy::ExecPolicyManager; use crate::image_preparation::prepare_response_items; use crate::parse_turn_item; @@ -53,7 +53,6 @@ use codex_app_server_protocol::McpServerElicitationRequestParams; use codex_config::types::AuthKeyringBackendKind; use codex_config::types::OAuthCredentialsStoreMode; use codex_exec_server::Environment; -use codex_exec_server::EnvironmentManager; use codex_exec_server::FileSystemSandboxContext; use codex_extension_api::ExtensionDataInit; use codex_extension_api::LoadedUserInstructions; @@ -408,7 +407,6 @@ pub(crate) struct CodexSpawnArgs { pub(crate) installation_id: String, pub(crate) auth_manager: Arc, pub(crate) models_manager: SharedModelsManager, - pub(crate) environment_manager: Arc, pub(crate) skills_manager: Arc, pub(crate) plugins_manager: Arc, pub(crate) mcp_manager: Arc, @@ -430,7 +428,7 @@ pub(crate) struct CodexSpawnArgs { pub(crate) parent_rollout_thread_trace: ThreadTraceContext, pub(crate) user_shell_override: Option, pub(crate) parent_trace: Option, - pub(crate) environment_selections: ResolvedTurnEnvironments, + pub(crate) turn_environments: TurnEnvironments, pub(crate) thread_extension_init: ExtensionDataInit, pub(crate) analytics_events_client: Option, pub(crate) thread_store: Arc, @@ -494,7 +492,6 @@ impl Codex { installation_id, auth_manager, models_manager, - environment_manager, skills_manager, plugins_manager, mcp_manager, @@ -512,7 +509,7 @@ impl Codex { inherited_exec_policy, parent_rollout_thread_trace, parent_trace: _, - environment_selections, + turn_environments, thread_extension_init, analytics_events_client, thread_store, @@ -531,8 +528,7 @@ impl Codex { .startup_warnings .extend(user_instruction_provider_warnings); let loaded_agents_md = - load_project_instructions(&mut config, user_instructions, &environment_selections) - .await; + load_project_instructions(&mut config, user_instructions, &turn_environments).await; let exec_policy = if crate::guardian::is_guardian_reviewer_source(&session_source) { // Guardian review should rely on the built-in shell safety checks, @@ -622,7 +618,7 @@ impl Codex { windows_sandbox_level: WindowsSandboxLevel::from_config(&config), environments: TurnEnvironmentSelections::new( config.cwd.clone(), - environment_selections.to_selections(), + turn_environments.to_selections(), ), workspace_roots: config.workspace_roots.clone(), codex_home: config.codex_home.clone(), @@ -661,7 +657,7 @@ impl Codex { extensions, thread_extension_init, agent_control, - environment_manager, + turn_environments, analytics_events_client, thread_store, parent_rollout_thread_trace, @@ -1431,6 +1427,7 @@ impl Session { &self, updates: SessionSettingsUpdate, ) -> ConstraintResult<()> { + let updated_turn_environments = self.turn_environments_for_update(&updates).await; let notify_config_contributors = !self.services.extensions.config_contributors().is_empty(); let ( previous_config, @@ -1463,6 +1460,9 @@ impl Session { let codex_home = updated.codex_home.clone(); let session_source = updated.session_source.clone(); state.session_configuration = updated; + if let Some(turn_environments) = updated_turn_environments { + state.turn_environments = turn_environments; + } ( previous_config, new_config, diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index be3b3e3a33dc..479209346d43 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -435,18 +435,11 @@ pub(crate) struct AppServerClientMetadata { async fn warm_plugins_and_skills_for_session_init( config: Arc, - environment_manager: Arc, plugins_manager: Arc, skills_manager: Arc, - environments: Vec, + turn_environments: TurnEnvironments, ) -> Vec { - let fs = crate::environment_selection::resolve_environment_selections( - environment_manager.as_ref(), - &environments, - ) - .await - .ok() - .and_then(|resolved| resolved.primary_filesystem()); + let fs = turn_environments.primary_filesystem(); let plugins_input = config.plugins_config_input(); let plugin_outcome = plugins_manager.plugins_for_config(&plugins_input).await; let effective_skill_roots = plugin_outcome.effective_plugin_skill_roots(); @@ -487,7 +480,7 @@ impl Session { extensions: Arc>, thread_extension_init: ExtensionDataInit, agent_control: AgentControl, - environment_manager: Arc, + turn_environments: TurnEnvironments, analytics_events_client: Option, thread_store: Arc, parent_rollout_thread_trace: ThreadTraceContext, @@ -627,10 +620,9 @@ impl Session { let plugin_and_skill_warmup_fut = warm_plugins_and_skills_for_session_init( Arc::clone(&config), - Arc::clone(&environment_manager), Arc::clone(&plugins_manager), Arc::clone(&skills_manager), - session_configuration.environment_selections().to_vec(), + turn_environments.clone(), ) .instrument(info_span!( "session_init.plugin_skill_warmup", @@ -855,7 +847,7 @@ impl Session { session_configuration.thread_name = thread_name.clone(); validate_config_lock_if_configured(&session_configuration).await?; export_config_lock_if_configured(&session_configuration, thread_id).await?; - let state = SessionState::new(session_configuration.clone()); + let state = SessionState::new(session_configuration.clone(), turn_environments); let managed_network_requirements_configured = config .config_layer_stack .requirements_toml() @@ -1034,7 +1026,6 @@ impl Session { ), code_mode_service: crate::tools::code_mode::CodeModeService::new(), tool_search_handler_cache: Default::default(), - environment_manager, }; let (out_of_band_elicitation_paused, _out_of_band_elicitation_paused_rx) = watch::channel(false); @@ -1116,28 +1107,17 @@ impl Session { *cancel_guard = cancel_token.clone(); cancel_token }; - let turn_environment = crate::environment_selection::resolve_environment_selections( - sess.services.environment_manager.as_ref(), - session_configuration.environment_selections(), - ) - .await - .map_err(|err| { - CodexErr::InvalidRequest(err.to_string().replace( - "unknown turn environment id", - "unknown stored MCP environment id", - )) - })? - .primary() - .cloned(); - let mcp_runtime_context = match turn_environment { - Some(turn_environment) => McpRuntimeContext::new( - Arc::clone(&sess.services.environment_manager), - turn_environment.cwd().to_path_buf(), - ), - None => McpRuntimeContext::new( - Arc::clone(&sess.services.environment_manager), - session_configuration.cwd().to_path_buf(), - ), + let mcp_runtime_context = { + let state = sess.state.lock().await; + let cwd = state + .turn_environments + .primary() + .map(|turn_environment| turn_environment.cwd().to_path_buf()) + .unwrap_or_else(|| session_configuration.cwd().to_path_buf()); + McpRuntimeContext::new( + Arc::clone(&state.turn_environments.environment_manager), + cwd, + ) }; let mcp_connection_manager = McpConnectionManager::new( &mcp_servers, diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index 9eee4da39cf7..ed69d262474b 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -3313,7 +3313,7 @@ async fn set_rate_limits_retains_previous_credits() { user_shell_override: None, }; - let mut state = SessionState::new(session_configuration); + let mut state = session_state_for_tests(session_configuration).await; let initial = RateLimitSnapshot { limit_id: None, limit_name: None, @@ -3420,7 +3420,7 @@ async fn set_rate_limits_updates_plan_type_when_present() { user_shell_override: None, }; - let mut state = SessionState::new(session_configuration); + let mut state = session_state_for_tests(session_configuration).await; let initial = RateLimitSnapshot { limit_id: None, limit_name: None, @@ -4032,18 +4032,20 @@ async fn emit_subagent_session_started_includes_fork_lineage_from_session_config ); } -fn turn_environments_for_tests( - environment: &Arc, - cwd: &codex_utils_absolute_path::AbsolutePathBuf, -) -> crate::environment_selection::ResolvedTurnEnvironments { - crate::environment_selection::ResolvedTurnEnvironments { - turn_environments: vec![TurnEnvironment::new( - codex_exec_server::LOCAL_ENVIRONMENT_ID.to_string(), - Arc::clone(environment), - cwd.clone(), - /*shell*/ None, - )], - } +async fn turn_environments_for_configuration( + session_configuration: &SessionConfiguration, +) -> TurnEnvironments { + TurnEnvironments::resolve( + Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), + session_configuration.environment_selections(), + ) + .await + .expect("environment selections should resolve") +} + +async fn session_state_for_tests(session_configuration: SessionConfiguration) -> SessionState { + let turn_environments = turn_environments_for_configuration(&session_configuration).await; + SessionState::new(session_configuration, turn_environments) } #[tokio::test] @@ -4405,7 +4407,7 @@ async fn active_profile_update_rebuilds_network_proxy_config() -> std::io::Resul #[cfg_attr(windows, ignore)] #[tokio::test] async fn new_default_turn_uses_config_aware_skills_for_role_overrides() { - let (session, _turn_context) = make_session_and_context().await; + let (session, turn_context) = make_session_and_context().await; let parent_config = session.get_config().await; let codex_home = parent_config.codex_home.clone(); let skill_dir = codex_home.join("skills").join("demo"); @@ -4417,11 +4419,9 @@ async fn new_default_turn_uses_config_aware_skills_for_role_overrides() { ) .expect("write skill"); - let skill_fs = session - .services - .environment_manager - .default_environment() - .map(|environment| environment.get_filesystem()) + let skill_fs = turn_context + .environments + .primary_filesystem() .unwrap_or_else(|| std::sync::Arc::clone(&codex_exec_server::LOCAL_FS)); let parent_outcome = session .services @@ -4658,8 +4658,7 @@ async fn session_update_settings_does_not_rewrite_sticky_environment_cwds() { async fn relative_cwd_update_without_environments_resolves_under_session_cwd() { let (session, _turn_context) = make_session_and_context().await; let original_cwd = { - let mut state = session.state.lock().await; - state.session_configuration.environments.environments = Vec::new(); + let state = session.state.lock().await; state.session_configuration.cwd().clone() }; let updated_cwd = original_cwd.join("project"); @@ -4819,6 +4818,7 @@ async fn session_new_fails_when_zsh_fork_enabled_without_packaged_zsh() { config.codex_home.clone(), /*bundled_skills_enabled*/ true, )); + let turn_environments = turn_environments_for_configuration(&session_configuration).await; let result = Session::new( session_configuration, Arc::clone(&config), @@ -4836,7 +4836,7 @@ async fn session_new_fails_when_zsh_fork_enabled_without_packaged_zsh() { Arc::new(codex_extension_api::ExtensionRegistryBuilder::new().build()), codex_extension_api::ExtensionDataInit::default(), AgentControl::default(), - Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), + turn_environments, /*analytics_events_client*/ None, Arc::new(codex_thread_store::LocalThreadStore::new( codex_thread_store::LocalThreadStoreConfig::from_config(config.as_ref()), @@ -4931,7 +4931,11 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) { session_configuration.session_source.clone(), ); - let state = SessionState::new(session_configuration.clone()); + let state = session_state_for_tests(session_configuration.clone()).await; + let turn_environments = state.turn_environments.clone(); + let environment = turn_environments + .primary_environment() + .expect("primary environment"); let plugins_manager = Arc::new(PluginsManager::new(config.codex_home.to_path_buf())); let mcp_manager = Arc::new(McpManager::new(Arc::clone(&plugins_manager))); let skills_manager = Arc::new(SkillsManager::new( @@ -4939,11 +4943,6 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) { /*bundled_skills_enabled*/ true, )); let network_approval = Arc::new(NetworkApprovalService::default()); - let environment = Arc::new( - codex_exec_server::Environment::create_for_tests(/*exec_server_url*/ None) - .expect("create environment"), - ); - let services = SessionServices { mcp_connection_manager: Arc::new(arc_swap::ArcSwap::from_pointee( McpConnectionManager::new_uninitialized_with_permission_profile( @@ -5013,7 +5012,6 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) { ), code_mode_service: crate::tools::code_mode::CodeModeService::new(), tool_search_handler_cache: Default::default(), - environment_manager: Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), }; let plugin_outcome = services @@ -5030,7 +5028,6 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) { .skills_for_config(&skills_input, Some(Arc::clone(&skill_fs))) .await, ); - let turn_environments = turn_environments_for_tests(&environment, session_configuration.cwd()); let turn_context = Session::make_turn_context( thread_id, SessionId::from(thread_id), @@ -5158,6 +5155,7 @@ async fn make_session_with_config_and_rx( config.codex_home.clone(), /*bundled_skills_enabled*/ true, )); + let turn_environments = turn_environments_for_configuration(&session_configuration).await; let session = Session::new( session_configuration, @@ -5176,7 +5174,7 @@ async fn make_session_with_config_and_rx( Arc::new(codex_extension_api::ExtensionRegistryBuilder::new().build()), codex_extension_api::ExtensionDataInit::default(), AgentControl::default(), - Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), + turn_environments, /*analytics_events_client*/ None, Arc::new(codex_thread_store::LocalThreadStore::new( codex_thread_store::LocalThreadStoreConfig::from_config(config.as_ref()), @@ -5260,6 +5258,7 @@ async fn make_session_with_history_source_and_agent_control_and_rx( config.codex_home.clone(), /*bundled_skills_enabled*/ true, )); + let turn_environments = turn_environments_for_configuration(&session_configuration).await; let session = Session::new( session_configuration, @@ -5278,7 +5277,7 @@ async fn make_session_with_history_source_and_agent_control_and_rx( Arc::new(codex_extension_api::ExtensionRegistryBuilder::new().build()), codex_extension_api::ExtensionDataInit::default(), agent_control, - Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), + turn_environments, /*analytics_events_client*/ None, Arc::new(codex_thread_store::LocalThreadStore::new( codex_thread_store::LocalThreadStoreConfig::from_config(config.as_ref()), @@ -6202,6 +6201,28 @@ async fn turn_environments_set_primary_environment() { let turn_cwd = turn_context.cwd.clone(); assert_eq!(turn_cwd.as_path(), selected_cwd.as_path()); assert_eq!(turn_context.config.cwd.as_path(), selected_cwd.as_path()); + + let stored_environment = { + let state = session.state.lock().await; + state + .turn_environments + .primary_environment() + .expect("stored primary environment") + }; + assert!(Arc::ptr_eq( + &stored_environment, + &turn_environment.environment + )); + + let default_turn = session.new_default_turn().await; + assert!(Arc::ptr_eq( + &stored_environment, + &default_turn + .environments + .primary() + .expect("default turn primary environment") + .environment + )); } #[tokio::test] @@ -6210,10 +6231,18 @@ async fn default_turn_does_not_overlay_legacy_fallback_cwd_onto_stored_thread_en let session_cwd = session.get_config().await.cwd.clone(); let selected_cwd = AbsolutePathBuf::try_from(session_cwd.as_path().join("selected")).expect("absolute path"); + let turn_environments = { + let state = session.state.lock().await; + state.turn_environments.clone() + } + .with_selections(&[local(selected_cwd.clone())]) + .await + .expect("environment selection should resolve"); { let mut state = session.state.lock().await; state.session_configuration.environments.environments = vec![local(selected_cwd.clone())]; + state.turn_environments = turn_environments; } let turn_context = session.new_default_turn().await; @@ -6242,6 +6271,7 @@ async fn default_turn_honors_empty_stored_thread_environments() { { let mut state = session.state.lock().await; state.session_configuration.environments.environments = Vec::new(); + state.turn_environments.turn_environments.clear(); } let turn_context = session.new_default_turn().await; @@ -6937,7 +6967,11 @@ where session_configuration.session_source.clone(), ); - let state = SessionState::new(session_configuration.clone()); + let state = session_state_for_tests(session_configuration.clone()).await; + let turn_environments = state.turn_environments.clone(); + let environment = turn_environments + .primary_environment() + .expect("primary environment"); let plugins_manager = Arc::new(PluginsManager::new(config.codex_home.to_path_buf())); let mcp_manager = Arc::new(McpManager::new(Arc::clone(&plugins_manager))); let skills_manager = Arc::new(SkillsManager::new( @@ -6945,11 +6979,6 @@ where /*bundled_skills_enabled*/ true, )); let network_approval = Arc::new(NetworkApprovalService::default()); - let environment = Arc::new( - codex_exec_server::Environment::create_for_tests(/*exec_server_url*/ None) - .expect("create environment"), - ); - let services = SessionServices { mcp_connection_manager: Arc::new(arc_swap::ArcSwap::from_pointee( McpConnectionManager::new_uninitialized_with_permission_profile( @@ -7019,7 +7048,6 @@ where ), code_mode_service: crate::tools::code_mode::CodeModeService::new(), tool_search_handler_cache: Default::default(), - environment_manager: Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), }; let plugin_outcome = services @@ -7036,7 +7064,6 @@ where .skills_for_config(&skills_input, Some(Arc::clone(&skill_fs))) .await, ); - let turn_environments = turn_environments_for_tests(&environment, session_configuration.cwd()); let turn_context = Arc::new(Session::make_turn_context( thread_id, SessionId::from(thread_id), diff --git a/codex-rs/core/src/session/tests/guardian_tests.rs b/codex-rs/core/src/session/tests/guardian_tests.rs index 0d3c9c033b8f..f60d78394e46 100644 --- a/codex-rs/core/src/session/tests/guardian_tests.rs +++ b/codex-rs/core/src/session/tests/guardian_tests.rs @@ -1,6 +1,6 @@ use super::*; use crate::compact::InitialContextInjection; -use crate::environment_selection::ResolvedTurnEnvironments; +use crate::environment_selection::TurnEnvironments; use crate::exec_policy::ExecPolicyManager; use crate::guardian::GUARDIAN_REVIEWER_NAME; use crate::sandboxing::SandboxPermissions; @@ -709,7 +709,6 @@ async fn guardian_subagent_does_not_inherit_parent_exec_policy_rules() { installation_id: "11111111-1111-4111-8111-111111111111".to_string(), auth_manager, models_manager, - environment_manager: Arc::new(EnvironmentManager::default_for_tests()), skills_manager, plugins_manager, mcp_manager, @@ -729,7 +728,8 @@ async fn guardian_subagent_does_not_inherit_parent_exec_policy_rules() { parent_rollout_thread_trace: codex_rollout_trace::ThreadTraceContext::disabled(), user_shell_override: None, parent_trace: None, - environment_selections: ResolvedTurnEnvironments { + turn_environments: TurnEnvironments { + environment_manager: Arc::new(EnvironmentManager::default_for_tests()), turn_environments: Vec::new(), }, thread_extension_init: codex_extension_api::ExtensionDataInit::default(), diff --git a/codex-rs/core/src/session/turn_context.rs b/codex-rs/core/src/session/turn_context.rs index 70b085fc5db5..7afd3c22f044 100644 --- a/codex-rs/core/src/session/turn_context.rs +++ b/codex-rs/core/src/session/turn_context.rs @@ -2,7 +2,7 @@ use super::*; use crate::SkillLoadOutcome; use crate::agents_md::LoadedAgentsMd; use crate::config::GhostSnapshotConfig; -use crate::environment_selection::ResolvedTurnEnvironments; +use crate::environment_selection::TurnEnvironments; use codex_core_skills::HostLoadedSkills; use codex_model_provider::SharedModelProvider; use codex_model_provider::create_model_provider; @@ -102,7 +102,7 @@ pub struct TurnContext { pub(crate) session_source: SessionSource, pub(crate) parent_thread_id: Option, pub(crate) thread_source: Option, - pub(crate) environments: ResolvedTurnEnvironments, + pub(crate) environments: TurnEnvironments, /// The session's absolute working directory. All relative paths provided /// by the model as well as sandbox policies are resolved against this path /// instead of `std::env::current_dir()`. @@ -502,7 +502,7 @@ impl Session { model_info: ModelInfo, models_manager: &SharedModelsManager, network: Option, - environments: ResolvedTurnEnvironments, + environments: TurnEnvironments, cwd: AbsolutePathBuf, sub_id: String, skills_outcome: Arc, @@ -612,11 +612,34 @@ impl Session { } } + pub(super) async fn turn_environments_for_update( + &self, + updates: &SessionSettingsUpdate, + ) -> Option { + let selections = updates.environments.as_ref()?.environments.as_slice(); + let current = { + let state = self.state.lock().await; + state.turn_environments.clone() + }; + let turn_environments = match current.with_selections(selections).await { + Ok(turn_environments) => turn_environments, + Err(err) => { + warn!("failed to resolve turn environments: {err}"); + current + .with_selections(&[]) + .await + .expect("empty turn environment selections should resolve") + } + }; + Some(turn_environments) + } + pub(crate) async fn new_turn_with_sub_id( &self, sub_id: String, updates: SessionSettingsUpdate, ) -> CodexResult> { + let updated_turn_environments = self.turn_environments_for_update(&updates).await; let notify_config_contributors = !self.services.extensions.config_contributors().is_empty(); let update_result: CodexResult<_> = { let mut state = self.state.lock().await; @@ -636,6 +659,9 @@ impl Session { let new_config = notify_config_contributors .then(|| Self::build_effective_session_config(&next)); state.session_configuration = next.clone(); + if let Some(turn_environments) = updated_turn_environments { + state.turn_environments = turn_environments; + } Ok(( next, permission_profile_changed, @@ -688,60 +714,28 @@ impl Session { } Ok(self - .new_turn_from_configuration( + .new_turn_context_from_state( sub_id, - session_configuration, updates.final_output_json_schema, + TurnMultiAgentRuntime::ResolveAndStore, ) .await) } - async fn new_turn_from_configuration( - &self, - sub_id: String, - session_configuration: SessionConfiguration, - final_output_json_schema: Option>, - ) -> Arc { - self.new_turn_context_from_configuration( - sub_id, - session_configuration, - final_output_json_schema, - TurnMultiAgentRuntime::ResolveAndStore, - ) - .await - } - - async fn new_startup_prewarm_turn_from_configuration( - &self, - sub_id: String, - session_configuration: SessionConfiguration, - ) -> Arc { - self.new_turn_context_from_configuration( - sub_id, - session_configuration, - /*final_output_json_schema*/ None, - TurnMultiAgentRuntime::Preview, - ) - .await - } - #[instrument(name = "turn_context.build", level = "trace", skip_all)] - async fn new_turn_context_from_configuration( + async fn new_turn_context_from_state( &self, sub_id: String, - session_configuration: SessionConfiguration, final_output_json_schema: Option>, multi_agent_runtime: TurnMultiAgentRuntime, ) -> Arc { - let turn_environments = crate::environment_selection::resolve_environment_selections( - self.services.environment_manager.as_ref(), - session_configuration.environment_selections(), - ) - .await - .unwrap_or_else(|err| { - warn!("failed to resolve turn environments: {err}"); - ResolvedTurnEnvironments::default() - }); + let (session_configuration, turn_environments) = { + let state = self.state.lock().await; + ( + state.session_configuration.clone(), + state.turn_environments.clone(), + ) + }; let primary_turn_environment = turn_environments.primary().cloned(); let cwd = primary_turn_environment .as_ref() @@ -854,11 +848,10 @@ impl Session { } pub(crate) async fn new_default_turn_with_sub_id(&self, sub_id: String) -> Arc { - let session_configuration = self.default_turn_configuration().await; - self.new_turn_from_configuration( + self.new_turn_context_from_state( sub_id, - session_configuration, /*final_output_json_schema*/ None, + TurnMultiAgentRuntime::ResolveAndStore, ) .await } @@ -867,13 +860,11 @@ impl Session { &self, sub_id: String, ) -> Arc { - let session_configuration = self.default_turn_configuration().await; - self.new_startup_prewarm_turn_from_configuration(sub_id, session_configuration) - .await - } - - async fn default_turn_configuration(&self) -> SessionConfiguration { - let state = self.state.lock().await; - state.session_configuration.clone() + self.new_turn_context_from_state( + sub_id, + /*final_output_json_schema*/ None, + TurnMultiAgentRuntime::Preview, + ) + .await } } diff --git a/codex-rs/core/src/state/service.rs b/codex-rs/core/src/state/service.rs index 30ef37048eac..dc95d2e24c2a 100644 --- a/codex-rs/core/src/state/service.rs +++ b/codex-rs/core/src/state/service.rs @@ -22,7 +22,6 @@ use arc_swap::ArcSwap; use arc_swap::ArcSwapOption; use codex_analytics::AnalyticsEventsClient; use codex_core_plugins::PluginsManager; -use codex_exec_server::EnvironmentManager; use codex_extension_api::ExtensionData; use codex_extension_api::ExtensionDataInit; use codex_extension_api::ExtensionRegistry; @@ -83,9 +82,6 @@ pub(crate) struct SessionServices { pub(crate) model_client: ModelClient, pub(crate) code_mode_service: CodeModeService, pub(crate) tool_search_handler_cache: ToolSearchHandlerCache, - /// Shared process-level environment registry. Sessions carry an `Arc` handle so they can pass - /// the same manager through child-thread spawn paths without reconstructing it. - pub(crate) environment_manager: Arc, } impl SessionServices { diff --git a/codex-rs/core/src/state/session.rs b/codex-rs/core/src/state/session.rs index 269d3e0f607e..2061d67bc53e 100644 --- a/codex-rs/core/src/state/session.rs +++ b/codex-rs/core/src/state/session.rs @@ -11,6 +11,7 @@ use super::AdditionalContextStore; use super::auto_compact_window::AutoCompactWindow; use super::auto_compact_window::AutoCompactWindowSnapshot; use crate::context_manager::ContextManager; +use crate::environment_selection::TurnEnvironments; use crate::session::PreviousTurnSettings; use crate::session::session::SessionConfiguration; use crate::session_startup_prewarm::SessionStartupPrewarmHandle; @@ -23,6 +24,8 @@ use codex_utils_output_truncation::TruncationPolicy; /// Persistent, session-scoped state previously stored directly on `Session`. pub(crate) struct SessionState { pub(crate) session_configuration: SessionConfiguration, + /// Resolved runtime environments matching `session_configuration.environments`. + pub(crate) turn_environments: TurnEnvironments, pub(crate) history: ContextManager, pub(crate) latest_rate_limits: Option, pub(crate) server_reasoning_included: bool, @@ -44,10 +47,14 @@ pub(crate) struct SessionState { impl SessionState { /// Create a new session state mirroring previous `State::default()` semantics. - pub(crate) fn new(session_configuration: SessionConfiguration) -> Self { + pub(crate) fn new( + session_configuration: SessionConfiguration, + turn_environments: TurnEnvironments, + ) -> Self { let history = ContextManager::new(); Self { session_configuration, + turn_environments, history, latest_rate_limits: None, server_reasoning_included: false, diff --git a/codex-rs/core/src/state/session_tests.rs b/codex-rs/core/src/state/session_tests.rs index 0fbb92b958f0..b3801e257c22 100644 --- a/codex-rs/core/src/state/session_tests.rs +++ b/codex-rs/core/src/state/session_tests.rs @@ -1,16 +1,27 @@ use super::*; +use crate::environment_selection::TurnEnvironments; use crate::session::tests::make_session_configuration_for_tests; use crate::state::AutoCompactWindowSnapshot; +use codex_exec_server::EnvironmentManager; use codex_protocol::protocol::CreditsSnapshot; use codex_protocol::protocol::RateLimitWindow; use codex_protocol::protocol::SpendControlLimitSnapshot; use pretty_assertions::assert_eq; +use std::sync::Arc; + +async fn make_session_state() -> SessionState { + let session_configuration = make_session_configuration_for_tests().await; + let turn_environments = + TurnEnvironments::resolve(Arc::new(EnvironmentManager::default_for_tests()), &[]) + .await + .expect("environment selections should resolve"); + SessionState::new(session_configuration, turn_environments) +} #[tokio::test] // Verifies connector merging deduplicates repeated IDs. async fn merge_connector_selection_deduplicates_entries() { - let session_configuration = make_session_configuration_for_tests().await; - let mut state = SessionState::new(session_configuration); + let mut state = make_session_state().await; let merged = state.merge_connector_selection([ "calendar".to_string(), "calendar".to_string(), @@ -26,8 +37,7 @@ async fn merge_connector_selection_deduplicates_entries() { #[tokio::test] // Verifies clearing connector selection removes all saved IDs. async fn clear_connector_selection_removes_entries() { - let session_configuration = make_session_configuration_for_tests().await; - let mut state = SessionState::new(session_configuration); + let mut state = make_session_state().await; state.merge_connector_selection(["calendar".to_string()]); state.clear_connector_selection(); @@ -37,8 +47,7 @@ async fn clear_connector_selection_removes_entries() { #[tokio::test] async fn set_rate_limits_defaults_limit_id_to_codex_when_missing() { - let session_configuration = make_session_configuration_for_tests().await; - let mut state = SessionState::new(session_configuration); + let mut state = make_session_state().await; state.set_rate_limits(RateLimitSnapshot { limit_id: None, @@ -66,8 +75,7 @@ async fn set_rate_limits_defaults_limit_id_to_codex_when_missing() { #[tokio::test] async fn replace_history_clears_auto_compact_window_prefill() { - let session_configuration = make_session_configuration_for_tests().await; - let mut state = SessionState::new(session_configuration); + let mut state = make_session_state().await; state.set_auto_compact_window_estimated_prefill(/*tokens*/ 100); state.replace_history(Vec::new(), /*reference_context_item*/ None); @@ -82,8 +90,7 @@ async fn replace_history_clears_auto_compact_window_prefill() { #[tokio::test] async fn set_rate_limits_defaults_to_codex_when_limit_id_missing_after_other_bucket() { - let session_configuration = make_session_configuration_for_tests().await; - let mut state = SessionState::new(session_configuration); + let mut state = make_session_state().await; state.set_rate_limits(RateLimitSnapshot { limit_id: Some("codex_other".to_string()), @@ -125,8 +132,7 @@ async fn set_rate_limits_defaults_to_codex_when_limit_id_missing_after_other_buc #[tokio::test] async fn set_rate_limits_carries_account_metadata_from_codex_to_codex_other() { - let session_configuration = make_session_configuration_for_tests().await; - let mut state = SessionState::new(session_configuration); + let mut state = make_session_state().await; state.set_rate_limits(RateLimitSnapshot { limit_id: Some("codex".to_string()), diff --git a/codex-rs/core/src/thread_manager.rs b/codex-rs/core/src/thread_manager.rs index 19070898ddb2..a565a30ae580 100644 --- a/codex-rs/core/src/thread_manager.rs +++ b/codex-rs/core/src/thread_manager.rs @@ -4,8 +4,8 @@ use crate::attestation::AttestationProvider; use crate::codex_thread::CodexThread; use crate::config::Config; use crate::config::ThreadStoreConfig; +use crate::environment_selection::TurnEnvironments; use crate::environment_selection::default_thread_environment_selections; -use crate::environment_selection::resolve_environment_selections; use crate::mcp::McpManager; use crate::rollout::truncation; use crate::session::Codex; @@ -1383,9 +1383,8 @@ impl ThreadManagerState { threads.remove(&resumed.conversation_id); } } - let environment_selections = - resolve_environment_selections(self.environment_manager.as_ref(), &environments) - .await?; + let turn_environments = + TurnEnvironments::resolve(Arc::clone(&self.environment_manager), &environments).await?; let user_instructions = self .user_instructions_for_spawn(&session_source, parent_thread_id, forked_from_thread_id) .await; @@ -1409,7 +1408,6 @@ impl ThreadManagerState { installation_id: self.installation_id.clone(), auth_manager, models_manager: Arc::clone(&self.models_manager), - environment_manager: Arc::clone(&self.environment_manager), skills_manager: Arc::clone(&self.skills_manager), plugins_manager: Arc::clone(&self.plugins_manager), mcp_manager: Arc::clone(&self.mcp_manager), @@ -1427,7 +1425,7 @@ impl ThreadManagerState { parent_rollout_thread_trace, user_shell_override, parent_trace, - environment_selections, + turn_environments, thread_extension_init, analytics_events_client: self.analytics_events_client.clone(), thread_store: Arc::clone(&self.thread_store), From f7c8157651ea3a2f589ec17ae4f367af90120b11 Mon Sep 17 00:00:00 2001 From: pakrym-oai Date: Fri, 12 Jun 2026 16:22:58 -0700 Subject: [PATCH 02/14] core: resolve turn environments inside codex --- codex-rs/core/src/codex_delegate.rs | 3 ++- codex-rs/core/src/session/mod.rs | 9 +++++++-- codex-rs/core/src/session/tests/guardian_tests.rs | 7 ++----- codex-rs/core/src/thread_manager.rs | 6 ++---- 4 files changed, 13 insertions(+), 12 deletions(-) diff --git a/codex-rs/core/src/codex_delegate.rs b/codex-rs/core/src/codex_delegate.rs index 2e38c4e7d342..08deb3bf3d47 100644 --- a/codex-rs/core/src/codex_delegate.rs +++ b/codex-rs/core/src/codex_delegate.rs @@ -107,7 +107,8 @@ pub(crate) async fn run_codex_thread_interactive( inherited_exec_policy: Some(Arc::clone(&parent_session.services.exec_policy)), parent_rollout_thread_trace: codex_rollout_trace::ThreadTraceContext::disabled(), parent_trace: None, - turn_environments: parent_ctx.environments.clone(), + environment_manager: Arc::clone(&parent_ctx.environments.environment_manager), + environments: parent_ctx.environments.to_selections(), thread_extension_init: codex_extension_api::ExtensionDataInit::default(), analytics_events_client: Some(parent_session.services.analytics_events_client.clone()), thread_store: Arc::clone(&parent_session.services.thread_store), diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 641f0c52ce93..008d81626ba7 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -53,6 +53,7 @@ use codex_app_server_protocol::McpServerElicitationRequestParams; use codex_config::types::AuthKeyringBackendKind; use codex_config::types::OAuthCredentialsStoreMode; use codex_exec_server::Environment; +use codex_exec_server::EnvironmentManager; use codex_exec_server::FileSystemSandboxContext; use codex_extension_api::ExtensionDataInit; use codex_extension_api::LoadedUserInstructions; @@ -428,7 +429,8 @@ pub(crate) struct CodexSpawnArgs { pub(crate) parent_rollout_thread_trace: ThreadTraceContext, pub(crate) user_shell_override: Option, pub(crate) parent_trace: Option, - pub(crate) turn_environments: TurnEnvironments, + pub(crate) environment_manager: Arc, + pub(crate) environments: Vec, pub(crate) thread_extension_init: ExtensionDataInit, pub(crate) analytics_events_client: Option, pub(crate) thread_store: Arc, @@ -509,13 +511,16 @@ impl Codex { inherited_exec_policy, parent_rollout_thread_trace, parent_trace: _, - turn_environments, + environment_manager, + environments, thread_extension_init, analytics_events_client, thread_store, attestation_provider, inherited_multi_agent_version, } = args; + let turn_environments = + TurnEnvironments::resolve(environment_manager, &environments).await?; let (tx_sub, rx_sub) = async_channel::bounded(SUBMISSION_CHANNEL_CAPACITY); let (tx_event, rx_event) = async_channel::unbounded(); diff --git a/codex-rs/core/src/session/tests/guardian_tests.rs b/codex-rs/core/src/session/tests/guardian_tests.rs index f60d78394e46..4bb0066efa6c 100644 --- a/codex-rs/core/src/session/tests/guardian_tests.rs +++ b/codex-rs/core/src/session/tests/guardian_tests.rs @@ -1,6 +1,5 @@ use super::*; use crate::compact::InitialContextInjection; -use crate::environment_selection::TurnEnvironments; use crate::exec_policy::ExecPolicyManager; use crate::guardian::GUARDIAN_REVIEWER_NAME; use crate::sandboxing::SandboxPermissions; @@ -728,10 +727,8 @@ async fn guardian_subagent_does_not_inherit_parent_exec_policy_rules() { parent_rollout_thread_trace: codex_rollout_trace::ThreadTraceContext::disabled(), user_shell_override: None, parent_trace: None, - turn_environments: TurnEnvironments { - environment_manager: Arc::new(EnvironmentManager::default_for_tests()), - turn_environments: Vec::new(), - }, + environment_manager: Arc::new(EnvironmentManager::default_for_tests()), + environments: Vec::new(), thread_extension_init: codex_extension_api::ExtensionDataInit::default(), analytics_events_client: None, thread_store, diff --git a/codex-rs/core/src/thread_manager.rs b/codex-rs/core/src/thread_manager.rs index a565a30ae580..adbfcda22ced 100644 --- a/codex-rs/core/src/thread_manager.rs +++ b/codex-rs/core/src/thread_manager.rs @@ -4,7 +4,6 @@ use crate::attestation::AttestationProvider; use crate::codex_thread::CodexThread; use crate::config::Config; use crate::config::ThreadStoreConfig; -use crate::environment_selection::TurnEnvironments; use crate::environment_selection::default_thread_environment_selections; use crate::mcp::McpManager; use crate::rollout::truncation; @@ -1383,8 +1382,6 @@ impl ThreadManagerState { threads.remove(&resumed.conversation_id); } } - let turn_environments = - TurnEnvironments::resolve(Arc::clone(&self.environment_manager), &environments).await?; let user_instructions = self .user_instructions_for_spawn(&session_source, parent_thread_id, forked_from_thread_id) .await; @@ -1425,7 +1422,8 @@ impl ThreadManagerState { parent_rollout_thread_trace, user_shell_override, parent_trace, - turn_environments, + environment_manager: Arc::clone(&self.environment_manager), + environments, thread_extension_init, analytics_events_client: self.analytics_events_client.clone(), thread_store: Arc::clone(&self.thread_store), From 1d4e81968f2b93c3f34ff296f5215bf1e8c7ffac Mon Sep 17 00:00:00 2001 From: pakrym-oai Date: Fri, 12 Jun 2026 18:07:55 -0700 Subject: [PATCH 03/14] core: skip unresolved turn environments --- codex-rs/core/src/environment_selection.rs | 36 +++++++++++++++++++++- 1 file changed, 35 insertions(+), 1 deletion(-) diff --git a/codex-rs/core/src/environment_selection.rs b/codex-rs/core/src/environment_selection.rs index 57c1fb90a8f5..7e50a352b4cb 100644 --- a/codex-rs/core/src/environment_selection.rs +++ b/codex-rs/core/src/environment_selection.rs @@ -63,7 +63,16 @@ impl TurnEnvironments { && environment.cwd_uri() == &selected_environment.cwd }) { Some(environment) => environment.clone(), - None => self.resolve_selection(selected_environment).await?, + None => match self.resolve_selection(selected_environment).await { + Ok(environment) => environment, + Err(err) => { + tracing::warn!( + "skipping unresolved turn environment `{}`: {err}", + selected_environment.environment_id + ); + continue; + } + }, }; turn_environments.push(turn_environment); } @@ -294,6 +303,31 @@ url = "ws://127.0.0.1:8765" ); } + #[tokio::test] + async fn unresolved_environment_selections_are_skipped() { + let cwd = AbsolutePathBuf::current_dir().expect("cwd"); + let manager = Arc::new(EnvironmentManager::default_for_tests()); + let local = TurnEnvironmentSelection { + environment_id: LOCAL_ENVIRONMENT_ID.to_string(), + cwd: cwd.clone(), + }; + + let resolved = TurnEnvironments::resolve( + manager, + &[ + TurnEnvironmentSelection { + environment_id: "missing".to_string(), + cwd, + }, + local.clone(), + ], + ) + .await + .expect("valid environment selections should resolve"); + + assert_eq!(resolved.to_selections(), vec![local]); + } + #[tokio::test] async fn matching_environment_id_and_cwd_reuse_resolved_environment() { let cwd = AbsolutePathBuf::current_dir().expect("cwd"); From 987bc42be7d9c8d9c4e1e952f20cf6c8856bcebb Mon Sep 17 00:00:00 2001 From: pakrym-oai Date: Fri, 12 Jun 2026 18:13:48 -0700 Subject: [PATCH 04/14] core: update turn environments in place --- codex-rs/core/src/environment_selection.rs | 79 +++++++++------------- codex-rs/core/src/session/mod.rs | 13 ++-- codex-rs/core/src/session/tests.rs | 14 ++-- codex-rs/core/src/session/turn_context.rs | 32 ++------- codex-rs/core/src/state/session_tests.rs | 4 +- 5 files changed, 51 insertions(+), 91 deletions(-) diff --git a/codex-rs/core/src/environment_selection.rs b/codex-rs/core/src/environment_selection.rs index 7e50a352b4cb..734d3b4dda82 100644 --- a/codex-rs/core/src/environment_selection.rs +++ b/codex-rs/core/src/environment_selection.rs @@ -36,27 +36,21 @@ impl TurnEnvironments { pub(crate) async fn resolve( environment_manager: Arc, environments: &[TurnEnvironmentSelection], - ) -> CodexResult { - Self { + ) -> Self { + let mut resolved = Self { environment_manager, turn_environments: Vec::new(), - } - .with_selections(environments) - .await + }; + resolved.update_selections(environments).await; + resolved } - pub(crate) async fn with_selections( - &self, - environments: &[TurnEnvironmentSelection], - ) -> CodexResult { + pub(crate) async fn update_selections(&mut self, environments: &[TurnEnvironmentSelection]) { let mut seen_environment_ids = HashSet::with_capacity(environments.len()); let mut turn_environments = Vec::with_capacity(environments.len()); for selected_environment in environments { if !seen_environment_ids.insert(selected_environment.environment_id.as_str()) { - return Err(CodexErr::InvalidRequest(format!( - "duplicate turn environment id `{}`", - selected_environment.environment_id - ))); + continue; } let turn_environment = match self.turn_environments.iter().find(|environment| { environment.environment_id == selected_environment.environment_id @@ -76,10 +70,7 @@ impl TurnEnvironments { }; turn_environments.push(turn_environment); } - Ok(Self { - environment_manager: Arc::clone(&self.environment_manager), - turn_environments, - }) + self.turn_environments = turn_environments; } async fn resolve_selection( @@ -238,28 +229,28 @@ url = "ws://127.0.0.1:8765" } #[tokio::test] - async fn resolve_environment_selections_rejects_duplicate_ids() { + async fn resolve_environment_selections_keeps_first_duplicate_id() { let cwd = AbsolutePathBuf::current_dir().expect("cwd"); let cwd_uri = PathUri::from_abs_path(&cwd); let manager = Arc::new(EnvironmentManager::default_for_tests()); + let first = TurnEnvironmentSelection { + environment_id: LOCAL_ENVIRONMENT_ID.to_string(), + cwd: cwd_uri.clone(), + }; - let err = TurnEnvironments::resolve( + let resolved = TurnEnvironments::resolve( manager, &[ + first.clone(), TurnEnvironmentSelection { - environment_id: "local".to_string(), - cwd: cwd_uri.clone(), - }, - TurnEnvironmentSelection { - environment_id: "local".to_string(), + environment_id: LOCAL_ENVIRONMENT_ID.to_string(), cwd: cwd_uri.join("other").expect("other cwd URI"), }, ], ) - .await - .expect_err("duplicate environment id should fail"); + .await; - assert!(err.to_string().contains("duplicate")); + assert_eq!(resolved.to_selections(), vec![first]); } #[tokio::test] @@ -276,8 +267,7 @@ url = "ws://127.0.0.1:8765" cwd: selected_cwd_uri, }], ) - .await - .expect("environment selections should resolve"); + .await; assert_eq!( resolved @@ -306,10 +296,11 @@ url = "ws://127.0.0.1:8765" #[tokio::test] async fn unresolved_environment_selections_are_skipped() { let cwd = AbsolutePathBuf::current_dir().expect("cwd"); + let cwd_uri = PathUri::from_abs_path(&cwd); let manager = Arc::new(EnvironmentManager::default_for_tests()); let local = TurnEnvironmentSelection { environment_id: LOCAL_ENVIRONMENT_ID.to_string(), - cwd: cwd.clone(), + cwd: cwd_uri.clone(), }; let resolved = TurnEnvironments::resolve( @@ -317,13 +308,12 @@ url = "ws://127.0.0.1:8765" &[ TurnEnvironmentSelection { environment_id: "missing".to_string(), - cwd, + cwd: cwd_uri, }, local.clone(), ], ) - .await - .expect("valid environment selections should resolve"); + .await; assert_eq!(resolved.to_selections(), vec![local]); } @@ -342,9 +332,7 @@ url = "ws://127.0.0.1:8765" environment_id: REMOTE_ENVIRONMENT_ID.to_string(), cwd: PathUri::from_abs_path(&cwd), }; - let initial = TurnEnvironments::resolve(Arc::clone(&manager), &[selection.clone()]) - .await - .expect("environment selection should resolve"); + let initial = TurnEnvironments::resolve(Arc::clone(&manager), &[selection.clone()]).await; manager .upsert_environment( REMOTE_ENVIRONMENT_ID.to_string(), @@ -352,17 +340,17 @@ url = "ws://127.0.0.1:8765" ) .expect("replace environment"); - let reused = initial - .with_selections(std::slice::from_ref(&selection)) - .await - .expect("matching environment selection should resolve"); - let changed = reused - .with_selections(&[TurnEnvironmentSelection { + let mut reused = initial.clone(); + reused + .update_selections(std::slice::from_ref(&selection)) + .await; + let mut changed = reused.clone(); + changed + .update_selections(&[TurnEnvironmentSelection { cwd: PathUri::from_abs_path(&cwd.join("changed")), ..selection }]) - .await - .expect("changed environment selection should resolve"); + .await; assert!(Arc::ptr_eq( &initial.primary().expect("initial environment").environment, @@ -386,8 +374,7 @@ url = "ws://127.0.0.1:8765" cwd: cwd_uri, }], ) - .await - .expect("local environment should resolve"); + .await; let remote_environment = Arc::new( Environment::create_for_tests(Some("ws://127.0.0.1:8765".to_string())) .expect("remote environment"), diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 008d81626ba7..23b3f3120781 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -519,8 +519,7 @@ impl Codex { attestation_provider, inherited_multi_agent_version, } = args; - let turn_environments = - TurnEnvironments::resolve(environment_manager, &environments).await?; + let turn_environments = TurnEnvironments::resolve(environment_manager, &environments).await; let (tx_sub, rx_sub) = async_channel::bounded(SUBMISSION_CHANNEL_CAPACITY); let (tx_event, rx_event) = async_channel::unbounded(); @@ -1432,7 +1431,6 @@ impl Session { &self, updates: SessionSettingsUpdate, ) -> ConstraintResult<()> { - let updated_turn_environments = self.turn_environments_for_update(&updates).await; let notify_config_contributors = !self.services.extensions.config_contributors().is_empty(); let ( previous_config, @@ -1464,10 +1462,13 @@ impl Session { let next_cwd = updated.cwd().clone(); let codex_home = updated.codex_home.clone(); let session_source = updated.session_source.clone(); - state.session_configuration = updated; - if let Some(turn_environments) = updated_turn_environments { - state.turn_environments = turn_environments; + if updates.environments.is_some() { + state + .turn_environments + .update_selections(updated.environment_selections()) + .await; } + state.session_configuration = updated; ( previous_config, new_config, diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index ed69d262474b..74428c7b8bc0 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -4040,7 +4040,6 @@ async fn turn_environments_for_configuration( session_configuration.environment_selections(), ) .await - .expect("environment selections should resolve") } async fn session_state_for_tests(session_configuration: SessionConfiguration) -> SessionState { @@ -6231,18 +6230,13 @@ async fn default_turn_does_not_overlay_legacy_fallback_cwd_onto_stored_thread_en let session_cwd = session.get_config().await.cwd.clone(); let selected_cwd = AbsolutePathBuf::try_from(session_cwd.as_path().join("selected")).expect("absolute path"); - let turn_environments = { - let state = session.state.lock().await; - state.turn_environments.clone() - } - .with_selections(&[local(selected_cwd.clone())]) - .await - .expect("environment selection should resolve"); - { let mut state = session.state.lock().await; + state + .turn_environments + .update_selections(&[local(selected_cwd.clone())]) + .await; state.session_configuration.environments.environments = vec![local(selected_cwd.clone())]; - state.turn_environments = turn_environments; } let turn_context = session.new_default_turn().await; diff --git a/codex-rs/core/src/session/turn_context.rs b/codex-rs/core/src/session/turn_context.rs index 7afd3c22f044..f404653f884d 100644 --- a/codex-rs/core/src/session/turn_context.rs +++ b/codex-rs/core/src/session/turn_context.rs @@ -612,34 +612,11 @@ impl Session { } } - pub(super) async fn turn_environments_for_update( - &self, - updates: &SessionSettingsUpdate, - ) -> Option { - let selections = updates.environments.as_ref()?.environments.as_slice(); - let current = { - let state = self.state.lock().await; - state.turn_environments.clone() - }; - let turn_environments = match current.with_selections(selections).await { - Ok(turn_environments) => turn_environments, - Err(err) => { - warn!("failed to resolve turn environments: {err}"); - current - .with_selections(&[]) - .await - .expect("empty turn environment selections should resolve") - } - }; - Some(turn_environments) - } - pub(crate) async fn new_turn_with_sub_id( &self, sub_id: String, updates: SessionSettingsUpdate, ) -> CodexResult> { - let updated_turn_environments = self.turn_environments_for_update(&updates).await; let notify_config_contributors = !self.services.extensions.config_contributors().is_empty(); let update_result: CodexResult<_> = { let mut state = self.state.lock().await; @@ -658,10 +635,13 @@ impl Session { }); let new_config = notify_config_contributors .then(|| Self::build_effective_session_config(&next)); - state.session_configuration = next.clone(); - if let Some(turn_environments) = updated_turn_environments { - state.turn_environments = turn_environments; + if updates.environments.is_some() { + state + .turn_environments + .update_selections(next.environment_selections()) + .await; } + state.session_configuration = next.clone(); Ok(( next, permission_profile_changed, diff --git a/codex-rs/core/src/state/session_tests.rs b/codex-rs/core/src/state/session_tests.rs index b3801e257c22..f4aa8df32772 100644 --- a/codex-rs/core/src/state/session_tests.rs +++ b/codex-rs/core/src/state/session_tests.rs @@ -12,9 +12,7 @@ use std::sync::Arc; async fn make_session_state() -> SessionState { let session_configuration = make_session_configuration_for_tests().await; let turn_environments = - TurnEnvironments::resolve(Arc::new(EnvironmentManager::default_for_tests()), &[]) - .await - .expect("environment selections should resolve"); + TurnEnvironments::resolve(Arc::new(EnvironmentManager::default_for_tests()), &[]).await; SessionState::new(session_configuration, turn_environments) } From 1a286807a98e081f00c5ba0030c93879dcf31ed0 Mon Sep 17 00:00:00 2001 From: pakrym-oai Date: Fri, 12 Jun 2026 23:44:49 -0700 Subject: [PATCH 05/14] codex: fix CI failure on PR #27955 --- codex-rs/core/src/session/mod.rs | 10 ++++------ codex-rs/core/src/session/turn_context.rs | 23 +++++++++++++++++------ 2 files changed, 21 insertions(+), 12 deletions(-) diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 23b3f3120781..dc53a0c9f4b2 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -1431,6 +1431,7 @@ impl Session { &self, updates: SessionSettingsUpdate, ) -> ConstraintResult<()> { + let updated_turn_environments = self.turn_environments_for_update(&updates).await; let notify_config_contributors = !self.services.extensions.config_contributors().is_empty(); let ( previous_config, @@ -1462,13 +1463,10 @@ impl Session { let next_cwd = updated.cwd().clone(); let codex_home = updated.codex_home.clone(); let session_source = updated.session_source.clone(); - if updates.environments.is_some() { - state - .turn_environments - .update_selections(updated.environment_selections()) - .await; - } state.session_configuration = updated; + if let Some(turn_environments) = updated_turn_environments { + state.turn_environments = turn_environments; + } ( previous_config, new_config, diff --git a/codex-rs/core/src/session/turn_context.rs b/codex-rs/core/src/session/turn_context.rs index f404653f884d..f76954d7aa9e 100644 --- a/codex-rs/core/src/session/turn_context.rs +++ b/codex-rs/core/src/session/turn_context.rs @@ -617,6 +617,7 @@ impl Session { sub_id: String, updates: SessionSettingsUpdate, ) -> CodexResult> { + let updated_turn_environments = self.turn_environments_for_update(&updates).await; let notify_config_contributors = !self.services.extensions.config_contributors().is_empty(); let update_result: CodexResult<_> = { let mut state = self.state.lock().await; @@ -635,13 +636,10 @@ impl Session { }); let new_config = notify_config_contributors .then(|| Self::build_effective_session_config(&next)); - if updates.environments.is_some() { - state - .turn_environments - .update_selections(next.environment_selections()) - .await; - } state.session_configuration = next.clone(); + if let Some(turn_environments) = updated_turn_environments { + state.turn_environments = turn_environments; + } Ok(( next, permission_profile_changed, @@ -702,6 +700,19 @@ impl Session { .await) } + pub(super) async fn turn_environments_for_update( + &self, + updates: &SessionSettingsUpdate, + ) -> Option { + let selections = updates.environments.as_ref()?.environments.as_slice(); + let mut turn_environments = { + let state = self.state.lock().await; + state.turn_environments.clone() + }; + turn_environments.update_selections(selections).await; + Some(turn_environments) + } + #[instrument(name = "turn_context.build", level = "trace", skip_all)] async fn new_turn_context_from_state( &self, From 4c94d20ae013b408a6f9851fd8741232660f16f1 Mon Sep 17 00:00:00 2001 From: pakrym-oai Date: Fri, 12 Jun 2026 23:55:53 -0700 Subject: [PATCH 06/14] codex: fix CI failure on PR #27955 --- codex-rs/core/src/environment_selection.rs | 3 ++- codex-rs/core/src/session/tests.rs | 12 ++++++++---- 2 files changed, 10 insertions(+), 5 deletions(-) diff --git a/codex-rs/core/src/environment_selection.rs b/codex-rs/core/src/environment_selection.rs index 734d3b4dda82..e975c8f133e9 100644 --- a/codex-rs/core/src/environment_selection.rs +++ b/codex-rs/core/src/environment_selection.rs @@ -332,7 +332,8 @@ url = "ws://127.0.0.1:8765" environment_id: REMOTE_ENVIRONMENT_ID.to_string(), cwd: PathUri::from_abs_path(&cwd), }; - let initial = TurnEnvironments::resolve(Arc::clone(&manager), &[selection.clone()]).await; + let initial = + TurnEnvironments::resolve(Arc::clone(&manager), std::slice::from_ref(&selection)).await; manager .upsert_environment( REMOTE_ENVIRONMENT_ID.to_string(), diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index 74428c7b8bc0..d9ddbe7e8846 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -6230,12 +6230,16 @@ async fn default_turn_does_not_overlay_legacy_fallback_cwd_onto_stored_thread_en let session_cwd = session.get_config().await.cwd.clone(); let selected_cwd = AbsolutePathBuf::try_from(session_cwd.as_path().join("selected")).expect("absolute path"); + let mut turn_environments = { + let state = session.state.lock().await; + state.turn_environments.clone() + }; + turn_environments + .update_selections(&[local(selected_cwd.clone())]) + .await; { let mut state = session.state.lock().await; - state - .turn_environments - .update_selections(&[local(selected_cwd.clone())]) - .await; + state.turn_environments = turn_environments; state.session_configuration.environments.environments = vec![local(selected_cwd.clone())]; } From c6481cb945ff550d6165890653522ae221b8635a Mon Sep 17 00:00:00 2001 From: pakrym-oai Date: Sat, 13 Jun 2026 00:04:52 -0700 Subject: [PATCH 07/14] codex: fix CI failure on PR #27955 --- codex-rs/core/src/thread_manager_tests.rs | 52 ----------------------- 1 file changed, 52 deletions(-) diff --git a/codex-rs/core/src/thread_manager_tests.rs b/codex-rs/core/src/thread_manager_tests.rs index a5670b05333a..778fc3669041 100644 --- a/codex-rs/core/src/thread_manager_tests.rs +++ b/codex-rs/core/src/thread_manager_tests.rs @@ -294,58 +294,6 @@ async fn shutdown_all_threads_bounded_submits_shutdown_to_every_thread() { assert!(manager.list_thread_ids().await.is_empty()); } -#[tokio::test] -async fn start_thread_rejects_explicit_local_environment_when_default_provider_is_disabled() { - let temp_dir = tempdir().expect("tempdir"); - let mut config = test_config().await; - config.codex_home = temp_dir.path().join("codex-home").abs(); - config.cwd = config.codex_home.abs(); - std::fs::create_dir_all(&config.codex_home).expect("create codex home"); - - let runtime_paths = codex_exec_server::ExecServerRuntimePaths::new( - std::env::current_exe().expect("current exe path"), - /*codex_linux_sandbox_exe*/ None, - ) - .expect("runtime paths"); - let environment_manager = Arc::new( - codex_exec_server::EnvironmentManager::create_for_tests( - Some("none".to_string()), - Some(runtime_paths), - ) - .await, - ); - let manager = ThreadManager::with_models_provider_and_home_for_tests( - CodexAuth::from_api_key("dummy"), - config.model_provider.clone(), - config.codex_home.to_path_buf(), - environment_manager, - ); - - let result = manager - .start_thread_with_options(StartThreadOptions { - config: config.clone(), - initial_history: InitialHistory::New, - session_source: None, - thread_source: None, - dynamic_tools: Vec::new(), - metrics_service_name: None, - parent_trace: None, - environments: vec![TurnEnvironmentSelection { - environment_id: "local".to_string(), - cwd: PathUri::from_abs_path(&config.cwd), - }], - thread_extension_init: Default::default(), - }) - .await; - let err = match result { - Ok(_) => panic!("explicit local environment should not resolve when provider is disabled"), - Err(err) => err, - }; - - assert_eq!(err.to_string(), "unknown turn environment id `local`"); - assert!(manager.list_thread_ids().await.is_empty()); -} - #[tokio::test] async fn start_thread_keeps_internal_threads_hidden_from_normal_lookups() { let temp_dir = tempdir().expect("tempdir"); From 7a0b15b72f8e29ce9356712d77cfe5521345a6bc Mon Sep 17 00:00:00 2001 From: pakrym-oai Date: Mon, 15 Jun 2026 14:00:50 -0700 Subject: [PATCH 08/14] core: move turn environments into session services --- codex-rs/core/src/agents_md.rs | 6 +- codex-rs/core/src/agents_md_tests.rs | 37 +++-- codex-rs/core/src/codex_delegate.rs | 5 +- codex-rs/core/src/codex_delegate_tests.rs | 2 +- .../core/src/context/environment_context.rs | 2 +- codex-rs/core/src/environment_selection.rs | 141 +++++++++++------- codex-rs/core/src/mcp_tool_call_tests.rs | 2 +- codex-rs/core/src/session/mcp.rs | 5 +- codex-rs/core/src/session/mod.rs | 17 ++- codex-rs/core/src/session/session.rs | 17 ++- codex-rs/core/src/session/tests.rs | 126 ++++++++-------- .../core/src/session/tests/guardian_tests.rs | 4 +- codex-rs/core/src/session/turn.rs | 3 +- codex-rs/core/src/session/turn_context.rs | 40 ++--- codex-rs/core/src/state/service.rs | 2 + codex-rs/core/src/state/session.rs | 9 +- codex-rs/core/src/state/session_tests.rs | 7 +- codex-rs/core/src/thread_manager_tests.rs | 24 +-- .../agent_jobs/spawn_agents_on_csv.rs | 2 +- .../src/tools/handlers/extension_tools.rs | 5 +- codex-rs/core/src/tools/handlers/mod.rs | 1 - .../core/src/tools/handlers/view_image.rs | 3 +- codex-rs/core/src/tools/spec_plan.rs | 1 - codex-rs/core/src/tools/spec_plan_tests.rs | 12 +- codex-rs/core/src/unified_exec/mod_tests.rs | 6 +- .../src/unified_exec/process_manager_tests.rs | 3 +- 26 files changed, 247 insertions(+), 235 deletions(-) diff --git a/codex-rs/core/src/agents_md.rs b/codex-rs/core/src/agents_md.rs index 107bc3145aa7..a2db1618e60a 100644 --- a/codex-rs/core/src/agents_md.rs +++ b/codex-rs/core/src/agents_md.rs @@ -18,7 +18,7 @@ use crate::config::Config; use crate::context::ContextualUserFragment; use crate::context::UserInstructions as ContextUserInstructions; -use crate::environment_selection::TurnEnvironments; +use crate::environment_selection::TurnEnvironmentsSnapshot; use codex_app_server_protocol::ConfigLayerSource; use codex_config::ConfigLayerStackOrdering; use codex_config::default_project_root_markers; @@ -48,10 +48,10 @@ const AGENTS_MD_SEPARATOR: &str = "\n\n--- project-doc ---\n\n"; pub(crate) async fn load_project_instructions( config: &mut Config, user_instructions: Option, - environments: &TurnEnvironments, + environments: &TurnEnvironmentsSnapshot, ) -> Option { let mut loaded = LoadedAgentsMd::from_user_instructions(user_instructions); - for turn_environment in &environments.turn_environments { + for turn_environment in environments.iter() { 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 1a49af30d761..0adb1a8b1f78 100644 --- a/codex-rs/core/src/agents_md_tests.rs +++ b/codex-rs/core/src/agents_md_tests.rs @@ -1,6 +1,6 @@ use super::*; use crate::config::ConfigBuilder; -use crate::environment_selection::TurnEnvironments; +use crate::environment_selection::TurnEnvironmentsSnapshot; use crate::session::turn_context::TurnEnvironment; use codex_config::ConfigLayerEntry; use codex_config::ConfigLayerStack; @@ -9,7 +9,6 @@ use codex_config::ConfigRequirementsToml; use codex_exec_server::CopyOptions; use codex_exec_server::CreateDirectoryOptions; use codex_exec_server::Environment; -use codex_exec_server::EnvironmentManager; use codex_exec_server::ExecutorFileSystemFuture; use codex_exec_server::FileMetadata; use codex_exec_server::FileSystemSandboxContext; @@ -255,24 +254,22 @@ async fn agents_md_paths(config: &TestConfig) -> std::io::Result( environments: [(&str, AbsolutePathBuf); N], -) -> TurnEnvironments { - TurnEnvironments { - environment_manager: Arc::new(EnvironmentManager::default_for_tests()), - turn_environments: environments - .into_iter() - .map(|(environment_id, cwd)| { - TurnEnvironment::new( - environment_id.to_string(), - Arc::new( - Environment::create_for_tests(/*exec_server_url*/ None) - .expect("local environment"), - ), - cwd, - /*shell*/ None, - ) - }) - .collect(), - } +) -> TurnEnvironmentsSnapshot { + environments + .into_iter() + .map(|(environment_id, cwd)| { + TurnEnvironment::new( + environment_id.to_string(), + Arc::new( + Environment::create_for_tests(/*exec_server_url*/ None) + .expect("local environment"), + ), + cwd, + /*shell*/ None, + ) + }) + .collect::>() + .into() } fn project_provenance(path: AbsolutePathBuf, cwd: AbsolutePathBuf) -> InstructionProvenance { diff --git a/codex-rs/core/src/codex_delegate.rs b/codex-rs/core/src/codex_delegate.rs index 08deb3bf3d47..31e756dfca30 100644 --- a/codex-rs/core/src/codex_delegate.rs +++ b/codex-rs/core/src/codex_delegate.rs @@ -107,7 +107,10 @@ pub(crate) async fn run_codex_thread_interactive( inherited_exec_policy: Some(Arc::clone(&parent_session.services.exec_policy)), parent_rollout_thread_trace: codex_rollout_trace::ThreadTraceContext::disabled(), parent_trace: None, - environment_manager: Arc::clone(&parent_ctx.environments.environment_manager), + environment_manager: parent_session + .services + .turn_environments + .environment_manager(), environments: parent_ctx.environments.to_selections(), thread_extension_init: codex_extension_api::ExtensionDataInit::default(), analytics_events_client: Some(parent_session.services.analytics_events_client.clone()), diff --git a/codex-rs/core/src/codex_delegate_tests.rs b/codex-rs/core/src/codex_delegate_tests.rs index c881823d8044..bd4babee646f 100644 --- a/codex-rs/core/src/codex_delegate_tests.rs +++ b/codex-rs/core/src/codex_delegate_tests.rs @@ -187,7 +187,7 @@ 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(); + parent_ctx_mut.environments.environments_mut()[0].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/environment_context.rs b/codex-rs/core/src/context/environment_context.rs index 8aa07d1ed5a5..9a8c9700cdac 100644 --- a/codex-rs/core/src/context/environment_context.rs +++ b/codex-rs/core/src/context/environment_context.rs @@ -421,7 +421,7 @@ impl EnvironmentContext { pub(crate) fn from_turn_context(turn_context: &TurnContext, shell: &Shell) -> Self { let mut context = Self::new( EnvironmentContextEnvironment::from_turn_environments( - &turn_context.environments.turn_environments, + &turn_context.environments, shell, ), turn_context.current_date.clone(), diff --git a/codex-rs/core/src/environment_selection.rs b/codex-rs/core/src/environment_selection.rs index e975c8f133e9..c501a4899393 100644 --- a/codex-rs/core/src/environment_selection.rs +++ b/codex-rs/core/src/environment_selection.rs @@ -1,4 +1,5 @@ use std::collections::HashSet; +use std::ops::Deref; use std::sync::Arc; use codex_exec_server::EnvironmentManager; @@ -8,6 +9,7 @@ use codex_protocol::error::Result as CodexResult; use codex_protocol::protocol::TurnEnvironmentSelection; use codex_utils_absolute_path::AbsolutePathBuf; use codex_utils_path_uri::PathUri; +use tokio::sync::Mutex; use crate::session::turn_context::TurnEnvironment; use crate::shell::Shell; @@ -26,33 +28,34 @@ pub(crate) fn default_thread_environment_selections( .collect() } -#[derive(Clone, Debug)] +#[derive(Debug)] pub(crate) struct TurnEnvironments { - pub(crate) environment_manager: Arc, - pub(crate) turn_environments: Vec, + environment_manager: Arc, + snapshot: Mutex, } impl TurnEnvironments { pub(crate) async fn resolve( environment_manager: Arc, environments: &[TurnEnvironmentSelection], - ) -> Self { - let mut resolved = Self { + ) -> Arc { + let resolved = Arc::new(Self { environment_manager, - turn_environments: Vec::new(), - }; + snapshot: Mutex::new(TurnEnvironmentsSnapshot::default()), + }); resolved.update_selections(environments).await; resolved } - pub(crate) async fn update_selections(&mut self, environments: &[TurnEnvironmentSelection]) { + pub(crate) async fn update_selections(&self, environments: &[TurnEnvironmentSelection]) { + let current = self.snapshot().await; let mut seen_environment_ids = HashSet::with_capacity(environments.len()); let mut turn_environments = Vec::with_capacity(environments.len()); for selected_environment in environments { if !seen_environment_ids.insert(selected_environment.environment_id.as_str()) { continue; } - let turn_environment = match self.turn_environments.iter().find(|environment| { + let turn_environment = match current.iter().find(|environment| { environment.environment_id == selected_environment.environment_id && environment.cwd_uri() == &selected_environment.cwd }) { @@ -70,7 +73,7 @@ impl TurnEnvironments { }; turn_environments.push(turn_environment); } - self.turn_environments = turn_environments; + *self.snapshot.lock().await = TurnEnvironmentsSnapshot(turn_environments); } async fn resolve_selection( @@ -112,21 +115,25 @@ impl TurnEnvironments { )) } - pub(crate) fn to_selections(&self) -> Vec { - self.turn_environments - .iter() - .map(TurnEnvironment::selection) - .collect() + pub(crate) async fn snapshot(&self) -> TurnEnvironmentsSnapshot { + self.snapshot.lock().await.clone() + } + + pub(crate) fn environment_manager(&self) -> Arc { + Arc::clone(&self.environment_manager) } +} + +#[derive(Clone, Debug, Default)] +pub(crate) struct TurnEnvironmentsSnapshot(Vec); +impl TurnEnvironmentsSnapshot { pub(crate) fn primary(&self) -> Option<&TurnEnvironment> { - self.turn_environments.first() + self.first() } - #[cfg(test)] - pub(crate) fn primary_environment(&self) -> Option> { - self.primary() - .map(|environment| Arc::clone(&environment.environment)) + pub(crate) fn to_selections(&self) -> Vec { + self.iter().map(TurnEnvironment::selection).collect() } pub(crate) fn primary_filesystem(&self) -> Option> { @@ -135,12 +142,35 @@ impl TurnEnvironments { } pub(crate) fn single_local_environment_cwd(&self) -> Option<&AbsolutePathBuf> { - let [environment] = self.turn_environments.as_slice() else { + let [environment] = self.as_slice() else { return None; }; (!environment.environment.is_remote()).then_some(environment.cwd()) } + + pub(crate) fn as_slice(&self) -> &[TurnEnvironment] { + &self.0 + } + + #[cfg(test)] + pub(crate) fn environments_mut(&mut self) -> &mut Vec { + &mut self.0 + } +} + +impl Deref for TurnEnvironmentsSnapshot { + type Target = [TurnEnvironment]; + + fn deref(&self) -> &Self::Target { + &self.0 + } +} + +impl From> for TurnEnvironmentsSnapshot { + fn from(environments: Vec) -> Self { + Self(environments) + } } #[cfg(test)] @@ -250,7 +280,7 @@ url = "ws://127.0.0.1:8765" ) .await; - assert_eq!(resolved.to_selections(), vec![first]); + assert_eq!(resolved.snapshot().await.to_selections(), vec![first]); } #[tokio::test] @@ -269,15 +299,16 @@ url = "ws://127.0.0.1:8765" ) .await; + let resolved = resolved.snapshot().await; assert_eq!( resolved - .primary() + .first() .expect("primary environment") .environment_id, "local" ); assert_eq!( - resolved.primary().expect("primary environment").shell, + resolved.first().expect("primary environment").shell, Some( Shell::from_environment_shell_info( manager @@ -315,7 +346,7 @@ url = "ws://127.0.0.1:8765" ) .await; - assert_eq!(resolved.to_selections(), vec![local]); + assert_eq!(resolved.snapshot().await.to_selections(), vec![local]); } #[tokio::test] @@ -341,25 +372,38 @@ url = "ws://127.0.0.1:8765" ) .expect("replace environment"); - let mut reused = initial.clone(); - reused + let initial_snapshot = initial.snapshot().await; + initial .update_selections(std::slice::from_ref(&selection)) .await; - let mut changed = reused.clone(); - changed + let reused_snapshot = initial.snapshot().await; + initial .update_selections(&[TurnEnvironmentSelection { cwd: PathUri::from_abs_path(&cwd.join("changed")), ..selection }]) .await; + let changed_snapshot = initial.snapshot().await; assert!(Arc::ptr_eq( - &initial.primary().expect("initial environment").environment, - &reused.primary().expect("reused environment").environment, + &initial_snapshot + .first() + .expect("initial environment") + .environment, + &reused_snapshot + .first() + .expect("reused environment") + .environment, )); assert!(!Arc::ptr_eq( - &reused.primary().expect("reused environment").environment, - &changed.primary().expect("changed environment").environment, + &reused_snapshot + .first() + .expect("reused environment") + .environment, + &changed_snapshot + .first() + .expect("changed environment") + .environment, )); } @@ -376,31 +420,26 @@ url = "ws://127.0.0.1:8765" }], ) .await; + let local = local.snapshot().await; let remote_environment = Arc::new( Environment::create_for_tests(Some("ws://127.0.0.1:8765".to_string())) .expect("remote environment"), ); - let remote = TurnEnvironments { - environment_manager: Arc::clone(&local_manager), - turn_environments: vec![TurnEnvironment::new( + let remote = TurnEnvironmentsSnapshot::from(vec![TurnEnvironment::new( + REMOTE_ENVIRONMENT_ID.to_string(), + remote_environment.clone(), + cwd.clone(), + /*shell*/ None, + )]); + let multiple = TurnEnvironmentsSnapshot::from(vec![ + local.first().expect("local environment").clone(), + TurnEnvironment::new( REMOTE_ENVIRONMENT_ID.to_string(), - remote_environment.clone(), + remote_environment, cwd.clone(), /*shell*/ None, - )], - }; - let multiple = TurnEnvironments { - environment_manager: local_manager, - turn_environments: vec![ - local.primary().expect("local environment").clone(), - TurnEnvironment::new( - REMOTE_ENVIRONMENT_ID.to_string(), - remote_environment, - cwd.clone(), - /*shell*/ None, - ), - ], - }; + ), + ]); assert_eq!(local.single_local_environment_cwd(), Some(&cwd)); assert_eq!(remote.single_local_environment_cwd(), None); diff --git a/codex-rs/core/src/mcp_tool_call_tests.rs b/codex-rs/core/src/mcp_tool_call_tests.rs index bb8d714b5baf..9a8a9174a12c 100644 --- a/codex-rs/core/src/mcp_tool_call_tests.rs +++ b/codex-rs/core/src/mcp_tool_call_tests.rs @@ -1268,7 +1268,7 @@ async fn install_host_owned_codex_apps_manager(session: &Session, turn_context: CancellationToken::new(), turn_context.permission_profile(), codex_mcp::McpRuntimeContext::new( - Arc::clone(&turn_context.environments.environment_manager), + session.services.turn_environments.environment_manager(), { #[allow(deprecated)] turn_context.cwd.to_path_buf() diff --git a/codex-rs/core/src/session/mcp.rs b/codex-rs/core/src/session/mcp.rs index 0427900024c2..cc6e42f85222 100644 --- a/codex-rs/core/src/session/mcp.rs +++ b/codex-rs/core/src/session/mcp.rs @@ -317,13 +317,14 @@ impl Session { auth.as_ref(), ) .await; + let environment_manager = self.services.turn_environments.environment_manager(); let mcp_runtime_context = match turn_context.environments.primary() { Some(turn_environment) => McpRuntimeContext::new( - Arc::clone(&turn_context.environments.environment_manager), + Arc::clone(&environment_manager), turn_environment.cwd().to_path_buf(), ), None => McpRuntimeContext::new( - Arc::clone(&turn_context.environments.environment_manager), + environment_manager, #[allow(deprecated)] turn_context.cwd.to_path_buf(), ), diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index dc53a0c9f4b2..bab851a403c6 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -520,6 +520,7 @@ impl Codex { inherited_multi_agent_version, } = args; let turn_environments = TurnEnvironments::resolve(environment_manager, &environments).await; + let resolved_environments = turn_environments.snapshot().await; let (tx_sub, rx_sub) = async_channel::bounded(SUBMISSION_CHANNEL_CAPACITY); let (tx_event, rx_event) = async_channel::unbounded(); @@ -532,7 +533,7 @@ impl Codex { .startup_warnings .extend(user_instruction_provider_warnings); let loaded_agents_md = - load_project_instructions(&mut config, user_instructions, &turn_environments).await; + load_project_instructions(&mut config, user_instructions, &resolved_environments).await; let exec_policy = if crate::guardian::is_guardian_reviewer_source(&session_source) { // Guardian review should rely on the built-in shell safety checks, @@ -622,7 +623,7 @@ impl Codex { windows_sandbox_level: WindowsSandboxLevel::from_config(&config), environments: TurnEnvironmentSelections::new( config.cwd.clone(), - turn_environments.to_selections(), + resolved_environments.to_selections(), ), workspace_roots: config.workspace_roots.clone(), codex_home: config.codex_home.clone(), @@ -1431,7 +1432,6 @@ impl Session { &self, updates: SessionSettingsUpdate, ) -> ConstraintResult<()> { - let updated_turn_environments = self.turn_environments_for_update(&updates).await; let notify_config_contributors = !self.services.extensions.config_contributors().is_empty(); let ( previous_config, @@ -1464,9 +1464,6 @@ impl Session { let codex_home = updated.codex_home.clone(); let session_source = updated.session_source.clone(); state.session_configuration = updated; - if let Some(turn_environments) = updated_turn_environments { - state.turn_environments = turn_environments; - } ( previous_config, new_config, @@ -1477,7 +1474,12 @@ impl Session { session_source, ) }; - + if let Some(environments) = &updates.environments { + self.services + .turn_environments + .update_selections(&environments.environments) + .await; + } self.emit_config_changed_contributors(previous_config.as_ref(), new_config.as_ref()); self.maybe_refresh_shell_snapshot_for_cwd( &previous_cwd, @@ -2356,7 +2358,6 @@ impl Session { let turn_environment = match args.environment_id.as_deref() { Some(environment_id) => turn_context .environments - .turn_environments .iter() .find(|environment| environment.environment_id == environment_id), None => turn_context.environments.primary(), diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index 479209346d43..af74742d6774 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -2,6 +2,7 @@ use super::input_queue::InputQueue; use super::*; use crate::agents_md::LoadedAgentsMd; use crate::config::ConstraintError; +use crate::environment_selection::TurnEnvironmentsSnapshot; use crate::skills::SkillError; use crate::state::ActiveTurn; use codex_extension_api::ExtensionDataInit; @@ -437,7 +438,7 @@ async fn warm_plugins_and_skills_for_session_init( config: Arc, plugins_manager: Arc, skills_manager: Arc, - turn_environments: TurnEnvironments, + turn_environments: TurnEnvironmentsSnapshot, ) -> Vec { let fs = turn_environments.primary_filesystem(); let plugins_input = config.plugins_config_input(); @@ -480,7 +481,7 @@ impl Session { extensions: Arc>, thread_extension_init: ExtensionDataInit, agent_control: AgentControl, - turn_environments: TurnEnvironments, + turn_environments: Arc, analytics_events_client: Option, thread_store: Arc, parent_rollout_thread_trace: ThreadTraceContext, @@ -622,7 +623,7 @@ impl Session { Arc::clone(&config), Arc::clone(&plugins_manager), Arc::clone(&skills_manager), - turn_environments.clone(), + turn_environments.snapshot().await, ) .instrument(info_span!( "session_init.plugin_skill_warmup", @@ -847,7 +848,7 @@ impl Session { session_configuration.thread_name = thread_name.clone(); validate_config_lock_if_configured(&session_configuration).await?; export_config_lock_if_configured(&session_configuration, thread_id).await?; - let state = SessionState::new(session_configuration.clone(), turn_environments); + let state = SessionState::new(session_configuration.clone()); let managed_network_requirements_configured = config .config_layer_stack .requirements_toml() @@ -962,6 +963,7 @@ impl Session { } let services = SessionServices { + turn_environments: Arc::clone(&turn_environments), // Initialize the MCP connection manager with an uninitialized // instance. It will be replaced with one created via // McpConnectionManager::new() once all its constructor args are @@ -1108,14 +1110,13 @@ impl Session { cancel_token }; let mcp_runtime_context = { - let state = sess.state.lock().await; - let cwd = state - .turn_environments + let turn_environments = sess.services.turn_environments.snapshot().await; + let cwd = turn_environments .primary() .map(|turn_environment| turn_environment.cwd().to_path_buf()) .unwrap_or_else(|| session_configuration.cwd().to_path_buf()); McpRuntimeContext::new( - Arc::clone(&state.turn_environments.environment_manager), + sess.services.turn_environments.environment_manager(), cwd, ) }; diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index d9ddbe7e8846..911854e9418d 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -3313,7 +3313,7 @@ async fn set_rate_limits_retains_previous_credits() { user_shell_override: None, }; - let mut state = session_state_for_tests(session_configuration).await; + let mut state = session_state_for_tests(session_configuration); let initial = RateLimitSnapshot { limit_id: None, limit_name: None, @@ -3420,7 +3420,7 @@ async fn set_rate_limits_updates_plan_type_when_present() { user_shell_override: None, }; - let mut state = session_state_for_tests(session_configuration).await; + let mut state = session_state_for_tests(session_configuration); let initial = RateLimitSnapshot { limit_id: None, limit_name: None, @@ -4034,7 +4034,7 @@ async fn emit_subagent_session_started_includes_fork_lineage_from_session_config async fn turn_environments_for_configuration( session_configuration: &SessionConfiguration, -) -> TurnEnvironments { +) -> Arc { TurnEnvironments::resolve( Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), session_configuration.environment_selections(), @@ -4042,9 +4042,8 @@ async fn turn_environments_for_configuration( .await } -async fn session_state_for_tests(session_configuration: SessionConfiguration) -> SessionState { - let turn_environments = turn_environments_for_configuration(&session_configuration).await; - SessionState::new(session_configuration, turn_environments) +fn session_state_for_tests(session_configuration: SessionConfiguration) -> SessionState { + SessionState::new(session_configuration) } #[tokio::test] @@ -4420,7 +4419,8 @@ async fn new_default_turn_uses_config_aware_skills_for_role_overrides() { let skill_fs = turn_context .environments - .primary_filesystem() + .first() + .map(|environment| environment.environment.get_filesystem()) .unwrap_or_else(|| std::sync::Arc::clone(&codex_exec_server::LOCAL_FS)); let parent_outcome = session .services @@ -4744,7 +4744,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.len(), 1); } #[tokio::test] @@ -4930,11 +4930,15 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) { session_configuration.session_source.clone(), ); - let state = session_state_for_tests(session_configuration.clone()).await; - let turn_environments = state.turn_environments.clone(); - let environment = turn_environments - .primary_environment() - .expect("primary environment"); + let state = session_state_for_tests(session_configuration.clone()); + let turn_environments = turn_environments_for_configuration(&session_configuration).await; + let resolved_turn_environments = turn_environments.snapshot().await; + let environment = Arc::clone( + &resolved_turn_environments + .first() + .expect("primary environment") + .environment, + ); let plugins_manager = Arc::new(PluginsManager::new(config.codex_home.to_path_buf())); let mcp_manager = Arc::new(McpManager::new(Arc::clone(&plugins_manager))); let skills_manager = Arc::new(SkillsManager::new( @@ -4943,6 +4947,7 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) { )); let network_approval = Arc::new(NetworkApprovalService::default()); let services = SessionServices { + turn_environments: Arc::clone(&turn_environments), mcp_connection_manager: Arc::new(arc_swap::ArcSwap::from_pointee( McpConnectionManager::new_uninitialized_with_permission_profile( &config.permissions.approval_policy, @@ -5042,7 +5047,7 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) { model_info, &models_manager, /*network*/ None, - turn_environments, + resolved_turn_environments, session_configuration.cwd().clone(), "turn_id".to_string(), skills_outcome, @@ -5591,7 +5596,7 @@ async fn request_permissions_emits_event_when_granular_policy_allows_requests() async move { let environment = turn_context .environments - .primary() + .first() .expect("primary environment") .selection(); session @@ -5664,8 +5669,8 @@ 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( + let current_environment = turn_context_mut.environments[0].clone(); + turn_context_mut.environments.environments_mut()[0] = TurnEnvironment::new( "remote".to_string(), current_environment.environment, environment_cwd.clone(), @@ -5824,7 +5829,7 @@ async fn request_permissions_response_materializes_session_cwd_grants_before_rec async move { let environment = turn_context .environments - .primary() + .first() .expect("primary environment") .selection(); session @@ -5915,7 +5920,7 @@ async fn request_permissions_is_auto_denied_when_granular_policy_blocks_tool_req let call_id = "call-1".to_string(); let environment = turn_context .environments - .primary() + .first() .expect("primary environment") .selection(); let response = session @@ -6186,26 +6191,29 @@ 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.len(), 1); let turn_environment = turn_context .environments - .primary() + .first() .expect("primary environment should be set"); assert!(std::sync::Arc::ptr_eq( &turn_environment.environment, - &turn_environments.turn_environments[0].environment + &turn_environments[0].environment )); - assert!(!turn_context.environments.turn_environments.is_empty()); + assert!(!turn_context.environments.is_empty()); #[allow(deprecated)] let turn_cwd = turn_context.cwd.clone(); assert_eq!(turn_cwd.as_path(), selected_cwd.as_path()); assert_eq!(turn_context.config.cwd.as_path(), selected_cwd.as_path()); let stored_environment = { - let state = session.state.lock().await; - state + session + .services .turn_environments - .primary_environment() + .snapshot() + .await + .first() + .map(|environment| Arc::clone(&environment.environment)) .expect("stored primary environment") }; assert!(Arc::ptr_eq( @@ -6218,7 +6226,7 @@ async fn turn_environments_set_primary_environment() { &stored_environment, &default_turn .environments - .primary() + .first() .expect("default turn primary environment") .environment )); @@ -6230,30 +6238,27 @@ async fn default_turn_does_not_overlay_legacy_fallback_cwd_onto_stored_thread_en let session_cwd = session.get_config().await.cwd.clone(); let selected_cwd = AbsolutePathBuf::try_from(session_cwd.as_path().join("selected")).expect("absolute path"); - let mut turn_environments = { - let state = session.state.lock().await; - state.turn_environments.clone() - }; - turn_environments + session + .services + .turn_environments .update_selections(&[local(selected_cwd.clone())]) .await; { let mut state = session.state.lock().await; - state.turn_environments = turn_environments; state.session_configuration.environments.environments = vec![local(selected_cwd.clone())]; } 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.len(), 1); let turn_environment = turn_context .environments - .primary() + .first() .expect("primary environment should be set"); assert!(std::sync::Arc::ptr_eq( &turn_environment.environment, - &turn_environments.turn_environments[0].environment + &turn_environments[0].environment )); #[allow(deprecated)] let turn_cwd = turn_context.cwd.clone(); @@ -6266,32 +6271,35 @@ async fn default_turn_honors_empty_stored_thread_environments() { let (session, _turn_context, _rx) = make_session_and_context_with_rx().await; let session_cwd = session.get_config().await.cwd.clone(); + session + .services + .turn_environments + .update_selections(&[]) + .await; { let mut state = session.state.lock().await; state.session_configuration.environments.environments = Vec::new(); - state.turn_environments.turn_environments.clear(); } 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.is_empty()); #[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.len(), 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[0].clone(); #[allow(deprecated)] let second_cwd = turn_context.cwd.join("second"); turn_context .environments - .turn_environments + .environments_mut() .push(TurnEnvironment::new( "second".to_string(), Arc::clone(&first_environment.environment), @@ -6302,7 +6310,7 @@ async fn primary_environment_uses_first_turn_environment() { assert_eq!( turn_context .environments - .primary() + .first() .expect("primary environment") .environment_id, first_environment.environment_id @@ -6310,18 +6318,14 @@ async fn primary_environment_uses_first_turn_environment() { assert_eq!( turn_context .environments - .turn_environments .iter() .find(|environment| environment.environment_id == "second") .expect("second environment") .cwd(), &second_cwd ); - assert_eq!(turn_context.environments.turn_environments.len(), 2); - assert_eq!( - turn_context.environments.turn_environments[1].cwd(), - &second_cwd - ); + assert_eq!(turn_context.environments.len(), 2); + assert_eq!(turn_context.environments[1].cwd(), &second_cwd); } #[tokio::test] @@ -6342,8 +6346,7 @@ async fn empty_turn_environments_clear_primary_environment() { .await .expect("turn should start"); - assert!(turn_context.environments.primary().is_none()); - assert!(turn_context.environments.turn_environments.is_empty()); + assert!(turn_context.environments.is_empty()); #[allow(deprecated)] let turn_cwd = turn_context.cwd.clone(); assert_eq!(turn_cwd, session.get_config().await.cwd); @@ -6965,11 +6968,15 @@ where session_configuration.session_source.clone(), ); - let state = session_state_for_tests(session_configuration.clone()).await; - let turn_environments = state.turn_environments.clone(); - let environment = turn_environments - .primary_environment() - .expect("primary environment"); + let state = session_state_for_tests(session_configuration.clone()); + let turn_environments = turn_environments_for_configuration(&session_configuration).await; + let resolved_turn_environments = turn_environments.snapshot().await; + let environment = Arc::clone( + &resolved_turn_environments + .first() + .expect("primary environment") + .environment, + ); let plugins_manager = Arc::new(PluginsManager::new(config.codex_home.to_path_buf())); let mcp_manager = Arc::new(McpManager::new(Arc::clone(&plugins_manager))); let skills_manager = Arc::new(SkillsManager::new( @@ -6978,6 +6985,7 @@ where )); let network_approval = Arc::new(NetworkApprovalService::default()); let services = SessionServices { + turn_environments: Arc::clone(&turn_environments), mcp_connection_manager: Arc::new(arc_swap::ArcSwap::from_pointee( McpConnectionManager::new_uninitialized_with_permission_profile( &config.permissions.approval_policy, @@ -7077,7 +7085,7 @@ where model_info, &models_manager, /*network*/ None, - turn_environments, + resolved_turn_environments, session_configuration.cwd().clone(), "turn_id".to_string(), skills_outcome, @@ -7271,7 +7279,7 @@ async fn environment_context_uses_session_shell_when_environment_shell_is_absent shell_type: crate::shell::ShellType::PowerShell, shell_path: PathBuf::from("powershell"), }); - for environment in &mut turn_context.environments.turn_environments { + for environment in turn_context.environments.environments_mut() { environment.shell = None; } @@ -7288,7 +7296,7 @@ async fn environment_context_uses_session_shell_when_environment_shell_is_absent let primary_environment = turn_context .environments - .turn_environments + .environments_mut() .first_mut() .expect("primary environment"); primary_environment.shell = Some(crate::shell::Shell { diff --git a/codex-rs/core/src/session/tests/guardian_tests.rs b/codex-rs/core/src/session/tests/guardian_tests.rs index 4bb0066efa6c..b63712ce0398 100644 --- a/codex-rs/core/src/session/tests/guardian_tests.rs +++ b/codex-rs/core/src/session/tests/guardian_tests.rs @@ -125,7 +125,7 @@ async fn request_permissions_routes_to_guardian_when_reviewer_is_enabled() { }; let environment = turn_context .environments - .primary() + .first() .expect("primary environment") .selection(); let response = tokio::time::timeout( @@ -221,7 +221,7 @@ async fn request_permissions_guardian_review_stops_when_cancelled() { async move { let environment = turn_context .environments - .primary() + .first() .expect("primary environment") .selection(); session diff --git a/codex-rs/core/src/session/turn.rs b/codex-rs/core/src/session/turn.rs index 3df795296030..b12b02fb879c 100644 --- a/codex-rs/core/src/session/turn.rs +++ b/codex-rs/core/src/session/turn.rs @@ -413,7 +413,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.iter() { let root = get_git_repo_root_with_fs( turn_environment.environment.get_filesystem().as_ref(), turn_environment.cwd(), @@ -631,7 +631,6 @@ async fn build_extension_turn_input_items( let environments = turn_context .environments - .turn_environments .iter() .enumerate() .map(|(index, environment)| TurnInputEnvironment { diff --git a/codex-rs/core/src/session/turn_context.rs b/codex-rs/core/src/session/turn_context.rs index f76954d7aa9e..1dff6af998ff 100644 --- a/codex-rs/core/src/session/turn_context.rs +++ b/codex-rs/core/src/session/turn_context.rs @@ -2,7 +2,7 @@ use super::*; use crate::SkillLoadOutcome; use crate::agents_md::LoadedAgentsMd; use crate::config::GhostSnapshotConfig; -use crate::environment_selection::TurnEnvironments; +use crate::environment_selection::TurnEnvironmentsSnapshot; use codex_core_skills::HostLoadedSkills; use codex_model_provider::SharedModelProvider; use codex_model_provider::create_model_provider; @@ -102,7 +102,7 @@ pub struct TurnContext { pub(crate) session_source: SessionSource, pub(crate) parent_thread_id: Option, pub(crate) thread_source: Option, - pub(crate) environments: TurnEnvironments, + pub(crate) environments: TurnEnvironmentsSnapshot, /// The session's absolute working directory. All relative paths provided /// by the model as well as sandbox policies are resolved against this path /// instead of `std::env::current_dir()`. @@ -199,7 +199,7 @@ impl TurnContext { } pub(crate) fn tool_environment_mode(&self) -> ToolEnvironmentMode { - ToolEnvironmentMode::from_count(self.environments.turn_environments.len()) + ToolEnvironmentMode::from_count(self.environments.len()) } pub(crate) async fn with_model( @@ -502,7 +502,7 @@ impl Session { model_info: ModelInfo, models_manager: &SharedModelsManager, network: Option, - environments: TurnEnvironments, + environments: TurnEnvironmentsSnapshot, cwd: AbsolutePathBuf, sub_id: String, skills_outcome: Arc, @@ -617,7 +617,6 @@ impl Session { sub_id: String, updates: SessionSettingsUpdate, ) -> CodexResult> { - let updated_turn_environments = self.turn_environments_for_update(&updates).await; let notify_config_contributors = !self.services.extensions.config_contributors().is_empty(); let update_result: CodexResult<_> = { let mut state = self.state.lock().await; @@ -637,9 +636,6 @@ impl Session { let new_config = notify_config_contributors .then(|| Self::build_effective_session_config(&next)); state.session_configuration = next.clone(); - if let Some(turn_environments) = updated_turn_environments { - state.turn_environments = turn_environments; - } Ok(( next, permission_profile_changed, @@ -677,7 +673,12 @@ impl Session { return Err(CodexErr::InvalidRequest(message)); } }; - + if let Some(environments) = &updates.environments { + self.services + .turn_environments + .update_selections(&environments.environments) + .await; + } self.emit_config_changed_contributors(previous_config.as_ref(), new_config.as_ref()); self.maybe_refresh_shell_snapshot_for_cwd( &previous_cwd, @@ -700,19 +701,6 @@ impl Session { .await) } - pub(super) async fn turn_environments_for_update( - &self, - updates: &SessionSettingsUpdate, - ) -> Option { - let selections = updates.environments.as_ref()?.environments.as_slice(); - let mut turn_environments = { - let state = self.state.lock().await; - state.turn_environments.clone() - }; - turn_environments.update_selections(selections).await; - Some(turn_environments) - } - #[instrument(name = "turn_context.build", level = "trace", skip_all)] async fn new_turn_context_from_state( &self, @@ -720,13 +708,11 @@ impl Session { final_output_json_schema: Option>, multi_agent_runtime: TurnMultiAgentRuntime, ) -> Arc { - let (session_configuration, turn_environments) = { + let session_configuration = { let state = self.state.lock().await; - ( - state.session_configuration.clone(), - state.turn_environments.clone(), - ) + state.session_configuration.clone() }; + let turn_environments = self.services.turn_environments.snapshot().await; let primary_turn_environment = turn_environments.primary().cloned(); let cwd = primary_turn_environment .as_ref() diff --git a/codex-rs/core/src/state/service.rs b/codex-rs/core/src/state/service.rs index dc95d2e24c2a..4a5de7287bb3 100644 --- a/codex-rs/core/src/state/service.rs +++ b/codex-rs/core/src/state/service.rs @@ -7,6 +7,7 @@ use crate::attestation::AttestationProvider; use crate::client::ModelClient; use crate::config::NetworkProxyAuditMetadata; use crate::config::StartedNetworkProxy; +use crate::environment_selection::TurnEnvironments; use crate::exec_policy::ExecPolicyManager; use crate::guardian::GuardianRejection; use crate::guardian::GuardianRejectionCircuitBreaker; @@ -40,6 +41,7 @@ use tokio::sync::Mutex; use tokio_util::sync::CancellationToken; pub(crate) struct SessionServices { + pub(crate) turn_environments: Arc, /// The latest manager; callers retain an owned handle while performing MCP I/O. pub(crate) mcp_connection_manager: Arc>, pub(crate) mcp_startup_cancellation_token: Mutex, diff --git a/codex-rs/core/src/state/session.rs b/codex-rs/core/src/state/session.rs index 2061d67bc53e..269d3e0f607e 100644 --- a/codex-rs/core/src/state/session.rs +++ b/codex-rs/core/src/state/session.rs @@ -11,7 +11,6 @@ use super::AdditionalContextStore; use super::auto_compact_window::AutoCompactWindow; use super::auto_compact_window::AutoCompactWindowSnapshot; use crate::context_manager::ContextManager; -use crate::environment_selection::TurnEnvironments; use crate::session::PreviousTurnSettings; use crate::session::session::SessionConfiguration; use crate::session_startup_prewarm::SessionStartupPrewarmHandle; @@ -24,8 +23,6 @@ use codex_utils_output_truncation::TruncationPolicy; /// Persistent, session-scoped state previously stored directly on `Session`. pub(crate) struct SessionState { pub(crate) session_configuration: SessionConfiguration, - /// Resolved runtime environments matching `session_configuration.environments`. - pub(crate) turn_environments: TurnEnvironments, pub(crate) history: ContextManager, pub(crate) latest_rate_limits: Option, pub(crate) server_reasoning_included: bool, @@ -47,14 +44,10 @@ pub(crate) struct SessionState { impl SessionState { /// Create a new session state mirroring previous `State::default()` semantics. - pub(crate) fn new( - session_configuration: SessionConfiguration, - turn_environments: TurnEnvironments, - ) -> Self { + pub(crate) fn new(session_configuration: SessionConfiguration) -> Self { let history = ContextManager::new(); Self { session_configuration, - turn_environments, history, latest_rate_limits: None, server_reasoning_included: false, diff --git a/codex-rs/core/src/state/session_tests.rs b/codex-rs/core/src/state/session_tests.rs index f4aa8df32772..00bf830a61bc 100644 --- a/codex-rs/core/src/state/session_tests.rs +++ b/codex-rs/core/src/state/session_tests.rs @@ -1,19 +1,14 @@ use super::*; -use crate::environment_selection::TurnEnvironments; use crate::session::tests::make_session_configuration_for_tests; use crate::state::AutoCompactWindowSnapshot; -use codex_exec_server::EnvironmentManager; use codex_protocol::protocol::CreditsSnapshot; use codex_protocol::protocol::RateLimitWindow; use codex_protocol::protocol::SpendControlLimitSnapshot; use pretty_assertions::assert_eq; -use std::sync::Arc; async fn make_session_state() -> SessionState { let session_configuration = make_session_configuration_for_tests().await; - let turn_environments = - TurnEnvironments::resolve(Arc::new(EnvironmentManager::default_for_tests()), &[]).await; - SessionState::new(session_configuration, turn_environments) + SessionState::new(session_configuration) } #[tokio::test] diff --git a/codex-rs/core/src/thread_manager_tests.rs b/codex-rs/core/src/thread_manager_tests.rs index 778fc3669041..c6fac14a321b 100644 --- a/codex-rs/core/src/thread_manager_tests.rs +++ b/codex-rs/core/src/thread_manager_tests.rs @@ -599,15 +599,9 @@ 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[0].cwd(), - &default_cwd - ); - assert_ne!( - resumed_turn.environments.turn_environments[0].cwd(), - &selected_cwd - ); + assert_eq!(resumed_turn.environments.len(), 1); + assert_eq!(resumed_turn.environments[0].cwd(), &default_cwd); + assert_ne!(resumed_turn.environments[0].cwd(), &selected_cwd); let forked = manager .fork_thread( @@ -626,15 +620,9 @@ 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[0].cwd(), - &default_cwd - ); - assert_ne!( - forked_turn.environments.turn_environments[0].cwd(), - &selected_cwd - ); + assert_eq!(forked_turn.environments.len(), 1); + assert_eq!(forked_turn.environments[0].cwd(), &default_cwd); + assert_ne!(forked_turn.environments[0].cwd(), &selected_cwd); } #[tokio::test] 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 2556989621cb..01ea1f61ff8c 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,7 @@ pub async fn handle( } fn single_local_environment_cwd(turn: &TurnContext) -> Result<&AbsolutePathBuf, FunctionCallError> { - let [turn_environment] = turn.environments.turn_environments.as_slice() else { + let [turn_environment] = turn.environments.as_slice() else { 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 33625317fe40..b5fa5702fc37 100644 --- a/codex-rs/core/src/tools/handlers/extension_tools.rs +++ b/codex-rs/core/src/tools/handlers/extension_tools.rs @@ -112,8 +112,8 @@ impl TurnItemEmitter for CoreTurnItemEmitter { async fn to_extension_call(invocation: &ToolInvocation) -> ExtensionToolCall { let conversation_history = ConversationHistory::new(invocation.session.clone_history().await.into_raw_items()); - let mut environments = Vec::with_capacity(invocation.turn.environments.turn_environments.len()); - for environment in &invocation.turn.environments.turn_environments { + let mut environments = Vec::with_capacity(invocation.turn.environments.len()); + for environment in invocation.turn.environments.iter() { let additional_permissions = apply_granted_turn_permissions( invocation.session.as_ref(), &environment.environment_id, @@ -313,7 +313,6 @@ mod tests { let truncation_policy = turn.truncation_policy; let expected_sandbox_cwds = turn .environments - .turn_environments .iter() .map(|environment| Some(environment.cwd_uri().clone())) .collect::>(); diff --git a/codex-rs/core/src/tools/handlers/mod.rs b/codex-rs/core/src/tools/handlers/mod.rs index 7791a819a530..a946014f59c4 100644 --- a/codex-rs/core/src/tools/handlers/mod.rs +++ b/codex-rs/core/src/tools/handlers/mod.rs @@ -156,7 +156,6 @@ fn resolve_tool_environment<'a>( || Ok(turn.environments.primary()), |environment_id| { turn.environments - .turn_environments .iter() .find(|environment| environment.environment_id == environment_id) .map(Some) diff --git a/codex-rs/core/src/tools/handlers/view_image.rs b/codex-rs/core/src/tools/handlers/view_image.rs index dc52e7e336c7..7cbb36d65fda 100644 --- a/codex-rs/core/src/tools/handlers/view_image.rs +++ b/codex-rs/core/src/tools/handlers/view_image.rs @@ -281,11 +281,10 @@ mod tests { fn replace_primary_environment_cwd(turn: &mut crate::TurnContext, cwd: AbsolutePathBuf) { let current = turn .environments - .turn_environments .first() .cloned() .expect("default local turn environment"); - turn.environments.turn_environments[0] = TurnEnvironment::new( + turn.environments.environments_mut()[0] = TurnEnvironment::new( current.environment_id, current.environment, cwd, diff --git a/codex-rs/core/src/tools/spec_plan.rs b/codex-rs/core/src/tools/spec_plan.rs index 872847611053..42f33475d297 100644 --- a/codex-rs/core/src/tools/spec_plan.rs +++ b/codex-rs/core/src/tools/spec_plan.rs @@ -628,7 +628,6 @@ fn unified_exec_should_include_shell_parameter(turn_context: &TurnContext) -> bo UnifiedExecShellMode::ZshFork(_) ) || turn_context .environments - .turn_environments .iter() .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 a1bf22310ae1..c5a0fb39f378 100644 --- a/codex-rs/core/src/tools/spec_plan_tests.rs +++ b/codex-rs/core/src/tools/spec_plan_tests.rs @@ -349,9 +349,11 @@ 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[0].clone(); second_environment.environment_id = "secondary".to_string(); - turn.environments.turn_environments.push(second_environment); + turn.environments + .environments_mut() + .push(second_environment); } fn mcp_tool(server: &str, namespace: &str, name: &str) -> ToolInfo { @@ -584,11 +586,11 @@ async fn zsh_fork_unified_exec_keeps_shell_parameter_when_remote_environment_ava codex_tools::UnifiedExecShellMode::ZshFork(zsh_fork_config_for_spec_plan_tests()); let remote_cwd = turn .environments - .primary() + .first() .expect("primary environment") .cwd() .clone(); - turn.environments.turn_environments.push( + turn.environments.environments_mut().push( crate::session::turn_context::TurnEnvironment::new( "remote".to_string(), Arc::new( @@ -615,7 +617,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_mut().clear(); set_feature(turn, Feature::ShellTool, /*enabled*/ true); turn.model_info.apply_patch_tool_type = Some(ApplyPatchToolType::Freeform); }) diff --git a/codex-rs/core/src/unified_exec/mod_tests.rs b/codex-rs/core/src/unified_exec/mod_tests.rs index ad39d62b9663..c4952b32c78f 100644 --- a/codex-rs/core/src/unified_exec/mod_tests.rs +++ b/codex-rs/core/src/unified_exec/mod_tests.rs @@ -112,7 +112,7 @@ async fn exec_command_with_tty( tty, Box::new(NoopSpawnLifecycle), turn.environments - .primary() + .first() .expect("turn environment") .environment .as_ref(), @@ -894,7 +894,7 @@ 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 = + turn.environments.environments_mut()[0].environment = Arc::new(remote_test_env.environment().clone()); #[allow(deprecated)] @@ -916,7 +916,7 @@ async fn remote_exec_server_rejects_inherited_fd_launches() -> anyhow::Result<() inherited_fds: vec![42], }), turn.environments - .primary() + .first() .expect("turn environment") .environment .as_ref(), diff --git a/codex-rs/core/src/unified_exec/process_manager_tests.rs b/codex-rs/core/src/unified_exec/process_manager_tests.rs index a84f6345ede6..96a9bb0391b6 100644 --- a/codex-rs/core/src/unified_exec/process_manager_tests.rs +++ b/codex-rs/core/src/unified_exec/process_manager_tests.rs @@ -205,7 +205,8 @@ async fn failed_initial_end_for_unstored_process_uses_fallback_output() { sandbox_cwd: turn.cwd.clone(), environment: turn .environments - .primary_environment() + .first() + .map(|environment| Arc::clone(&environment.environment)) .expect("primary environment"), shell_mode: codex_tools::UnifiedExecShellMode::Direct, network: None, From 56a40eecbb787783970d6dc00c4068f586b17cde Mon Sep 17 00:00:00 2001 From: pakrym-oai Date: Mon, 15 Jun 2026 14:16:47 -0700 Subject: [PATCH 09/14] core: reduce turn environment refactor churn --- codex-rs/core/src/agents_md.rs | 2 +- codex-rs/core/src/agents_md_tests.rs | 31 ++++--- codex-rs/core/src/codex_delegate.rs | 10 +- codex-rs/core/src/codex_delegate_tests.rs | 2 +- .../core/src/context/environment_context.rs | 2 +- codex-rs/core/src/environment_selection.rs | 87 ++++++++---------- codex-rs/core/src/session/mod.rs | 12 ++- codex-rs/core/src/session/session.rs | 2 +- codex-rs/core/src/session/tests.rs | 92 ++++++++++--------- .../core/src/session/tests/guardian_tests.rs | 8 +- codex-rs/core/src/session/turn.rs | 3 +- codex-rs/core/src/session/turn_context.rs | 61 ++++++++---- codex-rs/core/src/state/service.rs | 2 +- codex-rs/core/src/state/session_tests.rs | 23 ++--- codex-rs/core/src/thread_manager.rs | 4 +- codex-rs/core/src/thread_manager_tests.rs | 24 +++-- .../agent_jobs/spawn_agents_on_csv.rs | 2 +- .../src/tools/handlers/extension_tools.rs | 5 +- codex-rs/core/src/tools/handlers/mod.rs | 1 + .../core/src/tools/handlers/view_image.rs | 3 +- codex-rs/core/src/tools/spec_plan.rs | 1 + codex-rs/core/src/tools/spec_plan_tests.rs | 12 +-- codex-rs/core/src/unified_exec/mod_tests.rs | 6 +- .../src/unified_exec/process_manager_tests.rs | 3 +- 24 files changed, 220 insertions(+), 178 deletions(-) diff --git a/codex-rs/core/src/agents_md.rs b/codex-rs/core/src/agents_md.rs index a2db1618e60a..dfcba497ab3e 100644 --- a/codex-rs/core/src/agents_md.rs +++ b/codex-rs/core/src/agents_md.rs @@ -51,7 +51,7 @@ pub(crate) async fn load_project_instructions( environments: &TurnEnvironmentsSnapshot, ) -> Option { let mut loaded = LoadedAgentsMd::from_user_instructions(user_instructions); - for turn_environment in environments.iter() { + 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 0adb1a8b1f78..0f23fb014afc 100644 --- a/codex-rs/core/src/agents_md_tests.rs +++ b/codex-rs/core/src/agents_md_tests.rs @@ -255,21 +255,22 @@ async fn agents_md_paths(config: &TestConfig) -> std::io::Result( environments: [(&str, AbsolutePathBuf); N], ) -> TurnEnvironmentsSnapshot { - environments - .into_iter() - .map(|(environment_id, cwd)| { - TurnEnvironment::new( - environment_id.to_string(), - Arc::new( - Environment::create_for_tests(/*exec_server_url*/ None) - .expect("local environment"), - ), - cwd, - /*shell*/ None, - ) - }) - .collect::>() - .into() + TurnEnvironmentsSnapshot { + turn_environments: environments + .into_iter() + .map(|(environment_id, cwd)| { + TurnEnvironment::new( + environment_id.to_string(), + Arc::new( + Environment::create_for_tests(/*exec_server_url*/ None) + .expect("local environment"), + ), + cwd, + /*shell*/ None, + ) + }) + .collect(), + } } fn project_provenance(path: AbsolutePathBuf, cwd: AbsolutePathBuf) -> InstructionProvenance { diff --git a/codex-rs/core/src/codex_delegate.rs b/codex-rs/core/src/codex_delegate.rs index 31e756dfca30..b36522459181 100644 --- a/codex-rs/core/src/codex_delegate.rs +++ b/codex-rs/core/src/codex_delegate.rs @@ -90,6 +90,10 @@ pub(crate) async fn run_codex_thread_interactive( installation_id: parent_session.installation_id.clone(), auth_manager, models_manager, + environment_manager: parent_session + .services + .turn_environments + .environment_manager(), skills_manager: Arc::clone(&parent_session.services.skills_manager), plugins_manager: Arc::clone(&parent_session.services.plugins_manager), mcp_manager: Arc::clone(&parent_session.services.mcp_manager), @@ -107,11 +111,7 @@ pub(crate) async fn run_codex_thread_interactive( inherited_exec_policy: Some(Arc::clone(&parent_session.services.exec_policy)), parent_rollout_thread_trace: codex_rollout_trace::ThreadTraceContext::disabled(), parent_trace: None, - environment_manager: parent_session - .services - .turn_environments - .environment_manager(), - environments: parent_ctx.environments.to_selections(), + environment_selections: parent_ctx.environments.to_selections(), thread_extension_init: codex_extension_api::ExtensionDataInit::default(), analytics_events_client: Some(parent_session.services.analytics_events_client.clone()), thread_store: Arc::clone(&parent_session.services.thread_store), diff --git a/codex-rs/core/src/codex_delegate_tests.rs b/codex-rs/core/src/codex_delegate_tests.rs index bd4babee646f..c881823d8044 100644 --- a/codex-rs/core/src/codex_delegate_tests.rs +++ b/codex-rs/core/src/codex_delegate_tests.rs @@ -187,7 +187,7 @@ 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.environments_mut()[0].environment_id = "remote".to_string(); + parent_ctx_mut.environments.turn_environments[0].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/environment_context.rs b/codex-rs/core/src/context/environment_context.rs index 9a8c9700cdac..8aa07d1ed5a5 100644 --- a/codex-rs/core/src/context/environment_context.rs +++ b/codex-rs/core/src/context/environment_context.rs @@ -421,7 +421,7 @@ impl EnvironmentContext { pub(crate) fn from_turn_context(turn_context: &TurnContext, shell: &Shell) -> Self { let mut context = Self::new( EnvironmentContextEnvironment::from_turn_environments( - &turn_context.environments, + &turn_context.environments.turn_environments, shell, ), turn_context.current_date.clone(), diff --git a/codex-rs/core/src/environment_selection.rs b/codex-rs/core/src/environment_selection.rs index c501a4899393..c615193fc5ef 100644 --- a/codex-rs/core/src/environment_selection.rs +++ b/codex-rs/core/src/environment_selection.rs @@ -1,5 +1,4 @@ use std::collections::HashSet; -use std::ops::Deref; use std::sync::Arc; use codex_exec_server::EnvironmentManager; @@ -55,7 +54,7 @@ impl TurnEnvironments { if !seen_environment_ids.insert(selected_environment.environment_id.as_str()) { continue; } - let turn_environment = match current.iter().find(|environment| { + let turn_environment = match current.turn_environments.iter().find(|environment| { environment.environment_id == selected_environment.environment_id && environment.cwd_uri() == &selected_environment.cwd }) { @@ -73,7 +72,7 @@ impl TurnEnvironments { }; turn_environments.push(turn_environment); } - *self.snapshot.lock().await = TurnEnvironmentsSnapshot(turn_environments); + *self.snapshot.lock().await = TurnEnvironmentsSnapshot { turn_environments }; } async fn resolve_selection( @@ -125,15 +124,26 @@ impl TurnEnvironments { } #[derive(Clone, Debug, Default)] -pub(crate) struct TurnEnvironmentsSnapshot(Vec); +pub(crate) struct TurnEnvironmentsSnapshot { + pub(crate) turn_environments: Vec, +} impl TurnEnvironmentsSnapshot { pub(crate) fn primary(&self) -> Option<&TurnEnvironment> { - self.first() + self.turn_environments.first() + } + + #[cfg(test)] + pub(crate) fn primary_environment(&self) -> Option> { + self.primary() + .map(|environment| Arc::clone(&environment.environment)) } pub(crate) fn to_selections(&self) -> Vec { - self.iter().map(TurnEnvironment::selection).collect() + self.turn_environments + .iter() + .map(TurnEnvironment::selection) + .collect() } pub(crate) fn primary_filesystem(&self) -> Option> { @@ -142,35 +152,12 @@ impl TurnEnvironmentsSnapshot { } pub(crate) fn single_local_environment_cwd(&self) -> Option<&AbsolutePathBuf> { - let [environment] = self.as_slice() else { + let [environment] = self.turn_environments.as_slice() else { return None; }; (!environment.environment.is_remote()).then_some(environment.cwd()) } - - pub(crate) fn as_slice(&self) -> &[TurnEnvironment] { - &self.0 - } - - #[cfg(test)] - pub(crate) fn environments_mut(&mut self) -> &mut Vec { - &mut self.0 - } -} - -impl Deref for TurnEnvironmentsSnapshot { - type Target = [TurnEnvironment]; - - fn deref(&self) -> &Self::Target { - &self.0 - } -} - -impl From> for TurnEnvironmentsSnapshot { - fn from(environments: Vec) -> Self { - Self(environments) - } } #[cfg(test)] @@ -302,13 +289,13 @@ url = "ws://127.0.0.1:8765" let resolved = resolved.snapshot().await; assert_eq!( resolved - .first() + .primary() .expect("primary environment") .environment_id, "local" ); assert_eq!( - resolved.first().expect("primary environment").shell, + resolved.primary().expect("primary environment").shell, Some( Shell::from_environment_shell_info( manager @@ -387,21 +374,21 @@ url = "ws://127.0.0.1:8765" assert!(Arc::ptr_eq( &initial_snapshot - .first() + .primary() .expect("initial environment") .environment, &reused_snapshot - .first() + .primary() .expect("reused environment") .environment, )); assert!(!Arc::ptr_eq( &reused_snapshot - .first() + .primary() .expect("reused environment") .environment, &changed_snapshot - .first() + .primary() .expect("changed environment") .environment, )); @@ -425,21 +412,25 @@ url = "ws://127.0.0.1:8765" Environment::create_for_tests(Some("ws://127.0.0.1:8765".to_string())) .expect("remote environment"), ); - let remote = TurnEnvironmentsSnapshot::from(vec![TurnEnvironment::new( - REMOTE_ENVIRONMENT_ID.to_string(), - remote_environment.clone(), - cwd.clone(), - /*shell*/ None, - )]); - let multiple = TurnEnvironmentsSnapshot::from(vec![ - local.first().expect("local environment").clone(), - TurnEnvironment::new( + let remote = TurnEnvironmentsSnapshot { + turn_environments: vec![TurnEnvironment::new( REMOTE_ENVIRONMENT_ID.to_string(), - remote_environment, + remote_environment.clone(), cwd.clone(), /*shell*/ None, - ), - ]); + )], + }; + let multiple = TurnEnvironmentsSnapshot { + turn_environments: vec![ + local.primary().expect("local environment").clone(), + TurnEnvironment::new( + REMOTE_ENVIRONMENT_ID.to_string(), + remote_environment, + cwd.clone(), + /*shell*/ None, + ), + ], + }; assert_eq!(local.single_local_environment_cwd(), Some(&cwd)); assert_eq!(remote.single_local_environment_cwd(), None); diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index bab851a403c6..567afe77603e 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -408,6 +408,7 @@ pub(crate) struct CodexSpawnArgs { pub(crate) installation_id: String, pub(crate) auth_manager: Arc, pub(crate) models_manager: SharedModelsManager, + pub(crate) environment_manager: Arc, pub(crate) skills_manager: Arc, pub(crate) plugins_manager: Arc, pub(crate) mcp_manager: Arc, @@ -429,8 +430,7 @@ pub(crate) struct CodexSpawnArgs { pub(crate) parent_rollout_thread_trace: ThreadTraceContext, pub(crate) user_shell_override: Option, pub(crate) parent_trace: Option, - pub(crate) environment_manager: Arc, - pub(crate) environments: Vec, + pub(crate) environment_selections: Vec, pub(crate) thread_extension_init: ExtensionDataInit, pub(crate) analytics_events_client: Option, pub(crate) thread_store: Arc, @@ -494,6 +494,7 @@ impl Codex { installation_id, auth_manager, models_manager, + environment_manager, skills_manager, plugins_manager, mcp_manager, @@ -511,15 +512,15 @@ impl Codex { inherited_exec_policy, parent_rollout_thread_trace, parent_trace: _, - environment_manager, - environments, + environment_selections, thread_extension_init, analytics_events_client, thread_store, attestation_provider, inherited_multi_agent_version, } = args; - let turn_environments = TurnEnvironments::resolve(environment_manager, &environments).await; + let turn_environments = + TurnEnvironments::resolve(environment_manager, &environment_selections).await; let resolved_environments = turn_environments.snapshot().await; let (tx_sub, rx_sub) = async_channel::bounded(SUBMISSION_CHANNEL_CAPACITY); let (tx_event, rx_event) = async_channel::unbounded(); @@ -2358,6 +2359,7 @@ impl Session { let turn_environment = match args.environment_id.as_deref() { Some(environment_id) => turn_context .environments + .turn_environments .iter() .find(|environment| environment.environment_id == environment_id), None => turn_context.environments.primary(), diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index af74742d6774..bf5bcf085b3b 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -963,7 +963,6 @@ impl Session { } let services = SessionServices { - turn_environments: Arc::clone(&turn_environments), // Initialize the MCP connection manager with an uninitialized // instance. It will be replaced with one created via // McpConnectionManager::new() once all its constructor args are @@ -1028,6 +1027,7 @@ impl Session { ), code_mode_service: crate::tools::code_mode::CodeModeService::new(), tool_search_handler_cache: Default::default(), + turn_environments: Arc::clone(&turn_environments), }; let (out_of_band_elicitation_paused, _out_of_band_elicitation_paused_rx) = watch::channel(false); diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index 911854e9418d..004dd9eb0512 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -3313,7 +3313,7 @@ async fn set_rate_limits_retains_previous_credits() { user_shell_override: None, }; - let mut state = session_state_for_tests(session_configuration); + let mut state = SessionState::new(session_configuration); let initial = RateLimitSnapshot { limit_id: None, limit_name: None, @@ -3420,7 +3420,7 @@ async fn set_rate_limits_updates_plan_type_when_present() { user_shell_override: None, }; - let mut state = session_state_for_tests(session_configuration); + let mut state = SessionState::new(session_configuration); let initial = RateLimitSnapshot { limit_id: None, limit_name: None, @@ -4042,10 +4042,6 @@ async fn turn_environments_for_configuration( .await } -fn session_state_for_tests(session_configuration: SessionConfiguration) -> SessionState { - SessionState::new(session_configuration) -} - #[tokio::test] async fn session_configuration_apply_preserves_profile_file_system_policy_on_cwd_only_update() { let mut session_configuration = make_session_configuration_for_tests().await; @@ -4405,7 +4401,7 @@ async fn active_profile_update_rebuilds_network_proxy_config() -> std::io::Resul #[cfg_attr(windows, ignore)] #[tokio::test] async fn new_default_turn_uses_config_aware_skills_for_role_overrides() { - let (session, turn_context) = make_session_and_context().await; + let (session, _turn_context) = make_session_and_context().await; let parent_config = session.get_config().await; let codex_home = parent_config.codex_home.clone(); let skill_dir = codex_home.join("skills").join("demo"); @@ -4417,10 +4413,12 @@ async fn new_default_turn_uses_config_aware_skills_for_role_overrides() { ) .expect("write skill"); - let skill_fs = turn_context - .environments - .first() - .map(|environment| environment.environment.get_filesystem()) + let skill_fs = session + .services + .turn_environments + .environment_manager() + .default_environment() + .map(|environment| environment.get_filesystem()) .unwrap_or_else(|| std::sync::Arc::clone(&codex_exec_server::LOCAL_FS)); let parent_outcome = session .services @@ -4657,7 +4655,8 @@ async fn session_update_settings_does_not_rewrite_sticky_environment_cwds() { async fn relative_cwd_update_without_environments_resolves_under_session_cwd() { let (session, _turn_context) = make_session_and_context().await; let original_cwd = { - let state = session.state.lock().await; + let mut state = session.state.lock().await; + state.session_configuration.environments.environments = Vec::new(); state.session_configuration.cwd().clone() }; let updated_cwd = original_cwd.join("project"); @@ -4744,7 +4743,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.len(), 1); + assert_eq!(turn_context.environments.turn_environments.len(), 1); } #[tokio::test] @@ -4930,12 +4929,12 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) { session_configuration.session_source.clone(), ); - let state = session_state_for_tests(session_configuration.clone()); + let state = SessionState::new(session_configuration.clone()); let turn_environments = turn_environments_for_configuration(&session_configuration).await; let resolved_turn_environments = turn_environments.snapshot().await; let environment = Arc::clone( &resolved_turn_environments - .first() + .primary() .expect("primary environment") .environment, ); @@ -4947,7 +4946,6 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) { )); let network_approval = Arc::new(NetworkApprovalService::default()); let services = SessionServices { - turn_environments: Arc::clone(&turn_environments), mcp_connection_manager: Arc::new(arc_swap::ArcSwap::from_pointee( McpConnectionManager::new_uninitialized_with_permission_profile( &config.permissions.approval_policy, @@ -5016,6 +5014,7 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) { ), code_mode_service: crate::tools::code_mode::CodeModeService::new(), tool_search_handler_cache: Default::default(), + turn_environments: Arc::clone(&turn_environments), }; let plugin_outcome = services @@ -5596,7 +5595,7 @@ async fn request_permissions_emits_event_when_granular_policy_allows_requests() async move { let environment = turn_context .environments - .first() + .primary() .expect("primary environment") .selection(); session @@ -5669,8 +5668,8 @@ 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[0].clone(); - turn_context_mut.environments.environments_mut()[0] = TurnEnvironment::new( + 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, environment_cwd.clone(), @@ -5829,7 +5828,7 @@ async fn request_permissions_response_materializes_session_cwd_grants_before_rec async move { let environment = turn_context .environments - .first() + .primary() .expect("primary environment") .selection(); session @@ -5920,7 +5919,7 @@ async fn request_permissions_is_auto_denied_when_granular_policy_blocks_tool_req let call_id = "call-1".to_string(); let environment = turn_context .environments - .first() + .primary() .expect("primary environment") .selection(); let response = session @@ -6191,16 +6190,16 @@ async fn turn_environments_set_primary_environment() { .expect("turn should start"); let turn_environments = &turn_context.environments; - assert_eq!(turn_environments.len(), 1); + assert_eq!(turn_environments.turn_environments.len(), 1); let turn_environment = turn_context .environments - .first() + .primary() .expect("primary environment should be set"); assert!(std::sync::Arc::ptr_eq( &turn_environment.environment, - &turn_environments[0].environment + &turn_environments.turn_environments[0].environment )); - assert!(!turn_context.environments.is_empty()); + assert!(!turn_context.environments.turn_environments.is_empty()); #[allow(deprecated)] let turn_cwd = turn_context.cwd.clone(); assert_eq!(turn_cwd.as_path(), selected_cwd.as_path()); @@ -6212,8 +6211,7 @@ async fn turn_environments_set_primary_environment() { .turn_environments .snapshot() .await - .first() - .map(|environment| Arc::clone(&environment.environment)) + .primary_environment() .expect("stored primary environment") }; assert!(Arc::ptr_eq( @@ -6226,7 +6224,7 @@ async fn turn_environments_set_primary_environment() { &stored_environment, &default_turn .environments - .first() + .primary() .expect("default turn primary environment") .environment )); @@ -6251,14 +6249,14 @@ 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.len(), 1); + assert_eq!(turn_environments.turn_environments.len(), 1); let turn_environment = turn_context .environments - .first() + .primary() .expect("primary environment should be set"); assert!(std::sync::Arc::ptr_eq( &turn_environment.environment, - &turn_environments[0].environment + &turn_environments.turn_environments[0].environment )); #[allow(deprecated)] let turn_cwd = turn_context.cwd.clone(); @@ -6283,23 +6281,24 @@ async fn default_turn_honors_empty_stored_thread_environments() { let turn_context = session.new_default_turn().await; - assert!(turn_context.environments.is_empty()); + assert!(turn_context.environments.primary().is_none()); + assert!(turn_context.environments.turn_environments.is_empty()); #[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.len(), 0); + assert_eq!(turn_context.environments.turn_environments.len(), 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[0].clone(); + let first_environment = turn_context.environments.turn_environments[0].clone(); #[allow(deprecated)] let second_cwd = turn_context.cwd.join("second"); turn_context .environments - .environments_mut() + .turn_environments .push(TurnEnvironment::new( "second".to_string(), Arc::clone(&first_environment.environment), @@ -6310,7 +6309,7 @@ async fn primary_environment_uses_first_turn_environment() { assert_eq!( turn_context .environments - .first() + .primary() .expect("primary environment") .environment_id, first_environment.environment_id @@ -6318,14 +6317,18 @@ async fn primary_environment_uses_first_turn_environment() { assert_eq!( turn_context .environments + .turn_environments .iter() .find(|environment| environment.environment_id == "second") .expect("second environment") .cwd(), &second_cwd ); - assert_eq!(turn_context.environments.len(), 2); - assert_eq!(turn_context.environments[1].cwd(), &second_cwd); + assert_eq!(turn_context.environments.turn_environments.len(), 2); + assert_eq!( + turn_context.environments.turn_environments[1].cwd(), + &second_cwd + ); } #[tokio::test] @@ -6346,7 +6349,8 @@ async fn empty_turn_environments_clear_primary_environment() { .await .expect("turn should start"); - assert!(turn_context.environments.is_empty()); + assert!(turn_context.environments.primary().is_none()); + assert!(turn_context.environments.turn_environments.is_empty()); #[allow(deprecated)] let turn_cwd = turn_context.cwd.clone(); assert_eq!(turn_cwd, session.get_config().await.cwd); @@ -6968,12 +6972,12 @@ where session_configuration.session_source.clone(), ); - let state = session_state_for_tests(session_configuration.clone()); + let state = SessionState::new(session_configuration.clone()); let turn_environments = turn_environments_for_configuration(&session_configuration).await; let resolved_turn_environments = turn_environments.snapshot().await; let environment = Arc::clone( &resolved_turn_environments - .first() + .primary() .expect("primary environment") .environment, ); @@ -6985,7 +6989,6 @@ where )); let network_approval = Arc::new(NetworkApprovalService::default()); let services = SessionServices { - turn_environments: Arc::clone(&turn_environments), mcp_connection_manager: Arc::new(arc_swap::ArcSwap::from_pointee( McpConnectionManager::new_uninitialized_with_permission_profile( &config.permissions.approval_policy, @@ -7054,6 +7057,7 @@ where ), code_mode_service: crate::tools::code_mode::CodeModeService::new(), tool_search_handler_cache: Default::default(), + turn_environments: Arc::clone(&turn_environments), }; let plugin_outcome = services @@ -7279,7 +7283,7 @@ async fn environment_context_uses_session_shell_when_environment_shell_is_absent shell_type: crate::shell::ShellType::PowerShell, shell_path: PathBuf::from("powershell"), }); - for environment in turn_context.environments.environments_mut() { + for environment in &mut turn_context.environments.turn_environments { environment.shell = None; } @@ -7296,7 +7300,7 @@ async fn environment_context_uses_session_shell_when_environment_shell_is_absent let primary_environment = turn_context .environments - .environments_mut() + .turn_environments .first_mut() .expect("primary environment"); primary_environment.shell = Some(crate::shell::Shell { diff --git a/codex-rs/core/src/session/tests/guardian_tests.rs b/codex-rs/core/src/session/tests/guardian_tests.rs index b63712ce0398..c350afca8d12 100644 --- a/codex-rs/core/src/session/tests/guardian_tests.rs +++ b/codex-rs/core/src/session/tests/guardian_tests.rs @@ -125,7 +125,7 @@ async fn request_permissions_routes_to_guardian_when_reviewer_is_enabled() { }; let environment = turn_context .environments - .first() + .primary() .expect("primary environment") .selection(); let response = tokio::time::timeout( @@ -221,7 +221,7 @@ async fn request_permissions_guardian_review_stops_when_cancelled() { async move { let environment = turn_context .environments - .first() + .primary() .expect("primary environment") .selection(); session @@ -708,6 +708,7 @@ async fn guardian_subagent_does_not_inherit_parent_exec_policy_rules() { installation_id: "11111111-1111-4111-8111-111111111111".to_string(), auth_manager, models_manager, + environment_manager: Arc::new(EnvironmentManager::default_for_tests()), skills_manager, plugins_manager, mcp_manager, @@ -727,8 +728,7 @@ async fn guardian_subagent_does_not_inherit_parent_exec_policy_rules() { parent_rollout_thread_trace: codex_rollout_trace::ThreadTraceContext::disabled(), user_shell_override: None, parent_trace: None, - environment_manager: Arc::new(EnvironmentManager::default_for_tests()), - environments: Vec::new(), + environment_selections: Vec::new(), thread_extension_init: codex_extension_api::ExtensionDataInit::default(), analytics_events_client: None, thread_store, diff --git a/codex-rs/core/src/session/turn.rs b/codex-rs/core/src/session/turn.rs index b12b02fb879c..3df795296030 100644 --- a/codex-rs/core/src/session/turn.rs +++ b/codex-rs/core/src/session/turn.rs @@ -413,7 +413,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.iter() { + for turn_environment in &turn_context.environments.turn_environments { let root = get_git_repo_root_with_fs( turn_environment.environment.get_filesystem().as_ref(), turn_environment.cwd(), @@ -631,6 +631,7 @@ async fn build_extension_turn_input_items( let environments = turn_context .environments + .turn_environments .iter() .enumerate() .map(|(index, environment)| TurnInputEnvironment { diff --git a/codex-rs/core/src/session/turn_context.rs b/codex-rs/core/src/session/turn_context.rs index 1dff6af998ff..f7e13719337b 100644 --- a/codex-rs/core/src/session/turn_context.rs +++ b/codex-rs/core/src/session/turn_context.rs @@ -199,7 +199,7 @@ impl TurnContext { } pub(crate) fn tool_environment_mode(&self) -> ToolEnvironmentMode { - ToolEnvironmentMode::from_count(self.environments.len()) + ToolEnvironmentMode::from_count(self.environments.turn_environments.len()) } pub(crate) async fn with_model( @@ -693,25 +693,51 @@ impl Session { } Ok(self - .new_turn_context_from_state( + .new_turn_from_configuration( sub_id, + session_configuration, updates.final_output_json_schema, - TurnMultiAgentRuntime::ResolveAndStore, ) .await) } + async fn new_turn_from_configuration( + &self, + sub_id: String, + session_configuration: SessionConfiguration, + final_output_json_schema: Option>, + ) -> Arc { + self.new_turn_context_from_configuration( + sub_id, + session_configuration, + final_output_json_schema, + TurnMultiAgentRuntime::ResolveAndStore, + ) + .await + } + + async fn new_startup_prewarm_turn_from_configuration( + &self, + sub_id: String, + session_configuration: SessionConfiguration, + ) -> Arc { + self.new_turn_context_from_configuration( + sub_id, + session_configuration, + /*final_output_json_schema*/ None, + TurnMultiAgentRuntime::Preview, + ) + .await + } + #[instrument(name = "turn_context.build", level = "trace", skip_all)] - async fn new_turn_context_from_state( + async fn new_turn_context_from_configuration( &self, sub_id: String, + session_configuration: SessionConfiguration, final_output_json_schema: Option>, multi_agent_runtime: TurnMultiAgentRuntime, ) -> Arc { - let session_configuration = { - let state = self.state.lock().await; - state.session_configuration.clone() - }; let turn_environments = self.services.turn_environments.snapshot().await; let primary_turn_environment = turn_environments.primary().cloned(); let cwd = primary_turn_environment @@ -825,10 +851,11 @@ impl Session { } pub(crate) async fn new_default_turn_with_sub_id(&self, sub_id: String) -> Arc { - self.new_turn_context_from_state( + let session_configuration = self.default_turn_configuration().await; + self.new_turn_from_configuration( sub_id, + session_configuration, /*final_output_json_schema*/ None, - TurnMultiAgentRuntime::ResolveAndStore, ) .await } @@ -837,11 +864,13 @@ impl Session { &self, sub_id: String, ) -> Arc { - self.new_turn_context_from_state( - sub_id, - /*final_output_json_schema*/ None, - TurnMultiAgentRuntime::Preview, - ) - .await + let session_configuration = self.default_turn_configuration().await; + self.new_startup_prewarm_turn_from_configuration(sub_id, session_configuration) + .await + } + + async fn default_turn_configuration(&self) -> SessionConfiguration { + let state = self.state.lock().await; + state.session_configuration.clone() } } diff --git a/codex-rs/core/src/state/service.rs b/codex-rs/core/src/state/service.rs index 4a5de7287bb3..e6b9f152803d 100644 --- a/codex-rs/core/src/state/service.rs +++ b/codex-rs/core/src/state/service.rs @@ -41,7 +41,6 @@ use tokio::sync::Mutex; use tokio_util::sync::CancellationToken; pub(crate) struct SessionServices { - pub(crate) turn_environments: Arc, /// The latest manager; callers retain an owned handle while performing MCP I/O. pub(crate) mcp_connection_manager: Arc>, pub(crate) mcp_startup_cancellation_token: Mutex, @@ -84,6 +83,7 @@ pub(crate) struct SessionServices { pub(crate) model_client: ModelClient, pub(crate) code_mode_service: CodeModeService, pub(crate) tool_search_handler_cache: ToolSearchHandlerCache, + pub(crate) turn_environments: Arc, } impl SessionServices { diff --git a/codex-rs/core/src/state/session_tests.rs b/codex-rs/core/src/state/session_tests.rs index 00bf830a61bc..0fbb92b958f0 100644 --- a/codex-rs/core/src/state/session_tests.rs +++ b/codex-rs/core/src/state/session_tests.rs @@ -6,15 +6,11 @@ use codex_protocol::protocol::RateLimitWindow; use codex_protocol::protocol::SpendControlLimitSnapshot; use pretty_assertions::assert_eq; -async fn make_session_state() -> SessionState { - let session_configuration = make_session_configuration_for_tests().await; - SessionState::new(session_configuration) -} - #[tokio::test] // Verifies connector merging deduplicates repeated IDs. async fn merge_connector_selection_deduplicates_entries() { - let mut state = make_session_state().await; + let session_configuration = make_session_configuration_for_tests().await; + let mut state = SessionState::new(session_configuration); let merged = state.merge_connector_selection([ "calendar".to_string(), "calendar".to_string(), @@ -30,7 +26,8 @@ async fn merge_connector_selection_deduplicates_entries() { #[tokio::test] // Verifies clearing connector selection removes all saved IDs. async fn clear_connector_selection_removes_entries() { - let mut state = make_session_state().await; + let session_configuration = make_session_configuration_for_tests().await; + let mut state = SessionState::new(session_configuration); state.merge_connector_selection(["calendar".to_string()]); state.clear_connector_selection(); @@ -40,7 +37,8 @@ async fn clear_connector_selection_removes_entries() { #[tokio::test] async fn set_rate_limits_defaults_limit_id_to_codex_when_missing() { - let mut state = make_session_state().await; + let session_configuration = make_session_configuration_for_tests().await; + let mut state = SessionState::new(session_configuration); state.set_rate_limits(RateLimitSnapshot { limit_id: None, @@ -68,7 +66,8 @@ async fn set_rate_limits_defaults_limit_id_to_codex_when_missing() { #[tokio::test] async fn replace_history_clears_auto_compact_window_prefill() { - let mut state = make_session_state().await; + let session_configuration = make_session_configuration_for_tests().await; + let mut state = SessionState::new(session_configuration); state.set_auto_compact_window_estimated_prefill(/*tokens*/ 100); state.replace_history(Vec::new(), /*reference_context_item*/ None); @@ -83,7 +82,8 @@ async fn replace_history_clears_auto_compact_window_prefill() { #[tokio::test] async fn set_rate_limits_defaults_to_codex_when_limit_id_missing_after_other_bucket() { - let mut state = make_session_state().await; + let session_configuration = make_session_configuration_for_tests().await; + let mut state = SessionState::new(session_configuration); state.set_rate_limits(RateLimitSnapshot { limit_id: Some("codex_other".to_string()), @@ -125,7 +125,8 @@ async fn set_rate_limits_defaults_to_codex_when_limit_id_missing_after_other_buc #[tokio::test] async fn set_rate_limits_carries_account_metadata_from_codex_to_codex_other() { - let mut state = make_session_state().await; + let session_configuration = make_session_configuration_for_tests().await; + let mut state = SessionState::new(session_configuration); state.set_rate_limits(RateLimitSnapshot { limit_id: Some("codex".to_string()), diff --git a/codex-rs/core/src/thread_manager.rs b/codex-rs/core/src/thread_manager.rs index adbfcda22ced..dd4852c73227 100644 --- a/codex-rs/core/src/thread_manager.rs +++ b/codex-rs/core/src/thread_manager.rs @@ -1405,6 +1405,7 @@ impl ThreadManagerState { installation_id: self.installation_id.clone(), auth_manager, models_manager: Arc::clone(&self.models_manager), + environment_manager: Arc::clone(&self.environment_manager), skills_manager: Arc::clone(&self.skills_manager), plugins_manager: Arc::clone(&self.plugins_manager), mcp_manager: Arc::clone(&self.mcp_manager), @@ -1422,8 +1423,7 @@ impl ThreadManagerState { parent_rollout_thread_trace, user_shell_override, parent_trace, - environment_manager: Arc::clone(&self.environment_manager), - environments, + environment_selections: environments, thread_extension_init, analytics_events_client: self.analytics_events_client.clone(), thread_store: Arc::clone(&self.thread_store), diff --git a/codex-rs/core/src/thread_manager_tests.rs b/codex-rs/core/src/thread_manager_tests.rs index c6fac14a321b..778fc3669041 100644 --- a/codex-rs/core/src/thread_manager_tests.rs +++ b/codex-rs/core/src/thread_manager_tests.rs @@ -599,9 +599,15 @@ 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.len(), 1); - assert_eq!(resumed_turn.environments[0].cwd(), &default_cwd); - assert_ne!(resumed_turn.environments[0].cwd(), &selected_cwd); + assert_eq!(resumed_turn.environments.turn_environments.len(), 1); + assert_eq!( + resumed_turn.environments.turn_environments[0].cwd(), + &default_cwd + ); + assert_ne!( + resumed_turn.environments.turn_environments[0].cwd(), + &selected_cwd + ); let forked = manager .fork_thread( @@ -620,9 +626,15 @@ 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.len(), 1); - assert_eq!(forked_turn.environments[0].cwd(), &default_cwd); - assert_ne!(forked_turn.environments[0].cwd(), &selected_cwd); + assert_eq!(forked_turn.environments.turn_environments.len(), 1); + assert_eq!( + forked_turn.environments.turn_environments[0].cwd(), + &default_cwd + ); + assert_ne!( + forked_turn.environments.turn_environments[0].cwd(), + &selected_cwd + ); } #[tokio::test] 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 01ea1f61ff8c..2556989621cb 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,7 @@ pub async fn handle( } fn single_local_environment_cwd(turn: &TurnContext) -> Result<&AbsolutePathBuf, FunctionCallError> { - let [turn_environment] = turn.environments.as_slice() else { + let [turn_environment] = turn.environments.turn_environments.as_slice() else { 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 b5fa5702fc37..33625317fe40 100644 --- a/codex-rs/core/src/tools/handlers/extension_tools.rs +++ b/codex-rs/core/src/tools/handlers/extension_tools.rs @@ -112,8 +112,8 @@ impl TurnItemEmitter for CoreTurnItemEmitter { async fn to_extension_call(invocation: &ToolInvocation) -> ExtensionToolCall { let conversation_history = ConversationHistory::new(invocation.session.clone_history().await.into_raw_items()); - let mut environments = Vec::with_capacity(invocation.turn.environments.len()); - for environment in invocation.turn.environments.iter() { + let mut environments = Vec::with_capacity(invocation.turn.environments.turn_environments.len()); + for environment in &invocation.turn.environments.turn_environments { let additional_permissions = apply_granted_turn_permissions( invocation.session.as_ref(), &environment.environment_id, @@ -313,6 +313,7 @@ mod tests { let truncation_policy = turn.truncation_policy; let expected_sandbox_cwds = turn .environments + .turn_environments .iter() .map(|environment| Some(environment.cwd_uri().clone())) .collect::>(); diff --git a/codex-rs/core/src/tools/handlers/mod.rs b/codex-rs/core/src/tools/handlers/mod.rs index a946014f59c4..7791a819a530 100644 --- a/codex-rs/core/src/tools/handlers/mod.rs +++ b/codex-rs/core/src/tools/handlers/mod.rs @@ -156,6 +156,7 @@ fn resolve_tool_environment<'a>( || Ok(turn.environments.primary()), |environment_id| { turn.environments + .turn_environments .iter() .find(|environment| environment.environment_id == environment_id) .map(Some) diff --git a/codex-rs/core/src/tools/handlers/view_image.rs b/codex-rs/core/src/tools/handlers/view_image.rs index 7cbb36d65fda..dc52e7e336c7 100644 --- a/codex-rs/core/src/tools/handlers/view_image.rs +++ b/codex-rs/core/src/tools/handlers/view_image.rs @@ -281,10 +281,11 @@ mod tests { fn replace_primary_environment_cwd(turn: &mut crate::TurnContext, cwd: AbsolutePathBuf) { let current = turn .environments + .turn_environments .first() .cloned() .expect("default local turn environment"); - turn.environments.environments_mut()[0] = TurnEnvironment::new( + turn.environments.turn_environments[0] = TurnEnvironment::new( current.environment_id, current.environment, cwd, diff --git a/codex-rs/core/src/tools/spec_plan.rs b/codex-rs/core/src/tools/spec_plan.rs index 42f33475d297..872847611053 100644 --- a/codex-rs/core/src/tools/spec_plan.rs +++ b/codex-rs/core/src/tools/spec_plan.rs @@ -628,6 +628,7 @@ fn unified_exec_should_include_shell_parameter(turn_context: &TurnContext) -> bo UnifiedExecShellMode::ZshFork(_) ) || turn_context .environments + .turn_environments .iter() .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 c5a0fb39f378..a1bf22310ae1 100644 --- a/codex-rs/core/src/tools/spec_plan_tests.rs +++ b/codex-rs/core/src/tools/spec_plan_tests.rs @@ -349,11 +349,9 @@ impl ToolExecutor for DeferredExtensionTool { } fn duplicate_primary_environment(turn: &mut TurnContext) { - let mut second_environment = turn.environments[0].clone(); + let mut second_environment = turn.environments.turn_environments[0].clone(); second_environment.environment_id = "secondary".to_string(); - turn.environments - .environments_mut() - .push(second_environment); + turn.environments.turn_environments.push(second_environment); } fn mcp_tool(server: &str, namespace: &str, name: &str) -> ToolInfo { @@ -586,11 +584,11 @@ async fn zsh_fork_unified_exec_keeps_shell_parameter_when_remote_environment_ava codex_tools::UnifiedExecShellMode::ZshFork(zsh_fork_config_for_spec_plan_tests()); let remote_cwd = turn .environments - .first() + .primary() .expect("primary environment") .cwd() .clone(); - turn.environments.environments_mut().push( + turn.environments.turn_environments.push( crate::session::turn_context::TurnEnvironment::new( "remote".to_string(), Arc::new( @@ -617,7 +615,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.environments_mut().clear(); + turn.environments.turn_environments.clear(); set_feature(turn, Feature::ShellTool, /*enabled*/ true); turn.model_info.apply_patch_tool_type = Some(ApplyPatchToolType::Freeform); }) diff --git a/codex-rs/core/src/unified_exec/mod_tests.rs b/codex-rs/core/src/unified_exec/mod_tests.rs index c4952b32c78f..ad39d62b9663 100644 --- a/codex-rs/core/src/unified_exec/mod_tests.rs +++ b/codex-rs/core/src/unified_exec/mod_tests.rs @@ -112,7 +112,7 @@ async fn exec_command_with_tty( tty, Box::new(NoopSpawnLifecycle), turn.environments - .first() + .primary() .expect("turn environment") .environment .as_ref(), @@ -894,7 +894,7 @@ 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.environments_mut()[0].environment = + turn.environments.turn_environments[0].environment = Arc::new(remote_test_env.environment().clone()); #[allow(deprecated)] @@ -916,7 +916,7 @@ async fn remote_exec_server_rejects_inherited_fd_launches() -> anyhow::Result<() inherited_fds: vec![42], }), turn.environments - .first() + .primary() .expect("turn environment") .environment .as_ref(), diff --git a/codex-rs/core/src/unified_exec/process_manager_tests.rs b/codex-rs/core/src/unified_exec/process_manager_tests.rs index 96a9bb0391b6..a84f6345ede6 100644 --- a/codex-rs/core/src/unified_exec/process_manager_tests.rs +++ b/codex-rs/core/src/unified_exec/process_manager_tests.rs @@ -205,8 +205,7 @@ async fn failed_initial_end_for_unstored_process_uses_fallback_output() { sandbox_cwd: turn.cwd.clone(), environment: turn .environments - .first() - .map(|environment| Arc::clone(&environment.environment)) + .primary_environment() .expect("primary environment"), shell_mode: codex_tools::UnifiedExecShellMode::Direct, network: None, From f4c74a4c46a83ee25016477fe6c85a0f404140cb Mon Sep 17 00:00:00 2001 From: pakrym-oai Date: Mon, 15 Jun 2026 14:23:42 -0700 Subject: [PATCH 10/14] core: split thread environment initialization --- codex-rs/core/src/agents_md.rs | 4 +- codex-rs/core/src/agents_md_tests.rs | 6 +-- codex-rs/core/src/environment_selection.rs | 50 ++++++++++++---------- codex-rs/core/src/session/mod.rs | 8 ++-- codex-rs/core/src/session/session.rs | 6 +-- codex-rs/core/src/session/tests.rs | 14 +++--- codex-rs/core/src/session/turn_context.rs | 6 +-- codex-rs/core/src/state/service.rs | 4 +- 8 files changed, 53 insertions(+), 45 deletions(-) diff --git a/codex-rs/core/src/agents_md.rs b/codex-rs/core/src/agents_md.rs index dfcba497ab3e..52a386d694e9 100644 --- a/codex-rs/core/src/agents_md.rs +++ b/codex-rs/core/src/agents_md.rs @@ -18,7 +18,7 @@ use crate::config::Config; use crate::context::ContextualUserFragment; use crate::context::UserInstructions as ContextUserInstructions; -use crate::environment_selection::TurnEnvironmentsSnapshot; +use crate::environment_selection::TurnEnvironmentSnapshot; use codex_app_server_protocol::ConfigLayerSource; use codex_config::ConfigLayerStackOrdering; use codex_config::default_project_root_markers; @@ -48,7 +48,7 @@ const AGENTS_MD_SEPARATOR: &str = "\n\n--- project-doc ---\n\n"; pub(crate) async fn load_project_instructions( config: &mut Config, user_instructions: Option, - environments: &TurnEnvironmentsSnapshot, + environments: &TurnEnvironmentSnapshot, ) -> Option { let mut loaded = LoadedAgentsMd::from_user_instructions(user_instructions); for turn_environment in &environments.turn_environments { diff --git a/codex-rs/core/src/agents_md_tests.rs b/codex-rs/core/src/agents_md_tests.rs index 0f23fb014afc..565b7f4b999a 100644 --- a/codex-rs/core/src/agents_md_tests.rs +++ b/codex-rs/core/src/agents_md_tests.rs @@ -1,6 +1,6 @@ use super::*; use crate::config::ConfigBuilder; -use crate::environment_selection::TurnEnvironmentsSnapshot; +use crate::environment_selection::TurnEnvironmentSnapshot; use crate::session::turn_context::TurnEnvironment; use codex_config::ConfigLayerEntry; use codex_config::ConfigLayerStack; @@ -254,8 +254,8 @@ async fn agents_md_paths(config: &TestConfig) -> std::io::Result( environments: [(&str, AbsolutePathBuf); N], -) -> TurnEnvironmentsSnapshot { - TurnEnvironmentsSnapshot { +) -> TurnEnvironmentSnapshot { + TurnEnvironmentSnapshot { turn_environments: environments .into_iter() .map(|(environment_id, cwd)| { diff --git a/codex-rs/core/src/environment_selection.rs b/codex-rs/core/src/environment_selection.rs index c615193fc5ef..a6ad66dc52ea 100644 --- a/codex-rs/core/src/environment_selection.rs +++ b/codex-rs/core/src/environment_selection.rs @@ -28,22 +28,17 @@ pub(crate) fn default_thread_environment_selections( } #[derive(Debug)] -pub(crate) struct TurnEnvironments { +pub(crate) struct ThreadEnvironments { environment_manager: Arc, - snapshot: Mutex, + snapshot: Mutex, } -impl TurnEnvironments { - pub(crate) async fn resolve( - environment_manager: Arc, - environments: &[TurnEnvironmentSelection], - ) -> Arc { - let resolved = Arc::new(Self { +impl ThreadEnvironments { + pub(crate) fn new(environment_manager: Arc) -> Self { + Self { environment_manager, - snapshot: Mutex::new(TurnEnvironmentsSnapshot::default()), - }); - resolved.update_selections(environments).await; - resolved + snapshot: Mutex::new(TurnEnvironmentSnapshot::default()), + } } pub(crate) async fn update_selections(&self, environments: &[TurnEnvironmentSelection]) { @@ -72,7 +67,7 @@ impl TurnEnvironments { }; turn_environments.push(turn_environment); } - *self.snapshot.lock().await = TurnEnvironmentsSnapshot { turn_environments }; + *self.snapshot.lock().await = TurnEnvironmentSnapshot { turn_environments }; } async fn resolve_selection( @@ -114,7 +109,7 @@ impl TurnEnvironments { )) } - pub(crate) async fn snapshot(&self) -> TurnEnvironmentsSnapshot { + pub(crate) async fn snapshot(&self) -> TurnEnvironmentSnapshot { self.snapshot.lock().await.clone() } @@ -124,11 +119,11 @@ impl TurnEnvironments { } #[derive(Clone, Debug, Default)] -pub(crate) struct TurnEnvironmentsSnapshot { +pub(crate) struct TurnEnvironmentSnapshot { pub(crate) turn_environments: Vec, } -impl TurnEnvironmentsSnapshot { +impl TurnEnvironmentSnapshot { pub(crate) fn primary(&self) -> Option<&TurnEnvironment> { self.turn_environments.first() } @@ -173,6 +168,15 @@ mod tests { use super::*; + async fn resolve_turn_environments( + environment_manager: Arc, + selections: &[TurnEnvironmentSelection], + ) -> Arc { + let turn_environments = Arc::new(ThreadEnvironments::new(environment_manager)); + turn_environments.update_selections(selections).await; + turn_environments + } + fn test_runtime_paths() -> ExecServerRuntimePaths { ExecServerRuntimePaths::new( std::env::current_exe().expect("current exe"), @@ -255,7 +259,7 @@ url = "ws://127.0.0.1:8765" cwd: cwd_uri.clone(), }; - let resolved = TurnEnvironments::resolve( + let resolved = resolve_turn_environments( manager, &[ first.clone(), @@ -277,7 +281,7 @@ url = "ws://127.0.0.1:8765" let selected_cwd_uri = PathUri::from_abs_path(&selected_cwd); let manager = Arc::new(EnvironmentManager::default_for_tests()); - let resolved = TurnEnvironments::resolve( + let resolved = resolve_turn_environments( Arc::clone(&manager), &[TurnEnvironmentSelection { environment_id: "local".to_string(), @@ -321,7 +325,7 @@ url = "ws://127.0.0.1:8765" cwd: cwd_uri.clone(), }; - let resolved = TurnEnvironments::resolve( + let resolved = resolve_turn_environments( manager, &[ TurnEnvironmentSelection { @@ -351,7 +355,7 @@ url = "ws://127.0.0.1:8765" cwd: PathUri::from_abs_path(&cwd), }; let initial = - TurnEnvironments::resolve(Arc::clone(&manager), std::slice::from_ref(&selection)).await; + resolve_turn_environments(Arc::clone(&manager), std::slice::from_ref(&selection)).await; manager .upsert_environment( REMOTE_ENVIRONMENT_ID.to_string(), @@ -399,7 +403,7 @@ url = "ws://127.0.0.1:8765" let cwd = AbsolutePathBuf::current_dir().expect("cwd"); let cwd_uri = PathUri::from_abs_path(&cwd); let local_manager = Arc::new(EnvironmentManager::default_for_tests()); - let local = TurnEnvironments::resolve( + let local = resolve_turn_environments( Arc::clone(&local_manager), &[TurnEnvironmentSelection { environment_id: LOCAL_ENVIRONMENT_ID.to_string(), @@ -412,7 +416,7 @@ url = "ws://127.0.0.1:8765" Environment::create_for_tests(Some("ws://127.0.0.1:8765".to_string())) .expect("remote environment"), ); - let remote = TurnEnvironmentsSnapshot { + let remote = TurnEnvironmentSnapshot { turn_environments: vec![TurnEnvironment::new( REMOTE_ENVIRONMENT_ID.to_string(), remote_environment.clone(), @@ -420,7 +424,7 @@ url = "ws://127.0.0.1:8765" /*shell*/ None, )], }; - let multiple = TurnEnvironmentsSnapshot { + let multiple = TurnEnvironmentSnapshot { turn_environments: vec![ local.primary().expect("local environment").clone(), TurnEnvironment::new( diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 567afe77603e..14f4877d73ee 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -31,7 +31,7 @@ use crate::context::NetworkRuleSaved; use crate::context::PermissionsInstructions; use crate::context::PersonalitySpecInstructions; use crate::default_skill_metadata_budget; -use crate::environment_selection::TurnEnvironments; +use crate::environment_selection::ThreadEnvironments; use crate::exec_policy::ExecPolicyManager; use crate::image_preparation::prepare_response_items; use crate::parse_turn_item; @@ -519,8 +519,10 @@ impl Codex { attestation_provider, inherited_multi_agent_version, } = args; - let turn_environments = - TurnEnvironments::resolve(environment_manager, &environment_selections).await; + let turn_environments = Arc::new(ThreadEnvironments::new(environment_manager)); + turn_environments + .update_selections(&environment_selections) + .await; let resolved_environments = turn_environments.snapshot().await; let (tx_sub, rx_sub) = async_channel::bounded(SUBMISSION_CHANNEL_CAPACITY); let (tx_event, rx_event) = async_channel::unbounded(); diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index bf5bcf085b3b..79f5500377f4 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -2,7 +2,7 @@ use super::input_queue::InputQueue; use super::*; use crate::agents_md::LoadedAgentsMd; use crate::config::ConstraintError; -use crate::environment_selection::TurnEnvironmentsSnapshot; +use crate::environment_selection::TurnEnvironmentSnapshot; use crate::skills::SkillError; use crate::state::ActiveTurn; use codex_extension_api::ExtensionDataInit; @@ -438,7 +438,7 @@ async fn warm_plugins_and_skills_for_session_init( config: Arc, plugins_manager: Arc, skills_manager: Arc, - turn_environments: TurnEnvironmentsSnapshot, + turn_environments: TurnEnvironmentSnapshot, ) -> Vec { let fs = turn_environments.primary_filesystem(); let plugins_input = config.plugins_config_input(); @@ -481,7 +481,7 @@ impl Session { extensions: Arc>, thread_extension_init: ExtensionDataInit, agent_control: AgentControl, - turn_environments: Arc, + turn_environments: Arc, analytics_events_client: Option, thread_store: Arc, parent_rollout_thread_trace: ThreadTraceContext, diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index 004dd9eb0512..dfd3f122916f 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -4034,12 +4034,14 @@ async fn emit_subagent_session_started_includes_fork_lineage_from_session_config async fn turn_environments_for_configuration( session_configuration: &SessionConfiguration, -) -> Arc { - TurnEnvironments::resolve( - Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), - session_configuration.environment_selections(), - ) - .await +) -> Arc { + let turn_environments = Arc::new(ThreadEnvironments::new(Arc::new( + codex_exec_server::EnvironmentManager::default_for_tests(), + ))); + turn_environments + .update_selections(session_configuration.environment_selections()) + .await; + turn_environments } #[tokio::test] diff --git a/codex-rs/core/src/session/turn_context.rs b/codex-rs/core/src/session/turn_context.rs index f7e13719337b..ab0e4a34520c 100644 --- a/codex-rs/core/src/session/turn_context.rs +++ b/codex-rs/core/src/session/turn_context.rs @@ -2,7 +2,7 @@ use super::*; use crate::SkillLoadOutcome; use crate::agents_md::LoadedAgentsMd; use crate::config::GhostSnapshotConfig; -use crate::environment_selection::TurnEnvironmentsSnapshot; +use crate::environment_selection::TurnEnvironmentSnapshot; use codex_core_skills::HostLoadedSkills; use codex_model_provider::SharedModelProvider; use codex_model_provider::create_model_provider; @@ -102,7 +102,7 @@ pub struct TurnContext { pub(crate) session_source: SessionSource, pub(crate) parent_thread_id: Option, pub(crate) thread_source: Option, - pub(crate) environments: TurnEnvironmentsSnapshot, + pub(crate) environments: TurnEnvironmentSnapshot, /// The session's absolute working directory. All relative paths provided /// by the model as well as sandbox policies are resolved against this path /// instead of `std::env::current_dir()`. @@ -502,7 +502,7 @@ impl Session { model_info: ModelInfo, models_manager: &SharedModelsManager, network: Option, - environments: TurnEnvironmentsSnapshot, + environments: TurnEnvironmentSnapshot, cwd: AbsolutePathBuf, sub_id: String, skills_outcome: Arc, diff --git a/codex-rs/core/src/state/service.rs b/codex-rs/core/src/state/service.rs index e6b9f152803d..5d888bd83dc3 100644 --- a/codex-rs/core/src/state/service.rs +++ b/codex-rs/core/src/state/service.rs @@ -7,7 +7,7 @@ use crate::attestation::AttestationProvider; use crate::client::ModelClient; use crate::config::NetworkProxyAuditMetadata; use crate::config::StartedNetworkProxy; -use crate::environment_selection::TurnEnvironments; +use crate::environment_selection::ThreadEnvironments; use crate::exec_policy::ExecPolicyManager; use crate::guardian::GuardianRejection; use crate::guardian::GuardianRejectionCircuitBreaker; @@ -83,7 +83,7 @@ pub(crate) struct SessionServices { pub(crate) model_client: ModelClient, pub(crate) code_mode_service: CodeModeService, pub(crate) tool_search_handler_cache: ToolSearchHandlerCache, - pub(crate) turn_environments: Arc, + pub(crate) turn_environments: Arc, } impl SessionServices { From f5ea28cf215802c1050812ad2a9bde791b5f81e0 Mon Sep 17 00:00:00 2001 From: pakrym-oai Date: Mon, 15 Jun 2026 15:39:14 -0700 Subject: [PATCH 11/14] fix(core): resolve latest thread environments eagerly --- codex-rs/core/src/environment_selection.rs | 120 +++++++++++++++++---- codex-rs/core/src/session/mod.rs | 15 ++- codex-rs/core/src/session/tests.rs | 13 +-- codex-rs/core/src/session/turn_context.rs | 11 +- 4 files changed, 111 insertions(+), 48 deletions(-) diff --git a/codex-rs/core/src/environment_selection.rs b/codex-rs/core/src/environment_selection.rs index a6ad66dc52ea..1e014142c0d3 100644 --- a/codex-rs/core/src/environment_selection.rs +++ b/codex-rs/core/src/environment_selection.rs @@ -1,6 +1,7 @@ use std::collections::HashSet; use std::sync::Arc; +use arc_swap::ArcSwap; use codex_exec_server::EnvironmentManager; use codex_exec_server::ExecutorFileSystem; use codex_protocol::error::CodexErr; @@ -8,7 +9,10 @@ use codex_protocol::error::Result as CodexResult; use codex_protocol::protocol::TurnEnvironmentSelection; use codex_utils_absolute_path::AbsolutePathBuf; use codex_utils_path_uri::PathUri; -use tokio::sync::Mutex; +use futures::FutureExt; +use futures::future::BoxFuture; +use futures::future::Shared; +use tokio_util::task::AbortOnDropHandle; use crate::session::turn_context::TurnEnvironment; use crate::shell::Shell; @@ -27,25 +31,55 @@ pub(crate) fn default_thread_environment_selections( .collect() } -#[derive(Debug)] +type SnapshotTask = Shared>; + pub(crate) struct ThreadEnvironments { environment_manager: Arc, - snapshot: Mutex, + snapshot_task: ArcSwap, } impl ThreadEnvironments { pub(crate) fn new(environment_manager: Arc) -> Self { Self { environment_manager, - snapshot: Mutex::new(TurnEnvironmentSnapshot::default()), + snapshot_task: ArcSwap::from_pointee( + futures::future::ready(TurnEnvironmentSnapshot::default()) + .boxed() + .shared(), + ), + } + } + + pub(crate) fn update_selections(&self, environments: &[TurnEnvironmentSelection]) { + let previous = self + .snapshot_task + .load() + .peek() + .cloned() + .unwrap_or_default(); + let environment_manager = Arc::clone(&self.environment_manager); + let environments = environments.to_vec(); + let snapshot_task = AbortOnDropHandle::new(tokio::spawn(async move { + Self::resolve_snapshot(environment_manager, previous, environments).await + })); + let snapshot_task = async move { + snapshot_task + .await + .expect("environment resolution task should not panic") } + .boxed() + .shared(); + self.snapshot_task.store(Arc::new(snapshot_task)); } - pub(crate) async fn update_selections(&self, environments: &[TurnEnvironmentSelection]) { - let current = self.snapshot().await; + async fn resolve_snapshot( + environment_manager: Arc, + current: TurnEnvironmentSnapshot, + environments: Vec, + ) -> TurnEnvironmentSnapshot { let mut seen_environment_ids = HashSet::with_capacity(environments.len()); let mut turn_environments = Vec::with_capacity(environments.len()); - for selected_environment in environments { + for selected_environment in &environments { if !seen_environment_ids.insert(selected_environment.environment_id.as_str()) { continue; } @@ -54,7 +88,9 @@ impl ThreadEnvironments { && environment.cwd_uri() == &selected_environment.cwd }) { Some(environment) => environment.clone(), - None => match self.resolve_selection(selected_environment).await { + None => match Self::resolve_selection(&environment_manager, selected_environment) + .await + { Ok(environment) => environment, Err(err) => { tracing::warn!( @@ -67,16 +103,15 @@ impl ThreadEnvironments { }; turn_environments.push(turn_environment); } - *self.snapshot.lock().await = TurnEnvironmentSnapshot { turn_environments }; + TurnEnvironmentSnapshot { turn_environments } } async fn resolve_selection( - &self, + environment_manager: &EnvironmentManager, selected_environment: &TurnEnvironmentSelection, ) -> CodexResult { let environment_id = selected_environment.environment_id.clone(); - let environment = self - .environment_manager + let environment = environment_manager .get_environment(&environment_id) .ok_or_else(|| { CodexErr::InvalidRequest(format!("unknown turn environment id `{environment_id}`")) @@ -110,7 +145,7 @@ impl ThreadEnvironments { } pub(crate) async fn snapshot(&self) -> TurnEnvironmentSnapshot { - self.snapshot.lock().await.clone() + self.snapshot_task.load_full().as_ref().clone().await } pub(crate) fn environment_manager(&self) -> Arc { @@ -173,7 +208,8 @@ mod tests { selections: &[TurnEnvironmentSelection], ) -> Arc { let turn_environments = Arc::new(ThreadEnvironments::new(environment_manager)); - turn_environments.update_selections(selections).await; + turn_environments.update_selections(selections); + turn_environments.snapshot().await; turn_environments } @@ -340,6 +376,48 @@ url = "ws://127.0.0.1:8765" assert_eq!(resolved.snapshot().await.to_selections(), vec![local]); } + #[tokio::test] + async fn latest_environment_update_wins_while_previous_resolution_is_pending() { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0") + .await + .expect("bind websocket listener"); + let manager = Arc::new( + EnvironmentManager::create_for_tests_with_local( + Some(format!( + "ws://{}", + listener.local_addr().expect("listener address") + )), + test_runtime_paths(), + ) + .await, + ); + let cwd = AbsolutePathBuf::current_dir().expect("cwd"); + let turn_environments = Arc::new(ThreadEnvironments::new(manager)); + turn_environments.update_selections(&[TurnEnvironmentSelection { + environment_id: REMOTE_ENVIRONMENT_ID.to_string(), + cwd: PathUri::from_abs_path(&cwd), + }]); + let (_connection, _) = + tokio::time::timeout(std::time::Duration::from_secs(5), listener.accept()) + .await + .expect("remote resolution should connect") + .expect("accept remote resolution connection"); + let local = TurnEnvironmentSelection { + environment_id: LOCAL_ENVIRONMENT_ID.to_string(), + cwd: PathUri::from_abs_path(&cwd), + }; + + turn_environments.update_selections(std::slice::from_ref(&local)); + let snapshot = tokio::time::timeout( + std::time::Duration::from_secs(5), + turn_environments.snapshot(), + ) + .await + .expect("latest environment resolution should complete"); + + assert_eq!(snapshot.to_selections(), vec![local]); + } + #[tokio::test] async fn matching_environment_id_and_cwd_reuse_resolved_environment() { let cwd = AbsolutePathBuf::current_dir().expect("cwd"); @@ -364,16 +442,12 @@ url = "ws://127.0.0.1:8765" .expect("replace environment"); let initial_snapshot = initial.snapshot().await; - initial - .update_selections(std::slice::from_ref(&selection)) - .await; + initial.update_selections(std::slice::from_ref(&selection)); let reused_snapshot = initial.snapshot().await; - initial - .update_selections(&[TurnEnvironmentSelection { - cwd: PathUri::from_abs_path(&cwd.join("changed")), - ..selection - }]) - .await; + initial.update_selections(&[TurnEnvironmentSelection { + cwd: PathUri::from_abs_path(&cwd.join("changed")), + ..selection + }]); let changed_snapshot = initial.snapshot().await; assert!(Arc::ptr_eq( diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 14f4877d73ee..fb28987cba92 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -520,9 +520,7 @@ impl Codex { inherited_multi_agent_version, } = args; let turn_environments = Arc::new(ThreadEnvironments::new(environment_manager)); - turn_environments - .update_selections(&environment_selections) - .await; + turn_environments.update_selections(&environment_selections); let resolved_environments = turn_environments.snapshot().await; let (tx_sub, rx_sub) = async_channel::bounded(SUBMISSION_CHANNEL_CAPACITY); let (tx_event, rx_event) = async_channel::unbounded(); @@ -1466,6 +1464,11 @@ impl Session { let next_cwd = updated.cwd().clone(); let codex_home = updated.codex_home.clone(); let session_source = updated.session_source.clone(); + if updates.environments.is_some() { + self.services + .turn_environments + .update_selections(updated.environment_selections()); + } state.session_configuration = updated; ( previous_config, @@ -1477,12 +1480,6 @@ impl Session { session_source, ) }; - if let Some(environments) = &updates.environments { - self.services - .turn_environments - .update_selections(&environments.environments) - .await; - } self.emit_config_changed_contributors(previous_config.as_ref(), new_config.as_ref()); self.maybe_refresh_shell_snapshot_for_cwd( &previous_cwd, diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index dfd3f122916f..dd2634a148ef 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -4038,9 +4038,7 @@ async fn turn_environments_for_configuration( let turn_environments = Arc::new(ThreadEnvironments::new(Arc::new( codex_exec_server::EnvironmentManager::default_for_tests(), ))); - turn_environments - .update_selections(session_configuration.environment_selections()) - .await; + turn_environments.update_selections(session_configuration.environment_selections()); turn_environments } @@ -6241,8 +6239,7 @@ async fn default_turn_does_not_overlay_legacy_fallback_cwd_onto_stored_thread_en session .services .turn_environments - .update_selections(&[local(selected_cwd.clone())]) - .await; + .update_selections(&[local(selected_cwd.clone())]); { let mut state = session.state.lock().await; state.session_configuration.environments.environments = vec![local(selected_cwd.clone())]; @@ -6271,11 +6268,7 @@ async fn default_turn_honors_empty_stored_thread_environments() { let (session, _turn_context, _rx) = make_session_and_context_with_rx().await; let session_cwd = session.get_config().await.cwd.clone(); - session - .services - .turn_environments - .update_selections(&[]) - .await; + session.services.turn_environments.update_selections(&[]); { let mut state = session.state.lock().await; state.session_configuration.environments.environments = Vec::new(); diff --git a/codex-rs/core/src/session/turn_context.rs b/codex-rs/core/src/session/turn_context.rs index ab0e4a34520c..57f240ebb241 100644 --- a/codex-rs/core/src/session/turn_context.rs +++ b/codex-rs/core/src/session/turn_context.rs @@ -635,6 +635,11 @@ impl Session { }); let new_config = notify_config_contributors .then(|| Self::build_effective_session_config(&next)); + if updates.environments.is_some() { + self.services + .turn_environments + .update_selections(next.environment_selections()); + } state.session_configuration = next.clone(); Ok(( next, @@ -673,12 +678,6 @@ impl Session { return Err(CodexErr::InvalidRequest(message)); } }; - if let Some(environments) = &updates.environments { - self.services - .turn_environments - .update_selections(&environments.environments) - .await; - } self.emit_config_changed_contributors(previous_config.as_ref(), new_config.as_ref()); self.maybe_refresh_shell_snapshot_for_cwd( &previous_cwd, From b0970f98bd94800cab88d78822d45ba6fd1b3036 Mon Sep 17 00:00:00 2001 From: pakrym-oai Date: Mon, 15 Jun 2026 15:46:17 -0700 Subject: [PATCH 12/14] refactor(core): simplify environment snapshot task --- codex-rs/core/src/environment_selection.rs | 15 +++++---------- 1 file changed, 5 insertions(+), 10 deletions(-) diff --git a/codex-rs/core/src/environment_selection.rs b/codex-rs/core/src/environment_selection.rs index 1e014142c0d3..15807fd55837 100644 --- a/codex-rs/core/src/environment_selection.rs +++ b/codex-rs/core/src/environment_selection.rs @@ -12,7 +12,6 @@ use codex_utils_path_uri::PathUri; use futures::FutureExt; use futures::future::BoxFuture; use futures::future::Shared; -use tokio_util::task::AbortOnDropHandle; use crate::session::turn_context::TurnEnvironment; use crate::shell::Shell; @@ -59,17 +58,13 @@ impl ThreadEnvironments { .unwrap_or_default(); let environment_manager = Arc::clone(&self.environment_manager); let environments = environments.to_vec(); - let snapshot_task = AbortOnDropHandle::new(tokio::spawn(async move { + let (snapshot_task, snapshot) = async move { Self::resolve_snapshot(environment_manager, previous, environments).await - })); - let snapshot_task = async move { - snapshot_task - .await - .expect("environment resolution task should not panic") } - .boxed() - .shared(); - self.snapshot_task.store(Arc::new(snapshot_task)); + .remote_handle(); + drop(tokio::spawn(snapshot_task)); + self.snapshot_task + .store(Arc::new(snapshot.boxed().shared())); } async fn resolve_snapshot( From 23d04357c7870fea60111700ae4e29b6172c6c49 Mon Sep 17 00:00:00 2001 From: pakrym-oai Date: Mon, 15 Jun 2026 15:52:11 -0700 Subject: [PATCH 13/14] fix(core): publish environment task before spawning --- codex-rs/core/src/environment_selection.rs | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/codex-rs/core/src/environment_selection.rs b/codex-rs/core/src/environment_selection.rs index 15807fd55837..d9763a84fc83 100644 --- a/codex-rs/core/src/environment_selection.rs +++ b/codex-rs/core/src/environment_selection.rs @@ -1,7 +1,7 @@ use std::collections::HashSet; use std::sync::Arc; -use arc_swap::ArcSwap; +use arc_swap::ArcSwapAny; use codex_exec_server::EnvironmentManager; use codex_exec_server::ExecutorFileSystem; use codex_protocol::error::CodexErr; @@ -30,22 +30,22 @@ pub(crate) fn default_thread_environment_selections( .collect() } -type SnapshotTask = Shared>; +type SharedSnapshotTask = Arc>>; pub(crate) struct ThreadEnvironments { environment_manager: Arc, - snapshot_task: ArcSwap, + snapshot_task: ArcSwapAny, } impl ThreadEnvironments { pub(crate) fn new(environment_manager: Arc) -> Self { Self { environment_manager, - snapshot_task: ArcSwap::from_pointee( + snapshot_task: ArcSwapAny::new(Arc::new( futures::future::ready(TurnEnvironmentSnapshot::default()) .boxed() .shared(), - ), + )), } } @@ -62,9 +62,9 @@ impl ThreadEnvironments { Self::resolve_snapshot(environment_manager, previous, environments).await } .remote_handle(); - drop(tokio::spawn(snapshot_task)); self.snapshot_task .store(Arc::new(snapshot.boxed().shared())); + drop(tokio::spawn(snapshot_task)); } async fn resolve_snapshot( From ee626108cbb138883ab9390c02f555869f620ed4 Mon Sep 17 00:00:00 2001 From: pakrym-oai Date: Mon, 15 Jun 2026 15:54:36 -0700 Subject: [PATCH 14/14] refactor(core): simplify snapshot task storage --- codex-rs/core/src/environment_selection.rs | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/codex-rs/core/src/environment_selection.rs b/codex-rs/core/src/environment_selection.rs index d9763a84fc83..e1ab2ea93f25 100644 --- a/codex-rs/core/src/environment_selection.rs +++ b/codex-rs/core/src/environment_selection.rs @@ -1,7 +1,7 @@ use std::collections::HashSet; use std::sync::Arc; -use arc_swap::ArcSwapAny; +use arc_swap::ArcSwap; use codex_exec_server::EnvironmentManager; use codex_exec_server::ExecutorFileSystem; use codex_protocol::error::CodexErr; @@ -30,22 +30,22 @@ pub(crate) fn default_thread_environment_selections( .collect() } -type SharedSnapshotTask = Arc>>; +type SnapshotTask = Shared>; pub(crate) struct ThreadEnvironments { environment_manager: Arc, - snapshot_task: ArcSwapAny, + snapshot_task: ArcSwap, } impl ThreadEnvironments { pub(crate) fn new(environment_manager: Arc) -> Self { Self { environment_manager, - snapshot_task: ArcSwapAny::new(Arc::new( + snapshot_task: ArcSwap::from_pointee( futures::future::ready(TurnEnvironmentSnapshot::default()) .boxed() .shared(), - )), + ), } }