codex: address apply_patch tool spec review (#21048)

Co-authored-by: Codex <noreply@openai.com>

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<Arc<codex_exec_server::Environment>> {
         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());
+    }
 }

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",
     ],
 )

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<String>,
 }
 
+#[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(),

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()]));
 }

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;

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 <environment_context> block. If omitted, uses the primary environment.".to_string(),
-            )),
+            environment_id_schema(
+                "Optional environment id from the <environment_context> block. If omitted, uses the primary environment.",
+            ),
         );
     }
     properties.extend(create_approval_parameters(

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

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

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,
                 );

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,

codex-rs/tools/src/tool_spec.rs

@@ -80,6 +80,10 @@ impl From<LoadableToolSpec> 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 {}
 }

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 <image ...> tags)."