Support multi-environment apply_patch selection (#21617)

## Summary
- add multi-environment apply_patch routing for both freeform and
function-call tool flows
- parse and reconcile the optional environment selector in the main
apply_patch parser, then verify against the selected environment in the
handler
- carry environment_id through runtime and approval surfaces so
remote-targeted patches stay explicit end to end

## Testing
- just fmt
- remote exec-server e2e: `cargo test -p codex-core --test all
apply_patch_multi_environment_uses_remote_executor -- --nocapture` on
dev via `scripts/test-remote-env.sh`

---------

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

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

@@ -151,78 +151,87 @@ pub async fn maybe_parse_apply_patch_verified(
     }
 
     match maybe_parse_apply_patch(argv) {
-        MaybeApplyPatch::Body(ApplyPatchArgs {
-            patch,
-            hunks,
-            workdir,
-        }) => {
-            let effective_cwd = workdir
-                .as_ref()
-                .map(|dir| cwd.join(Path::new(dir)))
-                .unwrap_or_else(|| cwd.clone());
-            let mut changes = HashMap::new();
-            for hunk in hunks {
-                let path = hunk.resolve_path(&effective_cwd);
-                match hunk {
-                    Hunk::AddFile { contents, .. } => {
-                        changes.insert(
-                            path.into_path_buf(),
-                            ApplyPatchFileChange::Add { content: contents },
-                        );
-                    }
-                    Hunk::DeleteFile { .. } => {
-                        let content = match fs.read_file_text(&path, sandbox).await {
-                            Ok(content) => content,
-                            Err(e) => {
-                                return MaybeApplyPatchVerified::CorrectnessError(
-                                    ApplyPatchError::IoError(IoError {
-                                        context: format!("Failed to read {}", path.display()),
-                                        source: e,
-                                    }),
-                                );
-                            }
-                        };
-                        changes.insert(
-                            path.into_path_buf(),
-                            ApplyPatchFileChange::Delete { content },
-                        );
-                    }
-                    Hunk::UpdateFile {
-                        move_path, chunks, ..
-                    } => {
-                        let ApplyPatchFileUpdate {
-                            unified_diff,
-                            content: contents,
-                            ..
-                        } = match unified_diff_from_chunks(&path, &chunks, fs, sandbox).await {
-                            Ok(diff) => diff,
-                            Err(e) => {
-                                return MaybeApplyPatchVerified::CorrectnessError(e);
-                            }
-                        };
-                        changes.insert(
-                            path.into_path_buf(),
-                            ApplyPatchFileChange::Update {
-                                unified_diff,
-                                move_path: move_path.map(|p| effective_cwd.join(p).into_path_buf()),
-                                new_content: contents,
-                            },
-                        );
-                    }
-                }
-            }
-            MaybeApplyPatchVerified::Body(ApplyPatchAction {
-                changes,
-                patch,
-                cwd: effective_cwd,
-            })
-        }
+        MaybeApplyPatch::Body(args) => verify_apply_patch_args(args, cwd, fs, sandbox).await,
         MaybeApplyPatch::ShellParseError(e) => MaybeApplyPatchVerified::ShellParseError(e),
         MaybeApplyPatch::PatchParseError(e) => MaybeApplyPatchVerified::CorrectnessError(e.into()),
         MaybeApplyPatch::NotApplyPatch => MaybeApplyPatchVerified::NotApplyPatch,
     }
 }
 
+pub async fn verify_apply_patch_args(
+    args: ApplyPatchArgs,
+    cwd: &AbsolutePathBuf,
+    fs: &dyn ExecutorFileSystem,
+    sandbox: Option<&codex_exec_server::FileSystemSandboxContext>,
+) -> MaybeApplyPatchVerified {
+    let ApplyPatchArgs {
+        patch,
+        hunks,
+        workdir,
+        ..
+    } = args;
+    let effective_cwd = workdir
+        .as_ref()
+        .map(|dir| cwd.join(Path::new(dir)))
+        .unwrap_or_else(|| cwd.clone());
+    let mut changes = HashMap::new();
+    for hunk in hunks {
+        let path = hunk.resolve_path(&effective_cwd);
+        match hunk {
+            Hunk::AddFile { contents, .. } => {
+                changes.insert(
+                    path.into_path_buf(),
+                    ApplyPatchFileChange::Add { content: contents },
+                );
+            }
+            Hunk::DeleteFile { .. } => {
+                let content = match fs.read_file_text(&path, sandbox).await {
+                    Ok(content) => content,
+                    Err(e) => {
+                        return MaybeApplyPatchVerified::CorrectnessError(
+                            ApplyPatchError::IoError(IoError {
+                                context: format!("Failed to read {}", path.display()),
+                                source: e,
+                            }),
+                        );
+                    }
+                };
+                changes.insert(
+                    path.into_path_buf(),
+                    ApplyPatchFileChange::Delete { content },
+                );
+            }
+            Hunk::UpdateFile {
+                move_path, chunks, ..
+            } => {
+                let ApplyPatchFileUpdate {
+                    unified_diff,
+                    content: contents,
+                    ..
+                } = match unified_diff_from_chunks(&path, &chunks, fs, sandbox).await {
+                    Ok(diff) => diff,
+                    Err(e) => {
+                        return MaybeApplyPatchVerified::CorrectnessError(e);
+                    }
+                };
+                changes.insert(
+                    path.into_path_buf(),
+                    ApplyPatchFileChange::Update {
+                        unified_diff,
+                        move_path: move_path.map(|p| effective_cwd.join(p).into_path_buf()),
+                        new_content: contents,
+                    },
+                );
+            }
+        }
+    }
+    MaybeApplyPatchVerified::Body(ApplyPatchAction {
+        changes,
+        patch,
+        cwd: effective_cwd,
+    })
+}
+
 /// Extract the heredoc body (and optional `cd` workdir) from a `bash -lc` script
 /// that invokes the apply_patch tool using a heredoc.
 ///

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

@@ -26,6 +26,7 @@ pub use streaming_parser::StreamingPatchParser;
 use thiserror::Error;
 
 pub use invocation::maybe_parse_apply_patch_verified;
+pub use invocation::verify_apply_patch_args;
 pub use standalone_executable::main;
 
 use crate::invocation::ExtractHeredocError;
@@ -97,6 +98,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_id? hunk+ end_patch
 //! begin_patch: "*** Begin Patch" LF
+//! environment_id: "*** Environment ID: " filename LF
 //! end_patch: "*** End Patch" LF?
 //!
 //! hunk: add_hunk | delete_hunk | update_hunk
@@ -32,6 +33,7 @@ use std::path::PathBuf;
 use thiserror::Error;
 
 pub(crate) const BEGIN_PATCH_MARKER: &str = "*** Begin Patch";
+pub(crate) const ENVIRONMENT_ID_MARKER: &str = "*** Environment ID: ";
 pub(crate) const END_PATCH_MARKER: &str = "*** End Patch";
 pub(crate) const ADD_FILE_MARKER: &str = "*** Add File: ";
 pub(crate) const DELETE_FILE_MARKER: &str = "*** Delete File: ";
@@ -178,9 +180,9 @@ fn parse_patch_text(patch: &str, mode: ParseMode) -> Result<ApplyPatchArgs, Pars
         ParseMode::Lenient => check_patch_boundaries_lenient(&lines)?,
     };
 
+    let (environment_id, mut remaining_lines, mut line_number) =
+        parse_environment_id_preamble(hunk_lines)?;
     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,9 +194,28 @@ fn parse_patch_text(patch: &str, mode: ParseMode) -> Result<ApplyPatchArgs, Pars
         hunks,
         patch,
         workdir: None,
+        environment_id,
     })
 }
 
+fn parse_environment_id_preamble<'a>(
+    hunk_lines: &'a [&'a str],
+) -> Result<(Option<String>, &'a [&'a str], usize), ParseError> {
+    let Some(first_line) = hunk_lines.first() else {
+        return Ok((None, hunk_lines, 2));
+    };
+    let Some(environment_id) = first_line.trim_start().strip_prefix(ENVIRONMENT_ID_MARKER) else {
+        return Ok((None, hunk_lines, 2));
+    };
+    let environment_id = environment_id.trim();
+    if environment_id.is_empty() {
+        return Err(InvalidPatchError(
+            "apply_patch environment_id cannot be empty".to_string(),
+        ));
+    }
+    Ok((Some(environment_id.to_string()), &hunk_lines[1..], 3))
+}
+
 /// Checks the start and end lines of the patch text for `apply_patch`,
 /// returning an error if they do not match the expected markers.
 fn check_patch_boundaries_strict<'a>(
@@ -837,6 +858,7 @@ fn test_parse_patch_lenient() {
             hunks: expected_patch.clone(),
             patch: patch_text.to_string(),
             workdir: None,
+            environment_id: None,
         })
     );
 
@@ -851,6 +873,7 @@ fn test_parse_patch_lenient() {
             hunks: expected_patch.clone(),
             patch: patch_text.to_string(),
             workdir: None,
+            environment_id: None,
         })
     );
 
@@ -865,6 +888,7 @@ fn test_parse_patch_lenient() {
             hunks: expected_patch,
             patch: patch_text.to_string(),
             workdir: None,
+            environment_id: None,
         })
     );
 
@@ -891,3 +915,40 @@ fn test_parse_patch_lenient() {
         ))
     );
 }
+
+#[test]
+fn test_parse_patch_environment_id_preamble() {
+    assert_eq!(
+        parse_patch_text(
+            "*** Begin Patch\n\
+             *** Environment ID: remote\n\
+             *** Add File: hello.txt\n\
+             +hello\n\
+             *** End Patch",
+            ParseMode::Strict
+        ),
+        Ok(ApplyPatchArgs {
+            hunks: vec![AddFile {
+                path: PathBuf::from("hello.txt"),
+                contents: "hello\n".to_string(),
+            }],
+            patch: "*** Begin Patch\n*** Environment ID: remote\n*** Add File: hello.txt\n+hello\n*** End Patch".to_string(),
+            workdir: None,
+            environment_id: Some("remote".to_string()),
+        })
+    );
+
+    assert_eq!(
+        parse_patch_text(
+            "*** Begin Patch\n\
+             *** Environment ID:   \n\
+             *** Add File: hello.txt\n\
+             +hello\n\
+             *** End Patch",
+            ParseMode::Strict
+        ),
+        Err(InvalidPatchError(
+            "apply_patch environment_id cannot be empty".to_string()
+        ))
+    );
+}

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

@@ -16,6 +16,8 @@ use crate::parser::UpdateFileChunk;
 use Hunk::*;
 use ParseError::*;
 
+const ENVIRONMENT_ID_MARKER: &str = "*** Environment ID: ";
+
 #[derive(Debug, Default, Clone)]
 pub struct StreamingPatchParser {
     line_buffer: String,
@@ -29,7 +31,7 @@ struct StreamingParserState {
     hunks: Vec<Hunk>,
 }
 
-#[derive(Debug, Default, Clone)]
+#[derive(Debug, Default, Clone, Copy)]
 enum StreamingParserMode {
     #[default]
     NotStarted,
@@ -43,6 +45,13 @@ enum StreamingParserMode {
 }
 
 impl StreamingPatchParser {
+    // The live streaming parser only needs to keep the patch preview flowing.
+    // Environment selection and validation happen on the final tool invocation,
+    // so here we just tolerate and skip the optional preamble line.
+    fn is_environment_id_preamble_line(&self, line: &str) -> bool {
+        line.starts_with(ENVIRONMENT_ID_MARKER)
+    }
+
     fn ensure_update_hunk_is_not_empty(&self, line: &str) -> Result<(), ParseError> {
         if let Some(UpdateFile { path, chunks, .. }) = self.state.hunks.last() {
             if chunks.is_empty()
@@ -150,7 +159,7 @@ impl StreamingPatchParser {
 
     fn process_line(&mut self, line: &str) -> Result<(), ParseError> {
         let trimmed = line.trim();
-        match self.state.mode.clone() {
+        match self.state.mode {
             StreamingParserMode::NotStarted => {
                 if trimmed == BEGIN_PATCH_MARKER {
                     self.state.mode = StreamingParserMode::StartedPatch;
@@ -161,6 +170,9 @@ impl StreamingPatchParser {
                 ))
             }
             StreamingParserMode::StartedPatch => {
+                if self.is_environment_id_preamble_line(line) {
+                    return Ok(());
+                }
                 if self.handle_hunk_headers_and_end_patch(trimmed)? {
                     return Ok(());
                 }
@@ -431,6 +443,32 @@ mod tests {
         );
     }
 
+    #[test]
+    fn test_streaming_patch_parser_environment_id_mode() {
+        let patch = "\
+*** Begin Patch
+*** Environment ID: remote
+*** Add File: src/hello.txt
++hello
+*** End Patch
+";
+
+        let mut parser = StreamingPatchParser::default();
+        assert_eq!(
+            parser.push_delta(patch),
+            Ok(vec![AddFile {
+                path: PathBuf::from("src/hello.txt"),
+                contents: "hello\n".to_string(),
+            }])
+        );
+
+        let mut parser = StreamingPatchParser::default();
+        assert_eq!(
+            parser.push_delta("*** Begin Patch\n*** Environment ID:   \n"),
+            Ok(vec![])
+        );
+    }
+
     #[test]
     fn test_streaming_patch_parser_large_patch_split_by_character() {
         let patch = "\