Preserve environments when reloading V2 agents (#37847)

## What changed

- Restore a V2 agent's inherited environment selections when reloading its
  thread after residency eviction.
- Keep the restored execution environment and its tools available when the
  agent receives a follow-up task.

## Testing

- Add an integration test that evicts and reloads a V2 agent, then verifies
  its environment selection and `exec_command` tool are preserved.

GitOrigin-RevId: a5d4ad0fa465c717b32b7ae6807486429187a41d

codex-rs/core/src/thread_manager.rs

@@ -1585,10 +1585,15 @@ impl ThreadManagerState {
         } = options;
         let client_mcp_extensions = self.client_mcp_extensions_for_child(parent_thread_id).await;
         let thread_source = initial_history.get_resumed_thread_source();
+        let environments = inherited_environments
+            .as_ref()
+            .filter(|_| initial_history.get_multi_agent_version() == Some(MultiAgentVersion::V2))
+            .map(TurnEnvironmentSnapshot::to_selections);
         let options = StartThreadOptions {
             initial_history,
             session_source: Some(session_source),
             thread_source,
+            environments,
             client_mcp_extensions,
             ..StartThreadOptions::new(config)
         };

codex-rs/core/tests/suite/agent_execution.rs

@@ -1,5 +1,7 @@
 use anyhow::Result;
 use codex_features::Feature;
+use codex_protocol::protocol::EventMsg;
+use core_test_support::responses::ResponseMock;
 use core_test_support::responses::ev_assistant_message;
 use core_test_support::responses::ev_completed;
 use core_test_support::responses::ev_function_call_with_namespace;
@@ -8,6 +10,7 @@ use core_test_support::responses::mount_sse_once_match;
 use core_test_support::responses::sse;
 use core_test_support::responses::start_mock_server;
 use core_test_support::test_codex::test_codex;
+use core_test_support::wait_for_event;
 use pretty_assertions::assert_eq;
 use serde_json::json;
 use std::time::Duration;
@@ -36,6 +39,63 @@ fn has_function_call_output(request: &wiremock::Request, call_id: &str) -> bool
     })
 }
 
+async fn mount_root_collaboration_call(
+    server: &wiremock::MockServer,
+    prompt: &'static str,
+    call_id: &'static str,
+    tool_name: &'static str,
+    arguments: serde_json::Value,
+) {
+    let response_id = format!("resp-{call_id}");
+    mount_sse_once_match(
+        server,
+        move |request: &wiremock::Request| body_contains(request, prompt),
+        sse(vec![
+            ev_response_created(&response_id),
+            ev_function_call_with_namespace(
+                call_id,
+                MULTI_AGENT_V2_NAMESPACE,
+                tool_name,
+                &arguments.to_string(),
+            ),
+            ev_completed(&response_id),
+        ]),
+    )
+    .await;
+
+    let completion_id = format!("resp-{call_id}-complete");
+    mount_sse_once_match(
+        server,
+        move |request: &wiremock::Request| has_function_call_output(request, call_id),
+        sse(vec![
+            ev_response_created(&completion_id),
+            ev_assistant_message(&format!("msg-{call_id}"), "collaboration completed"),
+            ev_completed(&completion_id),
+        ]),
+    )
+    .await;
+}
+
+async fn mount_completed_worker(
+    server: &wiremock::MockServer,
+    task: &'static str,
+    parent_call_id: &'static str,
+) -> ResponseMock {
+    let response_id = format!("resp-worker-{parent_call_id}");
+    mount_sse_once_match(
+        server,
+        move |request: &wiremock::Request| {
+            body_contains(request, task) && !has_function_call_output(request, parent_call_id)
+        },
+        sse(vec![
+            ev_response_created(&response_id),
+            ev_assistant_message(&format!("msg-worker-{parent_call_id}"), "worker completed"),
+            ev_completed(&response_id),
+        ]),
+    )
+    .await
+}
+
 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
 async fn v2_nested_spawn_checks_shared_active_execution_capacity() -> Result<()> {
     let server = start_mock_server().await;
@@ -133,3 +193,140 @@ async fn v2_nested_spawn_checks_shared_active_execution_capacity() -> Result<()>
 
     Ok(())
 }
+
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn v2_residency_reload_preserves_inherited_environment_and_tools() -> Result<()> {
+    const EVICT_PROMPT: &str = "spawn the replacement worker";
+    const FOLLOWUP_PROMPT: &str = "continue the original worker";
+    const FOLLOWUP_TASK: &str = "continue work in the original environment";
+
+    let server = start_mock_server().await;
+    mount_root_collaboration_call(
+        &server,
+        FIRST_PROMPT,
+        "first-call",
+        "spawn_agent",
+        json!({ "message": FIRST_TASK, "task_name": "first", "fork_turns": "none" }),
+    )
+    .await;
+    mount_completed_worker(&server, FIRST_TASK, "first-call").await;
+
+    mount_root_collaboration_call(
+        &server,
+        EVICT_PROMPT,
+        "replacement-call",
+        "spawn_agent",
+        json!({ "message": SECOND_TASK, "task_name": "replacement", "fork_turns": "none" }),
+    )
+    .await;
+    mount_completed_worker(&server, SECOND_TASK, "replacement-call").await;
+
+    mount_root_collaboration_call(
+        &server,
+        FOLLOWUP_PROMPT,
+        "followup-call",
+        "followup_task",
+        json!({ "target": "first", "message": FOLLOWUP_TASK }),
+    )
+    .await;
+    let reloaded_worker_request =
+        mount_completed_worker(&server, FOLLOWUP_TASK, "followup-call").await;
+
+    let mut builder = test_codex()
+        .with_model("gpt-5.6-sol")
+        .with_exec_server_url("none")
+        .with_config(|config| {
+            for feature in [
+                Feature::Collab,
+                Feature::MultiAgentV2,
+                Feature::UnifiedExec,
+                Feature::DeferredExecutor,
+            ] {
+                config
+                    .features
+                    .enable(feature)
+                    .expect("test config should allow feature update");
+            }
+            config.use_experimental_unified_exec_tool = true;
+            config.multi_agent_v2.max_concurrent_threads_per_session = 2;
+        });
+    let test = builder.build_with_remote_and_local_env(&server).await?;
+    let child_environment = test.executor_environment().selection().clone();
+    if let Some(exec_server_url) = test.executor_environment().exec_server_url() {
+        test.thread_manager
+            .environment_manager()
+            .upsert_environment(
+                child_environment.environment_id.clone(),
+                exec_server_url.to_string(),
+                /*connect_timeout*/ None,
+            )?;
+    }
+    let mut created_threads = test.thread_manager.subscribe_thread_created();
+
+    test.submit_turn_with_environments(FIRST_PROMPT, Some(vec![child_environment.clone()]))
+        .await?;
+    let first_thread_id = created_threads.recv().await?;
+    let first_thread = test.thread_manager.get_thread(first_thread_id).await?;
+    wait_for_event(first_thread.as_ref(), |event| {
+        matches!(event, EventMsg::TurnComplete(_))
+    })
+    .await;
+
+    test.submit_text_turn(EVICT_PROMPT).await?;
+    let replacement_thread_id = created_threads.recv().await?;
+    let replacement_thread = test
+        .thread_manager
+        .get_thread(replacement_thread_id)
+        .await?;
+    wait_for_event(replacement_thread.as_ref(), |event| {
+        matches!(event, EventMsg::TurnComplete(_))
+    })
+    .await;
+    assert!(
+        test.thread_manager
+            .get_thread(first_thread_id)
+            .await
+            .is_err()
+    );
+
+    test.submit_text_turn(FOLLOWUP_PROMPT).await?;
+    let reloaded_worker = test.thread_manager.get_thread(first_thread_id).await?;
+    wait_for_event(reloaded_worker.as_ref(), |event| {
+        matches!(event, EventMsg::TurnComplete(_))
+    })
+    .await;
+    assert_eq!(
+        reloaded_worker
+            .config_snapshot()
+            .await
+            .environments
+            .environments,
+        vec![child_environment]
+    );
+
+    let worker_tools = |response_mock: &ResponseMock| {
+        response_mock
+            .requests()
+            .into_iter()
+            .find_map(|request| {
+                let body = request.body_json();
+                if body["client_metadata"]["thread_id"] != json!(first_thread_id) {
+                    return None;
+                }
+                body.get("tools")
+                    .or_else(|| {
+                        body["input"]
+                            .as_array()?
+                            .iter()
+                            .find(|item| item["type"] == "additional_tools")?
+                            .get("tools")
+                    })
+                    .cloned()
+            })
+            .expect("expected a model request for the original worker")
+    };
+    let reloaded_tools = worker_tools(&reloaded_worker_request);
+    assert!(reloaded_tools.to_string().contains("### `exec_command`"));
+
+    Ok(())
+}