codex: restore apply_patch registry test marker (#21048)

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

codex-rs/apply-patch/src/invocation.rs

@@ -155,6 +155,7 @@ pub async fn maybe_parse_apply_patch_verified(
             patch,
             hunks,
             workdir,
+            environment_id: _,
         }) => {
             let effective_cwd = workdir
                 .as_ref()

codex-rs/apply-patch/src/lib.rs

@@ -97,6 +97,7 @@ pub struct ApplyPatchArgs {
     pub patch: String,
     pub hunks: Vec<Hunk>,
     pub workdir: Option<String>,
+    pub environment_id: Option<String>,
 }
 
 #[derive(Debug, PartialEq)]

codex-rs/apply-patch/src/parser.rs

@@ -3,8 +3,9 @@
 //!
 //! The official Lark grammar for the apply-patch format is:
 //!
-//! start: begin_patch hunk+ end_patch
+//! 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
@@ -40,6 +41,7 @@ pub(crate) const MOVE_TO_MARKER: &str = "*** Move to: ";
 pub(crate) const EOF_MARKER: &str = "*** End of File";
 pub(crate) const CHANGE_CONTEXT_MARKER: &str = "@@ ";
 pub(crate) const EMPTY_CHANGE_CONTEXT_MARKER: &str = "@@";
+const ENVIRONMENT_ID_MARKER: &str = "*** Environment ID: ";
 
 /// Currently, the only OpenAI model that knowingly requires lenient parsing is
 /// gpt-4.1. While we could try to require everyone to pass in a strictness
@@ -177,10 +179,29 @@ fn parse_patch_text(patch: &str, mode: ParseMode) -> Result<ApplyPatchArgs, Pars
         ParseMode::Strict => check_patch_boundaries_strict(&lines)?,
         ParseMode::Lenient => check_patch_boundaries_lenient(&lines)?,
     };
+    let mut patch_lines = patch_lines.to_vec();
+    let mut hunk_lines = hunk_lines;
+    let mut line_number = 2;
+    let environment_id = match hunk_lines.first() {
+        Some(line) => match line.strip_prefix(ENVIRONMENT_ID_MARKER) {
+            Some("") => {
+                return Err(InvalidPatchError(
+                    "environment_id cannot be empty".to_string(),
+                ));
+            }
+            Some(environment_id) => {
+                patch_lines.remove(1);
+                hunk_lines = &hunk_lines[1..];
+                line_number += 1;
+                Some(environment_id.to_string())
+            }
+            None => None,
+        },
+        None => None,
+    };
 
     let mut hunks: Vec<Hunk> = Vec::new();
     let mut remaining_lines = hunk_lines;
-    let mut line_number = 2;
     while !remaining_lines.is_empty() {
         let (hunk, hunk_lines) = parse_one_hunk(remaining_lines, line_number)?;
         hunks.push(hunk);
@@ -192,6 +213,7 @@ fn parse_patch_text(patch: &str, mode: ParseMode) -> Result<ApplyPatchArgs, Pars
         hunks,
         patch,
         workdir: None,
+        environment_id,
     })
 }
 
@@ -749,6 +771,30 @@ fn test_parse_patch_accepts_relative_and_absolute_hunk_paths() {
     );
 }
 
+#[test]
+fn test_parse_patch_strips_environment_id_metadata() {
+    let patch_text = r#"*** Begin Patch
+*** Environment ID: selected
+*** Add File: hello.txt
++hello
+*** End Patch"#;
+
+    let parsed = parse_patch_text(patch_text, ParseMode::Strict).unwrap();
+
+    assert_eq!(
+        parsed,
+        ApplyPatchArgs {
+            hunks: vec![AddFile {
+                path: PathBuf::from("hello.txt"),
+                contents: "hello\n".to_string(),
+            }],
+            patch: "*** Begin Patch\n*** Add File: hello.txt\n+hello\n*** End Patch".to_string(),
+            workdir: None,
+            environment_id: Some("selected".to_string()),
+        }
+    );
+}
+
 #[test]
 fn test_hunk_resolve_path_accepts_relative_and_absolute_paths() {
     let cwd_dir = tempfile::tempdir().unwrap();
@@ -837,6 +883,7 @@ fn test_parse_patch_lenient() {
             hunks: expected_patch.clone(),
             patch: patch_text.to_string(),
             workdir: None,
+            environment_id: None,
         })
     );
 
@@ -851,6 +898,7 @@ fn test_parse_patch_lenient() {
             hunks: expected_patch.clone(),
             patch: patch_text.to_string(),
             workdir: None,
+            environment_id: None,
         })
     );
 
@@ -865,6 +913,7 @@ fn test_parse_patch_lenient() {
             hunks: expected_patch,
             patch: patch_text.to_string(),
             workdir: None,
+            environment_id: None,
         })
     );
 

codex-rs/core/src/apply_patch.rs

@@ -40,7 +40,7 @@ pub(crate) async fn apply_patch(
         turn_context.approval_policy.value(),
         &turn_context.permission_profile(),
         file_system_sandbox_policy,
-        &turn_context.cwd,
+        &action.cwd,
         turn_context.windows_sandbox_level,
     ) {
         SafetyCheck::AutoApprove {

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/core/src/tools/handlers/apply_patch.rs

@@ -121,6 +121,17 @@ impl ApplyPatchArgumentDiffConsumer {
     }
 }
 
+fn parse_apply_patch_input(
+    input: String,
+) -> Result<codex_apply_patch::ApplyPatchArgs, FunctionCallError> {
+    codex_apply_patch::parse_patch(&input).map_err(|err| {
+        FunctionCallError::RespondToModel(format!(
+            "apply_patch verification failed: {}",
+            codex_apply_patch::ApplyPatchError::from(err)
+        ))
+    })
+}
+
 fn convert_apply_patch_hunks_to_protocol(hunks: &[Hunk]) -> HashMap<PathBuf, FileChange> {
     hunks
         .iter()
@@ -259,6 +270,7 @@ fn apply_patch_payload_command(payload: &ToolPayload) -> Option<String> {
 async fn effective_patch_permissions(
     session: &Session,
     turn: &TurnContext,
+    cwd: &AbsolutePathBuf,
     action: &ApplyPatchAction,
 ) -> (
     Vec<AbsolutePathBuf>,
@@ -277,9 +289,9 @@ async fn effective_patch_permissions(
     );
     let effective_additional_permissions = apply_granted_turn_permissions(
         session,
-        turn.cwd.as_path(),
+        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;
 
@@ -351,12 +363,28 @@ impl ToolHandler for ApplyPatchHandler {
             ..
         } = invocation;
 
-        let patch_input = match payload {
+        let (patch_input, environment_id) = match payload {
             ToolPayload::Function { arguments } => {
                 let args: ApplyPatchToolArgs = parse_arguments(&arguments)?;
-                args.input
+                let parsed_input = parse_apply_patch_input(args.input)?;
+                let environment_id = match (args.environment_id, parsed_input.environment_id) {
+                    (Some(argument_environment_id), Some(header_environment_id))
+                        if argument_environment_id != header_environment_id =>
+                    {
+                        return Err(FunctionCallError::RespondToModel(
+                            "apply_patch environment_id argument conflicts with patch header metadata"
+                                .to_string(),
+                        ));
+                    }
+                    (Some(argument_environment_id), _) => Some(argument_environment_id),
+                    (None, header_environment_id) => header_environment_id,
+                };
+                (parsed_input.patch, environment_id)
+            }
+            ToolPayload::Custom { input } => {
+                let parsed_input = parse_apply_patch_input(input)?;
+                (parsed_input.patch, parsed_input.environment_id)
             }
-            ToolPayload::Custom { input } => input,
             _ => {
                 return Err(FunctionCallError::RespondToModel(
                     "apply_patch handler received unsupported payload".to_string(),
@@ -366,18 +394,25 @@ 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 command = vec!["apply_patch".to_string(), patch_input.clone()];
-        let Some(turn_environment) = turn.environments.primary() else {
+        let Some(turn_environment) = turn
+            .environments
+            .get_or_primary(environment_id.as_deref())
+            .cloned()
+        else {
             return Err(FunctionCallError::RespondToModel(
                 "apply_patch is unavailable in this session".to_string(),
             ));
         };
-        let fs = turn_environment.environment.get_filesystem();
-        let sandbox = turn_environment
-            .environment
-            .is_remote()
-            .then(|| turn.file_system_sandbox_context(/*additional_permissions*/ None));
+        let cwd = turn_environment.cwd.clone();
+        let command = vec!["apply_patch".to_string(), patch_input.clone()];
+        let environment = Arc::clone(&turn_environment.environment);
+        let fs = environment.get_filesystem();
+        let sandbox = turn_environment.environment.is_remote().then(|| {
+            let mut context =
+                turn.file_system_sandbox_context(/*additional_permissions*/ None);
+            context.cwd = Some(cwd.clone());
+            context
+        });
         match codex_apply_patch::maybe_parse_apply_patch_verified(
             &command,
             &cwd,
@@ -388,7 +423,13 @@ 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;
+                    effective_patch_permissions(
+                        session.as_ref(),
+                        turn.as_ref(),
+                        &changes.cwd,
+                        &changes,
+                    )
+                    .await;
                 match apply_patch::apply_patch(turn.as_ref(), &file_system_sandbox_policy, changes)
                     .await
                 {
@@ -410,6 +451,7 @@ impl ToolHandler for ApplyPatchHandler {
 
                         let req = ApplyPatchRequest {
                             action: apply.action,
+                            environment,
                             file_paths,
                             changes,
                             exec_approval_requirement: apply.exec_approval_requirement,
@@ -497,7 +539,13 @@ 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;
+                effective_patch_permissions(
+                    session.as_ref(),
+                    turn.as_ref(),
+                    &changes.cwd,
+                    &changes,
+                )
+                .await;
             match apply_patch::apply_patch(turn.as_ref(), &file_system_sandbox_policy, changes)
                 .await
             {
@@ -506,6 +554,11 @@ pub(crate) async fn intercept_apply_patch(
                     Ok(Some(FunctionToolOutput::from_text(content, Some(true))))
                 }
                 InternalApplyPatchInvocation::DelegateToRuntime(apply) => {
+                    let Some(turn_environment) = turn.environments.primary() else {
+                        return Err(FunctionCallError::RespondToModel(
+                            "apply_patch is unavailable in this session".to_string(),
+                        ));
+                    };
                     let changes = convert_apply_patch_to_protocol(&apply.action);
                     let emitter = ToolEmitter::apply_patch(changes.clone(), apply.auto_approved);
                     let event_ctx = ToolEventCtx::new(
@@ -518,6 +571,7 @@ pub(crate) async fn intercept_apply_patch(
 
                     let req = ApplyPatchRequest {
                         action: apply.action,
+                        environment: turn_environment.environment.clone(),
                         file_paths: approval_keys,
                         changes,
                         exec_approval_requirement: apply.exec_approval_requirement,

codex-rs/core/src/tools/runtimes/apply_patch.rs

@@ -18,6 +18,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;
@@ -33,11 +34,12 @@ use codex_sandboxing::policy_transforms::effective_permission_profile;
 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 action: ApplyPatchAction,
+    pub environment: Arc<Environment>,
     pub file_paths: Vec<AbsolutePathBuf>,
     pub changes: std::collections::HashMap<PathBuf, FileChange>,
     pub exec_approval_requirement: ExecApprovalRequirement,
@@ -77,7 +79,7 @@ impl ApplyPatchRuntime {
             effective_permission_profile(attempt.permissions, req.additional_permissions.as_ref());
         Some(FileSystemSandboxContext {
             permissions,
-            cwd: Some(attempt.sandbox_cwd.clone()),
+            cwd: Some(req.action.cwd.clone()),
             windows_sandbox_level: attempt.windows_sandbox_level,
             windows_sandbox_private_desktop: attempt.windows_sandbox_private_desktop,
             use_legacy_landlock: attempt.use_legacy_landlock,
@@ -185,17 +187,18 @@ impl Approvable<ApplyPatchRequest> for ApplyPatchRuntime {
 }
 
 impl ToolRuntime<ApplyPatchRequest, ExecToolCallOutput> for ApplyPatchRuntime {
+    fn sandbox_cwd<'a>(&self, req: &'a ApplyPatchRequest) -> Option<&'a AbsolutePathBuf> {
+        Some(&req.action.cwd)
+    }
+
     async fn run(
         &mut self,
         req: &ApplyPatchRequest,
         attempt: &SandboxAttempt<'_>,
-        ctx: &ToolCtx,
+        _ctx: &ToolCtx,
     ) -> Result<ExecToolCallOutput, ToolError> {
-        let turn_environment = ctx.turn.environments.primary().ok_or_else(|| {
-            ToolError::Rejected("apply_patch is unavailable in this session".to_string())
-        })?;
         let started_at = Instant::now();
-        let fs = turn_environment.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();

codex-rs/core/src/tools/runtimes/apply_patch_tests.rs

@@ -1,5 +1,6 @@
 use super::*;
 use crate::tools::sandboxing::SandboxAttempt;
+use codex_exec_server::Environment;
 use codex_protocol::config_types::WindowsSandboxLevel;
 use codex_protocol::models::AdditionalPermissionProfile;
 use codex_protocol::models::FileSystemPermissions;
@@ -15,6 +16,14 @@ use codex_sandboxing::policy_transforms::effective_network_sandbox_policy;
 use core_test_support::PathBufExt;
 use pretty_assertions::assert_eq;
 use std::collections::HashMap;
+use std::sync::Arc;
+
+fn test_environment() -> Arc<Environment> {
+    Arc::new(
+        Environment::create_for_tests(Some("ws://127.0.0.1:0".to_string()))
+            .expect("test environment"),
+    )
+}
 
 #[test]
 fn wants_no_sandbox_approval_granular_respects_sandbox_flag() {
@@ -40,8 +49,8 @@ fn wants_no_sandbox_approval_granular_respects_sandbox_flag() {
     );
 }
 
-#[test]
-fn guardian_review_request_includes_patch_context() {
+#[tokio::test]
+async fn guardian_review_request_includes_patch_context() {
     let path = std::env::temp_dir()
         .join("guardian-apply-patch-test.txt")
         .abs();
@@ -50,6 +59,7 @@ fn guardian_review_request_includes_patch_context() {
     let expected_patch = action.patch.clone();
     let request = ApplyPatchRequest {
         action,
+        environment: test_environment(),
         file_paths: vec![path.clone()],
         changes: HashMap::from([(
             path.to_path_buf(),
@@ -78,8 +88,8 @@ fn guardian_review_request_includes_patch_context() {
     );
 }
 
-#[test]
-fn permission_request_payload_uses_apply_patch_hook_name_and_aliases() {
+#[tokio::test]
+async fn permission_request_payload_uses_apply_patch_hook_name_and_aliases() {
     let runtime = ApplyPatchRuntime::new();
     let path = std::env::temp_dir()
         .join("apply-patch-permission-request-payload.txt")
@@ -88,6 +98,7 @@ fn permission_request_payload_uses_apply_patch_hook_name_and_aliases() {
     let expected_patch = action.patch.clone();
     let req = ApplyPatchRequest {
         action,
+        environment: test_environment(),
         file_paths: vec![path],
         changes: HashMap::new(),
         exec_approval_requirement: ExecApprovalRequirement::NeedsApproval {
@@ -113,8 +124,8 @@ fn permission_request_payload_uses_apply_patch_hook_name_and_aliases() {
     );
 }
 
-#[test]
-fn file_system_sandbox_context_uses_active_attempt() {
+#[tokio::test]
+async fn file_system_sandbox_context_uses_active_attempt() {
     let path = std::env::temp_dir()
         .join("apply-patch-runtime-attempt.txt")
         .abs();
@@ -127,6 +138,7 @@ fn file_system_sandbox_context_uses_active_attempt() {
     };
     let req = ApplyPatchRequest {
         action: ApplyPatchAction::new_add_for_test(&path, "hello".to_string()),
+        environment: test_environment(),
         file_paths: vec![path.clone()],
         changes: HashMap::new(),
         exec_approval_requirement: ExecApprovalRequirement::Skip {
@@ -168,7 +180,7 @@ fn file_system_sandbox_context_uses_active_attempt() {
     let expected_permissions =
         PermissionProfile::from_runtime_permissions(&file_system_policy, network_policy);
     assert_eq!(sandbox.permissions, expected_permissions);
-    assert_eq!(sandbox.cwd, Some(path.clone()));
+    assert_eq!(sandbox.cwd.as_ref(), Some(&req.action.cwd));
     assert_eq!(
         sandbox.windows_sandbox_level,
         WindowsSandboxLevel::RestrictedToken
@@ -177,13 +189,14 @@ fn file_system_sandbox_context_uses_active_attempt() {
     assert_eq!(sandbox.use_legacy_landlock, true);
 }
 
-#[test]
-fn no_sandbox_attempt_has_no_file_system_context() {
+#[tokio::test]
+async fn no_sandbox_attempt_has_no_file_system_context() {
     let path = std::env::temp_dir()
         .join("apply-patch-runtime-none.txt")
         .abs();
     let req = ApplyPatchRequest {
         action: ApplyPatchAction::new_add_for_test(&path, "hello".to_string()),
+        environment: test_environment(),
         file_paths: vec![path.clone()],
         changes: HashMap::new(),
         exec_approval_requirement: ExecApprovalRequirement::Skip {

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:
@@ -82,30 +85,50 @@ It is important to remember:
 #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
 pub struct ApplyPatchToolArgs {
     pub input: String,
+    #[serde(default)]
+    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_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,47 @@ fn exec_command_spec_includes_environment_id_only_for_multiple_selected_environm
     );
 }
 
+#[test]
+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 +2529,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 {}
 }