Use step environments for MCP file uploads (#35878)

## Why

An MCP tool call can run after its selected environment becomes ready, while the
turn context still contains the earlier starting-state snapshot. File argument
rewriting must use the environment available to the current step so it can
resolve and upload local files in that case.

## What changed

- Pass `StepContext` through MCP file argument rewriting.
- Resolve the primary environment from the step snapshot while retaining
  turn-scoped configuration and metadata.
- Add coverage for an environment that transitions from starting to ready
  during a turn.

GitOrigin-RevId: 935bea18ee425aabeacf4505eaff90565f9751e9

codex-rs/core/src/mcp_openai_file.rs

@@ -12,7 +12,7 @@
 //! module only handles uploading the files and rewriting the execution-time arguments.
 
 use crate::session::session::Session;
-use crate::session::turn_context::TurnContext;
+use crate::session::step_context::StepContext;
 use codex_api::OPENAI_FILE_UPLOAD_LIMIT_BYTES;
 use codex_api::upload_openai_file;
 use codex_http_client::RouteAwareClientPool;
@@ -23,7 +23,7 @@ use std::collections::HashMap;
 
 pub(crate) async fn rewrite_mcp_tool_arguments_for_openai_files(
     sess: &Session,
-    turn_context: &TurnContext,
+    step_context: &StepContext,
     arguments_value: Option<JsonValue>,
     openai_file_input_optional_fields: Option<&HashMap<String, Vec<String>>>,
 ) -> Result<Option<JsonValue>, String> {
@@ -45,7 +45,7 @@ pub(crate) async fn rewrite_mcp_tool_arguments_for_openai_files(
             continue;
         };
         let Some(uploaded_value) = rewrite_argument_value_for_openai_files(
-            turn_context,
+            step_context,
             &sess.services.openai_file_upload_client_pool,
             auth.as_ref(),
             field_name,
@@ -67,7 +67,7 @@ pub(crate) async fn rewrite_mcp_tool_arguments_for_openai_files(
 }
 
 async fn rewrite_argument_value_for_openai_files(
-    turn_context: &TurnContext,
+    step_context: &StepContext,
     client_pool: &RouteAwareClientPool,
     auth: Option<&CodexAuth>,
     field_name: &str,
@@ -77,7 +77,7 @@ async fn rewrite_argument_value_for_openai_files(
     match value {
         JsonValue::String(file_path) => {
             let rewritten = build_uploaded_argument_value(
-                turn_context,
+                step_context,
                 client_pool,
                 auth,
                 field_name,
@@ -95,7 +95,7 @@ async fn rewrite_argument_value_for_openai_files(
                     return Ok(None);
                 };
                 let rewritten = build_uploaded_argument_value(
-                    turn_context,
+                    step_context,
                     client_pool,
                     auth,
                     field_name,
@@ -113,7 +113,7 @@ async fn rewrite_argument_value_for_openai_files(
 }
 
 async fn build_uploaded_argument_value(
-    turn_context: &TurnContext,
+    step_context: &StepContext,
     client_pool: &RouteAwareClientPool,
     auth: Option<&CodexAuth>,
     field_name: &str,
@@ -133,7 +133,8 @@ async fn build_uploaded_argument_value(
     if !auth.uses_codex_backend() {
         return Err("ChatGPT auth is required to upload files for Codex Apps tools".to_string());
     }
-    let Some(turn_environment) = turn_context.environments.primary() else {
+    let turn_context = &step_context.turn;
+    let Some(turn_environment) = step_context.environments.primary() else {
         return Err(contextualize_error(
             "no primary turn environment is available".to_string(),
         ));
@@ -215,6 +216,7 @@ mod tests {
     use super::*;
     use crate::environment_selection::TurnEnvironmentState;
     use crate::session::tests::make_session_and_context;
+    use crate::session::turn_context::TurnContext;
     use crate::session::turn_context::TurnEnvironment;
     use codex_utils_absolute_path::AbsolutePathBuf;
     use codex_utils_path_uri::PathUri;
@@ -242,13 +244,14 @@ mod tests {
     #[tokio::test]
     async fn openai_file_argument_rewrite_requires_declared_file_params() {
         let (session, turn_context) = make_session_and_context().await;
+        let step_context = StepContext::for_test(Arc::new(turn_context));
         let arguments = Some(serde_json::json!({
             "file": "/tmp/codex-smoke-file.txt"
         }));
 
         let rewritten = rewrite_mcp_tool_arguments_for_openai_files(
             &session,
-            &Arc::new(turn_context),
+            &step_context,
             arguments.clone(),
             /*openai_file_input_optional_fields*/ None,
         )
@@ -259,7 +262,7 @@ mod tests {
     }
 
     #[tokio::test]
-    async fn build_uploaded_argument_value_includes_schema_declared_optional_fields() {
+    async fn build_uploaded_argument_value_uses_environment_that_becomes_ready_during_turn() {
         use wiremock::Mock;
         use wiremock::MockServer;
         use wiremock::ResponseTemplate;
@@ -311,13 +314,40 @@ mod tests {
             .await
             .expect("write local file");
         set_primary_environment_cwd(&mut turn_context, dir.path());
+        let selection = turn_context
+            .environments
+            .primary()
+            .expect("ready primary environment")
+            .selection();
+        let environments = crate::environment_selection::ThreadEnvironments::new(
+            session.services.turn_environments.environment_manager(),
+            crate::shell::default_user_shell(),
+            crate::shell_snapshot::ShellSnapshot::disabled(),
+            Default::default(),
+            /*non_blocking_snapshots*/ true,
+        );
+        environments.update_selections(std::slice::from_ref(&selection));
+        turn_context.environments = environments.snapshot().await;
+        turn_context
+            .environments
+            .starting()
+            .next()
+            .expect("environment should initially be starting")
+            .wait_until_ready()
+            .await
+            .expect("environment should become ready");
+        let step_environments = turn_context.environments.refresh_readiness();
 
         let mut config = (*turn_context.config).clone();
         config.chatgpt_base_url = format!("{}/backend-api", server.uri());
         turn_context.config = Arc::new(config);
+        let mut step_context = StepContext::for_test(Arc::new(turn_context));
+        Arc::get_mut(&mut step_context)
+            .expect("step context should be uniquely owned")
+            .environments = step_environments;
 
         let rewritten = build_uploaded_argument_value(
-            &turn_context,
+            &step_context,
             &session.services.openai_file_upload_client_pool,
             Some(&auth),
             "file",
@@ -349,9 +379,10 @@ mod tests {
         file.set_len(OPENAI_FILE_UPLOAD_LIMIT_BYTES + 1)
             .expect("size sparse file");
         set_primary_environment_cwd(&mut turn_context, dir.path());
+        let step_context = StepContext::for_test(Arc::new(turn_context));
 
         let error = build_uploaded_argument_value(
-            &turn_context,
+            &step_context,
             &session.services.openai_file_upload_client_pool,
             Some(&auth),
             "file",
@@ -423,8 +454,9 @@ mod tests {
         let mut config = (*turn_context.config).clone();
         config.chatgpt_base_url = format!("{}/backend-api", server.uri());
         turn_context.config = Arc::new(config);
+        let step_context = StepContext::for_test(Arc::new(turn_context));
         let rewritten = rewrite_argument_value_for_openai_files(
-            &turn_context,
+            &step_context,
             &session.services.openai_file_upload_client_pool,
             Some(&auth),
             "file",
@@ -535,8 +567,9 @@ mod tests {
         let mut config = (*turn_context.config).clone();
         config.chatgpt_base_url = format!("{}/backend-api", server.uri());
         turn_context.config = Arc::new(config);
+        let step_context = StepContext::for_test(Arc::new(turn_context));
         let rewritten = rewrite_argument_value_for_openai_files(
-            &turn_context,
+            &step_context,
             &session.services.openai_file_upload_client_pool,
             Some(&auth),
             "files",
@@ -567,9 +600,10 @@ mod tests {
         session.services.auth_manager = crate::test_support::auth_manager_from_auth(
             CodexAuth::create_dummy_chatgpt_auth_for_testing(),
         );
+        let step_context = StepContext::for_test(Arc::new(turn_context));
         let error = rewrite_mcp_tool_arguments_for_openai_files(
             &session,
-            &turn_context,
+            &step_context,
             Some(serde_json::json!({
                 "file": "/definitely/missing/file.csv",
             })),

codex-rs/core/src/mcp_tool_call.rs

@@ -428,7 +428,7 @@ async fn handle_approved_mcp_tool_call(
                         .await;
                     let rewritten_arguments = rewrite_mcp_tool_arguments_for_openai_files(
                         sess,
-                        turn_context,
+                        step_context,
                         arguments_value,
                         metadata.openai_file_input_optional_fields.as_ref(),
                     )

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

@@ -141,7 +141,6 @@ impl McpHandler {
         };
 
         let started = Instant::now();
-        // TODO(sayan): Use StepContext for MCP file arguments when MCP follows dynamic environments.
         let result = handle_mcp_tool_call(
             Arc::clone(&session),
             &step_context,