From 053ccb34f9ead8a2b4689375821098243c169742 Mon Sep 17 00:00:00 2001 From: starr-openai Date: Tue, 5 May 2026 16:07:27 -0700 Subject: [PATCH] codex: address apply_patch tool spec review (#21048) Co-authored-by: Codex --- codex-rs/core/src/environment_selection.rs | 44 +++++++++ codex-rs/tools/BUILD.bazel | 1 + codex-rs/tools/src/apply_patch_tool.rs | 29 +++++- codex-rs/tools/src/apply_patch_tool_tests.rs | 96 ++++++++++++------- codex-rs/tools/src/lib.rs | 1 + codex-rs/tools/src/local_tool.rs | 7 +- codex-rs/tools/src/tool_apply_patch.lark | 4 +- .../tool_apply_patch_with_environment.lark | 21 ++++ codex-rs/tools/src/tool_registry_plan.rs | 9 +- .../tools/src/tool_registry_plan_tests.rs | 60 +++++++++++- codex-rs/tools/src/tool_spec.rs | 4 + codex-rs/tools/src/view_image.rs | 1 - 12 files changed, 226 insertions(+), 51 deletions(-) create mode 100644 codex-rs/tools/src/tool_apply_patch_with_environment.lark diff --git a/codex-rs/core/src/environment_selection.rs b/codex-rs/core/src/environment_selection.rs index e9d617cbf4..013db32cf2 100644 --- a/codex-rs/core/src/environment_selection.rs +++ b/codex-rs/core/src/environment_selection.rs @@ -41,6 +41,19 @@ impl ResolvedTurnEnvironments { self.turn_environments.first() } + pub(crate) fn get_by_id(&self, environment_id: &str) -> Option<&TurnEnvironment> { + self.turn_environments + .iter() + .find(|environment| environment.environment_id == environment_id) + } + + pub(crate) fn get_or_primary(&self, environment_id: Option<&str>) -> Option<&TurnEnvironment> { + environment_id.map_or_else( + || self.primary(), + |environment_id| self.get_by_id(environment_id), + ) + } + pub(crate) fn primary_environment(&self) -> Option> { self.primary() .map(|environment| Arc::clone(&environment.environment)) @@ -177,4 +190,35 @@ mod tests { "local" ); } + + #[tokio::test] + async fn resolved_environment_selections_gets_by_id_or_primary() { + let cwd = AbsolutePathBuf::current_dir().expect("cwd"); + let manager = EnvironmentManager::default_for_tests(); + + let resolved = resolve_environment_selections( + &manager, + &[TurnEnvironmentSelection { + environment_id: "local".to_string(), + cwd, + }], + ) + .expect("environment selections should resolve"); + + assert_eq!( + resolved + .get_or_primary(/*environment_id*/ None) + .expect("primary environment") + .environment_id, + "local" + ); + assert_eq!( + resolved + .get_or_primary(Some("local")) + .expect("selected environment") + .environment_id, + "local" + ); + assert!(resolved.get_or_primary(Some("unknown")).is_none()); + } } diff --git a/codex-rs/tools/BUILD.bazel b/codex-rs/tools/BUILD.bazel index 7b1541e4e8..5b97df305b 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_with_environment.lark", ], ) diff --git a/codex-rs/tools/src/apply_patch_tool.rs b/codex-rs/tools/src/apply_patch_tool.rs index 8ba8eaad25..71aae54627 100644 --- a/codex-rs/tools/src/apply_patch_tool.rs +++ b/codex-rs/tools/src/apply_patch_tool.rs @@ -3,11 +3,14 @@ use crate::FreeformToolFormat; use crate::JsonSchema; use crate::ResponsesApiTool; use crate::ToolSpec; +use crate::tool_spec::environment_id_schema; use serde::Deserialize; use serde::Serialize; use std::collections::BTreeMap; const APPLY_PATCH_LARK_GRAMMAR: &str = include_str!("tool_apply_patch.lark"); +const APPLY_PATCH_LARK_GRAMMAR_WITH_ENVIRONMENT: &str = + include_str!("tool_apply_patch_with_environment.lark"); 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: @@ -86,28 +89,46 @@ pub struct ApplyPatchToolArgs { pub environment_id: Option, } +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct ApplyPatchToolOptions { + pub include_environment_id: bool, +} + /// 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 definition = if options.include_environment_id { + APPLY_PATCH_LARK_GRAMMAR_WITH_ENVIRONMENT + } 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(), 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.include_environment_id { + properties.insert( + "environment_id".to_string(), + environment_id_schema( + "Optional selected environment id to target. Omit this to use the primary environment.", + ), + ); + } 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..bb5a9d9cbd 100644 --- a/codex-rs/tools/src/apply_patch_tool_tests.rs +++ b/codex-rs/tools/src/apply_patch_tool_tests.rs @@ -1,46 +1,68 @@ use super::*; use crate::JsonSchema; use pretty_assertions::assert_eq; -use std::collections::BTreeMap; #[test] -fn create_apply_patch_freeform_tool_matches_expected_spec() { - assert_eq!( - create_apply_patch_freeform_tool(), - 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(), - format: FreeformToolFormat { - r#type: "grammar".to_string(), - syntax: "lark".to_string(), - definition: APPLY_PATCH_LARK_GRAMMAR.to_string(), - }, - }) - ); +fn freeform_tool_uses_single_environment_grammar_by_default() { + let ToolSpec::Freeform(tool) = create_apply_patch_freeform_tool(ApplyPatchToolOptions { + include_environment_id: false, + }) else { + panic!("expected freeform tool"); + }; + + assert_eq!(tool.name, "apply_patch"); + assert_eq!(tool.format.syntax, "lark"); + assert!(!tool.format.definition.contains("*** Environment ID: ")); } #[test] -fn create_apply_patch_json_tool_matches_expected_spec() { - assert_eq!( - create_apply_patch_json_tool(), - ToolSpec::Function(ResponsesApiTool { - name: "apply_patch".to_string(), - description: APPLY_PATCH_JSON_TOOL_DESCRIPTION.to_string(), - strict: false, - defer_loading: None, - parameters: JsonSchema::object( - BTreeMap::from([( - "input".to_string(), - JsonSchema::string(Some( - "The entire contents of the apply_patch command".to_string(), - ),), - )]), - Some(vec!["input".to_string()]), - Some(false.into()) - ), - output_schema: None, - }) - ); +fn freeform_tool_advertises_environment_metadata_in_multi_environment_mode() { + let ToolSpec::Freeform(tool) = create_apply_patch_freeform_tool(ApplyPatchToolOptions { + include_environment_id: true, + }) else { + panic!("expected freeform tool"); + }; + + assert!(tool.format.definition.contains("*** Environment ID: ")); +} + +#[test] +fn json_tool_omits_environment_id_by_default() { + let ToolSpec::Function(tool) = create_apply_patch_json_tool(ApplyPatchToolOptions { + include_environment_id: false, + }) else { + panic!("expected function tool"); + }; + + let properties = tool + .parameters + .properties + .as_ref() + .expect("expected properties"); + assert!(properties.contains_key("input")); + assert!(!properties.contains_key("environment_id")); + assert_eq!(tool.parameters.required, Some(vec!["input".to_string()])); +} + +#[test] +fn json_tool_advertises_environment_id_in_multi_environment_mode() { + let ToolSpec::Function(tool) = create_apply_patch_json_tool(ApplyPatchToolOptions { + include_environment_id: true, + }) else { + panic!("expected function tool"); + }; + + let properties = tool + .parameters + .properties + .as_ref() + .expect("expected properties"); + assert!(matches!( + properties.get("environment_id"), + Some(JsonSchema { + description: Some(_), + .. + }) + )); + assert_eq!(tool.parameters.required, Some(vec!["input".to_string()])); } diff --git a/codex-rs/tools/src/lib.rs b/codex-rs/tools/src/lib.rs index ebe54c382b..3600c0977e 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/local_tool.rs b/codex-rs/tools/src/local_tool.rs index aeabdbfa30..dab64d3fd2 100644 --- a/codex-rs/tools/src/local_tool.rs +++ b/codex-rs/tools/src/local_tool.rs @@ -1,6 +1,7 @@ use crate::JsonSchema; use crate::ResponsesApiTool; use crate::ToolSpec; +use crate::tool_spec::environment_id_schema; use serde_json::Value; use serde_json::json; use std::collections::BTreeMap; @@ -73,9 +74,9 @@ pub(crate) fn create_exec_command_tool_with_environment_id( if include_environment_id { properties.insert( "environment_id".to_string(), - JsonSchema::string(Some( - "Optional environment id from the block. If omitted, uses the primary environment.".to_string(), - )), + environment_id_schema( + "Optional environment id from the block. If omitted, uses the primary environment.", + ), ); } properties.extend(create_approval_parameters( diff --git a/codex-rs/tools/src/tool_apply_patch.lark b/codex-rs/tools/src/tool_apply_patch.lark index b2d82c7edd..5aa41b0af7 100644 --- a/codex-rs/tools/src/tool_apply_patch.lark +++ b/codex-rs/tools/src/tool_apply_patch.lark @@ -1,6 +1,5 @@ -start: begin_patch environment? hunk+ end_patch +start: begin_patch hunk+ end_patch begin_patch: "*** Begin Patch" LF -environment: "*** Environment ID: " environment_id LF end_patch: "*** End Patch" LF? hunk: add_hunk | delete_hunk | update_hunk @@ -9,7 +8,6 @@ delete_hunk: "*** Delete File: " filename LF update_hunk: "*** Update File: " filename LF change_move? change? filename: /(.+)/ -environment_id: /(.+)/ add_line: "+" /(.*)/ LF -> line change_move: "*** Move to: " filename LF diff --git a/codex-rs/tools/src/tool_apply_patch_with_environment.lark b/codex-rs/tools/src/tool_apply_patch_with_environment.lark new file mode 100644 index 0000000000..b2d82c7edd --- /dev/null +++ b/codex-rs/tools/src/tool_apply_patch_with_environment.lark @@ -0,0 +1,21 @@ +start: begin_patch environment? hunk+ end_patch +begin_patch: "*** Begin Patch" LF +environment: "*** Environment ID: " environment_id 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: /(.+)/ +environment_id: /(.+)/ +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_registry_plan.rs b/codex-rs/tools/src/tool_registry_plan.rs index c777bd4129..5e09de0e73 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_PLUGIN_INSTALL_TOOL_NAME; use crate::REQUEST_USER_INPUT_TOOL_NAME; @@ -75,6 +76,7 @@ pub fn build_tool_registry_plan( ) -> ToolRegistryPlan { let mut plan = ToolRegistryPlan::new(); let exec_permission_approvals_enabled = config.exec_permission_approvals_enabled; + let include_environment_id = matches!(config.environment_mode, ToolEnvironmentMode::Multiple); if config.code_mode_enabled { let namespace_descriptions = params @@ -337,17 +339,20 @@ pub fn build_tool_registry_plan( if config.environment_mode.has_environment() && let Some(apply_patch_tool_type) = &config.apply_patch_tool_type { + let apply_patch_options = ApplyPatchToolOptions { + include_environment_id, + }; match apply_patch_tool_type { ApplyPatchToolType::Freeform => { plan.push_spec( - create_apply_patch_freeform_tool(), + create_apply_patch_freeform_tool(apply_patch_options), /*supports_parallel_tool_calls*/ false, config.code_mode_enabled, ); } ApplyPatchToolType::Function => { plan.push_spec( - create_apply_patch_json_tool(), + create_apply_patch_json_tool(apply_patch_options), /*supports_parallel_tool_calls*/ false, 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 fe5507ddc7..2fef4af442 100644 --- a/codex-rs/tools/src/tool_registry_plan_tests.rs +++ b/codex-rs/tools/src/tool_registry_plan_tests.rs @@ -1,5 +1,6 @@ use super::*; use crate::AdditionalProperties; +use crate::ApplyPatchToolOptions; use crate::ConfiguredToolSpec; use crate::DiscoverablePluginInfo; use crate::DiscoverableTool; @@ -93,7 +94,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(&request_user_input_available_modes(&features)), - create_apply_patch_freeform_tool(), + create_apply_patch_freeform_tool(ApplyPatchToolOptions { + include_environment_id: false, + }), ToolSpec::WebSearch { external_web_access: Some(true), filters: None, @@ -203,6 +206,46 @@ fn exec_command_spec_includes_environment_id_only_for_multiple_selected_environm ); } +fn apply_patch_freeform_spec_includes_environment_metadata_only_for_multiple_selected_environments() +{ + let model_info = model_info(); + let available_models = Vec::new(); + let tools_config = ToolsConfig::new(&ToolsConfigParams { + model_info: &model_info, + available_models: &available_models, + features: &Features::with_defaults(), + image_generation_tool_auth_allowed: true, + web_search_mode: Some(WebSearchMode::Cached), + session_source: SessionSource::Cli, + permission_profile: &PermissionProfile::Disabled, + windows_sandbox_level: WindowsSandboxLevel::Disabled, + }); + + let (single_environment_tools, _) = build_specs( + &tools_config, + /*mcp_tools*/ None, + /*deferred_mcp_tools*/ None, + &[], + ); + assert_apply_patch_freeform_environment_metadata( + &single_environment_tools, + /*expected_present*/ false, + ); + + let multi_environment_config = + tools_config.with_environment_mode(ToolEnvironmentMode::Multiple); + let (multi_environment_tools, _) = build_specs( + &multi_environment_config, + /*mcp_tools*/ None, + /*deferred_mcp_tools*/ None, + &[], + ); + assert_apply_patch_freeform_environment_metadata( + &multi_environment_tools, + /*expected_present*/ true, + ); +} + #[test] fn test_build_specs_collab_tools_enabled() { let model_info = model_info(); @@ -2485,6 +2528,21 @@ fn assert_process_tool_environment_id( ); } +fn assert_apply_patch_freeform_environment_metadata( + tools: &[ConfiguredToolSpec], + expected_present: bool, +) { + let tool = find_tool(tools, "apply_patch"); + let ToolSpec::Freeform(FreeformTool { format, .. }) = &tool.spec else { + panic!("expected apply_patch freeform tool"); + }; + assert_eq!( + format.definition.contains("*** Environment ID: "), + expected_present, + "apply_patch freeform environment metadata presence" + ); +} + fn find_namespace_function_tool<'a>( tools: &'a [ConfiguredToolSpec], expected_namespace: &str, diff --git a/codex-rs/tools/src/tool_spec.rs b/codex-rs/tools/src/tool_spec.rs index 4236dcaa61..23a4348c41 100644 --- a/codex-rs/tools/src/tool_spec.rs +++ b/codex-rs/tools/src/tool_spec.rs @@ -80,6 +80,10 @@ impl From for ToolSpec { } } +pub(crate) fn environment_id_schema(description: &str) -> JsonSchema { + JsonSchema::string(Some(description.to_string())) +} + pub fn create_local_shell_tool() -> ToolSpec { ToolSpec::LocalShell {} } diff --git a/codex-rs/tools/src/view_image.rs b/codex-rs/tools/src/view_image.rs index 1d77ceadf3..142dd19ea3 100644 --- a/codex-rs/tools/src/view_image.rs +++ b/codex-rs/tools/src/view_image.rs @@ -24,7 +24,6 @@ pub fn create_view_image_tool(options: ViewImageToolOptions) -> ToolSpec { )), ); } - ToolSpec::Function(ResponsesApiTool { name: VIEW_IMAGE_TOOL_NAME.to_string(), description: "View a local image from the filesystem (only use if given a full filepath by the user, and the image isn't already attached to the thread context within tags)."