core: inherit step environments in subagents

codex-rs/core/src/tools/handlers/agent_jobs.rs

@@ -3,7 +3,7 @@ use crate::agent::status::is_final;
 use crate::config::Config;
 use crate::function_tool::FunctionCallError;
 use crate::session::session::Session;
-use crate::session::turn_context::TurnContext;
+use crate::session::step_context::StepContext;
 use crate::tools::handlers::multi_agents::build_agent_spawn_config;
 use crate::tools::handlers::parse_arguments;
 use codex_protocol::ThreadId;
@@ -12,6 +12,7 @@ use codex_protocol::protocol::AgentStatus;
 use codex_protocol::protocol::MultiAgentVersion;
 use codex_protocol::protocol::SessionSource;
 use codex_protocol::protocol::SubAgentSource;
+use codex_protocol::protocol::TurnEnvironmentSelection;
 use codex_protocol::user_input::UserInput;
 use codex_utils_absolute_path::AbsolutePathBuf;
 use futures::StreamExt;
@@ -88,6 +89,7 @@ struct ReportAgentJobResultToolResult {
 struct JobRunnerOptions {
     max_concurrency: usize,
     spawn_config: Config,
+    environment_selections: Vec<TurnEnvironmentSelection>,
 }
 
 #[derive(Debug, Clone)]
@@ -107,9 +109,10 @@ fn required_state_db(
 
 async fn build_runner_options(
     session: &Arc<Session>,
-    turn: &Arc<TurnContext>,
+    step_context: &StepContext,
     requested_concurrency: Option<usize>,
 ) -> Result<JobRunnerOptions, FunctionCallError> {
+    let turn = &step_context.turn;
     let multi_agent_version = turn.multi_agent_version;
     if multi_agent_version == MultiAgentVersion::Disabled {
         return Err(FunctionCallError::RespondToModel(
@@ -128,6 +131,7 @@ async fn build_runner_options(
     Ok(JobRunnerOptions {
         max_concurrency,
         spawn_config,
+        environment_selections: step_context.environments.to_selections(),
     })
 }
 
@@ -155,7 +159,6 @@ fn normalize_max_runtime_seconds(requested: Option<u64>) -> Result<Option<u64>,
 
 async fn run_agent_job_loop(
     session: Arc<Session>,
-    turn: Arc<TurnContext>,
     db: Arc<codex_state::StateRuntime>,
     job_id: String,
     options: JobRunnerOptions,
@@ -209,7 +212,7 @@ async fn run_agent_job_loop(
                         )))),
                         SpawnAgentOptions {
                             parent_thread_id: Some(session.thread_id),
-                            environments: Some(turn.environments.to_selections()),
+                            environments: Some(options.environment_selections.clone()),
                             ..Default::default()
                         },
                     )

codex-rs/core/src/tools/handlers/agent_jobs/spawn_agents_on_csv.rs

@@ -1,3 +1,4 @@
+use crate::environment_selection::TurnEnvironmentSnapshot;
 use crate::function_tool::FunctionCallError;
 use crate::tools::context::FunctionToolOutput;
 use crate::tools::context::ToolInvocation;
@@ -35,7 +36,7 @@ impl SpawnAgentsOnCsvHandler {
     ) -> Result<Box<dyn crate::tools::context::ToolOutput>, FunctionCallError> {
         let ToolInvocation {
             session,
-            turn,
+            step_context,
             payload,
             ..
         } = invocation;
@@ -49,7 +50,7 @@ impl SpawnAgentsOnCsvHandler {
             }
         };
 
-        handle(session, turn, arguments)
+        handle(session, step_context, arguments)
             .await
             .map(boxed_tool_output)
     }
@@ -68,9 +69,10 @@ impl CoreToolRuntime for SpawnAgentsOnCsvHandler {
 /// `report_agent_job_result`, then exported to CSV on completion.
 pub async fn handle(
     session: Arc<Session>,
-    turn: Arc<TurnContext>,
+    step_context: Arc<StepContext>,
     arguments: String,
 ) -> Result<FunctionToolOutput, FunctionCallError> {
+    let turn = &step_context.turn;
     let args: SpawnAgentsOnCsvArgs = parse_arguments(arguments.as_str())?;
     if args.instruction.trim().is_empty() {
         return Err(FunctionCallError::RespondToModel(
@@ -78,7 +80,7 @@ pub async fn handle(
         ));
     }
 
-    let cwd = single_local_environment_cwd(&turn)?;
+    let cwd = single_local_environment_cwd(&step_context.environments)?;
     let db = required_state_db(&session)?;
     let input_path = cwd.join(args.csv_path);
     let input_path_display = input_path.display().to_string();
@@ -182,7 +184,7 @@ pub async fn handle(
         })?;
 
     let requested_concurrency = args.max_concurrency.or(args.max_workers);
-    let options = match build_runner_options(&session, &turn, requested_concurrency).await {
+    let options = match build_runner_options(&session, &step_context, requested_concurrency).await {
         Ok(options) => options,
         Err(err) => {
             let error_message = err.to_string();
@@ -199,14 +201,7 @@ pub async fn handle(
                 "failed to transition agent job {job_id} to running: {err}"
             ))
         })?;
-    if let Err(err) = run_agent_job_loop(
-        session.clone(),
-        turn.clone(),
-        db.clone(),
-        job_id.clone(),
-        options,
-    )
-    .await
+    if let Err(err) = run_agent_job_loop(session.clone(), db.clone(), job_id.clone(), options).await
     {
         let error_message = format!("job runner failed: {err}");
         let _ = db
@@ -299,8 +294,15 @@ pub async fn handle(
     Ok(FunctionToolOutput::from_text(content, Some(true)))
 }
 
-fn single_local_environment_cwd(turn: &TurnContext) -> Result<AbsolutePathBuf, FunctionCallError> {
-    let [turn_environment] = turn.environments.turn_environments.as_slice() else {
+fn single_local_environment_cwd(
+    environments: &TurnEnvironmentSnapshot,
+) -> Result<AbsolutePathBuf, FunctionCallError> {
+    if !environments.starting.is_empty() {
+        return Err(FunctionCallError::RespondToModel(
+            "spawn_agents_on_csv requires exactly one local environment".to_string(),
+        ));
+    }
+    let [turn_environment] = environments.turn_environments.as_slice() else {
         return Err(FunctionCallError::RespondToModel(
             "spawn_agents_on_csv requires exactly one local environment".to_string(),
         ));

codex-rs/core/src/tools/handlers/multi_agents/spawn.rs

@@ -49,6 +49,7 @@ async fn handle_spawn_agent(
     let ToolInvocation {
         session,
         turn,
+        step_context,
         payload,
         call_id,
         ..
@@ -131,7 +132,7 @@ async fn handle_spawn_agent(
             fork_parent_spawn_call_id: args.fork_context.then(|| call_id.clone()),
             fork_mode: args.fork_context.then_some(SpawnAgentForkMode::FullHistory),
             parent_thread_id: Some(session.thread_id),
-            environments: Some(turn.environments.to_selections()),
+            environments: Some(step_context.environments.to_selections()),
         },
     ))
     .await

codex-rs/core/src/tools/handlers/multi_agents_tests.rs

@@ -7,6 +7,7 @@ use crate::init_state_db;
 use crate::local_agent_graph_store_from_state_db;
 use crate::session::step_context::StepContext;
 use crate::session::tests::make_session_and_context;
+use crate::session::turn_context::TurnEnvironment;
 use crate::session_prefix::format_inter_agent_completion_message;
 use crate::thread_manager::thread_store_from_config;
 use crate::tools::context::ToolOutput;
@@ -56,6 +57,7 @@ use codex_protocol::protocol::TurnAbortedEvent;
 use codex_protocol::protocol::TurnCompleteEvent;
 use codex_protocol::user_input::UserInput;
 use codex_state::DirectionalThreadSpawnEdgeStatus;
+use codex_utils_path_uri::PathUri;
 use core_test_support::TempDirExt;
 use pretty_assertions::assert_eq;
 use serde::Deserialize;
@@ -1042,6 +1044,60 @@ async fn spawn_agent_returns_agent_id_without_task_name() {
     assert_eq!(success, Some(true));
 }
 
+#[tokio::test]
+async fn spawn_agent_inherits_environments_from_step_context() {
+    let (mut session, turn) = make_session_and_context().await;
+    let manager = thread_manager();
+    session.services.agent_control = manager.agent_control();
+    let turn = Arc::new(turn);
+    let step_cwd = turn.config.cwd.join("step-environment");
+    tokio::fs::create_dir_all(&step_cwd)
+        .await
+        .expect("create step environment cwd");
+    let mut step_environments = turn.environments.clone();
+    let previous_environment = step_environments.turn_environments[0].clone();
+    step_environments.turn_environments[0] = TurnEnvironment::new(
+        previous_environment.environment_id,
+        previous_environment.environment,
+        PathUri::from_abs_path(&step_cwd),
+        previous_environment.shell,
+    );
+    let expected_environments = step_environments.to_selections();
+    let mut invocation = invocation(
+        Arc::new(session),
+        Arc::clone(&turn),
+        "spawn_agent",
+        function_payload(json!({"message": "inspect this repo"})),
+    );
+    invocation.step_context = Arc::new(StepContext::new(
+        turn,
+        step_environments,
+        /*loaded_agents_md*/ None,
+    ));
+
+    let output = SpawnAgentHandler::default()
+        .handle(invocation)
+        .await
+        .expect("spawn_agent should succeed");
+    let (content, _) = expect_text_output(output);
+    let result: serde_json::Value =
+        serde_json::from_str(&content).expect("spawn_agent result should be json");
+    let child_id = parse_agent_id(
+        result["agent_id"]
+            .as_str()
+            .expect("spawn_agent result should include agent_id"),
+    );
+    let child = manager
+        .get_thread(child_id)
+        .await
+        .expect("spawned agent thread should exist");
+
+    assert_eq!(
+        child.config_snapshot().await.environment_selections(),
+        expected_environments
+    );
+}
+
 #[tokio::test]
 async fn multi_agent_v2_spawn_requires_task_name() {
     let (mut session, mut turn) = make_session_and_context().await;

codex-rs/core/src/tools/handlers/multi_agents_v2/spawn.rs

@@ -42,6 +42,7 @@ async fn handle_spawn_agent(
     let ToolInvocation {
         session,
         turn,
+        step_context,
         payload,
         call_id,
         ..
@@ -128,7 +129,7 @@ async fn handle_spawn_agent(
                 fork_parent_spawn_call_id: fork_mode.as_ref().map(|_| call_id.clone()),
                 fork_mode,
                 parent_thread_id: Some(session.thread_id),
-                environments: Some(turn.environments.to_selections()),
+                environments: Some(step_context.environments.to_selections()),
             },
         ),
     )