Skip to content

feat: cost controls — pre-flight estimation, token budget, file priority - #21

Merged
brucearctor merged 3 commits into
mainfrom
feat/cost-controls
Jul 15, 2026
Merged

feat: cost controls — pre-flight estimation, token budget, file priority#21
brucearctor merged 3 commits into
mainfrom
feat/cost-controls

Conversation

@brucearctor

@brucearctor brucearctor commented Jul 15, 2026

Copy link
Copy Markdown
Contributor

What

Adds --max-tokens flag to cap total token usage per review.

Why

Enterprise teams can't adopt without cost predictability. A runaway review on a 500-file PR could blow the budget.

How

Pre-flight estimation

  • Estimates total tokens (input + ~25% output) before calling the model
  • If over budget, trims lowest-priority files to fit
  • Logs clear budget status

File priority sorting

  • Security-sensitive files (auth, crypto, secrets, config) → reviewed first
  • New files → prioritized over modifications
  • Generated code, pure renames, test files → deprioritized
  • When budget forces trimming, least important files are skipped

Runtime safety net

  • If actual token usage exceeds budget mid-review, stops cleanly
  • Catches cases where pre-flight estimate was too optimistic

Config

Method Example
Flag --max-tokens 50000
Env REVIEW_MAX_TOKENS=50000
YAML max_tokens: 50000

Test coverage

  • internal/diff: 82.7% (priority sorting tests)
  • internal/reviewer: 80.7% (budget estimation + trimming tests)

Summary by CodeRabbit

  • New Features
    • Added --max-tokens to cap total token usage per review (set 0 for unlimited).
    • Added token-budget pre-flight estimation and automatic trimming of lower-priority files when over budget.
    • Introduced per-file review priority to order diffs and improve what gets reviewed first.
    • Reviewer now stops once actual token usage reaches the configured limit.
  • Documentation
    • Updated README with --max-tokens, REVIEW_MAX_TOKENS, and example per-repo max_tokens configuration.
  • Tests
    • Added unit tests for file prioritization and budget estimation/trimming logic.

…riority

New --max-tokens flag caps total token usage per review (0 = unlimited).
Also available via REVIEW_MAX_TOKENS env var or max_tokens in yaml config.

Pre-flight estimation:
  - Estimates input + output tokens before calling the model
  - If over budget, trims lowest-priority files to fit
  - Logs clear budget status with file counts

File priority sorting:
  - Security-sensitive files (auth, crypto, secrets) reviewed first
  - New files prioritized over modifications
  - Generated code, renames, test files deprioritized
  - When budget forces trimming, least important files are skipped

Runtime safety net:
  - If actual token usage exceeds budget mid-review, stop cleanly
  - Catches cases where pre-flight estimate was too optimistic
@coderabbitai

coderabbitai Bot commented Jul 15, 2026

Copy link
Copy Markdown

Review Change Stack

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info
⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro Plus

Run ID: a03dc4ac-bc9b-44d8-87b2-5a43a5c924bf

📥 Commits

Reviewing files that changed from the base of the PR and between 59fa59e and 3b3266e.

📒 Files selected for processing (1)
  • internal/config/config.go

📝 Walkthrough

Walkthrough

Adds configurable review token budgets, priority-based diff ordering, pre-flight trimming, budget logging, and runtime stopping when accumulated usage reaches the configured limit.

Changes

Review budget enforcement

Layer / File(s) Summary
Token budget configuration
README.md, internal/config/config.go
Adds MaxTokens configuration through YAML, environment variables, and CLI flags; 0 represents unlimited usage.
File review prioritization
internal/diff/priority.go, internal/diff/priority_test.go
Scores files using security, size, file type, generated-code, test, new-file, and rename characteristics, then sorts them by descending priority.
Review budget enforcement
internal/reviewer/budget.go, internal/reviewer/budget_test.go, internal/reviewer/reviewer.go
Estimates token costs, removes lowest-priority files when needed, logs budget status, and stops chunk processing after the configured limit is reached.

Estimated code review effort: 3 (Moderate) | ~25 minutes

Possibly related PRs

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name Status Explanation Resolution
Docstring Coverage ⚠️ Warning Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%. Write docstrings for the functions missing them to satisfy the coverage threshold.
✅ Passed checks (4 passed)
Check name Status Explanation
Description Check ✅ Passed Check skipped - CodeRabbit’s high-level summary is enabled.
Title check ✅ Passed The title clearly reflects the main change: review cost controls with token budgeting, pre-flight estimation, and file prioritization.
Linked Issues check ✅ Passed Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check ✅ Passed Check skipped because no linked issues were found for this pull request.
✨ Finishing Touches
📝 Generate docstrings
  • Create stacked PR
  • Commit on current branch
🧪 Generate unit tests (beta)
  • Create PR with unit tests
  • Commit unit tests in branch feat/cost-controls

Comment @coderabbitai help to get the list of available commands.

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 5

🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@internal/config/config.go`:
- Around line 324-328: Update the REVIEW_MAX_TOKENS parsing in the configuration
loading flow to accept only valid positive integers. Treat malformed, zero, and
negative values as invalid and reject them without resetting MaxTokens to 0;
preserve the existing configured cap instead, so an environment override cannot
silently disable the token budget.
- Around line 424-426: Update the maxTokens handling in the configuration
construction flow to distinguish an explicitly provided --max-tokens=0 from an
unset flag, using the flag set’s visited-state tracking before assigning
c.MaxTokens. Preserve zero as the documented unlimited value, while validating
and rejecting negative values instead of ignoring them.

In `@internal/diff/priority.go`:
- Around line 70-81: The isSecuritySensitive function must stop using
unrestricted strings.Contains matches that classify unrelated paths as
sensitive. Replace the checks over securityPatterns and securityFiles with
path-component or filename-token matching that preserves legitimate security
matches while excluding collisions such as pkg/monkey.go and docs/authors.md,
and add negative tests for those cases.

In `@internal/reviewer/reviewer.go`:
- Around line 131-138: Update internal/reviewer/reviewer.go:131-138 so the
pre-flight budget check accounts for the complete request cost, including MR
metadata and related-code context, before invoking Review; trim or stop
processing when the resulting estimate reaches or exceeds MaxTokens. Update
internal/reviewer/reviewer.go:209-216 so the post-call guard also treats usage
equal to MaxTokens as exhausted, preventing another Review request at the limit.

In `@README.md`:
- Line 137: Update the README budget configuration documentation around the
--max-tokens option to document all supported sources: add REVIEW_MAX_TOKENS to
the environment-variable table and max_tokens to the YAML example/options,
preserving the existing CLI description and semantics.
🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

  • Push a commit to this branch (recommended)
  • Create a new PR with the fixes

ℹ️ Review info
⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro Plus

Run ID: ec676550-279c-453d-9443-845b6f0c6770

📥 Commits

Reviewing files that changed from the base of the PR and between 029e363 and 8fa4092.

📒 Files selected for processing (7)
  • README.md
  • internal/config/config.go
  • internal/diff/priority.go
  • internal/diff/priority_test.go
  • internal/reviewer/budget.go
  • internal/reviewer/budget_test.go
  • internal/reviewer/reviewer.go

Comment thread internal/config/config.go Outdated
Comment thread internal/diff/priority.go
Comment on lines +131 to +138
// Step 2d: Pre-flight budget check.
var skippedFiles []string
estimate := EstimateCost(diffs)
if r.cfg.MaxTokens > 0 && estimate.TotalEstimate > r.cfg.MaxTokens {
diffs, skippedFiles = TrimToBudget(diffs, r.cfg.MaxTokens)
estimate = EstimateCost(diffs) // Recalculate after trim.
}
LogBudgetStatus(estimate, r.cfg.MaxTokens, skippedFiles)

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🎯 Functional Correctness | 🟠 Major | 🏗️ Heavy lift

🧩 Analysis chain

🏁 Script executed:

#!/bin/bash
set -euo pipefail

printf '\n## File outline\n'
ast-grep outline internal/reviewer/reviewer.go --view expanded || true

printf '\n## Relevant line ranges\n'
sed -n '1,260p' internal/reviewer/reviewer.go | nl -ba | sed -n '110,240p'

printf '\n## Search MaxTokens and budget helpers\n'
rg -n "MaxTokens|EstimateCost|TrimToBudget|LogBudgetStatus|budget" internal/reviewer -S

printf '\n## Search for review call sites and prompt construction\n'
rg -n "Review\(|system prompt|MR metadata|related-code|chunks" internal/reviewer -S

Repository: OpticDiff/code-reviewer

Length of output: 1630


🏁 Script executed:

#!/bin/bash
set -euo pipefail

sed -n '1,260p' internal/reviewer/reviewer.go | nl -ba | sed -n '1,260p'

Repository: OpticDiff/code-reviewer

Length of output: 201


🏁 Script executed:

#!/bin/bash
set -euo pipefail
sed -n '1,260p' internal/reviewer/reviewer.go | nl -ba | sed -n '110,240p'
rg -n "MaxTokens|EstimateCost|TrimToBudget|LogBudgetStatus|budget" internal/reviewer -S

Repository: OpticDiff/code-reviewer

Length of output: 201


🏁 Script executed:

#!/bin/bash
set -euo pipefail

printf '\n## internal/reviewer/reviewer.go around Run\n'
awk 'NR>=66 && NR<=260 {printf "%4d %s\n", NR, $0}' internal/reviewer/reviewer.go

printf '\n## MaxTokens and budget helpers across repo\n'
rg -n "MaxTokens|EstimateCost|TrimToBudget|LogBudgetStatus|TotalEstimate|TotalTokens" . -S

printf '\n## Reviewer call construction\n'
rg -n "Review\(|GetDiffs\(|chunks|system prompt|MR metadata|related-code|context" internal/reviewer -S

Repository: OpticDiff/code-reviewer

Length of output: 22703


🏁 Script executed:

#!/bin/bash
set -euo pipefail

awk 'NR>=66 && NR<=260 {printf "%4d %s\n", NR, $0}' internal/reviewer/reviewer.go
rg -n "MaxTokens|EstimateCost|TrimToBudget|LogBudgetStatus|TotalEstimate|TotalTokens" . -S
rg -n "Review\(|GetDiffs\(|chunks|system prompt|MR metadata|related-code|context" internal/reviewer -S

Repository: OpticDiff/code-reviewer

Length of output: 22582


🏁 Script executed:

#!/bin/bash
set -euo pipefail

echo '## internal/reviewer/reviewer.go lines 66-260'
awk 'NR>=66 && NR<=260 {printf "%4d %s\n", NR, $0}' internal/reviewer/reviewer.go

echo
echo '## budget helpers'
rg -n "MaxTokens|EstimateCost|TrimToBudget|LogBudgetStatus|TotalEstimate|TotalTokens" internal/reviewer -S

echo
echo '## review call construction'
rg -n "Review\\(|GetDiffs\\(|chunks|system prompt|MR metadata|related-code|context" internal/reviewer -S

Repository: OpticDiff/code-reviewer

Length of output: 21324


🏁 Script executed:

#!/bin/bash
set -euo pipefail

echo '## internal/reviewer/budget.go'
awk 'NR>=1 && NR<=220 {printf "%4d %s\n", NR, $0}' internal/reviewer/budget.go

echo
echo '## prompt builders'
rg -n "BuildPromptWithCustom|BuildUserPromptWithContext|BuildUserPrompt|ContextSnippet" internal/model -S

Repository: OpticDiff/code-reviewer

Length of output: 6003


Enforce the limit before calling Review. The pre-flight estimate only counts diff size, while each request also adds MR metadata and related-code context, so trimming can still overshoot MaxTokens. The post-call guard is too late and only breaks on >, so it also allows another request when usage lands exactly on the limit.

  • internal/reviewer/reviewer.go#L131-L138
  • internal/reviewer/reviewer.go#L187-L216
📍 Affects 1 file
  • internal/reviewer/reviewer.go#L131-L138 (this comment)
  • internal/reviewer/reviewer.go#L209-L216
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@internal/reviewer/reviewer.go` around lines 131 - 138, Update
internal/reviewer/reviewer.go:131-138 so the pre-flight budget check accounts
for the complete request cost, including MR metadata and related-code context,
before invoking Review; trim or stop processing when the resulting estimate
reaches or exceeds MaxTokens. Update internal/reviewer/reviewer.go:209-216 so
the post-call guard also treats usage equal to MaxTokens as exhausted,
preventing another Review request at the limit.

Comment thread README.md
1. Config: use strconv.Atoi for REVIEW_MAX_TOKENS env parsing,
   reject invalid/negative values, don't silently reset to 0

2. Config: use fs.Visit to detect explicit --max-tokens=0 vs unset,
   reject negative values with warning

3. Priority: replace unrestricted strings.Contains with path-component
   matching to avoid false positives (monkey.go, authors.md, tokenizer.go)

4. Reviewer: use >= instead of > for runtime budget check so hitting
   exactly the limit also stops

5. README: document REVIEW_MAX_TOKENS env var and max_tokens yaml config

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 1

🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@internal/config/config.go`:
- Around line 326-335: Update the REVIEW_MAX_TOKENS parsing in the configuration
loader to accept zero as a valid value and assign it to c.MaxTokens, clearing
any YAML cap to represent unlimited. Retain the existing handling for positive
values and invalid or negative inputs, including the current warning behavior
for conversion errors.
🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

  • Push a commit to this branch (recommended)
  • Create a new PR with the fixes

ℹ️ Review info
⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro Plus

Run ID: 1ed465e0-b526-4f64-9f82-ca8743f32490

📥 Commits

Reviewing files that changed from the base of the PR and between 8fa4092 and 59fa59e.

📒 Files selected for processing (5)
  • README.md
  • internal/config/config.go
  • internal/diff/priority.go
  • internal/diff/priority_test.go
  • internal/reviewer/reviewer.go
🚧 Files skipped from review as they are similar to previous changes (4)
  • README.md
  • internal/diff/priority.go
  • internal/diff/priority_test.go
  • internal/reviewer/reviewer.go

Comment thread internal/config/config.go
Previously zero was silently ignored, meaning a YAML-configured cap
could not be overridden via env var. Now:
  - 0 = unlimited (explicitly clears any YAML cap)
  - >0 = token budget
  - <0 = rejected with warning
  - non-integer = rejected with warning
@brucearctor
brucearctor merged commit a9de473 into main Jul 15, 2026
4 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant