Support multi-environment apply_patch selection

Co-authored-by: Codex <noreply@openai.com>
This commit is contained in:
starr-openai
2026-05-07 16:43:32 -07:00
parent 99016ec732
commit ebf66ae577
27 changed files with 871 additions and 213 deletions

View File

@@ -14,7 +14,6 @@ use crate::ApplyPatchAction;
use crate::ApplyPatchArgs;
use crate::ApplyPatchError;
use crate::ApplyPatchFileChange;
use crate::ApplyPatchFileUpdate;
use crate::IoError;
use crate::MaybeApplyPatchVerified;
use crate::parser::Hunk;
@@ -154,6 +153,7 @@ pub async fn maybe_parse_apply_patch_verified(
MaybeApplyPatch::Body(ApplyPatchArgs {
patch,
hunks,
environment_id: _,
workdir,
}) => {
let effective_cwd = workdir
@@ -163,53 +163,44 @@ pub async fn maybe_parse_apply_patch_verified(
let mut changes = HashMap::new();
for hunk in hunks {
let path = hunk.resolve_path(&effective_cwd);
match hunk {
let change = match hunk {
Hunk::AddFile { contents, .. } => {
changes.insert(
path.into_path_buf(),
ApplyPatchFileChange::Add { content: contents },
);
ApplyPatchFileChange::Add { content: contents }
}
Hunk::DeleteFile { .. } => {
let content = match fs.read_file_text(&path, sandbox).await {
Ok(content) => content,
Err(e) => {
Err(source) => {
return MaybeApplyPatchVerified::CorrectnessError(
ApplyPatchError::IoError(IoError {
context: format!("Failed to read {}", path.display()),
source: e,
source,
}),
);
}
};
changes.insert(
path.into_path_buf(),
ApplyPatchFileChange::Delete { content },
);
ApplyPatchFileChange::Delete { content }
}
Hunk::UpdateFile {
move_path, chunks, ..
} => {
let ApplyPatchFileUpdate {
let update =
match unified_diff_from_chunks(&path, &chunks, fs, sandbox).await {
Ok(update) => update,
Err(err) => {
return MaybeApplyPatchVerified::CorrectnessError(err);
}
};
let (unified_diff, new_content) = update.into_parts();
ApplyPatchFileChange::Update {
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,
},
);
move_path: move_path
.map(|move_path| effective_cwd.join(move_path).into_path_buf()),
new_content,
}
}
}
};
changes.insert(path.into_path_buf(), change);
}
MaybeApplyPatchVerified::Body(ApplyPatchAction {
changes,
@@ -378,6 +369,7 @@ fn extract_apply_patch_from_bash(
#[cfg(test)]
mod tests {
use super::*;
use crate::ApplyPatchFileUpdate;
use crate::unified_diff_from_chunks;
use assert_matches::assert_matches;
use codex_exec_server::LOCAL_FS;

View File

@@ -19,8 +19,10 @@ use codex_utils_absolute_path::AbsolutePathBuf;
pub use parser::Hunk;
pub use parser::ParseError;
use parser::ParseError::*;
pub use parser::ParseOptions;
pub use parser::UpdateFileChunk;
pub use parser::parse_patch;
pub use parser::parse_patch_with_options;
use similar::TextDiff;
pub use streaming_parser::StreamingPatchParser;
use thiserror::Error;
@@ -96,6 +98,7 @@ impl PartialEq for IoError {
pub struct ApplyPatchArgs {
pub patch: String,
pub hunks: Vec<Hunk>,
pub environment_id: Option<String>,
pub workdir: Option<String>,
}
@@ -146,6 +149,18 @@ pub struct ApplyPatchAction {
}
impl ApplyPatchAction {
pub fn new(
changes: HashMap<PathBuf, ApplyPatchFileChange>,
patch: String,
cwd: AbsolutePathBuf,
) -> Self {
Self {
changes,
patch,
cwd,
}
}
pub fn is_empty(&self) -> bool {
self.changes.is_empty()
}
@@ -815,6 +830,12 @@ pub struct ApplyPatchFileUpdate {
content: String,
}
impl ApplyPatchFileUpdate {
pub fn into_parts(self) -> (String, String) {
(self.unified_diff, self.content)
}
}
pub async fn unified_diff_from_chunks(
path_abs: &AbsolutePathBuf,
chunks: &[UpdateFileChunk],

View File

@@ -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: " /(.+)/ 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: ";
@@ -124,12 +126,24 @@ pub struct UpdateFileChunk {
}
pub fn parse_patch(patch: &str) -> Result<ApplyPatchArgs, ParseError> {
parse_patch_with_options(patch, ParseOptions::default())
}
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct ParseOptions {
pub allow_environment_id: bool,
}
pub fn parse_patch_with_options(
patch: &str,
options: ParseOptions,
) -> Result<ApplyPatchArgs, ParseError> {
let mode = if PARSE_IN_STRICT_MODE {
ParseMode::Strict
} else {
ParseMode::Lenient
};
parse_patch_text(patch, mode)
parse_patch_text_with_options(patch, mode, options)
}
enum ParseMode {
@@ -171,30 +185,63 @@ enum ParseMode {
Lenient,
}
fn parse_patch_text(patch: &str, mode: ParseMode) -> Result<ApplyPatchArgs, ParseError> {
fn parse_patch_text_with_options(
patch: &str,
mode: ParseMode,
options: ParseOptions,
) -> Result<ApplyPatchArgs, ParseError> {
let lines: Vec<&str> = patch.trim().lines().collect();
let (patch_lines, hunk_lines) = match mode {
ParseMode::Strict => check_patch_boundaries_strict(&lines)?,
ParseMode::Lenient => check_patch_boundaries_lenient(&lines)?,
};
let (environment_id, hunk_lines, hunk_start_line_number) =
parse_environment_id(hunk_lines, options.allow_environment_id)?;
let mut hunks: Vec<Hunk> = Vec::new();
let mut remaining_lines = hunk_lines;
let mut line_number = 2;
let mut line_number = hunk_start_line_number;
while !remaining_lines.is_empty() {
let (hunk, hunk_lines) = parse_one_hunk(remaining_lines, line_number)?;
hunks.push(hunk);
line_number += hunk_lines;
remaining_lines = &remaining_lines[hunk_lines..]
}
let patch = patch_lines.join("\n");
let mut patch_lines_without_metadata = Vec::with_capacity(patch_lines.len());
patch_lines_without_metadata.push(patch_lines[0]);
patch_lines_without_metadata.extend_from_slice(hunk_lines);
patch_lines_without_metadata.push(patch_lines[patch_lines.len() - 1]);
let patch = patch_lines_without_metadata.join("\n");
Ok(ApplyPatchArgs {
hunks,
patch,
workdir: None,
environment_id,
})
}
fn parse_environment_id<'a>(
lines: &'a [&'a str],
allow_environment_id: bool,
) -> Result<(Option<String>, &'a [&'a str], usize), ParseError> {
let Some(first_line) = lines.first().map(|line| line.trim()) else {
return Ok((None, lines, 2));
};
if !allow_environment_id {
return Ok((None, lines, 2));
}
let Some(environment_id) = first_line.strip_prefix(ENVIRONMENT_ID_MARKER) else {
return Ok((None, lines, 2));
};
let environment_id = environment_id.trim();
if environment_id.is_empty() {
return Err(InvalidPatchError(
"Environment ID cannot be empty".to_string(),
));
}
Ok((Some(environment_id.to_string()), &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>(
@@ -557,20 +604,24 @@ fn test_update_file_chunk() {
#[test]
fn test_parse_patch() {
assert_eq!(
parse_patch_text("bad", ParseMode::Strict),
parse_patch_text_with_options("bad", ParseMode::Strict, ParseOptions::default()),
Err(InvalidPatchError(
"The first line of the patch must be '*** Begin Patch'".to_string()
))
);
assert_eq!(
parse_patch_text("*** Begin Patch\nbad", ParseMode::Strict),
parse_patch_text_with_options(
"*** Begin Patch\nbad",
ParseMode::Strict,
ParseOptions::default()
),
Err(InvalidPatchError(
"The last line of the patch must be '*** End Patch'".to_string()
))
);
assert_eq!(
parse_patch_text(
parse_patch_text_with_options(
concat!(
"*** Begin Patch",
" ",
@@ -578,7 +629,8 @@ fn test_parse_patch() {
" ",
"*** End Patch"
),
ParseMode::Strict
ParseMode::Strict,
ParseOptions::default()
)
.unwrap()
.hunks,
@@ -588,11 +640,12 @@ fn test_parse_patch() {
}]
);
assert_eq!(
parse_patch_text(
parse_patch_text_with_options(
"*** Begin Patch\n\
*** Update File: test.py\n\
*** End Patch",
ParseMode::Strict
ParseMode::Strict,
ParseOptions::default()
),
Err(InvalidHunkError {
message: "Update file hunk for path 'test.py' is empty".to_string(),
@@ -600,17 +653,18 @@ fn test_parse_patch() {
})
);
assert_eq!(
parse_patch_text(
parse_patch_text_with_options(
"*** Begin Patch\n\
*** End Patch",
ParseMode::Strict
ParseMode::Strict,
ParseOptions::default()
)
.unwrap()
.hunks,
Vec::new()
);
assert_eq!(
parse_patch_text(
parse_patch_text_with_options(
"*** Begin Patch\n\
*** Add File: path/add.py\n\
+abc\n\
@@ -622,7 +676,8 @@ fn test_parse_patch() {
- pass\n\
+ return 123\n\
*** End Patch",
ParseMode::Strict
ParseMode::Strict,
ParseOptions::default()
)
.unwrap()
.hunks,
@@ -648,7 +703,7 @@ fn test_parse_patch() {
);
// Update hunk followed by another hunk (Add File).
assert_eq!(
parse_patch_text(
parse_patch_text_with_options(
"*** Begin Patch\n\
*** Update File: file.py\n\
@@\n\
@@ -656,7 +711,8 @@ fn test_parse_patch() {
*** Add File: other.py\n\
+content\n\
*** End Patch",
ParseMode::Strict
ParseMode::Strict,
ParseOptions::default()
)
.unwrap()
.hunks,
@@ -681,13 +737,14 @@ fn test_parse_patch() {
// Update hunk without an explicit @@ header for the first chunk should parse.
// Use a raw string to preserve the leading space diff marker on the context line.
assert_eq!(
parse_patch_text(
parse_patch_text_with_options(
r#"*** Begin Patch
*** Update File: file2.py
import foo
+bar
*** End Patch"#,
ParseMode::Strict
ParseMode::Strict,
ParseOptions::default()
)
.unwrap()
.hunks,
@@ -724,7 +781,7 @@ fn test_parse_patch_accepts_relative_and_absolute_hunk_paths() {
);
assert_eq!(
parse_patch_text(&patch_text, ParseMode::Strict)
parse_patch_text_with_options(&patch_text, ParseMode::Strict, ParseOptions::default())
.unwrap()
.hunks,
vec![
@@ -828,64 +885,107 @@ fn test_parse_patch_lenient() {
let patch_text_in_heredoc = format!("<<EOF\n{patch_text}\nEOF\n");
assert_eq!(
parse_patch_text(&patch_text_in_heredoc, ParseMode::Strict),
parse_patch_text_with_options(
&patch_text_in_heredoc,
ParseMode::Strict,
ParseOptions::default()
),
Err(expected_error.clone())
);
assert_eq!(
parse_patch_text(&patch_text_in_heredoc, ParseMode::Lenient),
parse_patch_text_with_options(
&patch_text_in_heredoc,
ParseMode::Lenient,
ParseOptions::default()
),
Ok(ApplyPatchArgs {
hunks: expected_patch.clone(),
patch: patch_text.to_string(),
workdir: None,
environment_id: None,
})
);
let patch_text_in_single_quoted_heredoc = format!("<<'EOF'\n{patch_text}\nEOF\n");
assert_eq!(
parse_patch_text(&patch_text_in_single_quoted_heredoc, ParseMode::Strict),
parse_patch_text_with_options(
&patch_text_in_single_quoted_heredoc,
ParseMode::Strict,
ParseOptions::default()
),
Err(expected_error.clone())
);
assert_eq!(
parse_patch_text(&patch_text_in_single_quoted_heredoc, ParseMode::Lenient),
parse_patch_text_with_options(
&patch_text_in_single_quoted_heredoc,
ParseMode::Lenient,
ParseOptions::default()
),
Ok(ApplyPatchArgs {
hunks: expected_patch.clone(),
patch: patch_text.to_string(),
workdir: None,
environment_id: None,
})
);
let patch_text_in_double_quoted_heredoc = format!("<<\"EOF\"\n{patch_text}\nEOF\n");
assert_eq!(
parse_patch_text(&patch_text_in_double_quoted_heredoc, ParseMode::Strict),
parse_patch_text_with_options(
&patch_text_in_double_quoted_heredoc,
ParseMode::Strict,
ParseOptions::default()
),
Err(expected_error.clone())
);
assert_eq!(
parse_patch_text(&patch_text_in_double_quoted_heredoc, ParseMode::Lenient),
parse_patch_text_with_options(
&patch_text_in_double_quoted_heredoc,
ParseMode::Lenient,
ParseOptions::default()
),
Ok(ApplyPatchArgs {
hunks: expected_patch,
patch: patch_text.to_string(),
workdir: None,
environment_id: None,
})
);
let patch_text_in_mismatched_quotes_heredoc = format!("<<\"EOF'\n{patch_text}\nEOF\n");
assert_eq!(
parse_patch_text(&patch_text_in_mismatched_quotes_heredoc, ParseMode::Strict),
parse_patch_text_with_options(
&patch_text_in_mismatched_quotes_heredoc,
ParseMode::Strict,
ParseOptions::default()
),
Err(expected_error.clone())
);
assert_eq!(
parse_patch_text(&patch_text_in_mismatched_quotes_heredoc, ParseMode::Lenient),
parse_patch_text_with_options(
&patch_text_in_mismatched_quotes_heredoc,
ParseMode::Lenient,
ParseOptions::default()
),
Err(expected_error.clone())
);
let patch_text_with_missing_closing_heredoc =
"<<EOF\n*** Begin Patch\n*** Update File: file2.py\nEOF\n".to_string();
assert_eq!(
parse_patch_text(&patch_text_with_missing_closing_heredoc, ParseMode::Strict),
parse_patch_text_with_options(
&patch_text_with_missing_closing_heredoc,
ParseMode::Strict,
ParseOptions::default()
),
Err(expected_error)
);
assert_eq!(
parse_patch_text(&patch_text_with_missing_closing_heredoc, ParseMode::Lenient),
parse_patch_text_with_options(
&patch_text_with_missing_closing_heredoc,
ParseMode::Lenient,
ParseOptions::default()
),
Err(InvalidPatchError(
"The last line of the patch must be '*** End Patch'".to_string()
))

View File

@@ -6,6 +6,7 @@ use crate::parser::CHANGE_CONTEXT_MARKER;
use crate::parser::DELETE_FILE_MARKER;
use crate::parser::EMPTY_CHANGE_CONTEXT_MARKER;
use crate::parser::END_PATCH_MARKER;
use crate::parser::ENVIRONMENT_ID_MARKER;
use crate::parser::EOF_MARKER;
use crate::parser::Hunk;
use crate::parser::MOVE_TO_MARKER;
@@ -18,6 +19,8 @@ use ParseError::*;
#[derive(Debug, Default, Clone)]
pub struct StreamingPatchParser {
allow_environment_id: bool,
saw_environment_id: bool,
line_buffer: String,
state: StreamingParserState,
line_number: usize,
@@ -43,6 +46,10 @@ enum StreamingParserMode {
}
impl StreamingPatchParser {
pub fn set_allow_environment_id(&mut self, allow_environment_id: bool) {
self.allow_environment_id = allow_environment_id;
}
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()
@@ -161,6 +168,18 @@ impl StreamingPatchParser {
))
}
StreamingParserMode::StartedPatch => {
if self.allow_environment_id
&& !self.saw_environment_id
&& let Some(environment_id) = trimmed.strip_prefix(ENVIRONMENT_ID_MARKER)
{
if environment_id.trim().is_empty() {
return Err(InvalidPatchError(
"Environment ID cannot be empty".to_string(),
));
}
self.saw_environment_id = true;
return Ok(());
}
if self.handle_hunk_headers_and_end_patch(trimmed)? {
return Ok(());
}