# PR #2596: [apply-patch] Handle multiple context lines - URL: https://github.com/openai/codex/pull/2596 - Author: dylan-hurd-oai - Created: 2025-08-22 22:45:19 UTC - Updated: 2025-08-25 23:36:20 UTC - Changes: +218/-53, Files changed: 4, Commits: 1 ## Description ## Summary Fixes #2578 - the following example would be considered incorrect, despite it being valid according to our system instructions and our official Lark grammar. ``` *** Begin Patch *** Update File: src/app.py @@ class BaseClass @@ def method(): - old_line + new_line *** End Patch ``` ## Testing - [x] Added unit tests ## Full Diff ```diff diff --git a/codex-rs/apply-patch/src/lib.rs b/codex-rs/apply-patch/src/lib.rs index 84cb91201f..1f9301b121 100644 --- a/codex-rs/apply-patch/src/lib.rs +++ b/codex-rs/apply-patch/src/lib.rs @@ -532,32 +532,51 @@ fn compute_replacements( let mut line_index: usize = 0; for chunk in chunks { - // If a chunk has a `change_context`, we use seek_sequence to find it, then - // adjust our `line_index` to continue from there. - if let Some(ctx_line) = &chunk.change_context { - if let Some(idx) = seek_sequence::seek_sequence( - original_lines, - std::slice::from_ref(ctx_line), - line_index, - false, - ) { - line_index = idx + 1; - } else { - return Err(ApplyPatchError::ComputeReplacements(format!( - "Failed to find context '{}' in {}", - ctx_line, - path.display() - ))); + // If a chunk has context lines, we use seek_sequence to find each in order, + // then adjust our `line_index` to continue from there. + if !chunk.context_lines.is_empty() { + let total = chunk.context_lines.len(); + for (i, ctx_line) in chunk.context_lines.iter().enumerate() { + if let Some(idx) = seek_sequence::seek_sequence( + original_lines, + std::slice::from_ref(ctx_line), + line_index, + false, + ) { + line_index = idx + 1; + } else { + return Err(ApplyPatchError::ComputeReplacements(format!( + "Failed to find context {}/{}: '{}' in {}", + i + 1, + total, + ctx_line, + path.display() + ))); + } } } if chunk.old_lines.is_empty() { - // Pure addition (no old lines). We'll add them at the end or just - // before the final empty line if one exists. - let insertion_idx = if original_lines.last().is_some_and(|s| s.is_empty()) { - original_lines.len() - 1 + // Pure addition (no old lines). + // Prefer to insert at the matched context anchor if one exists and + // the hunk is not explicitly marked as end-of-file. + let insertion_idx = if chunk.is_end_of_file { + if original_lines.last().is_some_and(|s| s.is_empty()) { + original_lines.len() - 1 + } else { + original_lines.len() + } + } else if !chunk.context_lines.is_empty() { + // Insert immediately after the last matched context line. + line_index } else { - original_lines.len() + // No context provided: fall back to appending at the end (before + // the trailing empty line if present). + if original_lines.last().is_some_and(|s| s.is_empty()) { + original_lines.len() - 1 + } else { + original_lines.len() + } }; replacements.push((insertion_idx, 0, chunk.new_lines.clone())); continue; @@ -1270,6 +1289,57 @@ g ); } + #[test] + fn test_insert_addition_after_single_context_anchor() { + let dir = tempdir().unwrap(); + let path = dir.path().join("single_ctx.txt"); + fs::write(&path, "class BaseClass:\n def method():\nline1\nline2\n").unwrap(); + + let patch = wrap_patch(&format!( + r#"*** Update File: {} +@@ class BaseClass: ++INSERTED +"#, + path.display() + )); + + let mut stdout = Vec::new(); + let mut stderr = Vec::new(); + apply_patch(&patch, &mut stdout, &mut stderr).unwrap(); + + let contents = fs::read_to_string(path).unwrap(); + assert_eq!( + contents, + "class BaseClass:\nINSERTED\n def method():\nline1\nline2\n" + ); + } + + #[test] + fn test_insert_addition_after_multi_context_anchor() { + let dir = tempdir().unwrap(); + let path = dir.path().join("multi_ctx.txt"); + fs::write(&path, "class BaseClass:\n def method():\nline1\nline2\n").unwrap(); + + let patch = wrap_patch(&format!( + r#"*** Update File: {} +@@ class BaseClass: +@@ def method(): ++INSERTED +"#, + path.display() + )); + + let mut stdout = Vec::new(); + let mut stderr = Vec::new(); + apply_patch(&patch, &mut stdout, &mut stderr).unwrap(); + + let contents = fs::read_to_string(path).unwrap(); + assert_eq!( + contents, + "class BaseClass:\n def method():\nINSERTED\nline1\nline2\n" + ); + } + #[test] fn test_apply_patch_should_resolve_absolute_paths_in_cwd() { let session_dir = tempdir().unwrap(); diff --git a/codex-rs/apply-patch/src/parser.rs b/codex-rs/apply-patch/src/parser.rs index ff9dfd6f8f..1e2aeb946d 100644 --- a/codex-rs/apply-patch/src/parser.rs +++ b/codex-rs/apply-patch/src/parser.rs @@ -69,7 +69,7 @@ pub enum Hunk { path: PathBuf, move_path: Option, - /// Chunks should be in order, i.e. the `change_context` of one chunk + /// Chunks should be in order, i.e. the first context line of one chunk /// should occur later in the file than the previous chunk. chunks: Vec, }, @@ -89,12 +89,13 @@ use Hunk::*; #[derive(Debug, PartialEq, Clone)] pub struct UpdateFileChunk { - /// A single line of context used to narrow down the position of the chunk - /// (this is usually a class, method, or function definition.) - pub change_context: Option, + /// Context lines used to narrow down the position of the chunk. + /// Each entry is searched sequentially to progressively restrict the + /// search to the desired region (e.g. class → method). + pub context_lines: Vec, /// A contiguous block of lines that should be replaced with `new_lines`. - /// `old_lines` must occur strictly after `change_context`. + /// `old_lines` must occur strictly after the context. pub old_lines: Vec, pub new_lines: Vec, @@ -344,32 +345,38 @@ fn parse_update_file_chunk( line_number, }); } - // If we see an explicit context marker @@ or @@ , consume it; otherwise, optionally - // allow treating the chunk as starting directly with diff lines. - let (change_context, start_index) = if lines[0] == EMPTY_CHANGE_CONTEXT_MARKER { - (None, 1) - } else if let Some(context) = lines[0].strip_prefix(CHANGE_CONTEXT_MARKER) { - (Some(context.to_string()), 1) - } else { - if !allow_missing_context { - return Err(InvalidHunkError { - message: format!( - "Expected update hunk to start with a @@ context marker, got: '{}'", - lines[0] - ), - line_number, - }); + let mut context_lines = Vec::new(); + let mut start_index = 0; + let mut saw_context_marker = false; + while start_index < lines.len() { + if lines[start_index] == EMPTY_CHANGE_CONTEXT_MARKER { + saw_context_marker = true; + start_index += 1; + } else if let Some(context) = lines[start_index].strip_prefix(CHANGE_CONTEXT_MARKER) { + saw_context_marker = true; + context_lines.push(context.to_string()); + start_index += 1; + } else { + break; } - (None, 0) - }; + } + if !saw_context_marker && !allow_missing_context { + return Err(InvalidHunkError { + message: format!( + "Expected update hunk to start with a @@ context marker, got: '{}'", + lines[0] + ), + line_number, + }); + } if start_index >= lines.len() { return Err(InvalidHunkError { message: "Update hunk does not contain any lines".to_string(), - line_number: line_number + 1, + line_number: line_number + start_index, }); } let mut chunk = UpdateFileChunk { - change_context, + context_lines, old_lines: Vec::new(), new_lines: Vec::new(), is_end_of_file: false, @@ -381,7 +388,7 @@ fn parse_update_file_chunk( if parsed_lines == 0 { return Err(InvalidHunkError { message: "Update hunk does not contain any lines".to_string(), - line_number: line_number + 1, + line_number: line_number + start_index, }); } chunk.is_end_of_file = true; @@ -411,7 +418,7 @@ fn parse_update_file_chunk( message: format!( "Unexpected line found in update hunk: '{line_contents}'. Every line should start with ' ' (context line), '+' (added line), or '-' (removed line)" ), - line_number: line_number + 1, + line_number: line_number + start_index, }); } // Assume this is the start of the next hunk. @@ -491,7 +498,7 @@ fn test_parse_patch() { path: PathBuf::from("path/update.py"), move_path: Some(PathBuf::from("path/update2.py")), chunks: vec![UpdateFileChunk { - change_context: Some("def f():".to_string()), + context_lines: vec!["def f():".to_string()], old_lines: vec![" pass".to_string()], new_lines: vec![" return 123".to_string()], is_end_of_file: false @@ -518,7 +525,7 @@ fn test_parse_patch() { path: PathBuf::from("file.py"), move_path: None, chunks: vec![UpdateFileChunk { - change_context: None, + context_lines: Vec::new(), old_lines: vec![], new_lines: vec!["line".to_string()], is_end_of_file: false @@ -548,7 +555,7 @@ fn test_parse_patch() { path: PathBuf::from("file2.py"), move_path: None, chunks: vec![UpdateFileChunk { - change_context: None, + context_lines: Vec::new(), old_lines: vec!["import foo".to_string()], new_lines: vec!["import foo".to_string(), "bar".to_string()], is_end_of_file: false, @@ -568,7 +575,7 @@ fn test_parse_patch_lenient() { path: PathBuf::from("file2.py"), move_path: None, chunks: vec![UpdateFileChunk { - change_context: None, + context_lines: Vec::new(), old_lines: vec!["import foo".to_string()], new_lines: vec!["import foo".to_string(), "bar".to_string()], is_end_of_file: false, @@ -701,7 +708,7 @@ fn test_update_file_chunk() { ), Ok(( (UpdateFileChunk { - change_context: Some("change_context".to_string()), + context_lines: vec!["change_context".to_string()], old_lines: vec![ "".to_string(), "context".to_string(), @@ -723,7 +730,7 @@ fn test_update_file_chunk() { parse_update_file_chunk(&["@@", "+line", "*** End of File"], 123, false), Ok(( (UpdateFileChunk { - change_context: None, + context_lines: Vec::new(), old_lines: vec![], new_lines: vec!["line".to_string()], is_end_of_file: true @@ -731,4 +738,29 @@ fn test_update_file_chunk() { 3 )) ); + assert_eq!( + parse_update_file_chunk( + &[ + "@@ class BaseClass", + "@@ def method()", + " context", + "-old", + "+new", + ], + 123, + false + ), + Ok(( + (UpdateFileChunk { + context_lines: vec![ + "class BaseClass".to_string(), + " def method()".to_string() + ], + old_lines: vec!["context".to_string(), "old".to_string()], + new_lines: vec!["context".to_string(), "new".to_string()], + is_end_of_file: false + }), + 5 + )) + ); } diff --git a/codex-rs/exec/tests/fixtures/sse_apply_patch_context_update.json b/codex-rs/exec/tests/fixtures/sse_apply_patch_context_update.json new file mode 100644 index 0000000000..886fe64686 --- /dev/null +++ b/codex-rs/exec/tests/fixtures/sse_apply_patch_context_update.json @@ -0,0 +1,25 @@ +[ + { + "type": "response.output_item.done", + "item": { + "type": "custom_tool_call", + "name": "apply_patch", + "input": "*** Begin Patch\n*** Update File: app.py\n@@ class BaseClass:\n@@ def method():\n- return False\n+ return True\n*** End Patch", + "call_id": "__ID__" + } + }, + { + "type": "response.completed", + "response": { + "id": "__ID__", + "usage": { + "input_tokens": 0, + "input_tokens_details": null, + "output_tokens": 0, + "output_tokens_details": null, + "total_tokens": 0 + }, + "output": [] + } + } +] \ No newline at end of file diff --git a/codex-rs/exec/tests/suite/apply_patch.rs b/codex-rs/exec/tests/suite/apply_patch.rs index 5537853b02..a02f757f75 100644 --- a/codex-rs/exec/tests/suite/apply_patch.rs +++ b/codex-rs/exec/tests/suite/apply_patch.rs @@ -106,3 +106,41 @@ async fn test_apply_patch_freeform_tool() -> anyhow::Result<()> { ); Ok(()) } + +#[cfg(not(target_os = "windows"))] +#[tokio::test] +async fn test_apply_patch_context() -> anyhow::Result<()> { + use crate::suite::common::run_e2e_exec_test; + use codex_core::spawn::CODEX_SANDBOX_NETWORK_DISABLED_ENV_VAR; + + if std::env::var(CODEX_SANDBOX_NETWORK_DISABLED_ENV_VAR).is_ok() { + println!( + "Skipping test because it cannot execute when network is disabled in a Codex sandbox." + ); + return Ok(()); + } + + let tmp_cwd = tempdir().expect("failed to create temp dir"); + run_e2e_exec_test( + tmp_cwd.path(), + vec![ + include_str!("../fixtures/sse_apply_patch_freeform_add.json").to_string(), + include_str!("../fixtures/sse_apply_patch_context_update.json").to_string(), + include_str!("../fixtures/sse_response_completed.json").to_string(), + ], + ) + .await; + + // Verify final file contents + let final_path = tmp_cwd.path().join("app.py"); + let contents = std::fs::read_to_string(&final_path) + .unwrap_or_else(|e| panic!("failed reading {}: {e}", final_path.display())); + assert_eq!( + contents, + r#"class BaseClass: + def method(): + return True +"# + ); + Ok(()) +} ``` ## Review Comments ### codex-rs/apply-patch/src/lib.rs - Created: 2025-08-25 23:32:52 UTC | Link: https://github.com/openai/codex/pull/2596#discussion_r2299356487 ```diff @@ -532,32 +532,51 @@ fn compute_replacements( let mut line_index: usize = 0; for chunk in chunks { - // If a chunk has a `change_context`, we use seek_sequence to find it, then - // adjust our `line_index` to continue from there. - if let Some(ctx_line) = &chunk.change_context { - if let Some(idx) = seek_sequence::seek_sequence( - original_lines, - std::slice::from_ref(ctx_line), - line_index, - false, - ) { - line_index = idx + 1; - } else { - return Err(ApplyPatchError::ComputeReplacements(format!( - "Failed to find context '{}' in {}", - ctx_line, - path.display() - ))); + // If a chunk has context lines, we use seek_sequence to find each in order, + // then adjust our `line_index` to continue from there. + if !chunk.context_lines.is_empty() { + let total = chunk.context_lines.len(); + for (i, ctx_line) in chunk.context_lines.iter().enumerate() { + if let Some(idx) = seek_sequence::seek_sequence( + original_lines, + std::slice::from_ref(ctx_line), + line_index, + false, + ) { + line_index = idx + 1; + } else { + return Err(ApplyPatchError::ComputeReplacements(format!( + "Failed to find context {}/{}: '{}' in {}", + i + 1, + total, + ctx_line, + path.display() + ))); + } } } if chunk.old_lines.is_empty() { - // Pure addition (no old lines). We'll add them at the end or just - // before the final empty line if one exists. - let insertion_idx = if original_lines.last().is_some_and(|s| s.is_empty()) { - original_lines.len() - 1 + // Pure addition (no old lines). + // Prefer to insert at the matched context anchor if one exists and + // the hunk is not explicitly marked as end-of-file. + let insertion_idx = if chunk.is_end_of_file { + if original_lines.last().is_some_and(|s| s.is_empty()) { + original_lines.len() - 1 + } else { + original_lines.len() + } + } else if !chunk.context_lines.is_empty() { + // Insert immediately after the last matched context line. + line_index ``` > Is this correct? > > ```suggestion > let insertion_idx = if !chunk.is_end_of_file && !chunk.context_lines.is_empty() { > // Insert immediately after the last matched context line. > line_index > } else if original_lines.last().is_some_and(|s| s.is_empty()) { > original_lines.len() - 1 > } else { > original_lines.len() > }; > ``` ### codex-rs/exec/tests/fixtures/sse_apply_patch_context_update.json - Created: 2025-08-25 23:35:39 UTC | Link: https://github.com/openai/codex/pull/2596#discussion_r2299359432 ```diff @@ -0,0 +1,25 @@ +[ + { + "type": "response.output_item.done", + "item": { + "type": "custom_tool_call", + "name": "apply_patch", + "input": "*** Begin Patch\n*** Update File: app.py\n@@ class BaseClass:\n@@ def method():\n- return False\n+ return True\n*** End Patch", + "call_id": "__ID__" + } + }, + { + "type": "response.completed", + "response": { + "id": "__ID__", + "usage": { + "input_tokens": 0, + "input_tokens_details": null, + "output_tokens": 0, + "output_tokens_details": null, + "total_tokens": 0 + }, + "output": [] + } + } +] ``` > newline at eof? ### codex-rs/exec/tests/suite/apply_patch.rs - Created: 2025-08-25 23:35:53 UTC | Link: https://github.com/openai/codex/pull/2596#discussion_r2299359662 ```diff @@ -106,3 +106,41 @@ async fn test_apply_patch_freeform_tool() -> anyhow::Result<()> { ); Ok(()) } + +#[cfg(not(target_os = "windows"))] ``` > Are you sure these don't work on Windows?