feat(agent): add token/tool-call cost guardrails to the review path

[nitishagar]

What & why

Closes #409 (track 2). Adds runtime token-cost and tool-call budget guardrails to the diff-review path so a large MR stops itself before runaway cost, instead of timing out and losing all structured output. The #409 report quantified a single failed attempt at ≈90.4M tokens / ≈2,612 tool calls.

Scan already has --max-tokens-budget; this mirrors that proven gate onto the review path and adds the missing tool-call dimension plus the typed terminal state the report asked for.

Changes

New ocr review flags (mirror scan's flag wording):

• --max-tokens-budget N — cap aggregate token usage; dispatch stops once the running total + a per-file look-ahead would exceed it (0 = unlimited).
• --max-tool-calls N — cap aggregate tool-call count across the run (0 = unlimited).

Mechanism mirrors internal/scan/agent.go's dispatchBatch gate exactly: read the existing atomic Runner counters (no duplicate counter — a second one would be a drift bug) before acquiring the semaphore, and break the dispatch loop on exceed. In-flight workers are allowed to finish (overrun bounded by the in-flight count, ≤ concurrency), matching scan's documented contract.

Typed terminal state (INV-3): budget exhaustion returns the partial comments already produced with a nil error (returning an error would suppress all structured output via review_cmd.go's failure path) and signals budget-exceeded out-of-band. The output layer sets status: "budget_exceeded" (precedence over completed_with_warnings/completed_with_errors) and summary.budget_exceeded: true.

Structured usage on failure (INV-4): a non-budget failure now emits a structured usage record to stderr (token totals, tool-call count, elapsed, budget_exceeded=false) so the cost of a failed attempt is never lost. Carries only numeric tallies — no credentials or prompts.

Pre-review scale warning (INV-5): prints estimated scale (files, est. tokens) and configured budgets before any model spend; non-blocking, warn-only.

Scope decisions (what's NOT in this PR)

Per the issue's open product choices and to keep the change focused/mergeable:

• max_model_requests / max_elapsed_seconds / hard-refuse large_mr_policy are deferred follow-ups (each needs cross-package instrumentation; the pre-check is warn-only).
• Scan is unchanged except for a BudgetExceeded() bool { return false } accessor so the shared ResultProvider interface compiles; scan keeps its own token budget and its JSON output is unaffected.
• status remains a free-form string — only a new documented constant value budget_exceeded is added (no enum rewrite).

Additive JSON + zero-value defaults mean rolling back restores prior behavior with no migration. Detailed design/analysis and the invariant spec are in the linked plan bundle.

Testing

• New unit tests mirror internal/scan/budget_test.go: token-budget gate stops dispatch early and sets BudgetExceeded()==true with a token_budget_reached warning + partial comments; tool-call-budget variant; unlimited default runs all files; estimate helpers; CLI negative-value validation.
• New JSON-output tests assert status=="budget_exceeded" + summary.budget_exceeded==true, the precedence over warning/error statuses, the failure-path usage record (text + JSON), and that the success path omits the new field (omitempty).
• make test full -race suite green (23 packages, 0 failures, 0 race warnings).

Checklist

• make build / go vet ./... / gofmt clean
• make test green (-race)
• Tests added for new behavior
• No breaking CLI/JSON change (additive fields, 0/omitempty defaults)

[CLAassistant]

All committers have signed the CLA.

[github-actions]

🔍 OpenCodeReview found 5 issue(s) in this PR.

• ✅ Successfully posted inline: 5 comment(s)

[github-actions]

Bug: BudgetExceeded() always returns false even when scan's token budget is actually exceeded.

The scan agent has its own MaxTokensBudget mechanism (see Args.MaxTokensBudget at line 69 and the budget gate in dispatchBatch at lines 483-493). When the budget is hit, dispatchBatch returns budgetHit=true, causing dispatchSubtasks to stop scheduling remaining batches — exactly the same semantics as the diff-review agent's budgetExceeded flag.

However, this method hardcodes return false, so when a scan run is truncated by the token budget:

1. The JSON output status will be "success" or "completed_with_warnings" instead of "budget_exceeded"
2. The budget_exceeded field in the JSON summary will always be omitted/false
3. Callers relying on the typed budget_exceeded status (as documented in the ResultProvider interface) won't detect it

To fix this, track the budget-exceeded state in scan.Agent (similar to how internal/agent.Agent uses a.budgetExceeded) by propagating the budgetHit return value from dispatchBatch up to a field on the struct, then return that field here.

[nitishagar]

Not a bug — intentional scoped-out follow-up, not a defect in this PR.

This change targets the diff-review path only. Scan already has its own MaxTokensBudget gate (dispatchBatch returns budgetHit) and this PR deliberately does NOT wire that into a typed budget_exceeded status — it is explicitly listed under "What We're NOT Doing": scan-mode changes beyond the shared emitRunResult plumbing are out of scope. The BudgetExceeded() stub returning false exists for one reason: *scan.Agent must satisfy the shared ResultProvider interface, which gained the method so agent.Agent can signal budget-exceeded. Scan's JSON output is intentionally unchanged (status stays success / completed_with_).

Propagating scan's budgetHit into a typed status is a reasonable follow-up for review/scan parity — happy to open a separate issue. Doing it here would expand scope into scan core logic, which was scoped out to keep this PR focused and mergeable.

[github-actions]

This new file has no corresponding test file (internal/agent/estimate_test.go). The scan counterpart (internal/scan/estimate_test.go) has comprehensive tests covering humanTokens, estimateFileTokens, estimateCost, edge cases (empty/binary/deleted), and consistency between per-file and aggregate estimates.

Given that the comments explicitly state these functions "mirror" the scan versions and must be kept in sync, adding parallel tests here would:

1. Verify correctness of estimateDiffFileTokens (deleted/empty diffs → 0, positive for normal diffs).
2. Verify estimateDiffCost aggregates correctly and TotalTokens == InputTokens + OutputTokens.
3. Verify humanTokens output matches scan's (the sync invariant).
4. Verify sumToolCalls sums correctly.
5. Provide a regression safety net so future changes to either copy are caught if they diverge.

[github-actions]

The format string labels the second value as "next-file est" but actually prints projected, which is used + estimateDiffFileTokens(toDispatch[i]) — i.e., the total projected cost, not just the next file's estimate. This makes the log message misleading: a reader would interpret "used X + next-file est ≈ Y" as Y being the per-file estimate, when it's actually the sum.

Either print the per-file estimate separately (humanTokens(estimateDiffFileTokens(toDispatch[i]))) or relabel the field to "projected total".

Suggestion:

Suggested change

	fmt.Fprintf(stdout.Writer(), "[ocr] token budget reached (used %s + next-file est ≈ %s > budget %s) — skipping %s and remaining files\n",
	humanTokens(used), humanTokens(projected), humanTokens(a.args.MaxTokensBudget), toDispatch[i].NewPath)
	fmt.Fprintf(stdout.Writer(), "[ocr] token budget reached (used %s + next-file est ≈ %s = %s > budget %s) — skipping %s and remaining files\n",
	humanTokens(used), humanTokens(estimateDiffFileTokens(toDispatch[i])), humanTokens(projected), humanTokens(a.args.MaxTokensBudget), toDispatch[i].NewPath)

[lizhengfeng101]

Great work on the token budget guardrail — the design is clean and the pre-review estimate is a nice touch.

One thought on --max-tool-calls: I wonder if this flag pulls its weight given the existing controls. We already have:

• --max-tokens-budget (global) — directly caps cost, which is what users actually care about. When tokens are exhausted, tool calls stop naturally.
• --max-tools (per-file) — prevents any single subtask from spinning in an infinite tool-use loop, with a sensible default.

--max-tool-calls sits between these two but is harder to reason about in practice — there's no intuitive answer to "how many total tool calls should a 200-file review use?" and no good default we can recommend. It's also highly correlated with token usage, so in most scenarios the token budget will trip first anyway.

My suggestion: drop --max-tool-calls as a user-facing flag. The JSON output already reports tool_calls.total and tool_calls.by_tool distribution, so observability is covered without making it a configurable gate.

This keeps the CLI surface simple — one global knob (tokens = cost) and one local safety valve (per-file rounds) — without adding a third dial that users can't easily calibrate.

[nitishagar]

@lizhengfeng101 thanks for the thoughtful review — and glad the pre-review estimate landed well.

On dropping --max-tool-calls: I hear the calibration argument, and I considered it carefully, but I'd lean toward keeping it for this PR. The reasoning:

• It directly addresses the second headline from the 大 MR 场景下 open-code-review 可能在超时前产生极高 token 成本 #409 report. The failed attempt that motivated this issue was ~2,612 tool calls in a single run. --max-tokens-budget would eventually catch that too, but only after the tokens those tool calls generated have already been spent — the tool-call budget is the earlier, sharper signal for the specific failure mode the report describes (an agent looping on tool use).
• It's additive with a zero default, so it costs nothing when unused. Existing behavior and the common path are byte-identical (MaxToolCalls==0 skips the gate entirely). No new default to recommend because the default is "off."
• It pairs with the token budget as belt-and-suspenders, not a third independent dial. Tokens = cost (the user-facing knob, as you note); tool-calls = the runaway-loop signal. A pathological agent that issues cheap-but-numerous tool calls (small reads, repeated searches) can rack up calls faster than tokens, and the tool-call gate catches that the token gate wouldn't trip on as quickly.

That said — if after the above you'd still prefer it dropped, I'm happy to do it; it's a clean removal (the Args field, the gate branch, the flag, validation, help, and the tool-call-budget test). The token-budget gate and the typed budget_exceeded status stand on their own either way. Your call.

Everything else from the OCR pass I've addressed in e924d3b:

• The emitFailureUsage hardcoding budget_exceeded=false was a real edge (budget-trip + all-dispatched-fail) — now reports ag.BudgetExceeded() with a regression test.
• Misleading "next-file est" log label fixed (now prints used + next-file est = projected > budget).
• estimate_test.go added mirroring scan's, including a humanTokens parity table so the two copies can't silently diverge.
• stderr failure JSON now SetIndent'd for consistency.
• The scan.BudgetExceeded() stub returning false is intentional and documented — propagating scan's budgetHit into a typed status is explicitly scoped out (scan keeps its own budget + unchanged JSON); happy to open a follow-up for scan parity if there's appetite.

[lizhengfeng101]

Minor: misleading comment in review_cmd.go (lines 150–155)

The comment states:

A budget trip returns partial comments with a nil error and never reaches here, so budget_exceeded is always false in this record.

However, emitFailureUsage reads the agent's actual ag.BudgetExceeded() value (not a hardcoded false), and TestEmitFailureUsage_BudgetExceededPropagated explicitly covers the edge case where a budget gate trips and every dispatched file subsequently fails — reaching this failure path with BudgetExceeded()==true.

Suggested rewording:

// INV-4: emit a best-effort structured usage record on the failure path so
// the cost of the failed attempt is not lost. Budget exhaustion typically
// returns partial comments with a nil error and does not reach here, but a
// residual edge exists (budget trips, then all dispatched files error) —
// emitFailureUsage reports the agent's actual BudgetExceeded() state.

[lizhengfeng101]

Thanks for the detailed rationale on keeping --max-tool-calls — I appreciate the belt-and-suspenders framing, and you're right that the #409 report's 2,612 tool calls is a sharp signal.

That said, I'd still lean toward deferring it for now. My thinking:

• --max-tokens-budget already provides a strong guardrail for the failure mode 大 MR 场景下 open-code-review 可能在超时前产生极高 token 成本 #409 describes. Tool calls generate tokens, so a token budget will eventually catch a runaway loop — maybe not as early, but early enough to prevent the catastrophic ~90M token blowout.
• Every new flag adds cognitive load. Users now need to reason about two budget knobs, decide which to set, and understand the interaction between them. Keeping the surface area minimal until we have evidence that one budget isn't sufficient feels like the safer UX choice.
• We can always add it later with zero cost. Since the implementation already exists in this PR, promoting it to a user-facing flag in a follow-up is trivial if real-world usage of --max-tokens-budget reveals gaps.

Would you be open to shipping with just --max-tokens-budget for now, and re-evaluating --max-tool-calls once we see how the token budget performs in practice? If it turns out there are cases where cheap-but-numerous tool calls slip through, we'll have strong evidence to justify the second knob — and the implementation is already proven here.

Either way, the core design (budget gate, typed status, failure-path usage) is solid work. 👍

[nitishagar]

Sounds good — convinced by the UX-minimalism argument. I've dropped --max-tool-calls entirely in d149c0a (flag, validation, help, the Args.MaxToolCalls field, the dispatch gate's tool-call branch, the tool_budget_reached warning, the sumToolCalls helper, and the tool-call-budget test). No dead code left behind; --max-tokens-budget is the single budget knob this ships with.

If real-world usage reveals cheap-but-numerous tool-call loops slipping through the token budget, re-adding it is straightforward from this PR's history. Thanks for the review.

Also applied your comment wording fix in review_cmd.go (the stale "never reaches here" claim is gone).

[lizhengfeng101]

Minor: unconditional estimate output is a behavioral change for all text-mode users

The pre-run estimate line in agent.go:

est := estimateDiffCost(a.diffs)
fmt.Fprintf(stdout.Writer(), "[ocr] estimated cost: %s\n", est)

is printed unconditionally for every ocr review run in text mode (quiet handle suppresses it in JSON/agent modes, so those are fine). This means users who never set --max-tokens-budget still see a new [ocr] estimated cost: ... line they didn't ask for — and any script parsing text output may break on the unexpected line.

Suggestion: gate the estimate behind MaxTokensBudget > 0:

if a.args.MaxTokensBudget > 0 {
    est := estimateDiffCost(a.diffs)
    fmt.Fprintf(stdout.Writer(), "[ocr] estimated cost: %s\n", est)
    fmt.Fprintf(stdout.Writer(), "[ocr] token budget: %s (dispatch stops once exceeded)\n", humanTokens(a.args.MaxTokensBudget))
    if est.TotalTokens > a.args.MaxTokensBudget {
        fmt.Fprintf(stdout.Writer(), "[ocr] WARNING: estimate (%s) exceeds token budget (%s); review will stop partway\n",
            humanTokens(est.TotalTokens), humanTokens(a.args.MaxTokensBudget))
    }
}

Rationale: users who opt into a budget are exactly the audience who benefits from the pre-flight estimate ("my budget vs. projected cost"). Users who don't set a budget have already chosen "unlimited" and gain nothing from the extra line — it's noise. This keeps the INV-5 intent intact while eliminating the unconditional behavioral change for existing users.

[lizhengfeng101]

Suggestion: gate the estimate behind MaxTokensBudget > 0

[nitishagar]

Done in e944e5c — the pre-run estimate block is now gated behind MaxTokensBudget > 0, exactly as suggested. Users who never set a budget see no new output line (prior text-mode behavior unchanged); the estimate + budget + over-budget warning print only for the audience that opted into a cap. Tests green.

[lizhengfeng101]

CI is failing across all jobs — the root cause is an incorrect module path in the new test files.

The imports use github.com/open-code-review/open-code-review/internal/... but this repo's module path is github.com/alibaba/open-code-review. The go mod tidy step tries to resolve the non-existent repo and fails:

github.com/alibaba/open-code-review/internal/agent imports
    github.com/open-code-review/open-code-review/internal/llm: git ls-remote ... exit status 128:
    fatal: could not read Username for 'https://github.com': terminal prompts disabled

Affected files (grep for github.com/open-code-review/open-code-review):

• cmd/opencodereview/budget_output_test.go
• internal/agent/budget_test.go
• internal/agent/estimate.go

Fix: replace all occurrences of github.com/open-code-review/open-code-review with github.com/alibaba/open-code-review.

[lizhengfeng101]

@nitishagar

[nitishagar]

Updated @lizhengfeng101
Waiting for the CI pass.

[lizhengfeng101]

LGTM

[nitishagar]

Thanks @lizhengfeng101. I don't see a merge option. Assuming someone from the core team will move this forward.

[lizhengfeng101]

Thanks @lizhengfeng101. I don't see a merge option. Assuming someone from the core team will move this forward.

Running benchmark, please wait.