mirror of
https://github.com/alibaba/open-code-review.git
synced 2026-09-14 19:59:52 +08:00
chethanuk
75cb3d0c45
scan already detects the aggregate token-budget stop (it prints the
"[ocr] token budget reached" line and records a token_budget_reached
warning) but BudgetExceeded() was hard-coded false, so
summary.budget_exceeded never appeared in `ocr scan --format json`.
The write goes next to `budgetHit = true` in dispatchBatch's per-file
gate. That is the only site that sets budgetHit, and it covers all three
exits that carry the stop out of dispatchBatch: normal return, ctx-cancel
return, and the caller's `if budgetHit { break }`. Setting it at the
dispatchSubtasks break instead would lose it on the ctx-cancel path.
Plain bool, no mutex: dispatchBatch's loop is the only writer, it runs on
the caller's goroutine, and the value is read by emitRunResult after Run
returns. The spawned subtask goroutines never touch it. Matches the
existing internal/agent.Agent.budgetExceeded field.
Status and exit code are untouched — reaching the budget is a controlled
truncation, so out.Status stays the warning-derived value.
62 lines
2.0 KiB
Go
62 lines
2.0 KiB
Go
// SPDX-License-Identifier: Apache-2.0
|
|
// Copyright 2026 alibaba/open-code-review Contributors
|
|
|
|
package scan
|
|
|
|
import (
|
|
"context"
|
|
"testing"
|
|
|
|
"github.com/alibaba/open-code-review/internal/session"
|
|
"github.com/alibaba/open-code-review/internal/tool"
|
|
)
|
|
|
|
// TestBudgetExceededFlag pins BudgetExceeded() to the actual state of the
|
|
// token budget gate. Before this it was hard-coded false, so a budget stop
|
|
// never reached summary.budget_exceeded in `ocr scan --format json` (#771).
|
|
//
|
|
// Both cases run the same 50K-tokens-per-file fake through dispatchSubtasks;
|
|
// only MaxTokensBudget differs. MaxConcurrency=1 serializes dispatch so the
|
|
// gate trips deterministically.
|
|
//
|
|
// The zero-value case is not repeated here: TestScanGettersOnEmptyAgent in
|
|
// getters_test.go already asserts BudgetExceeded()==false on &Agent{}, and
|
|
// that assertion stops being vacuous now that the getter reads a field.
|
|
func TestBudgetExceededFlag(t *testing.T) {
|
|
cases := []struct {
|
|
name string
|
|
budget int64
|
|
items int
|
|
want bool
|
|
}{
|
|
{name: "gate trips", budget: 120_000, items: 10, want: true},
|
|
{name: "unlimited budget", budget: 0, items: 5, want: false},
|
|
}
|
|
|
|
for _, tc := range cases {
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
a := NewAgent(Args{
|
|
Template: budgetTestTemplate(),
|
|
LLMClient: &fakeBudgetClient{perCallTokens: 50_000},
|
|
CommentCollector: tool.NewCommentCollector(),
|
|
Tools: tool.NewRegistry(),
|
|
MaxConcurrency: 1,
|
|
MaxTokensBudget: tc.budget,
|
|
Session: session.New(t.TempDir(), "main", "test", session.SessionOptions{ReviewMode: session.ReviewModeFullScan}),
|
|
SkipPlan: true,
|
|
SkipDedup: true,
|
|
SkipSummary: true,
|
|
})
|
|
a.items = makeScanItems(tc.items)
|
|
a.args.Tools.Freeze()
|
|
|
|
if _, err := a.dispatchSubtasks(context.Background()); err != nil {
|
|
t.Fatalf("dispatchSubtasks: %v", err)
|
|
}
|
|
if got := a.BudgetExceeded(); got != tc.want {
|
|
t.Errorf("BudgetExceeded() = %v, want %v", got, tc.want)
|
|
}
|
|
})
|
|
}
|
|
}
|