fix: retry only failed chunks in large PR reviews

docs/docs/tools/review.md

@@ -273,6 +273,18 @@ for the authoritative default values.
     after chunking are still listed in the coverage footer. Every chunk is a separate model call,
     so a chunked review costs roughly `max_number_of_calls` times a normal one.
 
+    If a chunk fails or returns malformed output, successful chunks are retained and fallback
+    models retry only the pending work. Pending chunks can be split for a smaller model, within
+    the call limit; a larger model can also include previously omitted files. Chunks that still
+    exceed the model's budget are not sent to it.
+
+    If every fallback is exhausted, successful chunks are published as a partial review with a
+    failed-chunk coverage warning, even when `publish_output_no_suggestions = false`. Incomplete
+    reviews cannot resolve absent persistent findings. If no chunk succeeds, the review fails.
+    With `config.propagate_tool_errors = true`, an exhausted fallback chain still signals failure
+    to the caller after publishing the partial review and removing the progress comment.
+    Optional run details list all models that contributed to the merged review.
+
     Each chunk answers the same questions about a different part of the PR, so the answers are
     merged field by field:
 

pr_agent/algo/run_details.py

@@ -31,6 +31,8 @@ class RunDetails:
     # took over. Stays None when no prediction succeeded, which the renderer reads as
     # "nothing worth showing".
     model_used: Optional[str] = None
+    # Retain every model contributing to a merged review; keep scalar output for other tools.
+    models_used: list[str] = field(default_factory=list)
     # Sticky: once a fallback has won, a later success on the primary model must not
     # clear this, or the comment would hide that a fallback ran at all.
     fallback_used: bool = False

pr_agent/algo/utils.py

@@ -1901,7 +1901,10 @@ def show_run_details(gfm_supported: bool) -> str:
         return ""
 
     title = "⚙️ Agent run details"
-    lines = [f"- Model: {details.model_used}{' (fallback)' if details.fallback_used else ''}"]
+    if len(details.models_used) > 1:
+        lines = [f"- Models: {', '.join(details.models_used)}{' (includes fallback)' if details.fallback_used else ''}"]
+    else:
+        lines = [f"- Model: {details.model_used}{' (fallback)' if details.fallback_used else ''}"]
     if details.has_token_usage:
         # A counter still at zero after a successful call means the provider never
         # reported that component, so drop it instead of claiming it was zero.

pr_agent/tools/pr_reviewer.py

@@ -20,6 +20,7 @@
 )
 from pr_agent.algo.output_models import PRReview
 from pr_agent.algo.pr_processing import (
+    OUTPUT_BUFFER_TOKENS_SOFT_THRESHOLD,
     PreparedPRDiff,
     add_ai_metadata_to_diff_files,
     get_pr_diff,
@@ -34,7 +35,7 @@
     reconcile_review_findings,
 )
 from pr_agent.algo.review_merge import merge_review_chunks
-from pr_agent.algo.run_details import get_run_details, init_run_details
+from pr_agent.algo.run_details import get_run_details, init_run_details, record_model_used
 from pr_agent.algo.skills_loader import get_skills_context
 from pr_agent.algo.token_handler import TokenHandler
 from pr_agent.algo.utils import (
@@ -43,6 +44,7 @@
     PRReviewIdentity,
     add_pr_review_identity,
     convert_to_markdown_v2,
+    get_max_tokens,
     get_pr_review_comment_identifiers,
     github_action_output,
     hidden_marker_forms,
@@ -261,7 +263,11 @@ def parse_incremental(self, args: List[str]):
 
     async def run(self) -> None:
         init_run_details()
+        for name in ("_chunked_patches_diff_list", "_chunked_remaining_files_list", "_chunked_results",
+                     "_chunked_primary_model"):
+            self.__dict__.pop(name, None)
         progress_response = None
+        partial_review_error = None
         review_error = None
         review_failed = False
         persistent_write_failed = False
@@ -307,8 +313,14 @@ async def run(self) -> None:
             if get_settings().config.publish_output and not get_settings().config.get('is_auto_command', False):
                 progress_response = self.git_provider.publish_comment("Preparing review...", is_temporary=True)
 
-            await retry_with_fallback_models(self._prepare_prediction, model_type=ModelType.REGULAR,
-                                             git_provider=self.git_provider)
+            try:
+                await retry_with_fallback_models(self._prepare_prediction, model_type=ModelType.REGULAR,
+                                                 git_provider=self.git_provider)
+            except Exception as error:
+                if not self._merge_cached_review_chunks():
+                    raise
+                partial_review_error = error
+                get_logger().warning("Fallback models exhausted; publishing successful review chunks")
             if not self.prediction:
                 return None
 
@@ -325,6 +337,7 @@ async def run(self) -> None:
                 self._should_publish_review_no_suggestions(pr_review)
                 or state_changed
                 or state_blocked
+                or self.review_failed_chunk_count > 0
             )
             if not should_publish:
                 reason = "Review output is not published"
@@ -455,6 +468,9 @@ async def run(self) -> None:
                     self.git_provider.publish_comment(_review_failure_comment(review_error))
                 except Exception as e:
                     get_logger().exception(f"Failed to publish review failure result, error: {e}")
+            if (partial_review_error is not None and not review_failed
+                    and get_settings().config.get("propagate_tool_errors", False)):
+                raise partial_review_error
 
     def _review_finding_state_enabled(self) -> bool:
         settings = get_settings()
@@ -795,8 +811,10 @@ async def _prepare_prediction(self, model: str) -> None:
             self.patches_diff = output
             self.remaining_files_list = []
 
-        # a non-empty remaining_files_list means the token budget truncated the diff
-        if self.remaining_files_list and chunking_enabled:
+        # Resume an incomplete chunk plan even when a fallback model can fit the full diff.
+        # Otherwise the single-call path would bypass cached successful chunks.
+        has_incomplete_chunk_plan = hasattr(self, "_chunked_patches_diff_list")
+        if chunking_enabled and (self.remaining_files_list or has_incomplete_chunk_plan):
             prepared_diff = output if isinstance(output, PreparedPRDiff) else None
             if await self._prepare_chunked_prediction(model, prepared_diff):
                 return
@@ -816,56 +834,153 @@ async def _prepare_chunked_prediction(self, model: str,
 
         Returns False when chunking does not apply, leaving the single-call flow in place.
         """
-        multi_diff_kwargs = {
-            "max_calls": get_settings().pr_reviewer.get("max_number_of_calls", 3),
-            "add_line_numbers": True,
-            "return_remaining_files": True,
-        }
-        output_token_reserve = getattr(
-            getattr(self, "ai_handler", None), "get_output_token_reserve", None
-        )
-        if callable(output_token_reserve):
-            multi_diff_kwargs["output_token_reserve"] = output_token_reserve
-        if prepared_diff is not None:
-            multi_diff_kwargs["prepared_diff"] = prepared_diff
-        patches_diff_list, remaining_files_list = get_pr_multi_diffs(
-            self.git_provider,
-            self.token_handler,
-            model,
-            **multi_diff_kwargs)
+        patches_diff_list = getattr(self, "_chunked_patches_diff_list", None)
+        if patches_diff_list is not None:
+            self._resize_pending_review_chunks(model)
+            patches_diff_list = self._chunked_patches_diff_list
+        if patches_diff_list is not None and self._chunked_remaining_files_list:
+            self._include_newly_reviewable_files(model)
+        if patches_diff_list is None:
+            multi_diff_kwargs = {
+                "max_calls": get_settings().pr_reviewer.get("max_number_of_calls", 3),
+                "add_line_numbers": True,
+                "return_remaining_files": True,
+            }
+            output_token_reserve = getattr(
+                getattr(self, "ai_handler", None), "get_output_token_reserve", None
+            )
+            if callable(output_token_reserve):
+                multi_diff_kwargs["output_token_reserve"] = output_token_reserve
+            if prepared_diff is not None:
+                multi_diff_kwargs["prepared_diff"] = prepared_diff
+            patches_diff_list, remaining_files_list = get_pr_multi_diffs(
+                self.git_provider,
+                self.token_handler,
+                model,
+                **multi_diff_kwargs)
+            self._chunked_patches_diff_list = patches_diff_list
+            self._chunked_remaining_files_list = remaining_files_list
+            self._chunked_primary_model = model
         if len(patches_diff_list) < 2:
             get_logger().info("Large-diff chunking produced a single chunk, reviewing the PR in one call")
             return False
 
         get_logger().info(f"Number of PR chunk calls: {len(patches_diff_list)}")
         get_logger().debug("PR diff chunks", artifact=patches_diff_list)
+        chunk_results = getattr(self, "_chunked_results", {})
+        pending_indices = [index for index in range(len(patches_diff_list)) if index not in chunk_results]
         predictions = await asyncio.gather(
-            *[self._get_prediction(model, patches_diff) for patches_diff in patches_diff_list],
+            *[self._get_prediction(model, patches_diff_list[index]) for index in pending_indices],
             return_exceptions=True)
 
-        raw_predictions, chunk_outputs, chunk_errors = [], [], []
-        for chunk_index, prediction in enumerate(predictions):
+        chunk_errors = []
+        for chunk_index, prediction in zip(pending_indices, predictions, strict=True):
             if isinstance(prediction, Exception):
                 chunk_errors.append(prediction)
                 get_logger().warning(f"Failed to review chunk {chunk_index + 1}; retaining successful chunks",
                                      artifact={"error": prediction})
                 continue
             if isinstance(prediction, BaseException):
                 raise prediction
-            data = self._load_valid_review_yaml(prediction, source=f"review chunk {chunk_index + 1}")
+            try:
+                data = self._load_valid_review_yaml(prediction, source=f"review chunk {chunk_index + 1}")
+            except Exception as error:
+                chunk_errors.append(error)
+                get_logger().warning(f"Failed to parse review chunk {chunk_index + 1}; retrying it with fallback",
+                                     artifact={"error": error})
+                continue
             self._validate_review_schema(data)
-            raw_predictions.append(prediction)
-            chunk_outputs.append(data)
-
-        if not chunk_outputs:
-            raise chunk_errors[0]
+            chunk_results[chunk_index] = (prediction, data, model)
+        self._chunked_results = chunk_results
+
+        if len(chunk_results) < len(patches_diff_list):
+            if chunk_errors:
+                raise chunk_errors[0]
+            raise ValueError("No valid review output was produced for one or more chunks")
+
+        return self._merge_cached_review_chunks()
+
+    def _resize_pending_review_chunks(self, model: str) -> None:
+        """Split oversized pending chunks at file boundaries while preserving result order."""
+        chunks = self._chunked_patches_diff_list
+        results = getattr(self, "_chunked_results", {})
+        budget = get_max_tokens(model) - OUTPUT_BUFFER_TOKENS_SOFT_THRESHOLD - self.token_handler.prompt_tokens
+        max_calls = get_settings().pr_reviewer.get("max_number_of_calls", 3)
+        resized, retained = [], {}
+        for index, chunk in enumerate(chunks):
+            if index in results:
+                retained[len(resized)] = results[index]
+                resized.append(chunk)
+                continue
+            if self.token_handler.count_tokens(chunk) <= budget:
+                resized.append(chunk)
+                continue
+            sections = re.split(r"(?=^## File: ')", chunk, flags=re.MULTILINE)
+            parts, current = [], ""
+            for section in sections:
+                if current and self.token_handler.count_tokens(current + section) > budget:
+                    parts.append(current)
+                    current = ""
+                current += section
+            if current:
+                parts.append(current)
+            reserved = len(chunks) - index - 1
+            if (parts and len(resized) + len(parts) + reserved <= max_calls
+                    and all(self.token_handler.count_tokens(part) <= budget for part in parts)):
+                resized.extend(parts)
+            else:
+                # Retain unsplittable work for a later model and report it as failed if none can fit it.
+                resized.append(chunk)
+        self._chunked_patches_diff_list = resized
+        self._chunked_results = retained
+
+    def _include_newly_reviewable_files(self, model: str) -> None:
+        """Add newly fitting files to pending work without resending successful chunks."""
+        chunks = self._chunked_patches_diff_list
+        results = getattr(self, "_chunked_results", {})
+        remaining = self._chunked_remaining_files_list
+        pending = [index for index in range(len(chunks)) if index not in results]
+        budget = get_max_tokens(model) - OUTPUT_BUFFER_TOKENS_SOFT_THRESHOLD - self.token_handler.prompt_tokens
+        max_calls = get_settings().pr_reviewer.get("max_number_of_calls", 3)
+        included = set()
+        for section in re.split(r"(?=^## File: ')", self.patches_diff or "", flags=re.MULTILINE):
+            match = re.match(r"## File: '(.*)'\n", section)
+            if not match or match[1] not in remaining or match[1] in included:
+                continue
+            for index in pending:
+                combined = chunks[index] + "\n\n" + section
+                if self.token_handler.count_tokens(combined) <= budget:
+                    chunks[index] = combined
+                    included.add(match[1])
+                    break
+            else:
+                if len(chunks) < max_calls and self.token_handler.count_tokens(section) <= budget:
+                    pending.append(len(chunks))
+                    chunks.append(section)
+                    included.add(match[1])
+        self._chunked_remaining_files_list = [name for name in remaining if name not in included]
+
+    def _merge_cached_review_chunks(self) -> bool:
+        """Merge successful chunks in order, retaining incomplete coverage after exhausted retries."""
+        chunk_results = getattr(self, "_chunked_results", {})
+        if not chunk_results:
+            return False
 
-        # the raw text is kept for logging only; the merged verdict is in self.prediction_data
+        # Keep raw text for logging only; use the merged verdict from self.prediction_data.
+        indices = sorted(chunk_results)
+        raw_predictions = [chunk_results[index][0] for index in indices]
+        chunk_outputs = [chunk_results[index][1] for index in indices]
         self.prediction = "\n".join(raw_predictions)
         self.prediction_data = merge_review_chunks(chunk_outputs)
-        self.review_chunk_count = len(patches_diff_list)
-        self.review_failed_chunk_count = len(patches_diff_list) - len(chunk_outputs)
-        self.remaining_files_list = remaining_files_list
+        self.review_chunk_count = len(self._chunked_patches_diff_list)
+        self.review_failed_chunk_count = self.review_chunk_count - len(chunk_results)
+        self.remaining_files_list = self._chunked_remaining_files_list
+        models = list(dict.fromkeys(chunk_results[index][2] for index in indices))
+        details = get_run_details()
+        if details is not None:
+            details.models_used = models
+            for model in models:
+                record_model_used(model, is_fallback=model != self._chunked_primary_model)
         return True
 
     async def _get_prediction(self, model: str, patches_diff: Optional[str] = None) -> str:
@@ -880,6 +995,10 @@ async def _get_prediction(self, model: str, patches_diff: Optional[str] = None)
         Returns:
             A string representing the AI prediction for the pull request review.
         """
+        if patches_diff is not None:
+            budget = get_max_tokens(model) - OUTPUT_BUFFER_TOKENS_SOFT_THRESHOLD - self.token_handler.prompt_tokens
+            if self.token_handler.count_tokens(patches_diff) > budget:
+                raise ValueError("Review chunk exceeds the current model token budget")
         variables = copy.deepcopy(self.vars)
         variables["diff"] = self.patches_diff if patches_diff is None else patches_diff  # update diff