Check the first diff batch before loading every AI reviewable file

What does this MR do and why?

Duo Code Review's add_comments request loaded every file's patch just to check that one file is reviewable. It now checks the first 10 files and only falls back to the full load when none of them qualify, behind the duo_code_review_check_first_diff_batch flag. On GDK, a 1000-file MR went from ~60 MB to ~15-20 MB mem_total_bytes per request.

Detailed context for AI agents

Problem

POST /api/:version/ai/duo_workflows/code_review/add_comments runs on the ai-assisted Puma fleet, whose pods were being OOM-killed. Two earlier MRs cut the tail: the patch-only reviewable check (flag enabled 2026-10-02) and skipping the highlight cache (!259617 (merged), in production 2026-10-08 20:05 UTC). In Kibana, p99 mem_total_bytes (12 h buckets) went from ~100-115 MB on 2026-10-01 to ~60-75 MB, then to ~33-45 MB from 2026-10-09.

What is left in the request is ProcessReviewOutputService#valid?, which called FileExclusion#ai_reviewable_diff_files. That builds the full expanded collection (MergeRequest#ai_reviewable_diff_files, diffs(expanded: true).raw_diff_files), so every file's patch is read from the database or external diff storage and wrapped in a Gitlab::Diff::File, only for .empty?. On 2026-09-24 only 93 of ~13,500 requests returned 400, and that count covers every 400 reason, so the full load is almost always wasted. PublishCommentsWorker loads the collection again anyway.

The change

  • New FileExclusion#ai_reviewable_diff_files?: with the flag on, first_diff_batch_ai_reviewable? || ai_reviewable_diff_files.any?; with it off, ai_reviewable_diff_files.any? (master).
  • first_diff_batch_ai_reviewable? calls merge_request.merge_request_diff.diffs_in_batch(0, FIRST_DIFF_BATCH_SIZE, diff_options: { expanded: true }) (FIRST_DIFF_BATCH_SIZE = 10, below the batch loader's cap of 30), selects raw_diff_files with ai_reviewable_patch?, runs Ai::FileExclusionService on only those paths, and returns true if any is not excluded.
  • excluded_paths(file_paths) is extracted from excluded_files so both callers share the exclusion lookup. excluded_files behaves as before.
  • ProcessReviewOutputService#valid? uses the predicate. The worker (ProcessCommentsService) still uses the full ai_reviewable_diff_files and is unchanged.

Why the result is identical

  • Gitlab::Git::DiffCollection#each_serialized_patch applies max_files, max_lines and max_bytes cumulatively in order. A batch at offset 0 has no offset collection (PaginatedDiffs#diffs), so its files are exactly the first 10 files of the full collection, with the same overflow state.
  • Ai::FileExclusionService#check_paths evaluates each path on its own, so a path excluded in the batch is excluded in the full list and the other way round.
  • The batch uses merge_request_diff_files in association order (relative_order) after reorder_diff_files!, the same as load_diffs.
  • So a reviewable, non-excluded file in the batch is also one in the full collection (true on both sides). When the batch finds none, the old full check runs and decides (same result). The "nothing to review" message still comes from the full excluded_files.

Edge cases

  • without_files? diffs: diffs_in_batch goes through the Gitaly comparison and returns the full collection, so the result is exact. Cost is the same as today when a file is found, and two compares when none is (rare for an MR under review).
  • MR with no diff yet: merge_request_diff is an unsaved MergeRequestDiff.new. reorder_diff_files! returns early, the batch is empty, and the full check runs as before.
  • Legacy unsorted diff (sorted: false) on its first read: master's load_diffs reads the association loaded during reorder_diff_files!, while the batch re-queries in the new order. Only if the read-time limits also overflow could the batch return true where master returns false. Theoretical; the diff is sorted after that first read.
  • If the first 10 files have nothing reviewable (for example 10 leading binaries or excluded files), the request does the batch and then the full load, slightly more than today (+8% measured with a 30-file batch).
  • Extra queries on the common path: a COUNT for pagination_data and an OFFSET 0 LIMIT 10 read on merge_request_diff_files, the same queries the diffs batch endpoint already runs. The full read of every row is gone.

Feature flag

  • duo_code_review_check_first_diff_batch, gitlab_com_derisk, default off, actor merge_request.project (same as the earlier duo_code_review_skip_blob_loading). Rollout: #633895
  • Flag off: no batch is built and the full check runs as on master (7 queries on a 300-file MR, same as master). Flag on: one extra COUNT (8 queries in the test env, 1-3 more on GDK), same Gitaly calls (2).
  • The decision is meant to be identical either way; the flag guards against an error or a missed edge case in the batch load on the request path. No changelog entry, as the change is default-off.

Alternatives rejected

  • Walk every batch until a match. PaginatedDiffs#diffs reloads every earlier row for pages after the first (relation.limit(offset_index)) to keep the limits right, so this is quadratic in the worst case.
  • Paths first via pluck (0.5 MB). It returns nothing for without_files diffs, and it still needs patch content for ai_reviewable_patch?.
  • Move the check into PublishCommentsWorker. It changes the API contract (400 to 201) and needs a Duo Workflow Service change.

Measurements

GDK, real POST /api/v4/ai/duo_workflows/code_review/add_comments with a composite-identity token for the code review flow service account, 1000-file MR (40 lines each), mem_total_bytes and duration_s from log/api_json.log, warm requests:

Flag mem_total_bytes duration_s db_count Gitaly calls
off 59, 98, 63 MB 1.1-1.5 61-63 2
on 21, 20, 14, 14 MB 0.34-0.71 64-67 2
off again 60, 57, 61 MB 1.06-1.16 61-65 2

Earlier, with a 30-file batch: memory_profiler around the check, fresh MergeRequest per run, steady state (second run), local test env:

MR Master Branch
300 files x 60 lines 27.2 MB 6.1 MB
1000 files x 40 lines 80.0 MB 6.2 MB
40 leading binaries + 300 files (fallback) 30.1 MB 32.4 MB

Break-it pass

47 scenarios compared the branch against master's FileExclusion (loaded under another name) on the same MRs: basic inputs (binary, LFS on/off, empty, delete, rename by old/new path, unicode, -diff), batch boundaries (text at the last and first-past positions of the batch, all binary, whole first batch excluded, all excluded, include rules inside and past the batch), read-time max_files/max_lines overflow inside and past the batch, too-large patch, external diffs, without_files, second diff version, no diff rows, legacy unsorted diff, and the worker (ProcessCommentsService) draft notes and message. Every branch-vs-master comparison matched, with the flag on and off, at both the 30- and 10-file batch size. Removing the fallback or skipping batch exclusion each failed 4 scenarios (checked at 30 files), so the oracle catches breaks.

Specs

  • ee/spec/services/ai/duo_workflows/code_review/process_review_output_service_spec.rb: flag-off examples assert no batch is built and the full collection is loaded once, as on master. The stubbed examples now run under "when the first diff batch has no reviewable file". New examples with real diffs cover the fast path (no full load), the fallback when the first batch is excluded, and every file excluded (two DiffCollections, message lists * CHANGELOG). merge_request is now let_it_be_with_refind, since the memoised raw diffs on the shared object made the count order-dependent. The fast-path and every-file-excluded examples fail without the change.
  • Run locally and passing: that spec (21), process_comments_service_spec (50), create_comments_service_spec (66), ee/spec/requests/api/ai/duo_workflows/code_review_spec.rb (29). Rubocop is clean.

References

Screenshots or screen recordings

Not applicable, no UI change.

How to set up and validate locally

  1. Enable the flag: Feature.enable(:duo_code_review_check_first_diff_batch).
  2. Open a merge request with many changed files in a project with Duo Code Review (Agent Platform) enabled.
  3. Request a review from the Duo Code Review bot and wait for it to post comments.
  4. In log/api_json.log, find the POST /api/:version/ai/duo_workflows/code_review/add_comments line and compare mem_total_bytes with master. Comments should be posted as before.

MR acceptance checklist

Evaluate this MR against the MR acceptance checklist. It helps you analyze changes to reduce risks in quality, performance, reliability, security, and maintainability.

Edited by Marc Shaw

Merge request reports

Loading
Loading