From 905987c08f40bc8fb8f36d69e8137865e1a67393 Mon Sep 17 00:00:00 2001 From: starr-openai Date: Mon, 4 May 2026 10:55:49 -0700 Subject: [PATCH] Prepare selected environment plumbing (#20669) ## Why This is a prep PR in the multi-environment process-tool stack. It separates ownership/config cleanup from the behavior change that teaches process tools to route by selected environment, so the follow-up PR can focus on model-facing `environment_id` behavior. ## Stack 1. https://github.com/openai/codex/pull/20646 - `EnvironmentContext` rendering for selected environments 2. https://github.com/openai/codex/pull/20669 - selected-environment ownership and tool config prep (this PR) 3. https://github.com/openai/codex/pull/20647 - process-tool `environment_id` routing ## What Changed - keep the resolved turn environment list wrapped in `ResolvedTurnEnvironments` through `TurnContext` instead of unwrapping it back to a raw `Vec` - add `TurnContext::resolve_path_against` so cwd-relative path resolution has one shared helper - replace the old tool config boolean with `ToolEnvironmentMode::{None, Single, Multiple}` ## Testing - Tests not run locally; this prep refactor is covered by GitHub CI for the stack. Co-authored-by: Codex --- codex-rs/core/src/codex_delegate.rs | 5 +- .../core/src/context/environment_context.rs | 4 +- codex-rs/core/src/environment_selection.rs | 10 +-- codex-rs/core/src/session/mcp.rs | 2 +- codex-rs/core/src/session/mod.rs | 1 + codex-rs/core/src/session/review.rs | 2 +- codex-rs/core/src/session/session.rs | 2 +- codex-rs/core/src/session/tests.rs | 72 +++++++++++-------- codex-rs/core/src/session/turn_context.rs | 24 +++---- codex-rs/core/src/thread_manager.rs | 2 +- codex-rs/core/src/thread_manager_tests.rs | 24 +++++-- .../core/src/tools/handlers/agent_jobs.rs | 8 +-- .../core/src/tools/handlers/apply_patch.rs | 5 +- .../src/tools/handlers/multi_agents/spawn.rs | 40 +++++------ .../tools/handlers/multi_agents_v2/spawn.rs | 8 +-- codex-rs/core/src/tools/handlers/shell.rs | 2 +- .../core/src/tools/handlers/unified_exec.rs | 2 +- .../core/src/tools/handlers/view_image.rs | 2 +- .../core/src/tools/runtimes/apply_patch.rs | 2 +- .../core/src/tools/runtimes/unified_exec.rs | 7 +- codex-rs/core/src/unified_exec/mod_tests.rs | 9 ++- codex-rs/tools/src/lib.rs | 1 + codex-rs/tools/src/tool_config.rs | 29 ++++++-- codex-rs/tools/src/tool_registry_plan.rs | 12 ++-- .../tools/src/tool_registry_plan_tests.rs | 3 +- 25 files changed, 155 insertions(+), 123 deletions(-) diff --git a/codex-rs/core/src/codex_delegate.rs b/codex-rs/core/src/codex_delegate.rs index d142d33a2..1aeb12509 100644 --- a/codex-rs/core/src/codex_delegate.rs +++ b/codex-rs/core/src/codex_delegate.rs @@ -30,7 +30,6 @@ use tokio::time::timeout; use tokio_util::sync::CancellationToken; use crate::config::Config; -use crate::environment_selection::ResolvedTurnEnvironments; use crate::guardian::GuardianApprovalRequest; use crate::guardian::new_guardian_review_id; use crate::guardian::routes_approval_to_guardian; @@ -94,9 +93,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: ResolvedTurnEnvironments { - turn_environments: parent_ctx.environments.clone(), - }, + environment_selections: parent_ctx.environments.clone(), 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/context/environment_context.rs b/codex-rs/core/src/context/environment_context.rs index 1f7313de2..c647550ee 100644 --- a/codex-rs/core/src/context/environment_context.rs +++ b/codex-rs/core/src/context/environment_context.rs @@ -176,7 +176,9 @@ impl EnvironmentContext { pub(crate) fn from_turn_context(turn_context: &TurnContext) -> Self { Self::new( - EnvironmentContextEnvironment::from_turn_environments(&turn_context.environments), + EnvironmentContextEnvironment::from_turn_environments( + &turn_context.environments.turn_environments, + ), turn_context.current_date.clone(), turn_context.timezone.clone(), Self::network_from_turn_context(turn_context), diff --git a/codex-rs/core/src/environment_selection.rs b/codex-rs/core/src/environment_selection.rs index bbf059069..e9d617cbf 100644 --- a/codex-rs/core/src/environment_selection.rs +++ b/codex-rs/core/src/environment_selection.rs @@ -24,7 +24,7 @@ pub(crate) fn default_thread_environment_selections( .collect() } -#[derive(Clone, Debug)] +#[derive(Clone, Debug, Default)] pub(crate) struct ResolvedTurnEnvironments { pub(crate) turn_environments: Vec, } @@ -37,17 +37,17 @@ impl ResolvedTurnEnvironments { .collect() } - pub(crate) fn primary_turn_environment(&self) -> Option<&TurnEnvironment> { + pub(crate) fn primary(&self) -> Option<&TurnEnvironment> { self.turn_environments.first() } pub(crate) fn primary_environment(&self) -> Option> { - self.primary_turn_environment() + self.primary() .map(|environment| Arc::clone(&environment.environment)) } pub(crate) fn primary_filesystem(&self) -> Option> { - self.primary_turn_environment() + self.primary() .map(|environment| environment.environment.get_filesystem()) } } @@ -171,7 +171,7 @@ mod tests { assert_eq!( resolved - .primary_turn_environment() + .primary() .expect("primary environment") .environment_id, "local" diff --git a/codex-rs/core/src/session/mcp.rs b/codex-rs/core/src/session/mcp.rs index 2aa5adee2..5cdc1daa6 100644 --- a/codex-rs/core/src/session/mcp.rs +++ b/codex-rs/core/src/session/mcp.rs @@ -221,7 +221,7 @@ impl Session { let mcp_servers = with_codex_apps_mcp(mcp_servers, auth.as_ref(), &mcp_config); let auth_statuses = compute_auth_statuses(mcp_servers.iter(), store_mode, auth.as_ref()).await; - let mcp_runtime_environment = match turn_context.primary_environment() { + let mcp_runtime_environment = match turn_context.environments.primary() { Some(turn_environment) => McpRuntimeEnvironment::new( Arc::clone(&turn_environment.environment), turn_environment.cwd.to_path_buf(), diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index db297d642..cb8c21c4d 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -355,6 +355,7 @@ use codex_protocol::protocol::TokenUsage; use codex_protocol::protocol::TokenUsageInfo; use codex_protocol::protocol::WarningEvent; use codex_protocol::user_input::UserInput; +use codex_tools::ToolEnvironmentMode; use codex_tools::ToolsConfig; use codex_tools::ToolsConfigParams; use codex_utils_absolute_path::AbsolutePathBuf; diff --git a/codex-rs/core/src/session/review.rs b/codex-rs/core/src/session/review.rs index b879c78fa..67c1e3331 100644 --- a/codex-rs/core/src/session/review.rs +++ b/codex-rs/core/src/session/review.rs @@ -52,7 +52,7 @@ pub(super) async fn spawn_review_thread( ) .with_web_search_config(/*web_search_config*/ None) .with_allow_login_shell(config.permissions.allow_login_shell) - .with_has_environment(parent_turn_context.tools_config.has_environment) + .with_environment_mode(parent_turn_context.tools_config.environment_mode) .with_spawn_agent_usage_hint(config.multi_agent_v2.usage_hint_enabled) .with_spawn_agent_usage_hint_text(config.multi_agent_v2.usage_hint_text.clone()) .with_hide_spawn_agent_metadata(config.multi_agent_v2.hide_spawn_agent_metadata) diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index 50b3345d6..ac6f76675 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -968,7 +968,7 @@ impl Session { "unknown stored MCP environment id", )) })? - .primary_turn_environment() + .primary() .cloned(); let mcp_runtime_environment = match turn_environment { Some(turn_environment) => McpRuntimeEnvironment::new( diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index 89633daf3..b8554b25e 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -2944,13 +2944,15 @@ pub(crate) async fn make_session_configuration_for_tests() -> SessionConfigurati fn turn_environments_for_tests( environment: &Arc, cwd: &codex_utils_absolute_path::AbsolutePathBuf, -) -> Vec { - vec![TurnEnvironment { - environment_id: codex_exec_server::LOCAL_ENVIRONMENT_ID.to_string(), - environment: Arc::clone(environment), - cwd: cwd.clone(), - shell: "bash".to_string(), - }] +) -> crate::environment_selection::ResolvedTurnEnvironments { + crate::environment_selection::ResolvedTurnEnvironments { + turn_environments: vec![TurnEnvironment { + environment_id: codex_exec_server::LOCAL_ENVIRONMENT_ID.to_string(), + environment: Arc::clone(environment), + cwd: cwd.clone(), + shell: "bash".to_string(), + }], + } } #[tokio::test] @@ -3399,7 +3401,7 @@ async fn absolute_cwd_update_with_turn_environment_is_allowed() { assert_eq!(turn_context.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] @@ -4418,15 +4420,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 - .primary_environment() + .environments + .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()); assert_eq!(turn_context.cwd.as_path(), selected_cwd.as_path()); assert_eq!(turn_context.config.cwd.as_path(), selected_cwd.as_path()); } @@ -4449,13 +4452,14 @@ async fn default_turn_overlays_session_cwd_onto_stored_thread_environments() { 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 - .primary_environment() + .environments + .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_eq!(turn_context.cwd, session_cwd); assert_eq!(turn_context.config.cwd, session_cwd); @@ -4473,28 +4477,32 @@ async fn default_turn_honors_empty_stored_thread_environments() { let turn_context = session.new_default_turn().await; - assert!(turn_context.primary_environment().is_none()); - assert!(turn_context.environments.is_empty()); + assert!(turn_context.environments.primary().is_none()); + assert!(turn_context.environments.turn_environments.is_empty()); assert_eq!(turn_context.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(); let second_cwd = turn_context.cwd.join("second"); - turn_context.environments.push(TurnEnvironment { - environment_id: "second".to_string(), - environment: Arc::clone(&first_environment.environment), - cwd: second_cwd.clone(), - shell: first_environment.shell.clone(), - }); + turn_context + .environments + .turn_environments + .push(TurnEnvironment { + environment_id: "second".to_string(), + environment: Arc::clone(&first_environment.environment), + cwd: second_cwd.clone(), + shell: first_environment.shell.clone(), + }); assert_eq!( turn_context - .primary_environment() + .environments + .primary() .expect("primary environment") .environment_id, first_environment.environment_id @@ -4502,14 +4510,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] @@ -4527,8 +4539,8 @@ async fn empty_turn_environments_clear_primary_environment() { .await .expect("turn should start"); - assert!(turn_context.primary_environment().is_none()); - assert!(turn_context.environments.is_empty()); + assert!(turn_context.environments.primary().is_none()); + assert!(turn_context.environments.turn_environments.is_empty()); assert_eq!(turn_context.cwd, session.get_config().await.cwd); assert_eq!(turn_context.config.cwd, session.get_config().await.cwd); } diff --git a/codex-rs/core/src/session/turn_context.rs b/codex-rs/core/src/session/turn_context.rs index a87d2ab17..a68b5396a 100644 --- a/codex-rs/core/src/session/turn_context.rs +++ b/codex-rs/core/src/session/turn_context.rs @@ -1,5 +1,6 @@ use super::*; use crate::config::GhostSnapshotConfig; +use crate::environment_selection::ResolvedTurnEnvironments; use codex_model_provider::SharedModelProvider; use codex_model_provider::create_model_provider; use codex_protocol::models::AdditionalPermissionProfile; @@ -60,7 +61,7 @@ pub(crate) struct TurnContext { pub(crate) reasoning_effort: Option, pub(crate) reasoning_summary: ReasoningSummaryConfig, pub(crate) session_source: SessionSource, - pub(crate) environments: Vec, + pub(crate) environments: ResolvedTurnEnvironments, /// 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()`. @@ -106,10 +107,6 @@ impl TurnContext { self.permission_profile.network_sandbox_policy() } - pub(crate) fn primary_environment(&self) -> Option<&TurnEnvironment> { - self.environments.first() - } - pub(crate) fn sandbox_policy(&self) -> SandboxPolicy { let file_system_sandbox_policy = self.file_system_sandbox_policy(); let network_sandbox_policy = self.network_sandbox_policy(); @@ -198,7 +195,7 @@ impl TurnContext { .with_unified_exec_shell_mode(self.tools_config.unified_exec_shell_mode.clone()) .with_web_search_config(self.tools_config.web_search_config.clone()) .with_allow_login_shell(self.tools_config.allow_login_shell) - .with_has_environment(self.tools_config.has_environment) + .with_environment_mode(self.tools_config.environment_mode) .with_spawn_agent_usage_hint(config.multi_agent_v2.usage_hint_enabled) .with_spawn_agent_usage_hint_text(config.multi_agent_v2.usage_hint_text.clone()) .with_hide_spawn_agent_metadata(config.multi_agent_v2.hide_spawn_agent_metadata) @@ -435,7 +432,7 @@ impl Session { model_info: ModelInfo, models_manager: &SharedModelsManager, network: Option, - environments: Vec, + environments: ResolvedTurnEnvironments, cwd: AbsolutePathBuf, sub_id: String, skills_outcome: Arc, @@ -476,7 +473,9 @@ impl Session { ) .with_web_search_config(per_turn_config.web_search_config.clone()) .with_allow_login_shell(per_turn_config.permissions.allow_login_shell) - .with_has_environment(!environments.is_empty()) + .with_environment_mode(ToolEnvironmentMode::from_count( + environments.turn_environments.len(), + )) .with_spawn_agent_usage_hint(per_turn_config.multi_agent_v2.usage_hint_enabled) .with_spawn_agent_usage_hint_text(per_turn_config.multi_agent_v2.usage_hint_text.clone()) .with_hide_spawn_agent_metadata(per_turn_config.multi_agent_v2.hide_spawn_agent_metadata) @@ -647,12 +646,11 @@ impl Session { fn resolve_turn_environments( &self, environments: &[TurnEnvironmentSelection], - ) -> CodexResult> { + ) -> CodexResult { crate::environment_selection::resolve_environment_selections( self.services.environment_manager.as_ref(), environments, ) - .map(|resolved| resolved.turn_environments) } async fn new_turn_from_configuration( @@ -660,9 +658,9 @@ impl Session { sub_id: String, session_configuration: SessionConfiguration, final_output_json_schema: Option>, - turn_environments: Vec, + turn_environments: ResolvedTurnEnvironments, ) -> Arc { - let primary_turn_environment = turn_environments.first(); + let primary_turn_environment = turn_environments.primary(); let cwd = primary_turn_environment .map(|turn_environment| turn_environment.cwd.clone()) .unwrap_or_else(|| session_configuration.cwd.clone()); @@ -769,7 +767,7 @@ impl Session { Ok(turn_environments) => turn_environments, Err(err) => { warn!("failed to resolve stored session environments: {err}"); - Vec::new() + ResolvedTurnEnvironments::default() } }; diff --git a/codex-rs/core/src/thread_manager.rs b/codex-rs/core/src/thread_manager.rs index eb7419076..927f45279 100644 --- a/codex-rs/core/src/thread_manager.rs +++ b/codex-rs/core/src/thread_manager.rs @@ -1100,7 +1100,7 @@ impl ThreadManagerState { } let environment_selections = resolve_environment_selections(self.environment_manager.as_ref(), &environments)?; - let watch_registration = match environment_selections.primary_turn_environment() { + let watch_registration = match environment_selections.primary() { Some(turn_environment) if !turn_environment.environment.is_remote() => { self.skills_watcher .register_config( diff --git a/codex-rs/core/src/thread_manager_tests.rs b/codex-rs/core/src/thread_manager_tests.rs index 75dcc86ae..6c5345218 100644 --- a/codex-rs/core/src/thread_manager_tests.rs +++ b/codex-rs/core/src/thread_manager_tests.rs @@ -444,9 +444,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( @@ -465,9 +471,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.rs b/codex-rs/core/src/tools/handlers/agent_jobs.rs index d5f719feb..b3d9b481d 100644 --- a/codex-rs/core/src/tools/handlers/agent_jobs.rs +++ b/codex-rs/core/src/tools/handlers/agent_jobs.rs @@ -6,7 +6,6 @@ use crate::config::Config; use crate::function_tool::FunctionCallError; use crate::session::session::Session; use crate::session::turn_context::TurnContext; -use crate::session::turn_context::TurnEnvironment; use crate::tools::context::FunctionToolOutput; use crate::tools::context::ToolInvocation; use crate::tools::context::ToolPayload; @@ -541,12 +540,7 @@ async fn run_agent_job_loop( "agent_job:{job_id}" )))), SpawnAgentOptions { - environments: Some( - turn.environments - .iter() - .map(TurnEnvironment::selection) - .collect(), - ), + environments: Some(turn.environments.to_selections()), ..Default::default() }, ) diff --git a/codex-rs/core/src/tools/handlers/apply_patch.rs b/codex-rs/core/src/tools/handlers/apply_patch.rs index 294e16148..e2e020a96 100644 --- a/codex-rs/core/src/tools/handlers/apply_patch.rs +++ b/codex-rs/core/src/tools/handlers/apply_patch.rs @@ -363,7 +363,7 @@ impl ToolHandler for ApplyPatchHandler { // Avoid building temporary ExecParams/command vectors; derive directly from inputs. let cwd = turn.cwd.clone(); let command = vec!["apply_patch".to_string(), patch_input.clone()]; - let Some(turn_environment) = turn.primary_environment() else { + let Some(turn_environment) = turn.environments.primary() else { return Err(FunctionCallError::RespondToModel( "apply_patch is unavailable in this session".to_string(), )); @@ -475,7 +475,8 @@ pub(crate) async fn intercept_apply_patch( tool_name: &str, ) -> Result, FunctionCallError> { let sandbox = turn - .primary_environment() + .environments + .primary() .filter(|env| env.environment.is_remote()) .map(|_| turn.file_system_sandbox_context(/*additional_permissions*/ None)); match codex_apply_patch::maybe_parse_apply_patch_verified(command, cwd, fs, sandbox.as_ref()) diff --git a/codex-rs/core/src/tools/handlers/multi_agents/spawn.rs b/codex-rs/core/src/tools/handlers/multi_agents/spawn.rs index 777cb9be1..e36ee171d 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents/spawn.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents/spawn.rs @@ -6,7 +6,6 @@ use crate::agent::exceeds_thread_spawn_depth_limit; use crate::agent::next_thread_spawn_depth; use crate::agent::role::DEFAULT_ROLE_NAME; use crate::agent::role::apply_role_to_config; -use crate::session::turn_context::TurnEnvironment; pub(crate) struct Handler; @@ -83,29 +82,22 @@ impl ToolHandler for Handler { apply_spawn_agent_runtime_overrides(&mut config, turn.as_ref())?; apply_spawn_agent_overrides(&mut config, child_depth); - let result = Box::pin( - session.services.agent_control.spawn_agent_with_metadata( - config, - input_items, - Some(thread_spawn_source( - session.conversation_id, - &turn.session_source, - child_depth, - role_name, - /*task_name*/ None, - )?), - SpawnAgentOptions { - fork_parent_spawn_call_id: args.fork_context.then(|| call_id.clone()), - fork_mode: args.fork_context.then_some(SpawnAgentForkMode::FullHistory), - environments: Some( - turn.environments - .iter() - .map(TurnEnvironment::selection) - .collect(), - ), - }, - ), - ) + let result = Box::pin(session.services.agent_control.spawn_agent_with_metadata( + config, + input_items, + Some(thread_spawn_source( + session.conversation_id, + &turn.session_source, + child_depth, + role_name, + /*task_name*/ None, + )?), + SpawnAgentOptions { + fork_parent_spawn_call_id: args.fork_context.then(|| call_id.clone()), + fork_mode: args.fork_context.then_some(SpawnAgentForkMode::FullHistory), + environments: Some(turn.environments.to_selections()), + }, + )) .await .map_err(collab_spawn_error); let (new_thread_id, new_agent_metadata, status) = match &result { diff --git a/codex-rs/core/src/tools/handlers/multi_agents_v2/spawn.rs b/codex-rs/core/src/tools/handlers/multi_agents_v2/spawn.rs index 26b6750c4..509366de2 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents_v2/spawn.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents_v2/spawn.rs @@ -5,7 +5,6 @@ use crate::agent::control::render_input_preview; use crate::agent::next_thread_spawn_depth; use crate::agent::role::DEFAULT_ROLE_NAME; use crate::agent::role::apply_role_to_config; -use crate::session::turn_context::TurnEnvironment; use codex_protocol::AgentPath; use codex_protocol::protocol::InterAgentCommunication; use codex_protocol::protocol::Op; @@ -118,12 +117,7 @@ impl ToolHandler for Handler { SpawnAgentOptions { fork_parent_spawn_call_id: fork_mode.as_ref().map(|_| call_id.clone()), fork_mode, - environments: Some( - turn.environments - .iter() - .map(TurnEnvironment::selection) - .collect(), - ), + environments: Some(turn.environments.to_selections()), }, ) .await diff --git a/codex-rs/core/src/tools/handlers/shell.rs b/codex-rs/core/src/tools/handlers/shell.rs index fb80845bd..d498f3653 100644 --- a/codex-rs/core/src/tools/handlers/shell.rs +++ b/codex-rs/core/src/tools/handlers/shell.rs @@ -412,7 +412,7 @@ impl ShellHandler { } = args; let mut exec_params = exec_params; - let Some(turn_environment) = turn.primary_environment() else { + let Some(turn_environment) = turn.environments.primary() else { return Err(FunctionCallError::RespondToModel( "shell is unavailable in this session".to_string(), )); diff --git a/codex-rs/core/src/tools/handlers/unified_exec.rs b/codex-rs/core/src/tools/handlers/unified_exec.rs index 5aec8c8ba..563ee8d83 100644 --- a/codex-rs/core/src/tools/handlers/unified_exec.rs +++ b/codex-rs/core/src/tools/handlers/unified_exec.rs @@ -196,7 +196,7 @@ impl ToolHandler for UnifiedExecHandler { } }; - let Some(turn_environment) = turn.primary_environment() else { + let Some(turn_environment) = turn.environments.primary() else { return Err(FunctionCallError::RespondToModel( "unified exec is unavailable in this session".to_string(), )); diff --git a/codex-rs/core/src/tools/handlers/view_image.rs b/codex-rs/core/src/tools/handlers/view_image.rs index 43968c82f..d05807ef7 100644 --- a/codex-rs/core/src/tools/handlers/view_image.rs +++ b/codex-rs/core/src/tools/handlers/view_image.rs @@ -88,7 +88,7 @@ impl ToolHandler for ViewImageHandler { }; let abs_path = turn.resolve_path(Some(args.path)); - let Some(environment) = turn.primary_environment() else { + let Some(environment) = turn.environments.primary() else { return Err(FunctionCallError::RespondToModel( "view_image is unavailable in this session".to_string(), )); diff --git a/codex-rs/core/src/tools/runtimes/apply_patch.rs b/codex-rs/core/src/tools/runtimes/apply_patch.rs index e720243f2..fb330fa91 100644 --- a/codex-rs/core/src/tools/runtimes/apply_patch.rs +++ b/codex-rs/core/src/tools/runtimes/apply_patch.rs @@ -191,7 +191,7 @@ impl ToolRuntime for ApplyPatchRuntime { attempt: &SandboxAttempt<'_>, ctx: &ToolCtx, ) -> Result { - let turn_environment = ctx.turn.primary_environment().ok_or_else(|| { + let turn_environment = ctx.turn.environments.primary().ok_or_else(|| { ToolError::Rejected("apply_patch is unavailable in this session".to_string()) })?; let started_at = Instant::now(); diff --git a/codex-rs/core/src/tools/runtimes/unified_exec.rs b/codex-rs/core/src/tools/runtimes/unified_exec.rs index 520616823..1bc5d7212 100644 --- a/codex-rs/core/src/tools/runtimes/unified_exec.rs +++ b/codex-rs/core/src/tools/runtimes/unified_exec.rs @@ -254,7 +254,8 @@ impl<'a> ToolRuntime for UnifiedExecRunt } let environment_is_remote = ctx .turn - .primary_environment() + .environments + .primary() .is_some_and(|turn_environment| turn_environment.environment.is_remote()); let command = if environment_is_remote { base_command.to_vec() @@ -292,7 +293,7 @@ impl<'a> ToolRuntime for UnifiedExecRunt .await? { Some(prepared) => { - let Some(turn_environment) = ctx.turn.primary_environment() else { + let Some(turn_environment) = ctx.turn.environments.primary() else { return Err(ToolError::Rejected( "exec_command is unavailable in this session".to_string(), )); @@ -337,7 +338,7 @@ impl<'a> ToolRuntime for UnifiedExecRunt .env_for(command, options, managed_network) .map_err(|err| ToolError::Codex(err.into()))?; exec_env.exec_server_env_config = req.exec_server_env_config.clone(); - let Some(turn_environment) = ctx.turn.primary_environment() else { + let Some(turn_environment) = ctx.turn.environments.primary() else { return Err(ToolError::Rejected( "exec_command is unavailable in this session".to_string(), )); diff --git a/codex-rs/core/src/unified_exec/mod_tests.rs b/codex-rs/core/src/unified_exec/mod_tests.rs index 4420f11e3..f8a1480af 100644 --- a/codex-rs/core/src/unified_exec/mod_tests.rs +++ b/codex-rs/core/src/unified_exec/mod_tests.rs @@ -96,7 +96,8 @@ async fn exec_command_with_tty( &request, tty, Box::new(NoopSpawnLifecycle), - turn.primary_environment() + turn.environments + .primary() .expect("turn environment") .environment .as_ref(), @@ -594,7 +595,8 @@ 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[0].environment = Arc::new(remote_test_env.environment().clone()); + turn.environments.turn_environments[0].environment = + Arc::new(remote_test_env.environment().clone()); let request = test_exec_request( &turn, @@ -612,7 +614,8 @@ async fn remote_exec_server_rejects_inherited_fd_launches() -> anyhow::Result<() Box::new(TestSpawnLifecycle { inherited_fds: vec![42], }), - turn.primary_environment() + turn.environments + .primary() .expect("turn environment") .environment .as_ref(), diff --git a/codex-rs/tools/src/lib.rs b/codex-rs/tools/src/lib.rs index 40e5b0346..387d0394a 100644 --- a/codex-rs/tools/src/lib.rs +++ b/codex-rs/tools/src/lib.rs @@ -109,6 +109,7 @@ pub use responses_api::mcp_tool_to_deferred_responses_api_tool; pub use responses_api::mcp_tool_to_responses_api_tool; pub use responses_api::tool_definition_to_responses_api_tool; pub use tool_config::ShellCommandBackendConfig; +pub use tool_config::ToolEnvironmentMode; pub use tool_config::ToolUserShellType; pub use tool_config::ToolsConfig; pub use tool_config::ToolsConfigParams; diff --git a/codex-rs/tools/src/tool_config.rs b/codex-rs/tools/src/tool_config.rs index 32ee9e1e5..f2fc402cc 100644 --- a/codex-rs/tools/src/tool_config.rs +++ b/codex-rs/tools/src/tool_config.rs @@ -88,7 +88,7 @@ pub struct ToolsConfig { pub shell_type: ConfigShellToolType, pub shell_command_backend: ShellCommandBackendConfig, pub unified_exec_shell_mode: UnifiedExecShellMode, - pub has_environment: bool, + pub environment_mode: ToolEnvironmentMode, pub allow_login_shell: bool, pub apply_patch_tool_type: Option, pub web_search_mode: Option, @@ -129,6 +129,27 @@ pub struct ToolsConfigParams<'a> { pub windows_sandbox_level: WindowsSandboxLevel, } +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum ToolEnvironmentMode { + None, + Single, + Multiple, +} + +impl ToolEnvironmentMode { + pub fn from_count(count: usize) -> Self { + match count { + 0 => Self::None, + 1 => Self::Single, + _ => Self::Multiple, + } + } + + pub fn has_environment(self) -> bool { + !matches!(self, Self::None) + } +} + impl ToolsConfig { pub fn new(params: &ToolsConfigParams<'_>) -> Self { let ToolsConfigParams { @@ -205,7 +226,7 @@ impl ToolsConfig { shell_type, shell_command_backend, unified_exec_shell_mode: UnifiedExecShellMode::Direct, - has_environment: true, + environment_mode: ToolEnvironmentMode::Single, allow_login_shell: true, apply_patch_tool_type, web_search_mode: *web_search_mode, @@ -306,8 +327,8 @@ impl ToolsConfig { self } - pub fn with_has_environment(mut self, has_environment: bool) -> Self { - self.has_environment = has_environment; + pub fn with_environment_mode(mut self, environment_mode: ToolEnvironmentMode) -> Self { + self.environment_mode = environment_mode; self } diff --git a/codex-rs/tools/src/tool_registry_plan.rs b/codex-rs/tools/src/tool_registry_plan.rs index c00c7c8b1..594ab2ddc 100644 --- a/codex-rs/tools/src/tool_registry_plan.rs +++ b/codex-rs/tools/src/tool_registry_plan.rs @@ -135,7 +135,7 @@ pub fn build_tool_registry_plan( ); } - if config.has_environment { + if config.environment_mode.has_environment() { match &config.shell_type { ConfigShellToolType::Default => { plan.push_spec( @@ -184,7 +184,9 @@ pub fn build_tool_registry_plan( } } - if config.has_environment && config.shell_type != ConfigShellToolType::Disabled { + if config.environment_mode.has_environment() + && config.shell_type != ConfigShellToolType::Disabled + { plan.register_handler("shell", ToolHandlerKind::Shell); plan.register_handler("container.exec", ToolHandlerKind::Shell); plan.register_handler("local_shell", ToolHandlerKind::Shell); @@ -324,7 +326,7 @@ pub fn build_tool_registry_plan( ); } - if config.has_environment + if config.environment_mode.has_environment() && let Some(apply_patch_tool_type) = &config.apply_patch_tool_type { match apply_patch_tool_type { @@ -346,7 +348,7 @@ pub fn build_tool_registry_plan( plan.register_handler("apply_patch", ToolHandlerKind::ApplyPatch); } - if config.has_environment + if config.environment_mode.has_environment() && config .experimental_supported_tools .iter() @@ -393,7 +395,7 @@ pub fn build_tool_registry_plan( ); } - if config.has_environment { + if config.environment_mode.has_environment() { plan.push_spec( create_view_image_tool(ViewImageToolOptions { can_request_original_image_detail: config.can_request_original_image_detail, diff --git a/codex-rs/tools/src/tool_registry_plan_tests.rs b/codex-rs/tools/src/tool_registry_plan_tests.rs index 71aef8b1a..6ff356a61 100644 --- a/codex-rs/tools/src/tool_registry_plan_tests.rs +++ b/codex-rs/tools/src/tool_registry_plan_tests.rs @@ -11,6 +11,7 @@ use crate::ResponsesApiNamespaceTool; use crate::ResponsesApiTool; use crate::ResponsesApiWebSearchFilters; use crate::ResponsesApiWebSearchUserLocation; +use crate::ToolEnvironmentMode; use crate::ToolHandlerSpec; use crate::ToolName; use crate::ToolNamespace; @@ -544,7 +545,7 @@ fn disabled_environment_omits_environment_backed_tools() { permission_profile: &PermissionProfile::Disabled, windows_sandbox_level: WindowsSandboxLevel::Disabled, }) - .with_has_environment(/*has_environment*/ false); + .with_environment_mode(ToolEnvironmentMode::None); tools_config .experimental_supported_tools .push("list_dir".to_string());