From bbad6a804bd9520b5c22cef768d2f02e8e191280 Mon Sep 17 00:00:00 2001 From: starr-openai Date: Mon, 20 Apr 2026 15:23:54 -0700 Subject: [PATCH] Fix sticky environment CI fallout Shrink app.rs below the blob policy threshold, clean up dead-code fallout, and update stale exec-server listen URL test expectations. Co-authored-by: Codex --- codex-rs/core/src/apply_patch.rs | 3 +- .../core/src/context/environment_context.rs | 42 +++++- .../core/src/context_manager/history_tests.rs | 1 + codex-rs/core/src/session/mod.rs | 8 ++ .../session/rollout_reconstruction_tests.rs | 8 ++ codex-rs/core/src/session/tests.rs | 88 ++++++++++++ codex-rs/core/src/session/turn_context.rs | 30 +++- .../core/src/tools/handlers/apply_patch.rs | 131 +++++++++++++++--- .../src/tools/handlers/apply_patch_tests.rs | 25 ++++ .../core/src/tools/handlers/unified_exec.rs | 8 +- .../core/src/tools/handlers/view_image.rs | 37 +++-- .../core/src/tools/runtimes/apply_patch.rs | 39 ++++-- .../src/tools/runtimes/apply_patch_tests.rs | 17 +++ codex-rs/core/tests/suite/resume_warning.rs | 1 + .../exec-server/src/server/transport_tests.rs | 2 +- .../exec-server/tests/command_transport.rs | 50 +++++++ codex-rs/protocol/src/protocol.rs | 51 ++++++- codex-rs/rollout/src/recorder_tests.rs | 1 + codex-rs/state/src/extract.rs | 3 + codex-rs/tools/BUILD.bazel | 1 + codex-rs/tools/src/apply_patch_tool.rs | 37 ++++- codex-rs/tools/src/apply_patch_tool_tests.rs | 36 ++++- codex-rs/tools/src/lib.rs | 1 + .../src/tool_apply_patch_environment.lark | 21 +++ codex-rs/tools/src/tool_config.rs | 5 + codex-rs/tools/src/tool_registry_plan.rs | 10 +- .../tools/src/tool_registry_plan_tests.rs | 36 ++++- codex-rs/tools/src/view_image.rs | 9 ++ codex-rs/tui/src/lib.rs | 1 + 29 files changed, 631 insertions(+), 71 deletions(-) create mode 100644 codex-rs/exec-server/tests/command_transport.rs create mode 100644 codex-rs/tools/src/tool_apply_patch_environment.lark diff --git a/codex-rs/core/src/apply_patch.rs b/codex-rs/core/src/apply_patch.rs index c05a459049..0788ae216c 100644 --- a/codex-rs/core/src/apply_patch.rs +++ b/codex-rs/core/src/apply_patch.rs @@ -32,6 +32,7 @@ pub(crate) struct ApplyPatchRuntimeInvocation { pub(crate) async fn apply_patch( turn_context: &TurnContext, + cwd: &codex_utils_absolute_path::AbsolutePathBuf, file_system_sandbox_policy: &FileSystemSandboxPolicy, action: ApplyPatchAction, ) -> InternalApplyPatchInvocation { @@ -40,7 +41,7 @@ pub(crate) async fn apply_patch( turn_context.approval_policy.value(), turn_context.sandbox_policy.get(), file_system_sandbox_policy, - &turn_context.cwd, + cwd, turn_context.windows_sandbox_level, ) { SafetyCheck::AutoApprove { diff --git a/codex-rs/core/src/context/environment_context.rs b/codex-rs/core/src/context/environment_context.rs index d3105a03f2..046f172537 100644 --- a/codex-rs/core/src/context/environment_context.rs +++ b/codex-rs/core/src/context/environment_context.rs @@ -1,6 +1,5 @@ use crate::session::turn_context::TurnContext; use crate::shell::Shell; -use codex_features::Feature; use codex_protocol::protocol::TurnContextItem; use codex_protocol::protocol::TurnContextNetworkItem; use std::path::PathBuf; @@ -41,6 +40,7 @@ impl NetworkContext { } impl EnvironmentContext { + #[cfg(test)] pub(crate) fn new( cwd: Option, shell: String, @@ -117,7 +117,7 @@ impl EnvironmentContext { }; EnvironmentContext::new_with_environments( cwd, - after.environments.clone(), + Self::environments_for_diff(before, after), after.shell.clone(), after.current_date.clone(), after.timezone.clone(), @@ -142,8 +142,9 @@ impl EnvironmentContext { turn_context_item: &TurnContextItem, shell: String, ) -> Self { - Self::new( + Self::new_with_environments( Some(turn_context_item.cwd.clone()), + Self::environments_from_turn_context_item(turn_context_item), shell, turn_context_item.current_date.clone(), turn_context_item.timezone.clone(), @@ -197,7 +198,7 @@ impl EnvironmentContext { fn environments_from_turn_context( turn_context: &TurnContext, ) -> Option> { - if !turn_context.features.enabled(Feature::MultiEnvironmentTools) { + if !turn_context.tools_config.multi_environment_tools { return None; } let environments = turn_context.environments.as_ref()?; @@ -217,6 +218,39 @@ impl EnvironmentContext { .collect(), ) } + + fn environments_from_turn_context_item( + turn_context_item: &TurnContextItem, + ) -> Option> { + let environments = turn_context_item.environments.as_ref()?; + if environments.is_empty() { + return None; + } + + Some( + environments + .iter() + .enumerate() + .map(|(index, environment)| EnvironmentContextEnvironment { + id: environment.environment_id.clone(), + cwd: environment.cwd.to_path_buf(), + primary: index == 0, + }) + .collect(), + ) + } + + fn environments_for_diff( + before: &TurnContextItem, + after: &EnvironmentContext, + ) -> Option> { + let before_environments = Self::environments_from_turn_context_item(before); + if before_environments == after.environments { + None + } else { + after.environments.clone() + } + } } impl ContextualUserFragment for EnvironmentContext { diff --git a/codex-rs/core/src/context_manager/history_tests.rs b/codex-rs/core/src/context_manager/history_tests.rs index 1df14ca8bc..af7f620c28 100644 --- a/codex-rs/core/src/context_manager/history_tests.rs +++ b/codex-rs/core/src/context_manager/history_tests.rs @@ -129,6 +129,7 @@ fn reference_context_item() -> TurnContextItem { turn_id: Some("reference-turn".to_string()), trace_id: None, cwd: PathBuf::from("/tmp/reference-cwd"), + environments: None, current_date: Some("2026-03-23".to_string()), timezone: Some("America/Los_Angeles".to_string()), approval_policy: AskForApproval::OnRequest, diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index da8daf1f21..ac75741fff 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -1257,11 +1257,19 @@ impl Session { .reconstruct_history_from_rollout(turn_context, rollout_items) .await; let previous_turn_settings = reconstructed_rollout.previous_turn_settings.clone(); + let environment_selections = reconstructed_rollout + .reference_context_item + .as_ref() + .and_then(|context_item| context_item.environments.clone()); self.replace_history( reconstructed_rollout.history, reconstructed_rollout.reference_context_item, ) .await; + if let Some(environment_selections) = environment_selections { + let mut state = self.state.lock().await; + state.session_configuration.environment_selections = Some(environment_selections); + } self.set_previous_turn_settings(previous_turn_settings.clone()) .await; previous_turn_settings diff --git a/codex-rs/core/src/session/rollout_reconstruction_tests.rs b/codex-rs/core/src/session/rollout_reconstruction_tests.rs index 3dca7f0ce6..86a8f65bd9 100644 --- a/codex-rs/core/src/session/rollout_reconstruction_tests.rs +++ b/codex-rs/core/src/session/rollout_reconstruction_tests.rs @@ -64,6 +64,7 @@ async fn record_initial_history_resumed_bare_turn_context_does_not_hydrate_previ turn_id: Some(turn_context.sub_id.clone()), trace_id: turn_context.trace_id.clone(), cwd: turn_context.cwd.to_path_buf(), + environments: None, current_date: turn_context.current_date.clone(), timezone: turn_context.timezone.clone(), approval_policy: turn_context.approval_policy.value(), @@ -104,6 +105,7 @@ async fn record_initial_history_resumed_hydrates_previous_turn_settings_from_lif turn_id: Some(turn_context.sub_id.clone()), trace_id: turn_context.trace_id.clone(), cwd: turn_context.cwd.to_path_buf(), + environments: None, current_date: turn_context.current_date.clone(), timezone: turn_context.timezone.clone(), approval_policy: turn_context.approval_policy.value(), @@ -899,6 +901,7 @@ async fn record_initial_history_resumed_turn_context_after_compaction_reestablis turn_id: Some(turn_context.sub_id.clone()), trace_id: turn_context.trace_id.clone(), cwd: turn_context.cwd.to_path_buf(), + environments: None, current_date: turn_context.current_date.clone(), timezone: turn_context.timezone.clone(), approval_policy: turn_context.approval_policy.value(), @@ -975,6 +978,7 @@ async fn record_initial_history_resumed_turn_context_after_compaction_reestablis turn_id: Some(turn_context.sub_id.clone()), trace_id: turn_context.trace_id.clone(), cwd: turn_context.cwd.to_path_buf(), + environments: None, current_date: turn_context.current_date.clone(), timezone: turn_context.timezone.clone(), approval_policy: turn_context.approval_policy.value(), @@ -1005,6 +1009,7 @@ async fn record_initial_history_resumed_aborted_turn_without_id_clears_active_tu turn_id: Some(turn_context.sub_id.clone()), trace_id: turn_context.trace_id.clone(), cwd: turn_context.cwd.to_path_buf(), + environments: None, current_date: turn_context.current_date.clone(), timezone: turn_context.timezone.clone(), approval_policy: turn_context.approval_policy.value(), @@ -1118,6 +1123,7 @@ async fn record_initial_history_resumed_unmatched_abort_preserves_active_turn_fo turn_id: Some(current_turn_id.clone()), trace_id: turn_context.trace_id.clone(), cwd: turn_context.cwd.to_path_buf(), + environments: None, current_date: turn_context.current_date.clone(), timezone: turn_context.timezone.clone(), approval_policy: turn_context.approval_policy.value(), @@ -1229,6 +1235,7 @@ async fn record_initial_history_resumed_trailing_incomplete_turn_compaction_clea turn_id: Some(turn_context.sub_id.clone()), trace_id: turn_context.trace_id.clone(), cwd: turn_context.cwd.to_path_buf(), + environments: None, current_date: turn_context.current_date.clone(), timezone: turn_context.timezone.clone(), approval_policy: turn_context.approval_policy.value(), @@ -1379,6 +1386,7 @@ async fn record_initial_history_resumed_replaced_incomplete_compacted_turn_clear turn_id: Some(turn_context.sub_id.clone()), trace_id: turn_context.trace_id.clone(), cwd: turn_context.cwd.to_path_buf(), + environments: None, current_date: turn_context.current_date.clone(), timezone: turn_context.timezone.clone(), approval_policy: turn_context.approval_policy.value(), diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index d921d807b8..72c2ef82d8 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -1733,10 +1733,15 @@ async fn fork_startup_context_then_first_turn_diff_snapshot() -> anyhow::Result< async fn record_initial_history_forked_hydrates_previous_turn_settings() { let (session, turn_context) = make_session_and_context().await; let previous_model = "forked-rollout-model"; + let previous_environments = vec![codex_protocol::protocol::TurnEnvironmentSelection { + environment_id: "local".to_string(), + cwd: test_path_buf("/tmp/forked-local").abs(), + }]; let previous_context_item = TurnContextItem { turn_id: Some(turn_context.sub_id.clone()), trace_id: turn_context.trace_id.clone(), cwd: turn_context.cwd.to_path_buf(), + environments: Some(previous_environments.clone()), current_date: turn_context.current_date.clone(), timezone: turn_context.timezone.clone(), approval_policy: turn_context.approval_policy.value(), @@ -1798,6 +1803,15 @@ async fn record_initial_history_forked_hydrates_previous_turn_settings() { realtime_active: Some(turn_context.realtime_active), }) ); + assert_eq!( + session + .state + .lock() + .await + .session_configuration + .environment_selections, + Some(previous_environments) + ); assert_eq!(history.raw_items(), &[]); assert_eq!( serde_json::to_value(session.reference_context_item().await) @@ -3974,6 +3988,14 @@ async fn turn_environment_selection_sets_primary_environment() { .resolve_path_for_environment(&primary_environment, Some("relative/path".to_string())), selected_cwd.join("relative/path") ); + + assert_eq!( + turn_context.to_turn_context_item().environments, + Some(vec![codex_protocol::protocol::TurnEnvironmentSelection { + environment_id: "local".to_string(), + cwd: selected_cwd, + }]) + ); } #[tokio::test] @@ -4028,6 +4050,20 @@ async fn multiple_turn_environment_selections_use_first_as_primary_environment() .environment_for_tool(Some("local")) .expect("tool environment should resolve by id"); assert_eq!(tool_environment.cwd, first_cwd); + + assert_eq!( + turn_context.to_turn_context_item().environments, + Some(vec![ + codex_protocol::protocol::TurnEnvironmentSelection { + environment_id: "local".to_string(), + cwd: first_cwd, + }, + codex_protocol::protocol::TurnEnvironmentSelection { + environment_id: "local".to_string(), + cwd: second_cwd, + }, + ]) + ); } #[tokio::test] @@ -4096,6 +4132,58 @@ async fn primary_environment_falls_back_to_compatibility_projection() { ); } +#[tokio::test] +async fn multi_environment_feature_without_selections_uses_single_environment_behavior() { + let (_session, turn_context, _rx) = make_session_and_context_with_auth_and_config_and_rx( + CodexAuth::from_api_key("Test API Key"), + Vec::new(), + |config| config.features.enable(Feature::MultiEnvironmentTools), + ) + .await; + + assert!( + turn_context + .features + .enabled(Feature::MultiEnvironmentTools) + ); + assert!(turn_context.environments.is_none()); + assert!(!turn_context.tools_config.multi_environment_tools); + assert_eq!(turn_context.to_turn_context_item().environments, None); + + let primary_environment = turn_context + .primary_environment() + .expect("primary environment should resolve from compatibility fields"); + assert_eq!(primary_environment.environment_id, None); + assert_eq!(primary_environment.cwd, turn_context.cwd); +} + +#[tokio::test] +async fn multi_environment_tools_enable_only_with_feature_and_selection() { + let (session, _turn_context, _rx) = make_session_and_context_with_auth_and_config_and_rx( + CodexAuth::from_api_key("Test API Key"), + Vec::new(), + |config| config.features.enable(Feature::MultiEnvironmentTools), + ) + .await; + let selected_cwd = + AbsolutePathBuf::try_from(session.get_config().await.cwd.as_path().join("selected")) + .expect("absolute path"); + + let turn_context = session + .new_turn_with_sub_id( + "sub-1".to_string(), + SessionSettingsUpdate::default(), + Some(vec![codex_protocol::protocol::TurnEnvironmentSelection { + environment_id: "local".to_string(), + cwd: selected_cwd, + }]), + ) + .await + .expect("turn should start"); + + assert!(turn_context.tools_config.multi_environment_tools); +} + #[tokio::test] async fn unknown_turn_environment_selection_returns_error() { let (session, _turn_context, _rx) = make_session_and_context_with_rx().await; diff --git a/codex-rs/core/src/session/turn_context.rs b/codex-rs/core/src/session/turn_context.rs index 243f5b9327..6f15d9ba1f 100644 --- a/codex-rs/core/src/session/turn_context.rs +++ b/codex-rs/core/src/session/turn_context.rs @@ -10,6 +10,14 @@ pub(super) fn image_generation_tool_auth_allowed(auth_manager: Option<&AuthManag ) } +fn multi_environment_tools_enabled( + features: &codex_features::Features, + environments: Option<&[TurnEnvironment]>, +) -> bool { + features.enabled(codex_features::Feature::MultiEnvironmentTools) + && environments.is_some_and(|environments| !environments.is_empty()) +} + #[derive(Clone, Debug)] pub(crate) struct TurnSkillsContext { pub(crate) outcome: Arc, @@ -155,6 +163,8 @@ impl TurnContext { /*developer_instructions*/ None, ); let features = self.features.clone(); + let multi_environment_tools = + multi_environment_tools_enabled(&features, self.environments.as_deref()); let tools_config = ToolsConfig::new(&ToolsConfigParams { model_info: &model_info, available_models: &models_manager @@ -173,6 +183,7 @@ impl TurnContext { .with_web_search_config(self.tools_config.web_search_config.clone()) .with_allow_login_shell(self.tools_config.allow_login_shell) .with_has_environment(self.tools_config.has_environment) + .with_multi_environment_tools(multi_environment_tools) .with_spawn_agent_usage_hint(config.multi_agent_v2.usage_hint_enabled) .with_spawn_agent_usage_hint_text(config.multi_agent_v2.usage_hint_text.clone()) .with_hide_spawn_agent_metadata(config.multi_agent_v2.hide_spawn_agent_metadata) @@ -279,13 +290,6 @@ impl TurnContext { ) } - pub(crate) fn file_system_sandbox_context( - &self, - additional_permissions: Option, - ) -> FileSystemSandboxContext { - self.file_system_sandbox_context_for_cwd(&self.cwd, additional_permissions) - } - pub(crate) fn file_system_sandbox_context_for_cwd( &self, cwd: &AbsolutePathBuf, @@ -330,6 +334,15 @@ impl TurnContext { turn_id: Some(self.sub_id.clone()), trace_id: self.trace_id.clone(), cwd: self.cwd.to_path_buf(), + environments: self.environments.as_ref().map(|environments| { + environments + .iter() + .map(|environment| TurnEnvironmentSelection { + environment_id: environment.environment_id.clone(), + cwd: environment.cwd.clone(), + }) + .collect() + }), current_date: self.current_date.clone(), timezone: self.timezone.clone(), approval_policy: self.approval_policy.value(), @@ -449,6 +462,8 @@ impl Session { let session_source = session_configuration.session_source.clone(); let image_generation_tool_auth_allowed = image_generation_tool_auth_allowed(auth_manager.as_deref()); + let multi_environment_tools = + multi_environment_tools_enabled(&per_turn_config.features, environments.as_deref()); let auth_manager_for_context = auth_manager.clone(); let provider_for_context = create_model_provider(provider, auth_manager); let session_telemetry_for_context = session_telemetry; @@ -470,6 +485,7 @@ impl Session { .with_web_search_config(per_turn_config.web_search_config.clone()) .with_allow_login_shell(per_turn_config.permissions.allow_login_shell) .with_has_environment(environment.is_some()) + .with_multi_environment_tools(multi_environment_tools) .with_spawn_agent_usage_hint(per_turn_config.multi_agent_v2.usage_hint_enabled) .with_spawn_agent_usage_hint_text(per_turn_config.multi_agent_v2.usage_hint_text.clone()) .with_hide_spawn_agent_metadata(per_turn_config.multi_agent_v2.hide_spawn_agent_metadata) diff --git a/codex-rs/core/src/tools/handlers/apply_patch.rs b/codex-rs/core/src/tools/handlers/apply_patch.rs index 033315a694..c6c32e00bf 100644 --- a/codex-rs/core/src/tools/handlers/apply_patch.rs +++ b/codex-rs/core/src/tools/handlers/apply_patch.rs @@ -47,6 +47,7 @@ use codex_tools::ApplyPatchToolArgs; use codex_utils_absolute_path::AbsolutePathBuf; const APPLY_PATCH_ARGUMENT_DIFF_BUFFER_INTERVAL: Duration = Duration::from_millis(500); +const APPLY_PATCH_ENVIRONMENT_PREFIX: &str = "*** Environment: "; pub struct ApplyPatchHandler; @@ -239,9 +240,48 @@ fn write_permissions_for_paths( normalize_additional_permissions(permissions).ok() } +fn split_freeform_apply_patch_environment( + input: String, +) -> Result<(String, Option), FunctionCallError> { + let Some(rest) = input.strip_prefix(APPLY_PATCH_ENVIRONMENT_PREFIX) else { + return Ok((input, None)); + }; + let Some((environment_id, patch_input)) = rest.split_once('\n') else { + return Err(FunctionCallError::RespondToModel( + "apply_patch environment header must be followed by a patch body".to_string(), + )); + }; + let environment_id = environment_id.trim(); + if environment_id.is_empty() { + return Err(FunctionCallError::RespondToModel( + "apply_patch environment header must include an environment id".to_string(), + )); + } + Ok((patch_input.to_string(), Some(environment_id.to_string()))) +} + +fn merge_apply_patch_environment_ids( + argument_environment_id: Option, + header_environment_id: Option, +) -> Result, FunctionCallError> { + match (argument_environment_id, header_environment_id) { + (Some(argument_environment_id), Some(header_environment_id)) + if argument_environment_id != header_environment_id => + { + Err(FunctionCallError::RespondToModel(format!( + "apply_patch environment id mismatch: argument requested `{argument_environment_id}` but patch header requested `{header_environment_id}`" + ))) + } + (Some(environment_id), Some(_)) | (Some(environment_id), None) => Ok(Some(environment_id)), + (None, Some(environment_id)) => Ok(Some(environment_id)), + (None, None) => Ok(None), + } +} + async fn effective_patch_permissions( session: &Session, turn: &TurnContext, + cwd: &AbsolutePathBuf, action: &ApplyPatchAction, ) -> ( Vec, @@ -261,7 +301,7 @@ async fn effective_patch_permissions( session, turn.cwd.as_path(), crate::sandboxing::SandboxPermissions::UseDefault, - write_permissions_for_paths(&file_paths, &file_system_sandbox_policy, &turn.cwd), + write_permissions_for_paths(&file_paths, &file_system_sandbox_policy, cwd), ) .await; @@ -305,12 +345,29 @@ impl ToolHandler for ApplyPatchHandler { .. } = invocation; - let patch_input = match payload { + let multi_environment_tools = turn.tools_config.multi_environment_tools; + let (patch_input, requested_environment_id) = match payload { ToolPayload::Function { arguments } => { let args: ApplyPatchToolArgs = parse_arguments(&arguments)?; - args.input + if multi_environment_tools { + let (patch_input, header_environment_id) = + split_freeform_apply_patch_environment(args.input)?; + let environment_id = merge_apply_patch_environment_ids( + args.environment_id, + header_environment_id, + )?; + (patch_input, environment_id) + } else { + (args.input, None) + } + } + ToolPayload::Custom { input } => { + if multi_environment_tools { + split_freeform_apply_patch_environment(input)? + } else { + (input, None) + } } - ToolPayload::Custom { input } => input, _ => { return Err(FunctionCallError::RespondToModel( "apply_patch handler received unsupported payload".to_string(), @@ -320,17 +377,22 @@ impl ToolHandler for ApplyPatchHandler { // Re-parse and verify the patch so we can compute changes and approval. // Avoid building temporary ExecParams/command vectors; derive directly from inputs. - let cwd = turn.cwd.clone(); + let tool_environment = turn + .environment_for_tool(requested_environment_id.as_deref()) + .ok_or_else(|| { + FunctionCallError::RespondToModel(match requested_environment_id.as_deref() { + Some(environment_id) => { + format!("environment `{environment_id}` is unavailable in this session") + } + None => "apply_patch is unavailable in this session".to_string(), + }) + })?; + let cwd = tool_environment.cwd.clone(); let command = vec!["apply_patch".to_string(), patch_input.clone()]; - let Some(environment) = turn.environment.as_ref() else { - return Err(FunctionCallError::RespondToModel( - "apply_patch is unavailable in this session".to_string(), - )); - }; - let fs = environment.get_filesystem(); - let sandbox = environment - .is_remote() - .then(|| turn.file_system_sandbox_context(/*additional_permissions*/ None)); + let fs = tool_environment.environment.get_filesystem(); + let sandbox = tool_environment.environment.is_remote().then(|| { + turn.file_system_sandbox_context_for_cwd(&cwd, /*additional_permissions*/ None) + }); match codex_apply_patch::maybe_parse_apply_patch_verified( &command, &cwd, @@ -341,9 +403,15 @@ impl ToolHandler for ApplyPatchHandler { { codex_apply_patch::MaybeApplyPatchVerified::Body(changes) => { let (file_paths, effective_additional_permissions, file_system_sandbox_policy) = - effective_patch_permissions(session.as_ref(), turn.as_ref(), &changes).await; - match apply_patch::apply_patch(turn.as_ref(), &file_system_sandbox_policy, changes) - .await + effective_patch_permissions(session.as_ref(), turn.as_ref(), &cwd, &changes) + .await; + match apply_patch::apply_patch( + turn.as_ref(), + &cwd, + &file_system_sandbox_policy, + changes, + ) + .await { InternalApplyPatchInvocation::Output(item) => { let content = item?; @@ -362,6 +430,9 @@ impl ToolHandler for ApplyPatchHandler { emitter.begin(event_ctx).await; let req = ApplyPatchRequest { + environment_id: tool_environment.environment_id.clone(), + environment: tool_environment.environment, + cwd, action: apply.action, file_paths, changes, @@ -432,11 +503,13 @@ pub(crate) async fn intercept_apply_patch( call_id: &str, tool_name: &str, ) -> Result, FunctionCallError> { - let sandbox = turn - .environment + let tool_environment = turn.environment_for_tool(/*environment_id*/ None); + let sandbox = tool_environment .as_ref() - .filter(|env| env.is_remote()) - .map(|_| turn.file_system_sandbox_context(/*additional_permissions*/ None)); + .filter(|environment| environment.environment.is_remote()) + .map(|_| { + turn.file_system_sandbox_context_for_cwd(cwd, /*additional_permissions*/ None) + }); match codex_apply_patch::maybe_parse_apply_patch_verified(command, cwd, fs, sandbox.as_ref()) .await { @@ -450,8 +523,8 @@ pub(crate) async fn intercept_apply_patch( ) .await; let (approval_keys, effective_additional_permissions, file_system_sandbox_policy) = - effective_patch_permissions(session.as_ref(), turn.as_ref(), &changes).await; - match apply_patch::apply_patch(turn.as_ref(), &file_system_sandbox_policy, changes) + effective_patch_permissions(session.as_ref(), turn.as_ref(), cwd, &changes).await; + match apply_patch::apply_patch(turn.as_ref(), cwd, &file_system_sandbox_policy, changes) .await { InternalApplyPatchInvocation::Output(item) => { @@ -470,6 +543,18 @@ pub(crate) async fn intercept_apply_patch( emitter.begin(event_ctx).await; let req = ApplyPatchRequest { + environment_id: tool_environment + .as_ref() + .and_then(|environment| environment.environment_id.clone()), + environment: tool_environment + .as_ref() + .map(|environment| Arc::clone(&environment.environment)) + .ok_or_else(|| { + FunctionCallError::RespondToModel( + "apply_patch is unavailable in this session".to_string(), + ) + })?, + cwd: cwd.clone(), action: apply.action, file_paths: approval_keys, changes, 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 39c8029e67..d159c0b761 100644 --- a/codex-rs/core/src/tools/handlers/apply_patch_tests.rs +++ b/codex-rs/core/src/tools/handlers/apply_patch_tests.rs @@ -108,6 +108,31 @@ fn diff_consumer_sends_next_update_after_buffer_interval() { ); } +#[test] +fn split_freeform_apply_patch_environment_extracts_environment_header() { + let input = + "*** Environment: local\n*** Begin Patch\n*** Add File: hello.txt\n+hello\n*** End Patch\n" + .to_string(); + + let (patch, environment_id) = + split_freeform_apply_patch_environment(input).expect("header should parse"); + + assert_eq!(environment_id.as_deref(), Some("local")); + assert_eq!( + patch, + "*** Begin Patch\n*** Add File: hello.txt\n+hello\n*** End Patch\n" + ); +} + +#[test] +fn merge_apply_patch_environment_ids_rejects_mismatch() { + let err = + merge_apply_patch_environment_ids(Some("local".to_string()), Some("remote".to_string())) + .expect_err("mismatched environment ids should fail"); + + assert!(err.to_string().contains("environment id mismatch")); +} + #[tokio::test] async fn approval_keys_include_move_destination() { let tmp = TempDir::new().expect("tmp"); diff --git a/codex-rs/core/src/tools/handlers/unified_exec.rs b/codex-rs/core/src/tools/handlers/unified_exec.rs index 33f9bc6bf3..728f183ab0 100644 --- a/codex-rs/core/src/tools/handlers/unified_exec.rs +++ b/codex-rs/core/src/tools/handlers/unified_exec.rs @@ -199,10 +199,7 @@ impl ToolHandler for UnifiedExecHandler { let response = match tool_name.name.as_str() { "exec_command" => { let selector: ToolEnvironmentArgs = parse_arguments(&arguments)?; - let requested_environment_id = if context - .turn - .features - .enabled(Feature::MultiEnvironmentTools) + let requested_environment_id = if context.turn.tools_config.multi_environment_tools { selector.environment_id.as_deref() } else { @@ -243,8 +240,7 @@ impl ToolHandler for UnifiedExecHandler { turn.tools_config.allow_login_shell, ) .map_err(FunctionCallError::RespondToModel)?; - let command_for_display = - codex_shell_command::parse_command::shlex_join(&command); + let command_for_display = codex_shell_command::parse_command::shlex_join(&command); let ExecCommandArgs { workdir, diff --git a/codex-rs/core/src/tools/handlers/view_image.rs b/codex-rs/core/src/tools/handlers/view_image.rs index 8f3f69701f..b0c8355e99 100644 --- a/codex-rs/core/src/tools/handlers/view_image.rs +++ b/codex-rs/core/src/tools/handlers/view_image.rs @@ -29,6 +29,8 @@ const VIEW_IMAGE_UNSUPPORTED_MESSAGE: &str = struct ViewImageArgs { path: String, detail: Option, + #[serde(default, alias = "environmentId")] + environment_id: Option, } #[derive(Clone, Copy, Eq, PartialEq)] @@ -87,17 +89,31 @@ impl ToolHandler for ViewImageHandler { } }; - let abs_path = turn.resolve_path(Some(args.path)); - let Some(environment) = turn.environment.as_ref() else { - return Err(FunctionCallError::RespondToModel( - "view_image is unavailable in this session".to_string(), - )); + let requested_environment_id = if turn.tools_config.multi_environment_tools { + args.environment_id.as_deref() + } else { + None }; - let sandbox = environment - .is_remote() - .then(|| turn.file_system_sandbox_context(/*additional_permissions*/ None)); + let tool_environment = turn + .environment_for_tool(requested_environment_id) + .ok_or_else(|| { + FunctionCallError::RespondToModel(match requested_environment_id { + Some(environment_id) => { + format!("environment `{environment_id}` is unavailable in this session") + } + None => "view_image is unavailable in this session".to_string(), + }) + })?; + let abs_path = turn.resolve_path_for_environment(&tool_environment, Some(args.path)); + let sandbox = tool_environment.environment.is_remote().then(|| { + turn.file_system_sandbox_context_for_cwd( + &tool_environment.cwd, + /*additional_permissions*/ None, + ) + }); - let metadata = environment + let metadata = tool_environment + .environment .get_filesystem() .get_metadata(&abs_path, sandbox.as_ref()) .await @@ -114,7 +130,8 @@ impl ToolHandler for ViewImageHandler { abs_path.display() ))); } - let file_bytes = environment + let file_bytes = tool_environment + .environment .get_filesystem() .read_file(&abs_path, sandbox.as_ref()) .await diff --git a/codex-rs/core/src/tools/runtimes/apply_patch.rs b/codex-rs/core/src/tools/runtimes/apply_patch.rs index aac8e6bc13..f4de1cda44 100644 --- a/codex-rs/core/src/tools/runtimes/apply_patch.rs +++ b/codex-rs/core/src/tools/runtimes/apply_patch.rs @@ -16,6 +16,7 @@ use crate::tools::sandboxing::ToolError; use crate::tools::sandboxing::ToolRuntime; use crate::tools::sandboxing::with_cached_approval; use codex_apply_patch::ApplyPatchAction; +use codex_exec_server::Environment; use codex_exec_server::FileSystemSandboxContext; use codex_protocol::error::CodexErr; use codex_protocol::error::SandboxErr; @@ -35,10 +36,14 @@ use codex_sandboxing::SandboxablePreference; use codex_utils_absolute_path::AbsolutePathBuf; use futures::future::BoxFuture; use std::path::PathBuf; +use std::sync::Arc; use std::time::Instant; #[derive(Debug)] pub struct ApplyPatchRequest { + pub environment_id: Option, + pub environment: Arc, + pub cwd: AbsolutePathBuf, pub action: ApplyPatchAction, pub file_paths: Vec, pub changes: std::collections::HashMap, @@ -47,6 +52,12 @@ pub struct ApplyPatchRequest { pub permissions_preapproved: bool, } +#[derive(Clone, Debug, Eq, Hash, PartialEq, serde::Serialize)] +pub struct ApplyPatchApprovalKey { + pub environment_id: Option, + pub path: AbsolutePathBuf, +} + #[derive(Default)] pub struct ApplyPatchRuntime; @@ -61,7 +72,7 @@ impl ApplyPatchRuntime { ) -> GuardianApprovalRequest { GuardianApprovalRequest::ApplyPatch { id: call_id.to_string(), - cwd: req.action.cwd.clone(), + cwd: req.cwd.clone(), files: req.file_paths.clone(), patch: req.action.patch.clone(), } @@ -121,10 +132,17 @@ impl Sandboxable for ApplyPatchRuntime { } impl Approvable for ApplyPatchRuntime { - type ApprovalKey = AbsolutePathBuf; + type ApprovalKey = ApplyPatchApprovalKey; fn approval_keys(&self, req: &ApplyPatchRequest) -> Vec { - req.file_paths.clone() + req.file_paths + .iter() + .cloned() + .map(|path| ApplyPatchApprovalKey { + environment_id: req.environment_id.clone(), + path, + }) + .collect() } fn start_approval_async<'a>( @@ -201,23 +219,28 @@ impl Approvable for ApplyPatchRuntime { } impl ToolRuntime for ApplyPatchRuntime { + fn sandbox_cwd<'a>( + &self, + req: &'a ApplyPatchRequest, + _turn_ctx: &'a crate::session::turn_context::TurnContext, + ) -> &'a AbsolutePathBuf { + &req.cwd + } + async fn run( &mut self, req: &ApplyPatchRequest, attempt: &SandboxAttempt<'_>, ctx: &ToolCtx, ) -> Result { - let environment = ctx.turn.environment.as_ref().ok_or_else(|| { - ToolError::Rejected("apply_patch is unavailable in this session".to_string()) - })?; let started_at = Instant::now(); - let fs = environment.get_filesystem(); + let fs = req.environment.get_filesystem(); let sandbox = Self::file_system_sandbox_context_for_attempt(req, attempt); let mut stdout = Vec::new(); let mut stderr = Vec::new(); let result = codex_apply_patch::apply_patch( &req.action.patch, - &req.action.cwd, + &req.cwd, &mut stdout, &mut stderr, fs.as_ref(), diff --git a/codex-rs/core/src/tools/runtimes/apply_patch_tests.rs b/codex-rs/core/src/tools/runtimes/apply_patch_tests.rs index 7f3641c03c..c8e77a0a28 100644 --- a/codex-rs/core/src/tools/runtimes/apply_patch_tests.rs +++ b/codex-rs/core/src/tools/runtimes/apply_patch_tests.rs @@ -15,6 +15,11 @@ use codex_sandboxing::SandboxType; use core_test_support::PathBufExt; use pretty_assertions::assert_eq; use std::collections::HashMap; +use std::sync::Arc; + +fn local_environment() -> Arc { + Arc::new(Environment::create(/*exec_server_url*/ None).expect("local environment")) +} #[test] fn wants_no_sandbox_approval_granular_respects_sandbox_flag() { @@ -49,6 +54,9 @@ fn guardian_review_request_includes_patch_context() { let expected_cwd = action.cwd.clone(); let expected_patch = action.patch.clone(); let request = ApplyPatchRequest { + environment_id: None, + environment: local_environment(), + cwd: expected_cwd.clone(), action, file_paths: vec![path.clone()], changes: HashMap::from([( @@ -91,6 +99,9 @@ fn file_system_sandbox_context_uses_active_attempt() { )), }; let req = ApplyPatchRequest { + environment_id: None, + environment: local_environment(), + cwd: path.clone(), action: ApplyPatchAction::new_add_for_test(&path, "hello".to_string()), file_paths: vec![path.clone()], changes: HashMap::new(), @@ -147,6 +158,9 @@ fn file_system_sandbox_context_omits_legacy_equivalent_policy() { .join("apply-patch-runtime-legacy-equivalent.txt") .abs(); let req = ApplyPatchRequest { + environment_id: None, + environment: local_environment(), + cwd: path.clone(), action: ApplyPatchAction::new_add_for_test(&path, "hello".to_string()), file_paths: vec![path.clone()], changes: HashMap::new(), @@ -188,6 +202,9 @@ fn no_sandbox_attempt_has_no_file_system_context() { .join("apply-patch-runtime-none.txt") .abs(); let req = ApplyPatchRequest { + environment_id: None, + environment: local_environment(), + cwd: path.clone(), action: ApplyPatchAction::new_add_for_test(&path, "hello".to_string()), file_paths: vec![path.clone()], changes: HashMap::new(), diff --git a/codex-rs/core/tests/suite/resume_warning.rs b/codex-rs/core/tests/suite/resume_warning.rs index 8ea2be3147..2368b7c5e3 100644 --- a/codex-rs/core/tests/suite/resume_warning.rs +++ b/codex-rs/core/tests/suite/resume_warning.rs @@ -29,6 +29,7 @@ fn resume_history( turn_id: Some(turn_id.clone()), trace_id: None, cwd: config.cwd.to_path_buf(), + environments: None, current_date: None, timezone: None, approval_policy: config.permissions.approval_policy.value(), diff --git a/codex-rs/exec-server/src/server/transport_tests.rs b/codex-rs/exec-server/src/server/transport_tests.rs index bec91c936e..dc30421dac 100644 --- a/codex-rs/exec-server/src/server/transport_tests.rs +++ b/codex-rs/exec-server/src/server/transport_tests.rs @@ -45,6 +45,6 @@ fn parse_listen_url_rejects_unsupported_url() { parse_listen_url("http://127.0.0.1:1234").expect_err("unsupported scheme should fail"); assert_eq!( err.to_string(), - "unsupported --listen URL `http://127.0.0.1:1234`; expected `ws://IP:PORT`" + "unsupported --listen URL `http://127.0.0.1:1234`; expected `ws://IP:PORT` or `stdio://`" ); } diff --git a/codex-rs/exec-server/tests/command_transport.rs b/codex-rs/exec-server/tests/command_transport.rs new file mode 100644 index 0000000000..35eb73ca84 --- /dev/null +++ b/codex-rs/exec-server/tests/command_transport.rs @@ -0,0 +1,50 @@ +mod common; + +use codex_exec_server::ConfiguredEnvironmentManagerArgs; +use codex_exec_server::ConfiguredEnvironmentSpec; +use codex_exec_server::EnvironmentManager; +use codex_exec_server::RemoteExecServerTransport; +use codex_utils_absolute_path::AbsolutePathBuf; +use pretty_assertions::assert_eq; + +#[tokio::test] +async fn configured_command_environment_connects_lazily_over_stdio() -> anyhow::Result<()> { + let helper_paths = common::exec_server::test_codex_helper_paths()?; + let temp_dir = tempfile::tempdir()?; + let target_path = temp_dir.path().join("target.txt"); + let marker_path = temp_dir.path().join("spawned.txt"); + tokio::fs::write(&target_path, "ok").await?; + + let manager = EnvironmentManager::try_new(ConfiguredEnvironmentManagerArgs { + default_environment: Some("remote".to_string()), + environments: vec![ConfiguredEnvironmentSpec { + id: "remote".to_string(), + transport: RemoteExecServerTransport::Command { + command: format!( + "echo spawned > {marker_path:?}; exec {codex_exe:?} exec-server --listen stdio://", + marker_path = marker_path, + codex_exe = helper_paths.codex_exe, + ), + }, + }], + local_runtime_paths: None, + })?; + let environment = manager.default_environment().expect("default environment"); + + assert!( + tokio::fs::metadata(&marker_path).await.is_err(), + "command transport should not connect before the first remote operation" + ); + + let metadata = environment + .get_filesystem() + .get_metadata( + &AbsolutePathBuf::from_absolute_path(&target_path)?, + /*sandbox*/ None, + ) + .await?; + + assert_eq!(metadata.is_file, true); + assert_eq!(tokio::fs::read_to_string(&marker_path).await?, "spawned\n"); + Ok(()) +} diff --git a/codex-rs/protocol/src/protocol.rs b/codex-rs/protocol/src/protocol.rs index 6fe8e37e38..37504a8781 100644 --- a/codex-rs/protocol/src/protocol.rs +++ b/codex-rs/protocol/src/protocol.rs @@ -103,7 +103,7 @@ pub const REALTIME_CONVERSATION_OPEN_TAG: &str = ""; pub const REALTIME_CONVERSATION_CLOSE_TAG: &str = ""; pub const USER_MESSAGE_BEGIN: &str = "## My request for Codex:"; -#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, JsonSchema)] +#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, JsonSchema, TS)] pub struct TurnEnvironmentSelection { pub environment_id: String, pub cwd: AbsolutePathBuf, @@ -2932,6 +2932,8 @@ pub struct TurnContextItem { pub trace_id: Option, pub cwd: PathBuf, #[serde(default, skip_serializing_if = "Option::is_none")] + pub environments: Option>, + #[serde(default, skip_serializing_if = "Option::is_none")] pub current_date: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub timezone: Option, @@ -5082,17 +5084,64 @@ mod tests { }))?; assert_eq!(item.trace_id, None); + assert_eq!(item.environments, None); assert_eq!(item.network, None); assert_eq!(item.file_system_sandbox_policy, None); Ok(()) } + #[test] + fn turn_context_item_serializes_environments_when_present() -> Result<()> { + let item = TurnContextItem { + turn_id: None, + trace_id: None, + cwd: test_path_buf("/tmp/primary"), + environments: Some(vec![ + TurnEnvironmentSelection { + environment_id: "remote".to_string(), + cwd: test_path_buf("/tmp/primary").abs(), + }, + TurnEnvironmentSelection { + environment_id: "local".to_string(), + cwd: test_path_buf("/tmp/local").abs(), + }, + ]), + current_date: None, + timezone: None, + approval_policy: AskForApproval::Never, + sandbox_policy: SandboxPolicy::DangerFullAccess, + network: None, + file_system_sandbox_policy: None, + model: "gpt-5".to_string(), + personality: None, + collaboration_mode: None, + realtime_active: None, + effort: None, + summary: ReasoningSummaryConfig::Auto, + user_instructions: None, + developer_instructions: None, + final_output_json_schema: None, + truncation_policy: None, + }; + + let value = serde_json::to_value(item)?; + assert_eq!( + value["environments"], + json!([ + { "environment_id": "remote", "cwd": "/tmp/primary" }, + { "environment_id": "local", "cwd": "/tmp/local" }, + ]) + ); + Ok(()) + } + #[test] fn turn_context_item_serializes_network_when_present() -> Result<()> { let item = TurnContextItem { turn_id: None, trace_id: None, cwd: test_path_buf("/tmp"), + environments: None, current_date: None, timezone: None, approval_policy: AskForApproval::Never, diff --git a/codex-rs/rollout/src/recorder_tests.rs b/codex-rs/rollout/src/recorder_tests.rs index fa789d3a7e..01cda8f20a 100644 --- a/codex-rs/rollout/src/recorder_tests.rs +++ b/codex-rs/rollout/src/recorder_tests.rs @@ -543,6 +543,7 @@ async fn resume_candidate_matches_cwd_reads_latest_turn_context() -> std::io::Re turn_id: Some("turn-1".to_string()), trace_id: None, cwd: latest_cwd.clone(), + environments: None, current_date: None, timezone: None, approval_policy: AskForApproval::Never, diff --git a/codex-rs/state/src/extract.rs b/codex-rs/state/src/extract.rs index 64d7f5bbb7..586101d897 100644 --- a/codex-rs/state/src/extract.rs +++ b/codex-rs/state/src/extract.rs @@ -299,6 +299,7 @@ mod tests { turn_id: Some("turn-1".to_string()), trace_id: None, cwd: PathBuf::from("/parent/workspace"), + environments: None, current_date: None, timezone: None, approval_policy: AskForApproval::Never, @@ -338,6 +339,7 @@ mod tests { turn_id: Some("turn-1".to_string()), trace_id: None, cwd: PathBuf::from("/fallback/workspace"), + environments: None, current_date: None, timezone: None, approval_policy: AskForApproval::OnRequest, @@ -371,6 +373,7 @@ mod tests { turn_id: Some("turn-1".to_string()), trace_id: None, cwd: PathBuf::from("/fallback/workspace"), + environments: None, current_date: None, timezone: None, approval_policy: AskForApproval::OnRequest, diff --git a/codex-rs/tools/BUILD.bazel b/codex-rs/tools/BUILD.bazel index 7b1541e4e8..d37bbcfde0 100644 --- a/codex-rs/tools/BUILD.bazel +++ b/codex-rs/tools/BUILD.bazel @@ -5,5 +5,6 @@ codex_rust_crate( crate_name = "codex_tools", compile_data = [ "src/tool_apply_patch.lark", + "src/tool_apply_patch_environment.lark", ], ) diff --git a/codex-rs/tools/src/apply_patch_tool.rs b/codex-rs/tools/src/apply_patch_tool.rs index 469bb52367..3e5bc12a3c 100644 --- a/codex-rs/tools/src/apply_patch_tool.rs +++ b/codex-rs/tools/src/apply_patch_tool.rs @@ -8,6 +8,13 @@ use serde::Serialize; use std::collections::BTreeMap; const APPLY_PATCH_LARK_GRAMMAR: &str = include_str!("tool_apply_patch.lark"); +const APPLY_PATCH_ENVIRONMENT_LARK_GRAMMAR: &str = + include_str!("tool_apply_patch_environment.lark"); + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct ApplyPatchToolOptions { + pub multi_environment_tools: bool, +} const APPLY_PATCH_JSON_TOOL_DESCRIPTION: &str = r#"Use the `apply_patch` tool to edit files. Your patch language is a stripped‑down, file‑oriented diff format designed to be easy to parse and safe to apply. You can think of it as a high‑level envelope: @@ -82,30 +89,50 @@ It is important to remember: #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] pub struct ApplyPatchToolArgs { pub input: String, + #[serde(default, alias = "environmentId")] + pub environment_id: Option, } /// Returns a custom tool that can be used to edit files. Well-suited for GPT-5 models /// https://platform.openai.com/docs/guides/function-calling#custom-tools -pub fn create_apply_patch_freeform_tool() -> ToolSpec { +pub fn create_apply_patch_freeform_tool(options: ApplyPatchToolOptions) -> ToolSpec { + let description = if options.multi_environment_tools { + "Use the `apply_patch` tool to edit files. This is a FREEFORM tool, so do not wrap the patch in JSON. To target a non-primary selected environment, add a first line `*** Environment: ` before `*** Begin Patch`." + } else { + "Use the `apply_patch` tool to edit files. This is a FREEFORM tool, so do not wrap the patch in JSON." + }; + let definition = if options.multi_environment_tools { + APPLY_PATCH_ENVIRONMENT_LARK_GRAMMAR + } else { + APPLY_PATCH_LARK_GRAMMAR + }; ToolSpec::Freeform(FreeformTool { name: "apply_patch".to_string(), - description: "Use the `apply_patch` tool to edit files. This is a FREEFORM tool, so do not wrap the patch in JSON.".to_string(), + description: description.to_string(), format: FreeformToolFormat { r#type: "grammar".to_string(), syntax: "lark".to_string(), - definition: APPLY_PATCH_LARK_GRAMMAR.to_string(), + definition: definition.to_string(), }, }) } /// Returns a json tool that can be used to edit files. Should only be used with gpt-oss models -pub fn create_apply_patch_json_tool() -> ToolSpec { - let properties = BTreeMap::from([( +pub fn create_apply_patch_json_tool(options: ApplyPatchToolOptions) -> ToolSpec { + let mut properties = BTreeMap::from([( "input".to_string(), JsonSchema::string(Some( "The entire contents of the apply_patch command".to_string(), )), )]); + if options.multi_environment_tools { + properties.insert( + "environment_id".to_string(), + JsonSchema::string(Some( + "Optional selected environment id to apply the patch in; omit to use the primary environment.".to_string(), + )), + ); + } ToolSpec::Function(ResponsesApiTool { name: "apply_patch".to_string(), diff --git a/codex-rs/tools/src/apply_patch_tool_tests.rs b/codex-rs/tools/src/apply_patch_tool_tests.rs index c128594587..9c623b9c23 100644 --- a/codex-rs/tools/src/apply_patch_tool_tests.rs +++ b/codex-rs/tools/src/apply_patch_tool_tests.rs @@ -6,7 +6,9 @@ use std::collections::BTreeMap; #[test] fn create_apply_patch_freeform_tool_matches_expected_spec() { assert_eq!( - create_apply_patch_freeform_tool(), + create_apply_patch_freeform_tool(ApplyPatchToolOptions { + multi_environment_tools: false, + }), ToolSpec::Freeform(FreeformTool { name: "apply_patch".to_string(), description: @@ -24,7 +26,9 @@ fn create_apply_patch_freeform_tool_matches_expected_spec() { #[test] fn create_apply_patch_json_tool_matches_expected_spec() { assert_eq!( - create_apply_patch_json_tool(), + create_apply_patch_json_tool(ApplyPatchToolOptions { + multi_environment_tools: false, + }), ToolSpec::Function(ResponsesApiTool { name: "apply_patch".to_string(), description: APPLY_PATCH_JSON_TOOL_DESCRIPTION.to_string(), @@ -44,3 +48,31 @@ fn create_apply_patch_json_tool_matches_expected_spec() { }) ); } + +#[test] +fn create_apply_patch_freeform_tool_with_environment_matches_expected_spec() { + let ToolSpec::Freeform(tool) = create_apply_patch_freeform_tool(ApplyPatchToolOptions { + multi_environment_tools: true, + }) else { + panic!("apply_patch should be a freeform tool"); + }; + assert!( + tool.description + .contains("*** Environment: ") + ); + assert_eq!(tool.format.definition, APPLY_PATCH_ENVIRONMENT_LARK_GRAMMAR); +} + +#[test] +fn create_apply_patch_json_tool_with_environment_includes_environment_id() { + let ToolSpec::Function(tool) = create_apply_patch_json_tool(ApplyPatchToolOptions { + multi_environment_tools: true, + }) else { + panic!("apply_patch should be a function tool"); + }; + let properties = tool + .parameters + .properties + .expect("apply_patch parameters should include properties"); + assert!(properties.contains_key("environment_id")); +} diff --git a/codex-rs/tools/src/lib.rs b/codex-rs/tools/src/lib.rs index 2dc7c165d5..83789a8092 100644 --- a/codex-rs/tools/src/lib.rs +++ b/codex-rs/tools/src/lib.rs @@ -41,6 +41,7 @@ pub use agent_tool::create_spawn_agent_tool_v2; pub use agent_tool::create_wait_agent_tool_v1; pub use agent_tool::create_wait_agent_tool_v2; pub use apply_patch_tool::ApplyPatchToolArgs; +pub use apply_patch_tool::ApplyPatchToolOptions; pub use apply_patch_tool::create_apply_patch_freeform_tool; pub use apply_patch_tool::create_apply_patch_json_tool; pub use code_mode::augment_tool_spec_for_code_mode; diff --git a/codex-rs/tools/src/tool_apply_patch_environment.lark b/codex-rs/tools/src/tool_apply_patch_environment.lark new file mode 100644 index 0000000000..709e6b6ce9 --- /dev/null +++ b/codex-rs/tools/src/tool_apply_patch_environment.lark @@ -0,0 +1,21 @@ +start: environment? begin_patch hunk+ end_patch +environment: "*** Environment: " environment_id LF +environment_id: /[A-Za-z0-9_.:-]+/ +begin_patch: "*** Begin Patch" LF +end_patch: "*** End Patch" LF? + +hunk: add_hunk | delete_hunk | update_hunk +add_hunk: "*** Add File: " filename LF add_line+ +delete_hunk: "*** Delete File: " filename LF +update_hunk: "*** Update File: " filename LF change_move? change? + +filename: /(.+)/ +add_line: "+" /(.*)/ LF -> line + +change_move: "*** Move to: " filename LF +change: (change_context | change_line)+ eof_line? +change_context: ("@@" | "@@ " /(.+)/) LF +change_line: ("+" | "-" | " ") /(.*)/ LF +eof_line: "*** End of File" LF + +%import common.LF diff --git a/codex-rs/tools/src/tool_config.rs b/codex-rs/tools/src/tool_config.rs index 2ff0e064fe..28ca10e34f 100644 --- a/codex-rs/tools/src/tool_config.rs +++ b/codex-rs/tools/src/tool_config.rs @@ -272,6 +272,11 @@ impl ToolsConfig { self } + pub fn with_multi_environment_tools(mut self, multi_environment_tools: bool) -> Self { + self.multi_environment_tools = multi_environment_tools; + self + } + pub fn with_unified_exec_shell_mode( mut self, unified_exec_shell_mode: UnifiedExecShellMode, diff --git a/codex-rs/tools/src/tool_registry_plan.rs b/codex-rs/tools/src/tool_registry_plan.rs index 3a4b25679e..5ffa91896a 100644 --- a/codex-rs/tools/src/tool_registry_plan.rs +++ b/codex-rs/tools/src/tool_registry_plan.rs @@ -1,3 +1,4 @@ +use crate::ApplyPatchToolOptions; use crate::CommandToolOptions; use crate::REQUEST_USER_INPUT_TOOL_NAME; use crate::ResponsesApiNamespace; @@ -317,14 +318,18 @@ pub fn build_tool_registry_plan( match apply_patch_tool_type { ApplyPatchToolType::Freeform => { plan.push_spec( - create_apply_patch_freeform_tool(), + create_apply_patch_freeform_tool(ApplyPatchToolOptions { + multi_environment_tools: config.multi_environment_tools, + }), /*supports_parallel_tool_calls*/ false, config.code_mode_enabled, ); } ApplyPatchToolType::Function => { plan.push_spec( - create_apply_patch_json_tool(), + create_apply_patch_json_tool(ApplyPatchToolOptions { + multi_environment_tools: config.multi_environment_tools, + }), /*supports_parallel_tool_calls*/ false, config.code_mode_enabled, ); @@ -384,6 +389,7 @@ pub fn build_tool_registry_plan( plan.push_spec( create_view_image_tool(ViewImageToolOptions { can_request_original_image_detail: config.can_request_original_image_detail, + multi_environment_tools: config.multi_environment_tools, }), /*supports_parallel_tool_calls*/ true, config.code_mode_enabled, diff --git a/codex-rs/tools/src/tool_registry_plan_tests.rs b/codex-rs/tools/src/tool_registry_plan_tests.rs index 2b81dee155..c2b1f48045 100644 --- a/codex-rs/tools/src/tool_registry_plan_tests.rs +++ b/codex-rs/tools/src/tool_registry_plan_tests.rs @@ -90,7 +90,9 @@ fn test_full_toolset_specs_for_gpt5_codex_unified_exec_web_search() { create_write_stdin_tool(), create_update_plan_tool(), request_user_input_tool_spec(/*default_mode_request_user_input*/ false), - create_apply_patch_freeform_tool(), + create_apply_patch_freeform_tool(ApplyPatchToolOptions { + multi_environment_tools: false, + }), ToolSpec::WebSearch { external_web_access: Some(true), filters: None, @@ -101,6 +103,7 @@ fn test_full_toolset_specs_for_gpt5_codex_unified_exec_web_search() { create_image_generation_tool("png"), create_view_image_tool(ViewImageToolOptions { can_request_original_image_detail: config.can_request_original_image_detail, + multi_environment_tools: false, }), ] { expected.insert(spec.name().to_string(), spec); @@ -433,6 +436,37 @@ fn view_image_tool_includes_detail_with_original_detail_support() { assert!(description.contains("omit this field for default resized behavior")); } +#[test] +fn view_image_tool_includes_environment_id_when_multi_environment_tools_enabled() { + let mut model_info = model_info(); + model_info.supports_image_detail_original = false; + let mut features = Features::with_defaults(); + features.enable(Feature::MultiEnvironmentTools); + let available_models = Vec::new(); + let tools_config = ToolsConfig::new(&ToolsConfigParams { + model_info: &model_info, + available_models: &available_models, + features: &features, + image_generation_tool_auth_allowed: true, + web_search_mode: Some(WebSearchMode::Cached), + session_source: SessionSource::Cli, + sandbox_policy: &SandboxPolicy::DangerFullAccess, + windows_sandbox_level: WindowsSandboxLevel::Disabled, + }); + let (tools, _) = build_specs( + &tools_config, + /*mcp_tools*/ None, + /*deferred_mcp_tools*/ None, + &[], + ); + let view_image = find_tool(&tools, VIEW_IMAGE_TOOL_NAME); + let ToolSpec::Function(ResponsesApiTool { parameters, .. }) = &view_image.spec else { + panic!("view_image should be a function tool"); + }; + let (properties, _) = expect_object_schema(parameters); + assert!(properties.contains_key("environment_id")); +} + #[test] fn disabled_environment_omits_environment_backed_tools() { let model_info = model_info(); diff --git a/codex-rs/tools/src/view_image.rs b/codex-rs/tools/src/view_image.rs index 1d77ceadf3..017e09f732 100644 --- a/codex-rs/tools/src/view_image.rs +++ b/codex-rs/tools/src/view_image.rs @@ -9,6 +9,7 @@ use std::collections::BTreeMap; #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub struct ViewImageToolOptions { pub can_request_original_image_detail: bool, + pub multi_environment_tools: bool, } pub fn create_view_image_tool(options: ViewImageToolOptions) -> ToolSpec { @@ -24,6 +25,14 @@ pub fn create_view_image_tool(options: ViewImageToolOptions) -> ToolSpec { )), ); } + if options.multi_environment_tools { + properties.insert( + "environment_id".to_string(), + JsonSchema::string(Some( + "Optional selected environment id to load the image from; omit to use the primary environment.".to_string(), + )), + ); + } ToolSpec::Function(ResponsesApiTool { name: VIEW_IMAGE_TOOL_NAME.to_string(), diff --git a/codex-rs/tui/src/lib.rs b/codex-rs/tui/src/lib.rs index 96cba25785..37527bebc8 100644 --- a/codex-rs/tui/src/lib.rs +++ b/codex-rs/tui/src/lib.rs @@ -2190,6 +2190,7 @@ mod tests { turn_id: None, trace_id: None, cwd, + environments: None, current_date: None, timezone: None, approval_policy: config.permissions.approval_policy.value(),