fix(apply-patch) Support nested @@ context lines

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

@@ -780,9 +780,12 @@ 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 a chunk has one or more `change_context` lines, seek them in order
+        // to progressively narrow down the position of the chunk. This supports
+        // multiple @@ context headers such as:
+        // @@ class BaseClass:
+        // @@     def method():
+        for ctx_line in &chunk.change_context {
             if let Some(idx) = seek_sequence::seek_sequence(
                 original_lines,
                 std::slice::from_ref(ctx_line),
@@ -1545,6 +1548,34 @@ PATCH"#,
         assert_eq!(String::from_utf8(stderr).unwrap(), "");
     }
 
+    #[test]
+    fn test_apply_patch_multiple_change_contexts_success() {
+        let dir = tempdir().unwrap();
+        let path = dir.path().join("example.py");
+        let original =
+            include_str!("../tests/fixtures/scenarios/019_multiple_context_lines/input/example.py");
+        fs::write(&path, original).unwrap();
+
+        let patch = wrap_patch(&format!(
+            r#"*** Update File: {}
+@@ class BaseClass:
+@@     def method():
+-        # to_remove
++        # to_add"#,
+            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();
+        let expected = include_str!(
+            "../tests/fixtures/scenarios/019_multiple_context_lines/expected/example.py"
+        );
+        assert_eq!(contents, expected);
+    }
+
     #[test]
     fn test_unified_diff() {
         // Start with a file containing four lines.
@@ -1907,4 +1938,33 @@ g
         let result = apply_patch(&patch, &mut stdout, &mut stderr);
         assert!(result.is_err());
     }
+
+    #[test]
+    fn test_apply_patch_multiple_change_contexts_missing_context() {
+        let dir = tempdir().unwrap();
+        let path = dir.path().join("example.py");
+        let original =
+            include_str!("../tests/fixtures/scenarios/019_multiple_context_lines/input/example.py");
+        fs::write(&path, original).unwrap();
+
+        let patch = wrap_patch(&format!(
+            r#"*** Update File: {}
+@@ class BaseClass:
+@@     def missing():
+-        # to_remove
++        # to_add"#,
+            path.display()
+        ));
+
+        let mut stdout = Vec::new();
+        let mut stderr = Vec::new();
+        let result = apply_patch(&patch, &mut stdout, &mut stderr);
+
+        assert_matches!(
+            result,
+            Err(ApplyPatchError::IoError(IoError { context, .. }))
+                if context.contains("Failed to find context '    def missing():'")
+                    && context.contains(&path.display().to_string())
+        );
+    }
 }

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

@@ -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>,
+    /// Lines of context used to narrow down the position of the chunk.
+    /// These are usually class, method, or function definitions. If empty,
+    /// the chunk has no explicit context.
+    pub change_context: 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 last `change_context` line.
     pub old_lines: Vec<String>,
     pub new_lines: Vec<String>,
 
@@ -348,28 +349,42 @@ 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,
-            });
+    // Consume one or more explicit context markers (`@@` or `@@ <context>`). This
+    // supports multiple `@@` lines to narrow down the target location, e.g.:
+    // @@ class BaseClass:
+    // @@     def method():
+    // -old
+    // +new
+    let mut change_context: Vec<String> = Vec::new();
+    let mut start_index = 0;
+    while start_index < lines.len() {
+        let line = lines[start_index];
+        if line == EMPTY_CHANGE_CONTEXT_MARKER {
+            start_index += 1;
+            continue;
         }
-        (None, 0)
-    };
+        if let Some(context) = line.strip_prefix(CHANGE_CONTEXT_MARKER) {
+            change_context.push(context.to_string());
+            start_index += 1;
+            continue;
+        }
+        break;
+    }
+
+    if start_index == 0 && !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 {
@@ -495,7 +510,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()),
+                    change_context: vec!["def f():".to_string()],
                     old_lines: vec!["    pass".to_string()],
                     new_lines: vec!["    return 123".to_string()],
                     is_end_of_file: false
@@ -522,7 +537,7 @@ fn test_parse_patch() {
                 path: PathBuf::from("file.py"),
                 move_path: None,
                 chunks: vec![UpdateFileChunk {
-                    change_context: None,
+                    change_context: Vec::new(),
                     old_lines: vec![],
                     new_lines: vec!["line".to_string()],
                     is_end_of_file: false
@@ -552,7 +567,7 @@ fn test_parse_patch() {
             path: PathBuf::from("file2.py"),
             move_path: None,
             chunks: vec![UpdateFileChunk {
-                change_context: None,
+                change_context: Vec::new(),
                 old_lines: vec!["import foo".to_string()],
                 new_lines: vec!["import foo".to_string(), "bar".to_string()],
                 is_end_of_file: false,
@@ -572,7 +587,7 @@ fn test_parse_patch_lenient() {
         path: PathBuf::from("file2.py"),
         move_path: None,
         chunks: vec![UpdateFileChunk {
-            change_context: None,
+            change_context: Vec::new(),
             old_lines: vec!["import foo".to_string()],
             new_lines: vec!["import foo".to_string(), "bar".to_string()],
             is_end_of_file: false,
@@ -708,7 +723,7 @@ fn test_update_file_chunk() {
         ),
         Ok((
             (UpdateFileChunk {
-                change_context: Some("change_context".to_string()),
+                change_context: vec!["change_context".to_string()],
                 old_lines: vec![
                     "".to_string(),
                     "context".to_string(),
@@ -730,7 +745,7 @@ fn test_update_file_chunk() {
         parse_update_file_chunk(&["@@", "+line", "*** End of File"], 123, false),
         Ok((
             (UpdateFileChunk {
-                change_context: None,
+                change_context: Vec::new(),
                 old_lines: vec![],
                 new_lines: vec!["line".to_string()],
                 is_end_of_file: true
@@ -739,3 +754,31 @@ fn test_update_file_chunk() {
         ))
     );
 }
+
+#[test]
+fn test_update_file_chunk_multiple_change_context_lines() {
+    assert_eq!(
+        parse_update_file_chunk(
+            &[
+                "@@ class BaseClass:",
+                "@@     def method():",
+                "-        # to_remove",
+                "+        # to_add",
+            ],
+            200,
+            false
+        ),
+        Ok((
+            (UpdateFileChunk {
+                change_context: vec![
+                    "class BaseClass:".to_string(),
+                    "    def method():".to_string()
+                ],
+                old_lines: vec!["        # to_remove".to_string()],
+                new_lines: vec!["        # to_add".to_string()],
+                is_end_of_file: false
+            }),
+            4
+        ))
+    );
+}

codex-rs/apply-patch/tests/fixtures/scenarios/019_multiple_context_lines/expected/example.py

@@ -0,0 +1,10 @@
+class BaseClass:
+    def method():
+        # to_add
+        pass
+
+
+class OtherClass:
+    def method():
+        # untouched
+        pass

codex-rs/apply-patch/tests/fixtures/scenarios/019_multiple_context_lines/input/example.py

@@ -0,0 +1,11 @@
+class BaseClass:
+    def method():
+        # to_remove
+        pass
+
+
+class OtherClass:
+    def method():
+        # untouched
+        pass
+

codex-rs/apply-patch/tests/fixtures/scenarios/019_multiple_context_lines/patch.txt

@@ -0,0 +1,7 @@
+*** Begin Patch
+*** Update File: example.py
+@@ class BaseClass:
+@@     def method():
+-        # to_remove
++        # to_add
+*** End Patch

codex-rs/apply-patch/tests/fixtures/scenarios/020_multiple_context_lines_partial_success/expected/failure.py

@@ -0,0 +1,4 @@
+class Foo:
+    def foo():
+        # to_remove
+        pass

codex-rs/apply-patch/tests/fixtures/scenarios/020_multiple_context_lines_partial_success/expected/success.py

@@ -0,0 +1,10 @@
+class BaseClass:
+    def method():
+        # to_add
+        pass
+
+
+class OtherClass:
+    def method():
+        # untouched
+        pass

codex-rs/apply-patch/tests/fixtures/scenarios/020_multiple_context_lines_partial_success/input/failure.py

@@ -0,0 +1,4 @@
+class Foo:
+    def foo():
+        # to_remove
+        pass

codex-rs/apply-patch/tests/fixtures/scenarios/020_multiple_context_lines_partial_success/input/success.py

@@ -0,0 +1,10 @@
+class BaseClass:
+    def method():
+        # to_add
+        pass
+
+
+class OtherClass:
+    def method():
+        # untouched
+        pass

codex-rs/apply-patch/tests/fixtures/scenarios/020_multiple_context_lines_partial_success/patch.txt

@@ -0,0 +1,12 @@
+*** Begin Patch
+*** Update File: success.py
+@@ class BaseClass:
+@@     def method():
+-        # to_remove
++        # to_add
+*** Update File: failure.py
+@@ class Foo:
+@@     def missing():
+-        # to_remove
++        # to_add
+*** End Patch