fix: keep the current plan when LLM refinement drops hunks or does not reduce violations

Author: vitali87

pr_split/logs.py

@@ -59,6 +59,14 @@
 REFINEMENT_EXHAUSTED = (
     "Refinement iteration limit reached ({iterations}), {remaining} violation(s) remain"
 )
+REFINEMENT_REJECTED = (
+    "Refinement iteration {iteration} produced an invalid plan ({reason}); "
+    "keeping the current plan, {remaining} violation(s) remain"
+)
+REFINEMENT_NO_IMPROVEMENT = (
+    "Refinement iteration {iteration} did not reduce violations ({before} -> {after}); "
+    "keeping the current plan"
+)
 STACK_LINKED = "Linked stack for PRs {prs}"
 STACK_LINK_FAILED = "Could not link stack for PRs {prs}: {detail}"
 MERGE_NODE_NOT_STACKED = (

pr_split/planner/client.py

@@ -20,7 +20,7 @@
     Provider,
 )
 from ..diff_ops import ParsedDiff
-from ..exceptions import ErrorMsg, LLMError, PRSplitError
+from ..exceptions import ErrorMsg, LLMError, PlanValidationError, PRSplitError
 from ..schemas import Group, GroupAssignment
 from .chunker import (
     assign_uncovered_hunks,
@@ -42,7 +42,7 @@
     build_user_prompt,
 )
 from .scoring import score_plan
-from .validator import detect_loc_bound_violations
+from .validator import detect_loc_bound_violations, validate_coverage
 
 _ANTHROPIC_TOOL_DEF = anthropic.types.ToolParam(
     name=SPLIT_TOOL_NAME,
@@ -379,6 +379,29 @@ def _refine_plan_with_llm(
             raw = _call_llm(system=system, user=user, settings=settings)
             refined = _parse_groups(raw)
             recompute_estimated_loc(refined, parsed_diff)
+            # The incoming plan only had LOC warnings; never trade it for one
+            # that drops or duplicates hunks, or that is no better.
+            try:
+                validate_coverage(refined, parsed_diff)
+            except PlanValidationError as exc:
+                logger.warning(
+                    logs.REFINEMENT_REJECTED.format(
+                        iteration=iteration, reason=exc, remaining=len(violations)
+                    )
+                )
+                return groups
+            refined_violations = detect_loc_bound_violations(
+                refined, settings.max_loc, settings.min_loc
+            )
+            if len(refined_violations) >= len(violations):
+                logger.warning(
+                    logs.REFINEMENT_NO_IMPROVEMENT.format(
+                        iteration=iteration,
+                        before=len(violations),
+                        after=len(refined_violations),
+                    )
+                )
+                return groups
             groups = refined
         except LLMError:
             logger.warning(

tests/test_client.py

@@ -533,6 +533,73 @@ def test_refinement_falls_back_on_malformed_response(
         assert result == groups
         mock_call_llm.assert_called_once()
 
+    @patch("pr_split.planner.client._call_llm")
+    def test_refinement_that_drops_hunks_is_rejected(
+        self, mock_call_llm: MagicMock, monkeypatch: pytest.MonkeyPatch
+    ) -> None:
+        monkeypatch.delenv("ANTHROPIC_API_KEY", raising=False)
+        monkeypatch.delenv("OPENAI_API_KEY", raising=False)
+        parsed = parse_diff(_TWO_FILE_DIFF)
+        settings = Settings(
+            partition_strategy=PartitionStrategy.GRAPH,
+            min_loc=5,
+            max_loc=10,
+            max_refinement_iterations=2,
+        )
+        # Only a.py survives: b.py's hunk is gone from the plan.
+        mock_call_llm.return_value = RawToolOutput(
+            groups=[
+                {
+                    "id": "pr-1",
+                    "title": "feat: add a",
+                    "description": "Only a",
+                    "depends_on": [],
+                    "assignments": [
+                        {"file_path": "a.py", "assignment_type": "whole_file", "hunk_indices": [0]}
+                    ],
+                    "estimated_loc": 3,
+                }
+            ]
+        )
+
+        groups = _undersized_groups()
+        with patch("pr_split.planner.client.logger") as mock_logger:
+            result = _refine_plan_with_llm(groups, parsed, settings, system="system")
+
+        assert result is groups
+        assert [g.id for g in result] == ["pr-1", "pr-2"]
+        mock_call_llm.assert_called_once()
+        warning = mock_logger.warning.call_args[0][0]
+        assert "produced an invalid plan" in warning
+        assert "b.py[0] not assigned to any group" in warning
+
+    @patch("pr_split.planner.client._call_llm")
+    def test_refinement_that_does_not_improve_is_rejected(
+        self, mock_call_llm: MagicMock, monkeypatch: pytest.MonkeyPatch
+    ) -> None:
+        monkeypatch.delenv("ANTHROPIC_API_KEY", raising=False)
+        monkeypatch.delenv("OPENAI_API_KEY", raising=False)
+        parsed = parse_diff(_TWO_FILE_DIFF)
+        settings = Settings(
+            partition_strategy=PartitionStrategy.GRAPH,
+            min_loc=5,
+            max_loc=10,
+            max_refinement_iterations=3,
+        )
+        # Same two undersized groups handed straight back.
+        mock_call_llm.return_value = RawToolOutput(
+            groups=_groups_to_raw_dicts(_undersized_groups())  # type: ignore[typeddict-item]
+        )
+
+        groups = _undersized_groups()
+        with patch("pr_split.planner.client.logger") as mock_logger:
+            result = _refine_plan_with_llm(groups, parsed, settings, system="system")
+
+        assert result is groups
+        mock_call_llm.assert_called_once()
+        warning = mock_logger.warning.call_args[0][0]
+        assert "did not reduce violations (2 -> 2)" in warning
+
     def test_no_refinement_when_min_loc_is_none(self, monkeypatch: pytest.MonkeyPatch) -> None:
         monkeypatch.delenv("ANTHROPIC_API_KEY", raising=False)
         monkeypatch.delenv("OPENAI_API_KEY", raising=False)