From 611847e794b2058c03b51e9a1671da9de8119511 Mon Sep 17 00:00:00 2001 From: ChethanUK Date: Fri, 7 Aug 2026 20:01:30 +0200 Subject: [PATCH] feat(scan): report token budget stop in JSON summary.budget_exceeded MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- cmd/opencodereview/scan_budget_json_test.go | 139 ++++++++++++++++++++ internal/scan/agent.go | 17 ++- internal/scan/budget_exceeded_test.go | 61 +++++++++ internal/scan/getters_test.go | 7 +- 4 files changed, 215 insertions(+), 9 deletions(-) create mode 100644 cmd/opencodereview/scan_budget_json_test.go create mode 100644 internal/scan/budget_exceeded_test.go diff --git a/cmd/opencodereview/scan_budget_json_test.go b/cmd/opencodereview/scan_budget_json_test.go new file mode 100644 index 00000000..0897aa1f --- /dev/null +++ b/cmd/opencodereview/scan_budget_json_test.go @@ -0,0 +1,139 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright 2026 alibaba/open-code-review Contributors + +package main + +import ( + "context" + "encoding/json" + "os" + "path/filepath" + "strings" + "testing" + "time" + + "github.com/alibaba/open-code-review/internal/config/template" + "github.com/alibaba/open-code-review/internal/llm" + "github.com/alibaba/open-code-review/internal/scan" + "github.com/alibaba/open-code-review/internal/session" + "github.com/alibaba/open-code-review/internal/tool" +) + +// fakeScanBudgetClient finishes every file in one round and reports a fixed +// 50K token usage, so the aggregate budget gate trips deterministically. +type fakeScanBudgetClient struct{} + +func (fakeScanBudgetClient) CompletionsWithCtx(_ context.Context, _ llm.ChatRequest) (*llm.ChatResponse, error) { + return &llm.ChatResponse{ + Choices: []llm.Choice{{ + Message: llm.ResponseMessage{ + Role: "assistant", + ToolCalls: []llm.ToolCall{{ + ID: "1", + Type: "function", + Function: llm.FunctionCall{Name: "task_done", Arguments: "{}"}, + }}, + }, + FinishReason: "tool_calls", + }}, + Usage: &llm.UsageInfo{PromptTokens: 50_000, TotalTokens: 50_000}, + }, nil +} + +// TestScanBudgetJSON is the end-to-end contract for #771: a real *scan.Agent, +// real enumeration, real budget gate, through the shared emitRunResult +// pipeline. BudgetExceeded() was hard-coded false, so summary.budget_exceeded +// never appeared however early the gate cut the run short. +// +// The unlimited case asserts on the RAW capture, not the decoded struct: +// budget_exceeded is omitempty, so false and absent are indistinguishable +// after Unmarshal. +func TestScanBudgetJSON(t *testing.T) { + cases := []struct { + name string + budget int64 + want bool + wantStatus string + }{ + // The budget stop must NOT invent a typed status: it stays the + // ordinary warning-derived one (output.go leaves out.Status alone). + {name: "budget stop sets budget_exceeded", budget: 120_000, want: true, wantStatus: "completed_with_warnings"}, + {name: "unlimited budget omits the key", budget: 0, want: false, wantStatus: "success"}, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + t.Setenv("HOME", t.TempDir()) + + repoDir := t.TempDir() + // Zero-padded so lexical walk order is stable (f10 vs f2). + for _, name := range []string{"f01.go", "f02.go", "f03.go", "f04.go", "f05.go", "f06.go", "f07.go", "f08.go"} { + if err := os.WriteFile(filepath.Join(repoDir, name), []byte("package x\n"), 0o600); err != nil { + t.Fatalf("write %s: %v", name, err) + } + } + + ag := scan.NewAgent(scan.Args{ + RepoDir: repoDir, + Template: template.ScanTemplate{ + MaxTokens: 100000, + MaxToolRequestTimes: 5, + MainTask: template.LlmConversation{ + Messages: []template.ChatMessage{ + {Role: "system", Content: "scan"}, + {Role: "user", Content: "review {{file_content}}"}, + }, + }, + }, + LLMClient: fakeScanBudgetClient{}, + Tools: tool.NewRegistry(), + CommentCollector: tool.NewCommentCollector(), + MaxConcurrency: 1, // serialize so the gate is deterministic + MaxTokensBudget: tc.budget, + Session: session.New(t.TempDir(), "main", "test", session.SessionOptions{ReviewMode: session.ReviewModeFullScan}), + SkipPlan: true, + SkipDedup: true, + SkipSummary: true, + }) + + comments, err := ag.Run(context.Background()) + if err != nil { + t.Fatalf("scan Run: %v", err) + } + + raw := captureStdout(t, func() { + if err := emitRunResult(context.Background(), ag, comments, time.Now(), "json", "developer", nil, nil, nil); err != nil { + t.Fatalf("emitRunResult: %v", err) + } + }) + + var out jsonOutput + if err := json.Unmarshal([]byte(raw), &out); err != nil { + t.Fatalf("unmarshal %q: %v", raw, err) + } + if out.Status != tc.wantStatus { + t.Errorf("status = %q, want %q", out.Status, tc.wantStatus) + } + if out.Summary == nil { + t.Fatal("summary must be present") + } + if got := out.Summary.BudgetExceeded; got != tc.want { + t.Errorf("summary.budget_exceeded = %v, want %v", got, tc.want) + } + if !tc.want && strings.Contains(raw, "budget_exceeded") { + t.Errorf("unlimited budget must omit budget_exceeded entirely, got %s", raw) + } + + var warned bool + for _, w := range out.Warnings { + if w.Type == "token_budget_reached" { + warned = true + break + } + } + if warned != tc.want { + t.Errorf("token_budget_reached warning present = %v, want %v", warned, tc.want) + } + }) + } +} diff --git a/internal/scan/agent.go b/internal/scan/agent.go index 7e879a7e..66262ae7 100644 --- a/internal/scan/agent.go +++ b/internal/scan/agent.go @@ -106,6 +106,7 @@ type Agent struct { resumeInfo *session.ResumeInfo scanFingerprints map[string]string projectSummary string // populated post-run by maybeRunProjectSummary + budgetExceeded bool // set when the token budget gate stopped dispatch; written only by dispatchBatch's loop } // ProjectSummary returns the markdown project-level summary produced after @@ -223,13 +224,12 @@ func (a *Agent) Warnings() []llmloop.AgentWarning { return a.runner.Warnings() } // ToolCalls returns per-tool call counts accumulated during scan. func (a *Agent) ToolCalls() map[string]int64 { return a.runner.ToolCalls() } -// BudgetExceeded always returns false for scan. Scan self-limits via its own -// token budget gate and MaxToolRequestTimes; the typed budget_exceeded status -// and tool-call-budget plumbing are diff-review-path features (see -// internal/agent). This method exists only so *scan.Agent satisfies the -// cmd/opencodereview.ResultProvider interface, keeping scan's JSON output -// unchanged (status stays success / completed_with_*). -func (a *Agent) BudgetExceeded() bool { return false } +// BudgetExceeded reports whether the aggregate token budget gate stopped +// dispatch before every file was reviewed. Diagnostic only: scan still returns +// its partial comments and a nil error, and the value reaches output solely as +// summary.budget_exceeded. Per-file MaxToolRequestTimes exhaustion does NOT set +// it — that is an item-level outcome, not a run-level budget stop. +func (a *Agent) BudgetExceeded() bool { return a.budgetExceeded } func (a *Agent) recordWarning(warningType, file, message string) { a.runner.RecordWarning(warningType, file, message) @@ -636,6 +636,9 @@ func (a *Agent) dispatchBatch(ctx context.Context, batchIdx int, batch []model.S a.recordWarning("token_budget_reached", it.Path, fmt.Sprintf("stopped in batch #%d: used %d tokens + next-file estimate exceeds budget %d", batchIdx, used, a.args.MaxTokensBudget)) budgetHit = true + // budgetHit is per-batch and dies with this call; the field is + // the run-level signal emitRunResult reads after Run returns. + a.budgetExceeded = true break } } diff --git a/internal/scan/budget_exceeded_test.go b/internal/scan/budget_exceeded_test.go new file mode 100644 index 00000000..09281a56 --- /dev/null +++ b/internal/scan/budget_exceeded_test.go @@ -0,0 +1,61 @@ +// 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) + } + }) + } +} diff --git a/internal/scan/getters_test.go b/internal/scan/getters_test.go index d4aeec83..fc619ab1 100644 --- a/internal/scan/getters_test.go +++ b/internal/scan/getters_test.go @@ -6,8 +6,11 @@ package scan import "testing" // TestScanGettersOnEmptyAgent exercises the small ResultProvider getters that -// scan implements as constants or nil-safe guards, so review and scan can share -// the output pipeline. +// scan implements as zero values or nil-safe guards, so review and scan can +// share the output pipeline. BudgetExceeded now reads a field rather than +// returning a constant, so the assertion below pins the zero value that +// summary.budget_exceeded's omitempty depends on; TestBudgetExceededFlag +// covers the set case. func TestScanGettersOnEmptyAgent(t *testing.T) { a := &Agent{}