From a24d9864e4eb90d911933a661b47cde8a47a61b3 Mon Sep 17 00:00:00 2001 From: Sayan Sisodiya Date: Mon, 22 Jun 2026 20:41:58 -0700 Subject: [PATCH] core: use current step environments for tools --- codex-rs/core/src/compact_remote.rs | 2 + codex-rs/core/src/compact_remote_v2.rs | 2 + codex-rs/core/src/prompt_debug.rs | 9 ++- codex-rs/core/src/session/mod.rs | 19 ++++++- codex-rs/core/src/session/step_context.rs | 9 +++ codex-rs/core/src/session/tests.rs | 11 +++- .../core/src/session/tests/guardian_tests.rs | 5 ++ codex-rs/core/src/session/turn.rs | 18 ++++-- codex-rs/core/src/session/turn_context.rs | 4 -- codex-rs/core/src/session_startup_prewarm.rs | 4 ++ .../core/src/stream_events_utils_tests.rs | 4 +- codex-rs/core/src/tools/context.rs | 2 + .../core/src/tools/handlers/apply_patch.rs | 7 ++- .../src/tools/handlers/apply_patch_tests.rs | 2 + .../src/tools/handlers/extension_tools.rs | 9 ++- codex-rs/core/src/tools/handlers/mcp.rs | 5 ++ codex-rs/core/src/tools/handlers/mod.rs | 8 +-- .../src/tools/handlers/multi_agents_tests.rs | 2 + .../src/tools/handlers/request_permissions.rs | 7 ++- .../handlers/request_user_input_tests.rs | 2 + codex-rs/core/src/tools/handlers/shell.rs | 8 +-- .../src/tools/handlers/shell/shell_command.rs | 8 +++ .../core/src/tools/handlers/shell_tests.rs | 3 + .../handlers/unified_exec/exec_command.rs | 7 ++- .../src/tools/handlers/unified_exec_tests.rs | 4 ++ .../core/src/tools/handlers/view_image.rs | 7 ++- codex-rs/core/src/tools/parallel.rs | 3 + codex-rs/core/src/tools/registry_tests.rs | 2 + codex-rs/core/src/tools/router.rs | 20 ++++++- codex-rs/core/src/tools/router_tests.rs | 16 ++++-- codex-rs/core/src/tools/spec_plan.rs | 35 +++++++++--- codex-rs/core/src/tools/spec_plan_tests.rs | 35 +++++++++++- .../src/tools/tool_dispatch_trace_tests.rs | 2 + codex-rs/core/tests/suite/remote_env.rs | 56 +++++++++++++++++-- 34 files changed, 281 insertions(+), 56 deletions(-) diff --git a/codex-rs/core/src/compact_remote.rs b/codex-rs/core/src/compact_remote.rs index a8f6abc0ca..84c85f5697 100644 --- a/codex-rs/core/src/compact_remote.rs +++ b/codex-rs/core/src/compact_remote.rs @@ -18,6 +18,7 @@ use crate::hook_runtime::run_pre_compact_hooks; use crate::responses_metadata::CodexResponsesRequestKind; use crate::responses_metadata::CompactionTurnMetadata; use crate::session::session::Session; +use crate::session::step_context::StepContext; use crate::session::turn::built_tools; use crate::session::turn_context::TurnContext; use codex_analytics::CompactionImplementation; @@ -222,6 +223,7 @@ async fn run_remote_compact_task_inner_impl( let tool_router = built_tools( sess.as_ref(), turn_context.as_ref(), + Arc::new(StepContext::from_turn_context(turn_context.as_ref())), &CancellationToken::new(), ) .await?; diff --git a/codex-rs/core/src/compact_remote_v2.rs b/codex-rs/core/src/compact_remote_v2.rs index 9cd7a456be..947a3f1e2c 100644 --- a/codex-rs/core/src/compact_remote_v2.rs +++ b/codex-rs/core/src/compact_remote_v2.rs @@ -21,6 +21,7 @@ use crate::responses_metadata::CompactionTurnMetadata; use crate::responses_retry::ResponsesStreamRequest; use crate::responses_retry::handle_retryable_response_stream_error; use crate::session::session::Session; +use crate::session::step_context::StepContext; use crate::session::turn::built_tools; use crate::session::turn_context::TurnContext; use codex_analytics::CompactionImplementation; @@ -230,6 +231,7 @@ async fn run_remote_compact_task_inner_impl( let tool_router = built_tools( sess.as_ref(), turn_context.as_ref(), + Arc::new(StepContext::from_turn_context(turn_context.as_ref())), &CancellationToken::new(), ) .await?; diff --git a/codex-rs/core/src/prompt_debug.rs b/codex-rs/core/src/prompt_debug.rs index 307115c505..a64573751e 100644 --- a/codex-rs/core/src/prompt_debug.rs +++ b/codex-rs/core/src/prompt_debug.rs @@ -14,6 +14,7 @@ use tokio_util::sync::CancellationToken; use crate::config::Config; use crate::resolve_installation_id; use crate::session::session::Session; +use crate::session::step_context::StepContext; use crate::session::turn::build_prompt; use crate::session::turn::built_tools; use crate::state_db_bridge::StateDbHandle; @@ -90,7 +91,13 @@ pub(crate) async fn build_prompt_input_from_session( .clone_history() .await .for_prompt(&turn_context.model_info.input_modalities); - let router = built_tools(sess, turn_context.as_ref(), &CancellationToken::new()).await?; + let router = built_tools( + sess, + turn_context.as_ref(), + Arc::new(StepContext::from_turn_context(turn_context.as_ref())), + &CancellationToken::new(), + ) + .await?; let base_instructions = sess.get_base_instructions().await; let prompt = build_prompt( prompt_input, diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index a51cadd576..06c3fb1dcd 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -387,7 +387,6 @@ use codex_protocol::protocol::TokenUsageInfo; use codex_protocol::protocol::TurnModerationMetadataEvent; use codex_protocol::protocol::WarningEvent; use codex_protocol::user_input::UserInput; -use codex_tools::ToolEnvironmentMode; use codex_tools::UnifiedExecShellMode; use codex_utils_absolute_path::AbsolutePathBuf; #[cfg(test)] @@ -2788,6 +2787,24 @@ impl Session { world_state } + pub(crate) async fn capture_step_context( + &self, + turn_context: &TurnContext, + ) -> Arc { + let step_context = if turn_context + .config + .features + .enabled(Feature::DeferredExecutor) + { + StepContext { + environments: self.services.turn_environments.snapshot().await, + } + } else { + StepContext::from_turn_context(turn_context) + }; + Arc::new(step_context) + } + pub(crate) async fn record_inter_agent_communication( &self, turn_context: &TurnContext, diff --git a/codex-rs/core/src/session/step_context.rs b/codex-rs/core/src/session/step_context.rs index 17243ced67..b180850061 100644 --- a/codex-rs/core/src/session/step_context.rs +++ b/codex-rs/core/src/session/step_context.rs @@ -1,7 +1,16 @@ use crate::environment_selection::TurnEnvironmentSnapshot; +use crate::session::turn_context::TurnContext; /// Request-scoped state that may change between model sampling requests. #[derive(Debug)] pub(crate) struct StepContext { pub(crate) environments: TurnEnvironmentSnapshot, } + +impl StepContext { + pub(crate) fn from_turn_context(turn_context: &TurnContext) -> Self { + Self { + environments: turn_context.environments.clone(), + } + } +} diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index e715509a69..92fc60d4c1 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -8,6 +8,7 @@ use crate::context::ContextualUserFragment; use crate::context::TurnAborted; use crate::environment_selection::ThreadEnvironments; use crate::function_tool::FunctionCallError; +use crate::session::step_context::StepContext; use crate::shell::default_user_shell; use crate::shell_snapshot::ShellSnapshot; use crate::skills::SkillRenderSideEffects; @@ -595,8 +596,9 @@ async fn preview_session_start_hooks( } fn test_tool_runtime(session: Arc, turn_context: Arc) -> ToolCallRuntime { - let router = Arc::new(ToolRouter::from_turn_context( + let router = Arc::new(ToolRouter::from_contexts( &turn_context, + Arc::new(StepContext::from_turn_context(turn_context.as_ref())), crate::tools::router::ToolRouterParams { tool_suggest_candidates: None, mcp_tools: None, @@ -5899,6 +5901,7 @@ async fn request_permissions_tool_resolves_relative_paths_against_selected_envir handler .handle(ToolInvocation { session, + step_context: Arc::new(StepContext::from_turn_context(turn_context.as_ref())), turn: turn_context, cancellation_token: CancellationToken::new(), tracker, @@ -5973,6 +5976,7 @@ async fn request_permissions_tool_rejects_unknown_environment_id() { let result = RequestPermissionsHandler .handle(ToolInvocation { session: Arc::new(session), + step_context: Arc::new(StepContext::from_turn_context(&turn_context)), turn: Arc::new(turn_context), cancellation_token: CancellationToken::new(), tracker: Arc::new(tokio::sync::Mutex::new(TurnDiffTracker::new())), @@ -9898,8 +9902,9 @@ async fn fatal_tool_error_stops_turn_and_reports_error() { .await }; let deferred_mcp_tools = Some(tools.clone()); - let router = ToolRouter::from_turn_context( + let router = ToolRouter::from_contexts( &turn_context, + Arc::new(StepContext::from_turn_context(turn_context.as_ref())), crate::tools::router::ToolRouterParams { tool_suggest_candidates: None, deferred_mcp_tools, @@ -10155,6 +10160,7 @@ async fn rejects_escalated_permissions_when_policy_not_on_request() { .handle(ToolInvocation { session: Arc::clone(&session), turn: Arc::clone(&turn_context), + step_context: Arc::new(StepContext::from_turn_context(turn_context.as_ref())), cancellation_token: CancellationToken::new(), tracker: Arc::clone(&turn_diff_tracker), call_id, @@ -10308,6 +10314,7 @@ async fn unified_exec_rejects_escalated_permissions_when_policy_not_on_request() .handle(ToolInvocation { session: Arc::clone(&session), turn: Arc::clone(&turn_context), + step_context: Arc::new(StepContext::from_turn_context(turn_context.as_ref())), cancellation_token: CancellationToken::new(), tracker: Arc::clone(&tracker), call_id: "exec-call".to_string(), diff --git a/codex-rs/core/src/session/tests/guardian_tests.rs b/codex-rs/core/src/session/tests/guardian_tests.rs index 2f86cc1a05..ef53024cab 100644 --- a/codex-rs/core/src/session/tests/guardian_tests.rs +++ b/codex-rs/core/src/session/tests/guardian_tests.rs @@ -3,6 +3,7 @@ use crate::compact::InitialContextInjection; use crate::exec_policy::ExecPolicyManager; use crate::guardian::GUARDIAN_REVIEWER_NAME; use crate::sandboxing::SandboxPermissions; +use crate::session::step_context::StepContext; use crate::test_support::models_manager_with_provider; use crate::tools::context::ToolCallSource; use crate::tools::context::ToolOutput; @@ -333,6 +334,7 @@ async fn guardian_allows_shell_command_additional_permissions_requests_past_poli .handle(ToolInvocation { session: Arc::clone(&session), turn: Arc::clone(&turn_context), + step_context: Arc::new(StepContext::from_turn_context(turn_context.as_ref())), cancellation_token: CancellationToken::new(), tracker: Arc::new(tokio::sync::Mutex::new(TurnDiffTracker::new())), call_id: "test-call".to_string(), @@ -437,6 +439,7 @@ async fn strict_auto_review_turn_grant_forces_guardian_for_shell_command_policy_ .handle(ToolInvocation { session: Arc::clone(&session), turn: Arc::clone(&turn_context), + step_context: Arc::new(StepContext::from_turn_context(turn_context.as_ref())), cancellation_token: CancellationToken::new(), tracker: Arc::new(tokio::sync::Mutex::new(TurnDiffTracker::new())), call_id: "strict-shell-command-call".to_string(), @@ -484,6 +487,7 @@ async fn guardian_allows_unified_exec_additional_permissions_requests_past_polic .handle(ToolInvocation { session: Arc::clone(&session), turn: Arc::clone(&turn_context), + step_context: Arc::new(StepContext::from_turn_context(turn_context.as_ref())), cancellation_token: CancellationToken::new(), tracker: Arc::clone(&tracker), call_id: "exec-call".to_string(), @@ -612,6 +616,7 @@ async fn shell_command_allows_sticky_turn_permissions_without_inline_request_per .handle(ToolInvocation { session: Arc::clone(&session), turn: Arc::clone(&turn_context), + step_context: Arc::new(StepContext::from_turn_context(turn_context.as_ref())), cancellation_token: CancellationToken::new(), tracker: Arc::new(tokio::sync::Mutex::new(TurnDiffTracker::new())), call_id: "sticky-turn-grant".to_string(), diff --git a/codex-rs/core/src/session/turn.rs b/codex-rs/core/src/session/turn.rs index 607e2ad79e..646d1a461e 100644 --- a/codex-rs/core/src/session/turn.rs +++ b/codex-rs/core/src/session/turn.rs @@ -244,14 +244,12 @@ pub(crate) async fn run_turn( ) .await?; + let step_context = sess.capture_step_context(turn_context.as_ref()).await; if turn_context .config .features .enabled(Feature::DeferredExecutor) { - let step_context = StepContext { - environments: sess.services.turn_environments.snapshot().await, - }; world_state = sess .record_step_environment_context_if_changed( turn_context.as_ref(), @@ -278,6 +276,7 @@ pub(crate) async fn run_turn( run_sampling_request( Arc::clone(&sess), Arc::clone(&turn_context), + step_context, Arc::clone(&turn_extension_data), Arc::clone(&turn_diff_tracker), &mut client_session, @@ -1130,6 +1129,7 @@ pub(crate) fn build_prompt( async fn run_sampling_request( sess: Arc, turn_context: Arc, + step_context: Arc, turn_store: Arc, turn_diff_tracker: SharedTurnDiffTracker, client_session: &mut ModelClientSession, @@ -1137,7 +1137,13 @@ async fn run_sampling_request( input: Vec, cancellation_token: CancellationToken, ) -> CodexResult<(SamplingRequestResult, Vec)> { - let router = built_tools(sess.as_ref(), turn_context.as_ref(), &cancellation_token).await?; + let router = built_tools( + sess.as_ref(), + turn_context.as_ref(), + step_context, + &cancellation_token, + ) + .await?; let base_instructions = sess.get_base_instructions().await; @@ -1234,6 +1240,7 @@ async fn run_sampling_request( pub(crate) async fn built_tools( sess: &Session, turn_context: &TurnContext, + step_context: Arc, cancellation_token: &CancellationToken, ) -> CodexResult> { let mcp_connection_manager = sess.services.mcp_connection_manager.load_full(); @@ -1351,8 +1358,9 @@ pub(crate) async fn built_tools( ); let mcp_tools = has_mcp_servers.then_some(mcp_tool_exposure.direct_tools); let deferred_mcp_tools = mcp_tool_exposure.deferred_tools; - Ok(Arc::new(ToolRouter::from_turn_context( + Ok(Arc::new(ToolRouter::from_contexts( turn_context, + step_context, ToolRouterParams { mcp_tools, deferred_mcp_tools, diff --git a/codex-rs/core/src/session/turn_context.rs b/codex-rs/core/src/session/turn_context.rs index b546355f7e..7ff1a5b2cf 100644 --- a/codex-rs/core/src/session/turn_context.rs +++ b/codex-rs/core/src/session/turn_context.rs @@ -207,10 +207,6 @@ impl TurnContext { && self.config.orchestrator_mcp_enabled } - pub(crate) fn tool_environment_mode(&self) -> ToolEnvironmentMode { - ToolEnvironmentMode::from_count(self.environments.turn_environments.len()) - } - pub(crate) async fn with_model( &self, model: String, diff --git a/codex-rs/core/src/session_startup_prewarm.rs b/codex-rs/core/src/session_startup_prewarm.rs index 270134981b..1a3649d5ec 100644 --- a/codex-rs/core/src/session_startup_prewarm.rs +++ b/codex-rs/core/src/session_startup_prewarm.rs @@ -14,6 +14,7 @@ use crate::guardian::routes_approval_to_guardian; use crate::responses_metadata::CodexResponsesRequestKind; use crate::session::INITIAL_SUBMIT_ID; use crate::session::session::Session; +use crate::session::step_context::StepContext; use crate::session::turn::build_prompt; use crate::session::turn::built_tools; use codex_otel::STARTUP_PREWARM_AGE_AT_FIRST_TURN_METRIC; @@ -261,6 +262,9 @@ async fn schedule_startup_prewarm_inner( let startup_router = built_tools( session.as_ref(), startup_turn_context.as_ref(), + Arc::new(StepContext::from_turn_context( + startup_turn_context.as_ref(), + )), &startup_cancellation_token, ) .await?; diff --git a/codex-rs/core/src/stream_events_utils_tests.rs b/codex-rs/core/src/stream_events_utils_tests.rs index b01ed4c4ca..4580b94bc2 100644 --- a/codex-rs/core/src/stream_events_utils_tests.rs +++ b/codex-rs/core/src/stream_events_utils_tests.rs @@ -8,6 +8,7 @@ use super::image_generation_artifact_path; use super::last_assistant_message_from_item; use super::response_item_may_include_external_context; use super::save_image_generation_result; +use crate::session::step_context::StepContext; use crate::session::tests::make_session_and_context; use crate::tools::ToolRouter; use crate::tools::parallel::ToolCallRuntime; @@ -279,8 +280,9 @@ async fn handle_output_item_done_returns_contributed_last_agent_message() { session.services.extensions = Arc::new(builder.build()); let session = Arc::new(session); let turn_context = Arc::new(turn_context); - let router = Arc::new(ToolRouter::from_turn_context( + let router = Arc::new(ToolRouter::from_contexts( &turn_context, + Arc::new(StepContext::from_turn_context(turn_context.as_ref())), crate::tools::router::ToolRouterParams { tool_suggest_candidates: None, mcp_tools: None, diff --git a/codex-rs/core/src/tools/context.rs b/codex-rs/core/src/tools/context.rs index 57cf0d80b3..9268f66621 100644 --- a/codex-rs/core/src/tools/context.rs +++ b/codex-rs/core/src/tools/context.rs @@ -1,6 +1,7 @@ use crate::context_manager::truncate_function_output_payload; use crate::original_image_detail::sanitize_original_image_detail; use crate::session::session::Session; +use crate::session::step_context::StepContext; use crate::session::turn_context::TurnContext; use crate::tools::TELEMETRY_PREVIEW_MAX_BYTES; use crate::tools::TELEMETRY_PREVIEW_MAX_LINES; @@ -54,6 +55,7 @@ pub enum ToolCallSource { pub struct ToolInvocation { pub session: Arc, pub turn: Arc, + pub(crate) step_context: Arc, pub cancellation_token: CancellationToken, pub tracker: SharedTurnDiffTracker, pub call_id: String, diff --git a/codex-rs/core/src/tools/handlers/apply_patch.rs b/codex-rs/core/src/tools/handlers/apply_patch.rs index 5bb2b013c3..6326dbc1ad 100644 --- a/codex-rs/core/src/tools/handlers/apply_patch.rs +++ b/codex-rs/core/src/tools/handlers/apply_patch.rs @@ -348,6 +348,7 @@ impl ApplyPatchHandler { let ToolInvocation { session, turn, + step_context, tracker, call_id, tool_name, @@ -372,8 +373,10 @@ impl ApplyPatchHandler { require_environment_id(args.environment_id.as_deref(), self.multi_environment)?; // Verify the parsed patch against the selected environment filesystem. - let Some(turn_environment) = - resolve_tool_environment(turn.as_ref(), selected_environment_id.as_deref())? + let Some(turn_environment) = resolve_tool_environment( + &step_context.environments, + selected_environment_id.as_deref(), + )? else { return Err(FunctionCallError::RespondToModel( "apply_patch is unavailable in this session".to_string(), diff --git a/codex-rs/core/src/tools/handlers/apply_patch_tests.rs b/codex-rs/core/src/tools/handlers/apply_patch_tests.rs index f8d5f9710e..dd6e4f93f0 100644 --- a/codex-rs/core/src/tools/handlers/apply_patch_tests.rs +++ b/codex-rs/core/src/tools/handlers/apply_patch_tests.rs @@ -13,6 +13,7 @@ use std::sync::Arc; use tempfile::TempDir; use tokio::sync::Mutex; +use crate::session::step_context::StepContext; use crate::session::tests::make_session_and_context; use crate::tools::context::ToolInvocation; use crate::tools::hook_names::HookToolName; @@ -31,6 +32,7 @@ async fn invocation_for_payload(payload: ToolPayload) -> ToolInvocation { let (session, turn) = make_session_and_context().await; ToolInvocation { session: session.into(), + step_context: Arc::new(StepContext::from_turn_context(&turn)), turn: turn.into(), cancellation_token: tokio_util::sync::CancellationToken::new(), tracker: Arc::new(Mutex::new(TurnDiffTracker::new())), diff --git a/codex-rs/core/src/tools/handlers/extension_tools.rs b/codex-rs/core/src/tools/handlers/extension_tools.rs index 025cde231c..4d21c0cb7c 100644 --- a/codex-rs/core/src/tools/handlers/extension_tools.rs +++ b/codex-rs/core/src/tools/handlers/extension_tools.rs @@ -112,8 +112,9 @@ impl TurnItemEmitter for CoreTurnItemEmitter { async fn to_extension_call(invocation: &ToolInvocation) -> ExtensionToolCall { let conversation_history = ConversationHistory::new(invocation.session.clone_history().await.into_raw_items()); - let mut environments = Vec::with_capacity(invocation.turn.environments.turn_environments.len()); - for environment in &invocation.turn.environments.turn_environments { + let mut environments = + Vec::with_capacity(invocation.step_context.environments.turn_environments.len()); + for environment in &invocation.step_context.environments.turn_environments { // TODO(anp): Migrate extension ToolEnvironment and granted-permission lookup to PathUri // so extensions can receive foreign environment cwd values. let Ok(native_cwd) = environment.cwd().to_abs_path() else { @@ -175,6 +176,7 @@ mod tests { use super::CoreTurnItemEmitter; use super::ExtensionToolAdapter; + use crate::session::step_context::StepContext; use crate::tools::context::ToolCallSource; use crate::tools::context::ToolInvocation; use crate::tools::context::ToolPayload; @@ -274,6 +276,7 @@ mod tests { let (session, turn) = crate::session::tests::make_session_and_context().await; let invocation = ToolInvocation { session: session.into(), + step_context: Arc::new(StepContext::from_turn_context(&turn)), turn: turn.into(), cancellation_token: tokio_util::sync::CancellationToken::new(), tracker: Arc::new(tokio::sync::Mutex::new(TurnDiffTracker::new())), @@ -343,6 +346,7 @@ mod tests { assert_eq!(raw_history_item.item, expected_history_item); let invocation = ToolInvocation { session, + step_context: Arc::new(StepContext::from_turn_context(turn.as_ref())), turn, cancellation_token: tokio_util::sync::CancellationToken::new(), tracker: Arc::new(tokio::sync::Mutex::new(TurnDiffTracker::new())), @@ -544,6 +548,7 @@ mod tests { ); let invocation = ToolInvocation { session, + step_context: Arc::new(StepContext::from_turn_context(turn.as_ref())), turn, cancellation_token: tokio_util::sync::CancellationToken::new(), tracker: Arc::new(tokio::sync::Mutex::new(TurnDiffTracker::new())), diff --git a/codex-rs/core/src/tools/handlers/mcp.rs b/codex-rs/core/src/tools/handlers/mcp.rs index 016a0344cd..781d1a3ef0 100644 --- a/codex-rs/core/src/tools/handlers/mcp.rs +++ b/codex-rs/core/src/tools/handlers/mcp.rs @@ -317,6 +317,7 @@ mod search_tests; #[cfg(test)] mod tests { use super::*; + use crate::session::step_context::StepContext; use crate::session::tests::make_session_and_context; use crate::tools::context::ToolCallSource; use crate::tools::hook_names::HookToolName; @@ -345,6 +346,7 @@ mod tests { assert_eq!( handler.pre_tool_use_payload(&ToolInvocation { session: session.into(), + step_context: Arc::new(StepContext::from_turn_context(&turn)), turn: turn.into(), cancellation_token: tokio_util::sync::CancellationToken::new(), tracker: Arc::new(Mutex::new(TurnDiffTracker::new())), @@ -377,6 +379,7 @@ mod tests { assert_eq!( handler.pre_tool_use_payload(&ToolInvocation { session: session.into(), + step_context: Arc::new(StepContext::from_turn_context(&turn)), turn: turn.into(), cancellation_token: tokio_util::sync::CancellationToken::new(), tracker: Arc::new(Mutex::new(TurnDiffTracker::new())), @@ -405,6 +408,7 @@ mod tests { .with_updated_hook_input( ToolInvocation { session: session.into(), + step_context: Arc::new(StepContext::from_turn_context(&turn)), turn: turn.into(), cancellation_token: tokio_util::sync::CancellationToken::new(), tracker: Arc::new(Mutex::new(TurnDiffTracker::new())), @@ -452,6 +456,7 @@ mod tests { .expect("MCP tool spec should build"); let invocation = ToolInvocation { session: session.into(), + step_context: Arc::new(StepContext::from_turn_context(&turn)), turn: turn.into(), cancellation_token: tokio_util::sync::CancellationToken::new(), tracker: Arc::new(Mutex::new(TurnDiffTracker::new())), diff --git a/codex-rs/core/src/tools/handlers/mod.rs b/codex-rs/core/src/tools/handlers/mod.rs index d1e931aa95..47c81eabc6 100644 --- a/codex-rs/core/src/tools/handlers/mod.rs +++ b/codex-rs/core/src/tools/handlers/mod.rs @@ -46,10 +46,10 @@ use serde_json::Map; use serde_json::Value; use std::path::Path; +use crate::environment_selection::TurnEnvironmentSnapshot; use crate::function_tool::FunctionCallError; use crate::sandboxing::SandboxPermissions; use crate::session::session::Session; -use crate::session::turn_context::TurnContext; use crate::session::turn_context::TurnEnvironment; pub(crate) use crate::tools::code_mode::CodeModeExecuteHandler; pub(crate) use crate::tools::code_mode::CodeModeWaitHandler; @@ -153,13 +153,13 @@ fn resolve_workdir_base_path( } fn resolve_tool_environment<'a>( - turn: &'a TurnContext, + environments: &'a TurnEnvironmentSnapshot, environment_id: Option<&str>, ) -> Result, FunctionCallError> { environment_id.map_or_else( - || Ok(turn.environments.primary()), + || Ok(environments.primary()), |environment_id| { - turn.environments + environments .turn_environments .iter() .find(|environment| environment.environment_id == environment_id) diff --git a/codex-rs/core/src/tools/handlers/multi_agents_tests.rs b/codex-rs/core/src/tools/handlers/multi_agents_tests.rs index 2561cd7ded..47cfaa15e5 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents_tests.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents_tests.rs @@ -4,6 +4,7 @@ use crate::config::AgentRoleConfig; use crate::config::DEFAULT_AGENT_MAX_DEPTH; use crate::function_tool::FunctionCallError; use crate::init_state_db; +use crate::session::step_context::StepContext; use crate::session::tests::make_session_and_context; use crate::session_prefix::format_inter_agent_completion_message; use crate::thread_manager::thread_store_from_config; @@ -75,6 +76,7 @@ fn invocation( ) -> ToolInvocation { ToolInvocation { session, + step_context: Arc::new(StepContext::from_turn_context(turn.as_ref())), turn, cancellation_token: CancellationToken::new(), tracker: Arc::new(Mutex::new(TurnDiffTracker::default())), diff --git a/codex-rs/core/src/tools/handlers/request_permissions.rs b/codex-rs/core/src/tools/handlers/request_permissions.rs index afc15f9659..9fcb1c642c 100644 --- a/codex-rs/core/src/tools/handlers/request_permissions.rs +++ b/codex-rs/core/src/tools/handlers/request_permissions.rs @@ -47,6 +47,7 @@ impl RequestPermissionsHandler { let ToolInvocation { session, turn, + step_context, cancellation_token, call_id, payload, @@ -63,8 +64,10 @@ impl RequestPermissionsHandler { }; let environment_args: RequestPermissionsEnvironmentArgs = parse_arguments(&arguments)?; - let Some(turn_environment) = - resolve_tool_environment(turn.as_ref(), environment_args.environment_id.as_deref())? + let Some(turn_environment) = resolve_tool_environment( + &step_context.environments, + environment_args.environment_id.as_deref(), + )? else { return Err(FunctionCallError::RespondToModel( "request_permissions requires a primary environment".to_string(), diff --git a/codex-rs/core/src/tools/handlers/request_user_input_tests.rs b/codex-rs/core/src/tools/handlers/request_user_input_tests.rs index 7c577c54be..2851e30343 100644 --- a/codex-rs/core/src/tools/handlers/request_user_input_tests.rs +++ b/codex-rs/core/src/tools/handlers/request_user_input_tests.rs @@ -1,4 +1,5 @@ use super::*; +use crate::session::step_context::StepContext; use crate::session::tests::make_session_and_context; use crate::tools::context::ToolInvocation; use crate::tools::context::ToolPayload; @@ -27,6 +28,7 @@ async fn multi_agent_v2_request_user_input_rejects_subagent_threads() { } .handle(ToolInvocation { session: Arc::new(session), + step_context: Arc::new(StepContext::from_turn_context(&turn)), turn: Arc::new(turn), cancellation_token: tokio_util::sync::CancellationToken::new(), tracker: Arc::new(Mutex::new(TurnDiffTracker::default())), diff --git a/codex-rs/core/src/tools/handlers/shell.rs b/codex-rs/core/src/tools/handlers/shell.rs index 131bcdba95..87a9301093 100644 --- a/codex-rs/core/src/tools/handlers/shell.rs +++ b/codex-rs/core/src/tools/handlers/shell.rs @@ -8,6 +8,7 @@ use crate::exec::ExecParams; use crate::exec_policy::ExecApprovalRequest; use crate::function_tool::FunctionCallError; use crate::session::turn_context::TurnContext; +use crate::session::turn_context::TurnEnvironment; use crate::shell::ShellType; use crate::tools::context::FunctionToolOutput; use crate::tools::context::ToolPayload; @@ -53,6 +54,7 @@ struct RunExecLikeArgs { prefix_rule: Option>, session: Arc, turn: Arc, + turn_environment: TurnEnvironment, tracker: crate::tools::context::SharedTurnDiffTracker, call_id: String, shell_runtime_backend: ShellRuntimeBackend, @@ -69,16 +71,12 @@ async fn run_exec_like(args: RunExecLikeArgs) -> Result ToolInvocation { ToolInvocation { session, + step_context: Arc::new(StepContext::from_turn_context(turn.as_ref())), turn, cancellation_token: tokio_util::sync::CancellationToken::new(), tracker: Arc::new(tokio::sync::Mutex::new( diff --git a/codex-rs/core/src/tools/router.rs b/codex-rs/core/src/tools/router.rs index b488448398..7379212485 100644 --- a/codex-rs/core/src/tools/router.rs +++ b/codex-rs/core/src/tools/router.rs @@ -1,5 +1,6 @@ use crate::function_tool::FunctionCallError; use crate::session::session::Session; +use crate::session::step_context::StepContext; use crate::session::turn_context::TurnContext; use crate::tools::context::SharedTurnDiffTracker; use crate::tools::context::ToolInvocation; @@ -35,6 +36,7 @@ pub struct ToolCall { pub struct ToolRouter { registry: ToolRegistry, model_visible_specs: Vec, + step_context: Arc, } pub(crate) struct ToolRouterParams<'a> { @@ -58,18 +60,29 @@ pub(crate) struct ToolSuggestCandidates { } impl ToolRouter { - pub(crate) fn from_turn_context( + pub(crate) fn from_contexts( turn_context: &TurnContext, + step_context: Arc, params: ToolRouterParams<'_>, tool_search_handler_cache: &ToolSearchHandlerCache, ) -> Self { - build_tool_router(turn_context, params, tool_search_handler_cache) + build_tool_router( + turn_context, + step_context, + params, + tool_search_handler_cache, + ) } - pub(crate) fn from_parts(registry: ToolRegistry, model_visible_specs: Vec) -> Self { + pub(crate) fn from_parts( + registry: ToolRegistry, + model_visible_specs: Vec, + step_context: Arc, + ) -> Self { Self { registry, model_visible_specs, + step_context, } } @@ -226,6 +239,7 @@ impl ToolRouter { let invocation = ToolInvocation { session, turn, + step_context: Arc::clone(&self.step_context), cancellation_token, tracker, call_id, diff --git a/codex-rs/core/src/tools/router_tests.rs b/codex-rs/core/src/tools/router_tests.rs index 77ddfe48d5..0ebf052160 100644 --- a/codex-rs/core/src/tools/router_tests.rs +++ b/codex-rs/core/src/tools/router_tests.rs @@ -1,6 +1,7 @@ use std::sync::Arc; use crate::config::Config; +use crate::session::step_context::StepContext; use crate::session::tests::make_session_and_context; use crate::tools::context::ToolPayload; use crate::turn_diff_tracker::TurnDiffTracker; @@ -111,8 +112,9 @@ async fn parallel_support_does_not_match_namespaced_local_tool_names() -> anyhow .load_full() .list_all_tools() .await; - let router = ToolRouter::from_turn_context( + let router = ToolRouter::from_contexts( &turn, + Arc::new(StepContext::from_turn_context(&turn)), ToolRouterParams { tool_suggest_candidates: None, deferred_mcp_tools: None, @@ -179,8 +181,9 @@ async fn build_tool_call_uses_namespace_for_registry_name() -> anyhow::Result<() #[tokio::test] async fn mcp_parallel_support_uses_handler_data() -> anyhow::Result<()> { let (_, turn) = make_session_and_context().await; - let router = ToolRouter::from_turn_context( + let router = ToolRouter::from_contexts( &turn, + Arc::new(StepContext::from_turn_context(&turn)), ToolRouterParams { tool_suggest_candidates: None, deferred_mcp_tools: None, @@ -228,8 +231,9 @@ async fn mcp_parallel_support_uses_handler_data() -> anyhow::Result<()> { #[tokio::test] async fn tools_without_handlers_do_not_support_parallel() -> anyhow::Result<()> { let (_, turn) = make_session_and_context().await; - let router = ToolRouter::from_turn_context( + let router = ToolRouter::from_contexts( &turn, + Arc::new(StepContext::from_turn_context(&turn)), ToolRouterParams { tool_suggest_candidates: None, deferred_mcp_tools: None, @@ -283,8 +287,9 @@ async fn specs_filter_deferred_dynamic_tools() -> anyhow::Result<()> { ], })]; - let router = ToolRouter::from_turn_context( + let router = ToolRouter::from_contexts( &turn, + Arc::new(StepContext::from_turn_context(&turn)), ToolRouterParams { tool_suggest_candidates: None, deferred_mcp_tools: None, @@ -348,8 +353,9 @@ async fn extension_tool_executors_are_model_visible_and_dispatchable() -> anyhow let mut expected_history_item = history_item.clone(); expected_history_item.set_turn_id_if_missing(&turn.sub_id); - let router = ToolRouter::from_turn_context( + let router = ToolRouter::from_contexts( &turn, + Arc::new(StepContext::from_turn_context(&turn)), ToolRouterParams { tool_suggest_candidates: None, deferred_mcp_tools: None, diff --git a/codex-rs/core/src/tools/spec_plan.rs b/codex-rs/core/src/tools/spec_plan.rs index 51cd436503..51da1e8a94 100644 --- a/codex-rs/core/src/tools/spec_plan.rs +++ b/codex-rs/core/src/tools/spec_plan.rs @@ -1,5 +1,6 @@ use crate::agent::exceeds_thread_spawn_depth_limit; use crate::agent::next_thread_spawn_depth; +use crate::session::step_context::StepContext; use crate::session::turn_context::TurnContext; use crate::tools::code_mode::execute_spec::create_code_mode_tool; use crate::tools::context::ToolInvocation; @@ -144,6 +145,7 @@ impl PlannedTools { #[derive(Clone, Copy)] struct CoreToolPlanContext<'a> { turn_context: &'a TurnContext, + step_context: &'a StepContext, mcp_tools: Option<&'a [ToolInfo]>, deferred_mcp_tools: Option<&'a [ToolInfo]>, tool_suggest_candidates: Option<&'a crate::tools::router::ToolSuggestCandidates>, @@ -157,17 +159,23 @@ struct CoreToolPlanContext<'a> { #[instrument(level = "trace", skip_all)] pub(crate) fn build_tool_router( turn_context: &TurnContext, + step_context: Arc, params: ToolRouterParams<'_>, tool_search_handler_cache: &ToolSearchHandlerCache, ) -> ToolRouter { - let (model_visible_specs, registry) = - build_tool_specs_and_registry(turn_context, params, tool_search_handler_cache); - ToolRouter::from_parts(registry, model_visible_specs) + let (model_visible_specs, registry) = build_tool_specs_and_registry( + turn_context, + &step_context, + params, + tool_search_handler_cache, + ); + ToolRouter::from_parts(registry, model_visible_specs, step_context) } #[instrument(level = "trace", skip_all)] fn build_tool_specs_and_registry( turn_context: &TurnContext, + step_context: &StepContext, params: ToolRouterParams<'_>, tool_search_handler_cache: &ToolSearchHandlerCache, ) -> (Vec, ToolRegistry) { @@ -182,6 +190,7 @@ fn build_tool_specs_and_registry( crate::agent::role::spawn_tool_spec::build(&std::collections::BTreeMap::new()); let context = CoreToolPlanContext { turn_context, + step_context, mcp_tools: mcp_tools.as_deref(), deferred_mcp_tools: deferred_mcp_tools.as_deref(), tool_suggest_candidates: tool_suggest_candidates.as_ref(), @@ -625,11 +634,15 @@ fn standalone_web_search_enabled(turn_context: &TurnContext) -> bool { .enabled(Feature::StandaloneWebSearch)) } +fn tool_environment_mode(step_context: &StepContext) -> ToolEnvironmentMode { + ToolEnvironmentMode::from_count(step_context.environments.turn_environments.len()) +} + #[instrument(level = "trace", skip_all)] fn add_shell_tools(context: &CoreToolPlanContext<'_>, planned_tools: &mut PlannedTools) { let turn_context = context.turn_context; let features = turn_context.config.features.get(); - let environment_mode = turn_context.tool_environment_mode(); + let environment_mode = tool_environment_mode(context.step_context); if !environment_mode.has_environment() { return; } @@ -649,7 +662,10 @@ fn add_shell_tools(context: &CoreToolPlanContext<'_>, planned_tools: &mut Planne allow_login_shell, exec_permission_approvals_enabled, include_environment_id, - include_shell_parameter: unified_exec_should_include_shell_parameter(turn_context), + include_shell_parameter: unified_exec_should_include_shell_parameter( + turn_context, + context.step_context, + ), })); planned_tools.add(WriteStdinHandler); @@ -666,11 +682,14 @@ fn add_shell_tools(context: &CoreToolPlanContext<'_>, planned_tools: &mut Planne } } -fn unified_exec_should_include_shell_parameter(turn_context: &TurnContext) -> bool { +fn unified_exec_should_include_shell_parameter( + turn_context: &TurnContext, + step_context: &StepContext, +) -> bool { !matches!( &turn_context.unified_exec_shell_mode, UnifiedExecShellMode::ZshFork(_) - ) || turn_context + ) || step_context .environments .turn_environments .iter() @@ -690,7 +709,7 @@ fn add_mcp_resource_tools(context: &CoreToolPlanContext<'_>, planned_tools: &mut fn add_core_utility_tools(context: &CoreToolPlanContext<'_>, planned_tools: &mut PlannedTools) { let turn_context = context.turn_context; let features = turn_context.config.features.get(); - let environment_mode = turn_context.tool_environment_mode(); + let environment_mode = tool_environment_mode(context.step_context); planned_tools.add(PlanHandler); diff --git a/codex-rs/core/src/tools/spec_plan_tests.rs b/codex-rs/core/src/tools/spec_plan_tests.rs index 5d70fb8f17..eb2189189b 100644 --- a/codex-rs/core/src/tools/spec_plan_tests.rs +++ b/codex-rs/core/src/tools/spec_plan_tests.rs @@ -30,6 +30,7 @@ use codex_tools::ToolSpec; use pretty_assertions::assert_eq; use serde_json::json; +use crate::session::step_context::StepContext; use crate::session::tests::make_session_and_context; use crate::session::turn_context::TurnContext; use crate::tools::handlers::ToolSearchHandlerCache; @@ -178,8 +179,9 @@ async fn probe_with( ) -> ToolPlanProbe { let (_session, mut turn) = make_session_and_context().await; configure_turn(&mut turn); - let router = ToolRouter::from_turn_context( + let router = ToolRouter::from_contexts( &turn, + Arc::new(StepContext::from_turn_context(&turn)), ToolRouterParams { tool_suggest_candidates: inputs.tool_suggest_candidates, mcp_tools: inputs.mcp_tools, @@ -639,6 +641,31 @@ async fn environment_count_controls_environment_backed_tools() { )); } +#[tokio::test] +async fn environment_tools_follow_the_step_context() { + let (_session, mut turn) = make_session_and_context().await; + set_feature(&mut turn, Feature::UnifiedExec, /*enabled*/ true); + turn.model_info.apply_patch_tool_type = Some(ApplyPatchToolType::Freeform); + + let step_context = Arc::new(StepContext::from_turn_context(&turn)); + turn.environments.turn_environments.clear(); + + let plan = ToolPlanProbe::from_router(ToolRouter::from_contexts( + &turn, + step_context, + ToolRouterParams { + mcp_tools: None, + deferred_mcp_tools: None, + tool_suggest_candidates: None, + extension_tool_executors: Vec::new(), + dynamic_tools: &[], + }, + &Default::default(), + )); + + plan.assert_visible_contains(&["exec_command", "apply_patch", "view_image"]); +} + #[tokio::test] async fn host_context_gates_agent_job_tools() { let normal_agent_job = probe(|turn| { @@ -773,8 +800,9 @@ async fn tool_search_cache_rebuilds_when_deferred_sources_change() { let (_session, mut first_turn) = make_session_and_context().await; first_turn.model_info.supports_search_tool = true; - let first_router = ToolRouter::from_turn_context( + let first_router = ToolRouter::from_contexts( &first_turn, + Arc::new(StepContext::from_turn_context(&first_turn)), ToolRouterParams { mcp_tools: None, deferred_mcp_tools: Some(vec![mcp_tool("first", "mcp__first", "lookup")]), @@ -788,8 +816,9 @@ async fn tool_search_cache_rebuilds_when_deferred_sources_change() { let (_session, mut second_turn) = make_session_and_context().await; second_turn.model_info.supports_search_tool = true; - let second_router = ToolRouter::from_turn_context( + let second_router = ToolRouter::from_contexts( &second_turn, + Arc::new(StepContext::from_turn_context(&second_turn)), ToolRouterParams { mcp_tools: None, deferred_mcp_tools: Some(vec![mcp_tool("second", "mcp__second", "lookup")]), diff --git a/codex-rs/core/src/tools/tool_dispatch_trace_tests.rs b/codex-rs/core/src/tools/tool_dispatch_trace_tests.rs index b09996872e..349ed16d2c 100644 --- a/codex-rs/core/src/tools/tool_dispatch_trace_tests.rs +++ b/codex-rs/core/src/tools/tool_dispatch_trace_tests.rs @@ -13,6 +13,7 @@ use tokio_util::sync::CancellationToken; use crate::function_tool::FunctionCallError; use crate::session::session::Session; +use crate::session::step_context::StepContext; use crate::session::tests::make_session_and_context; use crate::session::turn_context::TurnContext; use crate::tools::code_mode::CodeModeWaitHandler; @@ -271,6 +272,7 @@ fn test_invocation_with_payload( ) -> ToolInvocation { ToolInvocation { session, + step_context: Arc::new(StepContext::from_turn_context(turn.as_ref())), turn, cancellation_token: CancellationToken::new(), tracker: Arc::new(tokio::sync::Mutex::new(TurnDiffTracker::new())), diff --git a/codex-rs/core/tests/suite/remote_env.rs b/codex-rs/core/tests/suite/remote_env.rs index 4f02001d8c..52ddb61eaf 100644 --- a/codex-rs/core/tests/suite/remote_env.rs +++ b/codex-rs/core/tests/suite/remote_env.rs @@ -547,7 +547,7 @@ async fn serve_environment_info(listener: TcpListener) { } #[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn deferred_executor_updates_model_context_after_startup() -> Result<()> { +async fn deferred_executor_updates_context_and_tools_after_startup() -> Result<()> { let listener = TcpListener::bind("127.0.0.1:0").await?; let server = start_mock_server().await; let user_input_call_id = "wait-for-startup"; @@ -580,11 +580,13 @@ async fn deferred_executor_updates_model_context_after_startup() -> Result<()> { sse(vec![ ev_response_created("resp-2"), ev_function_call( - "update-plan", - "update_plan", + "request-permissions", + "request_permissions", &json!({ - "explanation": "Continue after startup.", - "plan": [{"step": "Finish", "status": "completed"}] + "reason": "Verify that the ready environment is used.", + "permissions": { + "network": { "enabled": true } + } }) .to_string(), ), @@ -601,7 +603,17 @@ async fn deferred_executor_updates_model_context_after_startup() -> Result<()> { let mut builder = test_codex() .with_exec_server_url(format!("ws://{}", listener.local_addr()?)) .with_config(|config| { + config.use_experimental_unified_exec_tool = true; + config.permissions.approval_policy = Constrained::allow_any(AskForApproval::OnRequest); + config.approvals_reviewer = ApprovalsReviewer::User; assert!(config.features.enable(Feature::DeferredExecutor).is_ok()); + assert!(config.features.enable(Feature::UnifiedExec).is_ok()); + assert!( + config + .features + .enable(Feature::RequestPermissionsTool) + .is_ok() + ); assert!( config .features @@ -645,6 +657,30 @@ async fn deferred_executor_updates_model_context_after_startup() -> Result<()> { }, }) .await?; + let event = wait_for_event(&test.codex, |event| { + matches!( + event, + EventMsg::RequestPermissions(_) | EventMsg::TurnComplete(_) + ) + }) + .await; + let EventMsg::RequestPermissions(permission_request) = event else { + panic!("ready environment should be available to request_permissions: {event:?}"); + }; + assert_eq!( + permission_request.environment_id.as_deref(), + Some(REMOTE_ENVIRONMENT_ID) + ); + test.codex + .submit(Op::RequestPermissionsResponse { + id: permission_request.call_id, + response: RequestPermissionsResponse { + permissions: RequestPermissionProfile::default(), + scope: PermissionGrantScope::Turn, + strict_auto_review: false, + }, + }) + .await?; wait_for_event(&test.codex, |event| { matches!(event, EventMsg::TurnComplete(_)) }) @@ -652,6 +688,16 @@ async fn deferred_executor_updates_model_context_after_startup() -> Result<()> { let requests = response_mock.requests(); assert_eq!(requests.len(), 3); + let tool_names = |request_index: usize| { + requests[request_index].body_json()["tools"] + .as_array() + .expect("tools array") + .iter() + .filter_map(|tool| tool.get("name").and_then(Value::as_str).map(str::to_owned)) + .collect::>() + }; + assert!(!tool_names(0).contains(&"exec_command".to_string())); + assert!(tool_names(1).contains(&"exec_command".to_string())); assert!( requests[0] .message_input_texts("user")