[apply-patch] Handle multiple context lines

Author: dylan-hurd-oai

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

@@ -529,32 +529,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;
@@ -1267,6 +1286,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();

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

@@ -69,7 +69,7 @@ pub enum Hunk {
         path: PathBuf,
         move_path: Option<PathBuf>,
 
-        /// 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<UpdateFileChunk>,
     },
@@ -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<String>,
+    /// 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<String>,
 
     /// 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<String>,
     pub new_lines: Vec<String>,
 
@@ -344,32 +345,38 @@ fn parse_update_file_chunk(
             line_number,
         });
     }
-    // If we see an explicit context marker @@ or @@ <context>, 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,12 +730,37 @@ 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
             }),
             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
+        ))
+    );
 }

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": []
+    }
+  }
+]

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(())
+}