From 9728992fab1aaba5d1ca801f22402b7c8c3e5760 Mon Sep 17 00:00:00 2001 From: pakrym-oai Date: Mon, 15 Jun 2026 16:15:07 -0700 Subject: [PATCH] [codex] retain resolved environments across turns (#27955) ## Why Selected execution environments are thread-scoped resources, but startup and turn construction repeatedly resolved their IDs and working directories. That discarded existing environment handles and shell metadata even when a selection had not changed. Session configuration updates also need to affect future turns without changing the resolved environment set already captured by a running turn. ## What changed - Create a `ThreadEnvironments` service inside `Codex` from the spawned `EnvironmentManager` and raw environment selections, then store it on `SessionServices`. - Split service construction from `update_selections`, allowing session configuration updates to mutate the resolved set in place. - Retain an existing `TurnEnvironment` when its environment ID and working directory match; resolve only added or changed selections and remove selections that are no longer present. - Normalize duplicate IDs by keeping the first selection and skip individual selections that fail to resolve instead of rejecting the entire update. - Give each `TurnContext` a cloned `TurnEnvironmentSnapshot`, so later session configuration updates affect future turns without rewriting an active turn. - Reuse the service-owned environment manager and resolved snapshot for startup work, MCP initialization, and child-thread spawning instead of flowing resolved environments through spawn arguments. ## Test plan - `cargo check -p codex-core --tests` - `just test -p codex-core environment_selection` - `just test -p codex-core turn_environments` - `just test -p codex-core session_update_settings_does_not_rewrite_sticky_environment_cwds` - `just test -p codex-core default_turn_does_not_overlay_legacy_fallback_cwd_onto_stored_thread_environments` --- codex-rs/core/src/agents_md.rs | 4 +- codex-rs/core/src/agents_md_tests.rs | 6 +- codex-rs/core/src/codex_delegate.rs | 7 +- codex-rs/core/src/environment_selection.rs | 352 ++++++++++++++---- codex-rs/core/src/mcp_tool_call_tests.rs | 11 +- codex-rs/core/src/session/mcp.rs | 5 +- codex-rs/core/src/session/mod.rs | 20 +- codex-rs/core/src/session/session.rs | 51 +-- codex-rs/core/src/session/tests.rs | 98 +++-- .../core/src/session/tests/guardian_tests.rs | 5 +- codex-rs/core/src/session/turn_context.rs | 22 +- codex-rs/core/src/state/service.rs | 6 +- codex-rs/core/src/thread_manager.rs | 6 +- codex-rs/core/src/thread_manager_tests.rs | 52 --- 14 files changed, 405 insertions(+), 240 deletions(-) diff --git a/codex-rs/core/src/agents_md.rs b/codex-rs/core/src/agents_md.rs index d63342bdd..52a386d69 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::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: &ResolvedTurnEnvironments, + 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 8d0987ca4..565b7f4b9 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::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], -) -> ResolvedTurnEnvironments { - ResolvedTurnEnvironments { +) -> TurnEnvironmentSnapshot { + TurnEnvironmentSnapshot { 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 135909bf0..b36522459 100644 --- a/codex-rs/core/src/codex_delegate.rs +++ b/codex-rs/core/src/codex_delegate.rs @@ -90,7 +90,10 @@ 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), + 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), @@ -108,7 +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_selections: parent_ctx.environments.clone(), + 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/environment_selection.rs b/codex-rs/core/src/environment_selection.rs index 31c7494ca..e1ab2ea93 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,6 +9,9 @@ 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 futures::FutureExt; +use futures::future::BoxFuture; +use futures::future::Shared; use crate::session::turn_context::TurnEnvironment; use crate::shell::Shell; @@ -26,56 +30,81 @@ pub(crate) fn default_thread_environment_selections( .collect() } -#[derive(Clone, Debug, Default)] -pub(crate) struct ResolvedTurnEnvironments { - pub(crate) turn_environments: Vec, +type SnapshotTask = Shared>; + +pub(crate) struct ThreadEnvironments { + environment_manager: Arc, + snapshot_task: ArcSwap, } -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()) - } - - 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 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 - ))); +impl ThreadEnvironments { + pub(crate) fn new(environment_manager: Arc) -> Self { + Self { + environment_manager, + 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, snapshot) = async move { + Self::resolve_snapshot(environment_manager, previous, environments).await + } + .remote_handle(); + self.snapshot_task + .store(Arc::new(snapshot.boxed().shared())); + drop(tokio::spawn(snapshot_task)); + } + + 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 { + if !seen_environment_ids.insert(selected_environment.environment_id.as_str()) { + continue; + } + let turn_environment = match current.turn_environments.iter().find(|environment| { + environment.environment_id == selected_environment.environment_id + && environment.cwd_uri() == &selected_environment.cwd + }) { + Some(environment) => environment.clone(), + None => match Self::resolve_selection(&environment_manager, 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); + } + TurnEnvironmentSnapshot { turn_environments } + } + + async fn resolve_selection( + environment_manager: &EnvironmentManager, + selected_environment: &TurnEnvironmentSelection, + ) -> CodexResult { let environment_id = selected_environment.environment_id.clone(); let environment = environment_manager .get_environment(&environment_id) @@ -97,7 +126,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 +136,53 @@ pub(crate) async fn resolve_environment_selections( )) })?, shell, - )); + )) + } + + pub(crate) async fn snapshot(&self) -> TurnEnvironmentSnapshot { + self.snapshot_task.load_full().as_ref().clone().await + } + + pub(crate) fn environment_manager(&self) -> Arc { + Arc::clone(&self.environment_manager) + } +} + +#[derive(Clone, Debug, Default)] +pub(crate) struct TurnEnvironmentSnapshot { + pub(crate) turn_environments: Vec, +} + +impl TurnEnvironmentSnapshot { + 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 to_selections(&self) -> Vec { + self.turn_environments + .iter() + .map(TurnEnvironment::selection) + .collect() + } + + 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)] @@ -125,6 +198,16 @@ 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); + turn_environments.snapshot().await; + turn_environments + } + fn test_runtime_paths() -> ExecServerRuntimePaths { ExecServerRuntimePaths::new( std::env::current_exe().expect("current exe"), @@ -198,28 +281,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 = EnvironmentManager::default_for_tests(); + let manager = Arc::new(EnvironmentManager::default_for_tests()); + let first = TurnEnvironmentSelection { + environment_id: LOCAL_ENVIRONMENT_ID.to_string(), + cwd: cwd_uri.clone(), + }; - let err = resolve_environment_selections( - &manager, + let resolved = resolve_turn_environments( + 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.snapshot().await.to_selections(), vec![first]); } #[tokio::test] @@ -227,18 +310,18 @@ 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 = resolve_turn_environments( + Arc::clone(&manager), &[TurnEnvironmentSelection { environment_id: "local".to_string(), cwd: selected_cwd_uri, }], ) - .await - .expect("environment selections should resolve"); + .await; + let resolved = resolved.snapshot().await; assert_eq!( resolved .primary() @@ -263,25 +346,146 @@ 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_uri.clone(), + }; + + let resolved = resolve_turn_environments( + manager, + &[ + TurnEnvironmentSelection { + environment_id: "missing".to_string(), + cwd: cwd_uri, + }, + local.clone(), + ], + ) + .await; + + 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"); + 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 = + resolve_turn_environments(Arc::clone(&manager), std::slice::from_ref(&selection)).await; + manager + .upsert_environment( + REMOTE_ENVIRONMENT_ID.to_string(), + "ws://127.0.0.1:9876".to_string(), + ) + .expect("replace environment"); + + let initial_snapshot = initial.snapshot().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 + }]); + let changed_snapshot = initial.snapshot().await; + + assert!(Arc::ptr_eq( + &initial_snapshot + .primary() + .expect("initial environment") + .environment, + &reused_snapshot + .primary() + .expect("reused environment") + .environment, + )); + assert!(!Arc::ptr_eq( + &reused_snapshot + .primary() + .expect("reused environment") + .environment, + &changed_snapshot + .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 = resolve_turn_environments( + Arc::clone(&local_manager), &[TurnEnvironmentSelection { environment_id: LOCAL_ENVIRONMENT_ID.to_string(), cwd: cwd_uri, }], ) - .await - .expect("local environment should resolve"); + .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 = ResolvedTurnEnvironments { + let remote = TurnEnvironmentSnapshot { turn_environments: vec![TurnEnvironment::new( REMOTE_ENVIRONMENT_ID.to_string(), remote_environment.clone(), @@ -289,7 +493,7 @@ url = "ws://127.0.0.1:8765" /*shell*/ None, )], }; - let multiple = ResolvedTurnEnvironments { + let multiple = TurnEnvironmentSnapshot { 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 b413704cd..9a8a9174a 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( + session.services.turn_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 c033511a7..cc6e42f85 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(&self.services.environment_manager), + Arc::clone(&environment_manager), turn_environment.cwd().to_path_buf(), ), None => McpRuntimeContext::new( - Arc::clone(&self.services.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 d73d476e0..7843dd515 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::ThreadEnvironments; use crate::exec_policy::ExecPolicyManager; use crate::image_preparation::prepare_response_items; use crate::parse_turn_item; @@ -431,7 +431,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) environment_selections: Vec, pub(crate) thread_extension_init: ExtensionDataInit, pub(crate) analytics_events_client: Option, pub(crate) thread_store: Arc, @@ -520,6 +520,9 @@ impl Codex { attestation_provider, inherited_multi_agent_version, } = args; + let turn_environments = Arc::new(ThreadEnvironments::new(environment_manager)); + 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(); @@ -532,8 +535,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, &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, @@ -623,7 +625,7 @@ impl Codex { windows_sandbox_level: WindowsSandboxLevel::from_config(&config), environments: TurnEnvironmentSelections::new( config.cwd.clone(), - environment_selections.to_selections(), + resolved_environments.to_selections(), ), workspace_roots: config.workspace_roots.clone(), codex_home: config.codex_home.clone(), @@ -662,7 +664,7 @@ impl Codex { extensions, thread_extension_init, agent_control, - environment_manager, + turn_environments, analytics_events_client, thread_store, parent_rollout_thread_trace, @@ -1463,6 +1465,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, @@ -1474,7 +1481,6 @@ impl Session { session_source, ) }; - 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/session.rs b/codex-rs/core/src/session/session.rs index be3b3e3a3..79f550037 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::TurnEnvironmentSnapshot; use crate::skills::SkillError; use crate::state::ActiveTurn; use codex_extension_api::ExtensionDataInit; @@ -435,18 +436,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: TurnEnvironmentSnapshot, ) -> 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 +481,7 @@ impl Session { extensions: Arc>, thread_extension_init: ExtensionDataInit, agent_control: AgentControl, - environment_manager: Arc, + turn_environments: Arc, analytics_events_client: Option, thread_store: Arc, parent_rollout_thread_trace: ThreadTraceContext, @@ -627,10 +621,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.snapshot().await, ) .instrument(info_span!( "session_init.plugin_skill_warmup", @@ -1034,7 +1027,7 @@ impl Session { ), code_mode_service: crate::tools::code_mode::CodeModeService::new(), tool_search_handler_cache: Default::default(), - environment_manager, + turn_environments: Arc::clone(&turn_environments), }; let (out_of_band_elicitation_paused, _out_of_band_elicitation_paused_rx) = watch::channel(false); @@ -1116,28 +1109,16 @@ 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 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( + sess.services.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 80037ce49..dd887b421 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -4046,18 +4046,14 @@ 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, +) -> 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()); + turn_environments } #[tokio::test] @@ -4433,7 +4429,8 @@ async fn new_default_turn_uses_config_aware_skills_for_role_overrides() { let skill_fs = session .services - .environment_manager + .turn_environments + .environment_manager() .default_environment() .map(|environment| environment.get_filesystem()) .unwrap_or_else(|| std::sync::Arc::clone(&codex_exec_server::LOCAL_FS)); @@ -4833,6 +4830,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), @@ -4850,7 +4848,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()), @@ -4946,6 +4944,14 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) { ); 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 + .primary() + .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( @@ -4953,11 +4959,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( @@ -5027,7 +5028,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(), - environment_manager: Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), + turn_environments: Arc::clone(&turn_environments), }; let plugin_outcome = services @@ -5044,7 +5045,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), @@ -5060,7 +5060,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, @@ -5172,6 +5172,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, @@ -5190,7 +5191,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()), @@ -5274,6 +5275,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, @@ -5292,7 +5294,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()), @@ -6216,6 +6218,30 @@ 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 = { + session + .services + .turn_environments + .snapshot() + .await + .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] @@ -6224,7 +6250,10 @@ 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"); - + session + .services + .turn_environments + .update_selections(&[local(selected_cwd.clone())]); { let mut state = session.state.lock().await; state.session_configuration.environments.environments = vec![local(selected_cwd.clone())]; @@ -6253,6 +6282,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(&[]); { let mut state = session.state.lock().await; state.session_configuration.environments.environments = Vec::new(); @@ -6952,6 +6982,14 @@ where ); 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 + .primary() + .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( @@ -6959,11 +6997,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( @@ -7033,7 +7066,7 @@ 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()), + turn_environments: Arc::clone(&turn_environments), }; let plugin_outcome = services @@ -7050,7 +7083,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), @@ -7066,7 +7098,7 @@ where model_info, &models_manager, /*network*/ None, - turn_environments, + resolved_turn_environments, session_configuration.cwd().clone(), "turn_id".to_string(), skills_outcome, diff --git a/codex-rs/core/src/session/tests/guardian_tests.rs b/codex-rs/core/src/session/tests/guardian_tests.rs index a2cffe2de..184472084 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::ResolvedTurnEnvironments; use crate::exec_policy::ExecPolicyManager; use crate::guardian::GUARDIAN_REVIEWER_NAME; use crate::sandboxing::SandboxPermissions; @@ -731,9 +730,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_selections: ResolvedTurnEnvironments { - turn_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_context.rs b/codex-rs/core/src/session/turn_context.rs index 70b085fc5..57f240ebb 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::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: ResolvedTurnEnvironments, + 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: ResolvedTurnEnvironments, + environments: TurnEnvironmentSnapshot, cwd: AbsolutePathBuf, sub_id: String, skills_outcome: Arc, @@ -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,7 +678,6 @@ impl Session { return Err(CodexErr::InvalidRequest(message)); } }; - self.emit_config_changed_contributors(previous_config.as_ref(), new_config.as_ref()); self.maybe_refresh_shell_snapshot_for_cwd( &previous_cwd, @@ -733,15 +737,7 @@ impl Session { 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 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 30ef37048..5d888bd83 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::ThreadEnvironments; use crate::exec_policy::ExecPolicyManager; use crate::guardian::GuardianRejection; use crate::guardian::GuardianRejectionCircuitBreaker; @@ -22,7 +23,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 +83,7 @@ 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, + pub(crate) turn_environments: Arc, } impl SessionServices { diff --git a/codex-rs/core/src/thread_manager.rs b/codex-rs/core/src/thread_manager.rs index 19070898d..dd4852c73 100644 --- a/codex-rs/core/src/thread_manager.rs +++ b/codex-rs/core/src/thread_manager.rs @@ -5,7 +5,6 @@ use crate::codex_thread::CodexThread; use crate::config::Config; use crate::config::ThreadStoreConfig; 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 +1382,6 @@ impl ThreadManagerState { threads.remove(&resumed.conversation_id); } } - let environment_selections = - resolve_environment_selections(self.environment_manager.as_ref(), &environments) - .await?; let user_instructions = self .user_instructions_for_spawn(&session_source, parent_thread_id, forked_from_thread_id) .await; @@ -1427,7 +1423,7 @@ impl ThreadManagerState { parent_rollout_thread_trace, user_shell_override, parent_trace, - environment_selections, + 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 4a7890511..ddfa5e701 100644 --- a/codex-rs/core/src/thread_manager_tests.rs +++ b/codex-rs/core/src/thread_manager_tests.rs @@ -298,58 +298,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");