feat: cost controls — pre-flight estimation, token budget, file priority - #21
Conversation
…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
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (1)
📝 WalkthroughWalkthroughAdds configurable review token budgets, priority-based diff ordering, pre-flight trimming, budget logging, and runtime stopping when accumulated usage reaches the configured limit. ChangesReview budget enforcement
Estimated code review effort: 3 (Moderate) | ~25 minutes Possibly related PRs
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Comment |
There was a problem hiding this comment.
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
📒 Files selected for processing (7)
README.mdinternal/config/config.gointernal/diff/priority.gointernal/diff/priority_test.gointernal/reviewer/budget.gointernal/reviewer/budget_test.gointernal/reviewer/reviewer.go
| // 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) |
There was a problem hiding this comment.
🎯 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 -SRepository: 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 -SRepository: 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 -SRepository: 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 -SRepository: 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 -SRepository: 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 -SRepository: 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-L138internal/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.
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
There was a problem hiding this comment.
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
📒 Files selected for processing (5)
README.mdinternal/config/config.gointernal/diff/priority.gointernal/diff/priority_test.gointernal/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
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
What
Adds
--max-tokensflag 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
File priority sorting
Runtime safety net
Config
--max-tokens 50000REVIEW_MAX_TOKENS=50000max_tokens: 50000Test coverage
internal/diff: 82.7% (priority sorting tests)internal/reviewer: 80.7% (budget estimation + trimming tests)Summary by CodeRabbit
--max-tokensto cap total token usage per review (set0for unlimited).--max-tokens,REVIEW_MAX_TOKENS, and example per-repomax_tokensconfiguration.