fix: keep one assignment per file when merging chunk groups

pr_split/diff_ops/reconstructor.py

@@ -124,35 +124,43 @@ def apply_hunks(base_content: str, patch_file: PatchedFile, assigned_indices: li
     return "".join(lines)
 
 
+def _assigned_hunk_indices(
+    patch_file: PatchedFile, assignments: list[GroupAssignment]
+) -> list[int]:
+    """Union of the hunks every assignment for one file claims."""
+    covered: set[int] = set()
+    for assignment in assignments:
+        if assignment.assignment_type is AssignmentType.WHOLE_FILE:
+            covered.update(range(len(patch_file)))
+        else:
+            covered.update(assignment.hunk_indices)
+    return sorted(covered)
+
+
 def materialize_group_files(
     parsed_diff: ParsedDiff, group: Group, ref: str
 ) -> dict[str, str | None]:
-    logger.info(logs.MATERIALIZING_FILES.format(count=len(group.assignments), group=group.id))
     pf_map = {pf.path: pf for pf in parsed_diff.patch_set}
-    result: dict[str, str | None] = {}
+    # Several assignments may name the same file (e.g. merged across diff
+    # chunks); each file is written once from the union of their hunks.
+    assignments_by_path: dict[str, list[GroupAssignment]] = {}
     for assignment in group.assignments:
-        patch_file = pf_map.get(assignment.file_path)
+        assignments_by_path.setdefault(assignment.file_path, []).append(assignment)
+    logger.info(logs.MATERIALIZING_FILES.format(count=len(assignments_by_path), group=group.id))
+    result: dict[str, str | None] = {}
+    for file_path, assignments in assignments_by_path.items():
+        patch_file = pf_map.get(file_path)
         if patch_file is None:
             continue
         if patch_file.is_removed_file:
-            result[assignment.file_path] = None
+            result[file_path] = None
             continue
+        indices = _assigned_hunk_indices(patch_file, assignments)
         if patch_file.is_added_file:
-            all_indices = list(range(len(patch_file)))
-            match assignment.assignment_type:
-                case AssignmentType.WHOLE_FILE:
-                    indices = all_indices
-                case AssignmentType.PARTIAL_HUNKS:
-                    indices = assignment.hunk_indices
-            result[assignment.file_path] = "".join(
+            result[file_path] = "".join(
                 "".join(_hunk_target_lines(patch_file[idx])) for idx in indices
             )
             continue
-        base_content = _get_base_file_content(assignment.file_path, ref)
-        match assignment.assignment_type:
-            case AssignmentType.WHOLE_FILE:
-                indices = list(range(len(patch_file)))
-            case AssignmentType.PARTIAL_HUNKS:
-                indices = assignment.hunk_indices
-        result[assignment.file_path] = apply_hunks(base_content, patch_file, indices)
+        base_content = _get_base_file_content(file_path, ref)
+        result[file_path] = apply_hunks(base_content, patch_file, indices)
     return result

pr_split/planner/client.py

@@ -227,12 +227,38 @@ def _parse_groups(raw: RawToolOutput) -> list[Group]:
     return groups
 
 
+def _merge_assignment(existing: GroupAssignment, incoming: GroupAssignment) -> GroupAssignment:
+    if (
+        existing.assignment_type is AssignmentType.WHOLE_FILE
+        or incoming.assignment_type is AssignmentType.WHOLE_FILE
+    ):
+        assignment_type = AssignmentType.WHOLE_FILE
+    else:
+        assignment_type = AssignmentType.PARTIAL_HUNKS
+    return GroupAssignment(
+        file_path=existing.file_path,
+        assignment_type=assignment_type,
+        hunk_indices=sorted(set(existing.hunk_indices) | set(incoming.hunk_indices)),
+    )
+
+
 def _merge_chunk_groups(accumulated: list[Group], chunk_groups: list[Group]) -> list[Group]:
     acc_map = {g.id: g for g in accumulated}
     for cg in chunk_groups:
         if cg.id in acc_map:
             existing = acc_map[cg.id]
-            existing.assignments.extend(cg.assignments)
+            # A later chunk may assign more hunks of a file this group already
+            # holds. Keep one assignment per path: materialization writes a
+            # file once per assignment, so duplicates would drop hunks.
+            by_path: dict[str, GroupAssignment] = {}
+            for assignment in [*existing.assignments, *cg.assignments]:
+                if assignment.file_path in by_path:
+                    by_path[assignment.file_path] = _merge_assignment(
+                        by_path[assignment.file_path], assignment
+                    )
+                else:
+                    by_path[assignment.file_path] = assignment
+            existing.assignments = list(by_path.values())
             for dep in cg.depends_on:
                 if dep not in existing.depends_on:
                     existing.depends_on.append(dep)

tests/test_client.py

@@ -986,3 +986,105 @@ def test_large_diff_uses_chunking(
         result = _plan_split_with_llm(parsed, settings)
         assert len(result) == 1
         mock_chunked.assert_called_once()
+
+
+class TestMergeChunkGroupsSameFile:
+    def test_same_file_in_later_chunk_is_merged_into_one_assignment(self) -> None:
+        first = Group(
+            id="pr-1",
+            title="t",
+            description="d",
+            assignments=[
+                GroupAssignment(
+                    file_path="f.py",
+                    assignment_type=AssignmentType.PARTIAL_HUNKS,
+                    hunk_indices=[0],
+                )
+            ],
+        )
+        later = Group(
+            id="pr-1",
+            title="t",
+            description="d",
+            assignments=[
+                GroupAssignment(
+                    file_path="f.py",
+                    assignment_type=AssignmentType.PARTIAL_HUNKS,
+                    hunk_indices=[1],
+                )
+            ],
+        )
+        (merged,) = _merge_chunk_groups([first], [later])
+        assert len(merged.assignments) == 1
+        assert merged.assignments[0].hunk_indices == [0, 1]
+        assert merged.assignments[0].assignment_type == AssignmentType.PARTIAL_HUNKS
+
+    def test_whole_file_wins_when_merging(self) -> None:
+        partial = GroupAssignment(
+            file_path="f.py", assignment_type=AssignmentType.PARTIAL_HUNKS, hunk_indices=[0]
+        )
+        whole = GroupAssignment(
+            file_path="f.py", assignment_type=AssignmentType.WHOLE_FILE, hunk_indices=[]
+        )
+        first = Group(id="pr-1", title="t", description="d", assignments=[partial])
+        later = Group(id="pr-1", title="t", description="d", assignments=[whole])
+        (merged,) = _merge_chunk_groups([first], [later])
+        assert len(merged.assignments) == 1
+        assert merged.assignments[0].assignment_type == AssignmentType.WHOLE_FILE
+        assert merged.assignments[0].hunk_indices == [0]
+
+    def test_other_files_untouched(self) -> None:
+        first = Group(
+            id="pr-1",
+            title="t",
+            description="d",
+            assignments=[
+                GroupAssignment(
+                    file_path="a.py",
+                    assignment_type=AssignmentType.PARTIAL_HUNKS,
+                    hunk_indices=[0],
+                )
+            ],
+        )
+        later = Group(
+            id="pr-1",
+            title="t",
+            description="d",
+            assignments=[
+                GroupAssignment(
+                    file_path="b.py",
+                    assignment_type=AssignmentType.PARTIAL_HUNKS,
+                    hunk_indices=[2],
+                )
+            ],
+        )
+        (merged,) = _merge_chunk_groups([first], [later])
+        assert [(a.file_path, a.hunk_indices) for a in merged.assignments] == [
+            ("a.py", [0]),
+            ("b.py", [2]),
+        ]
+
+    def test_pre_existing_duplicates_in_accumulated_group_are_kept(self) -> None:
+        def _pa(path: str, hunks: list[int]) -> GroupAssignment:
+            return GroupAssignment(
+                file_path=path,
+                assignment_type=AssignmentType.PARTIAL_HUNKS,
+                hunk_indices=hunks,
+            )
+
+        first = Group(
+            id="pr-1", title="t", description="d", assignments=[_pa("f.py", [0]), _pa("f.py", [1])]
+        )
+        later = Group(id="pr-1", title="t", description="d", assignments=[_pa("g.py", [0])])
+        (merged,) = _merge_chunk_groups([first], [later])
+        assert [(a.file_path, a.hunk_indices) for a in merged.assignments] == [
+            ("f.py", [0, 1]),
+            ("g.py", [0]),
+        ]
+
+        first = Group(
+            id="pr-1", title="t", description="d", assignments=[_pa("f.py", [0]), _pa("f.py", [1])]
+        )
+        later = Group(id="pr-1", title="t", description="d", assignments=[_pa("f.py", [2])])
+        (merged,) = _merge_chunk_groups([first], [later])
+        assert [(a.file_path, a.hunk_indices) for a in merged.assignments] == [("f.py", [0, 1, 2])]

tests/test_reconstructor.py

@@ -382,6 +382,66 @@ def test_added_file_content_is_not_double_spaced(self) -> None:
         assert result["new_file.py"] == 'def hello():\n    return "world"\n\n'
 
 
+class TestMaterializeDuplicateAssignments:
+    @patch("pr_split.diff_ops.reconstructor._get_base_file_content")
+    def test_two_assignments_for_one_file_apply_both_hunks(self, mock_base: MagicMock) -> None:
+        mock_base.return_value = _base_content()
+        parsed = parse_diff(PATCH_TEXT)
+        group = Group(
+            id="pr-1",
+            title="t",
+            description="d",
+            assignments=[
+                GroupAssignment(
+                    file_path="example.py",
+                    assignment_type=AssignmentType.PARTIAL_HUNKS,
+                    hunk_indices=[0],
+                ),
+                GroupAssignment(
+                    file_path="example.py",
+                    assignment_type=AssignmentType.PARTIAL_HUNKS,
+                    hunk_indices=[1],
+                ),
+            ],
+        )
+        with patch("pr_split.diff_ops.reconstructor.logger.info") as mock_log:
+            result = materialize_group_files(parsed, group, "main")
+        content = result["example.py"]
+        assert content is not None
+        assert "inserted_after_1" in content
+        assert "inserted_after_11" in content
+        assert mock_base.call_count == 1
+        assert "Materializing 1 file" in mock_log.call_args[0][0]
+
+    def test_duplicate_assignments_on_new_file_apply_both_hunks(self) -> None:
+        # Two hunks in a new file (unidiff splits them when the context gap
+        # is large enough), each claimed by a separate PARTIAL assignment.
+        parsed = parse_diff(
+            "diff --git a/n.py b/n.py\nnew file mode 100644\n--- /dev/null\n+++ b/n.py\n"
+            "@@ -0,0 +1,2 @@\n+one\n+two\n"
+            "@@ -0,0 +10,1 @@\n+ten\n"
+        )
+        assert len(parsed.patch_set[0]) == 2
+        group = Group(
+            id="pr-1",
+            title="t",
+            description="d",
+            assignments=[
+                GroupAssignment(
+                    file_path="n.py",
+                    assignment_type=AssignmentType.PARTIAL_HUNKS,
+                    hunk_indices=[0],
+                ),
+                GroupAssignment(
+                    file_path="n.py",
+                    assignment_type=AssignmentType.PARTIAL_HUNKS,
+                    hunk_indices=[1],
+                ),
+            ],
+        )
+        assert materialize_group_files(parsed, group, "main")["n.py"] == "one\ntwo\nten\n"
+
+
 class TestMissingTrailingNewline:
     NO_EOL_PATCH = """\
 --- a/f.txt