From 250de82bfb51a210325e88bfe1f7c30b0fa514f0 Mon Sep 17 00:00:00 2001 From: sayan-oai Date: Wed, 29 Jul 2026 03:06:36 +0000 Subject: [PATCH] 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 | 64 +++++++++++++++++++------ codex-rs/core/src/mcp_tool_call.rs | 2 +- codex-rs/core/src/tools/handlers/mcp.rs | 1 - 3 files changed, 50 insertions(+), 17 deletions(-) diff --git a/codex-rs/core/src/mcp_openai_file.rs b/codex-rs/core/src/mcp_openai_file.rs index 9decd133f1..837a9af655 100644 --- a/codex-rs/core/src/mcp_openai_file.rs +++ b/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, openai_file_input_optional_fields: Option<&HashMap>>, ) -> Result, 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", })), diff --git a/codex-rs/core/src/mcp_tool_call.rs b/codex-rs/core/src/mcp_tool_call.rs index 9cebe69cf5..669edf35a1 100644 --- a/codex-rs/core/src/mcp_tool_call.rs +++ b/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(), ) diff --git a/codex-rs/core/src/tools/handlers/mcp.rs b/codex-rs/core/src/tools/handlers/mcp.rs index 7cadbcf0c3..8515920290 100644 --- a/codex-rs/core/src/tools/handlers/mcp.rs +++ b/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,