diff --git a/cmd/opencodereview/budget_output_test.go b/cmd/opencodereview/budget_output_test.go new file mode 100644 index 00000000..10cbe10a --- /dev/null +++ b/cmd/opencodereview/budget_output_test.go @@ -0,0 +1,223 @@ +package main + +import ( + "context" + "encoding/json" + "strings" + "testing" + "time" + + "github.com/alibaba/open-code-review/internal/agent" + "github.com/alibaba/open-code-review/internal/model" +) + +// TestEmitRunResult_JSONBudgetExceededStatus verifies that a provider signaling +// BudgetExceeded()==true produces JSON with status=="budget_exceeded" AND +// summary.budget_exceeded==true (INV-3 typed status), and that it takes +// precedence over completed_with_warnings. +func TestEmitRunResult_JSONBudgetExceededStatus(t *testing.T) { + ag := &mockResultProvider{ + filesReviewed: 3, + inputTokens: 100, + outputTokens: 50, + totalTokens: 150, + warnings: []agent.AgentWarning{{Type: "token_budget_reached", File: "big.go", Message: "stopped"}}, + toolCalls: map[string]int64{"file_read": 2}, + budgetExceeded: true, + } + got := captureStdout(t, func() { + err := emitRunResult(context.Background(), ag, nil, time.Now(), "json", "developer", nil) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + }) + var out jsonOutput + if err := json.Unmarshal([]byte(got), &out); err != nil { + t.Fatalf("unmarshal: %v", err) + } + if out.Status != "budget_exceeded" { + t.Errorf("status = %q, want budget_exceeded (must take precedence over completed_with_warnings)", out.Status) + } + if out.Summary == nil || !out.Summary.BudgetExceeded { + t.Errorf("summary.budget_exceeded = %v, want true", out.Summary) + } + // The token_budget_reached warning must still be present in the output so + // the reason is observable alongside the typed status. + var foundBudgetWarn bool + for _, w := range out.Warnings { + if w.Type == "token_budget_reached" { + foundBudgetWarn = true + break + } + } + if !foundBudgetWarn { + t.Error("expected token_budget_reached warning preserved in output") + } +} + +// TestEmitRunResult_JSONBudgetExceededPrecedenceOverErrors verifies that +// budget_exceeded takes precedence over completed_with_errors too — a budget +// trip is a distinct typed terminal state (INV-3 lists completed_with_errors +// as a status it must be distinct from). +func TestEmitRunResult_JSONBudgetExceededPrecedenceOverErrors(t *testing.T) { + ag := &mockResultProvider{ + filesReviewed: 1, + warnings: []agent.AgentWarning{{Type: "subtask_error", File: "x.go", Message: "boom"}}, + budgetExceeded: true, + } + got := captureStdout(t, func() { + err := emitRunResult(context.Background(), ag, nil, time.Now(), "json", "developer", nil) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + }) + var out jsonOutput + if err := json.Unmarshal([]byte(got), &out); err != nil { + t.Fatalf("unmarshal: %v", err) + } + if out.Status != "budget_exceeded" { + t.Errorf("status = %q, want budget_exceeded (must take precedence over completed_with_errors)", out.Status) + } +} + +// TestEmitRunResult_JSONNoBudgetIsSuccess verifies the default path is +// unchanged: BudgetExceeded()==false yields the normal status (regression +// guard that the new field/param didn't disturb existing behavior). +func TestEmitRunResult_JSONNoBudgetIsSuccess(t *testing.T) { + ag := &mockResultProvider{ + filesReviewed: 1, + inputTokens: 10, + outputTokens: 5, + totalTokens: 15, + } + got := captureStdout(t, func() { + err := emitRunResult(context.Background(), ag, nil, time.Now(), "json", "developer", nil) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + }) + var out jsonOutput + if err := json.Unmarshal([]byte(got), &out); err != nil { + t.Fatalf("unmarshal: %v", err) + } + if out.Status != "success" { + t.Errorf("status = %q, want success", out.Status) + } + if out.Summary != nil && out.Summary.BudgetExceeded { + t.Error("budget_exceeded should be false/omitted on a normal success run") + } +} + +// TestEmitFailureUsage_TextEmitsStructuredRecord verifies the non-budget +// failure path emits a structured usage record to stderr with the token totals +// and budget_exceeded=false (INV-4). Text format. +func TestEmitFailureUsage_TextEmitsStructuredRecord(t *testing.T) { + ag := &mockResultProvider{ + filesReviewed: 4, + inputTokens: 1000, + outputTokens: 500, + totalTokens: 1500, + toolCalls: map[string]int64{"file_read": 3, "code_comment": 2}, + sessionID: "sess-fail-1", + } + got := captureStderr(t, func() { + emitFailureUsage(ag, 42*time.Second, "text") + }) + for _, want := range []string{"usage on failure", "1500 total tokens", "5 tool calls", "budget_exceeded=false", "sess-fail-1"} { + if !strings.Contains(got, want) { + t.Errorf("stderr missing %q; got %q", want, got) + } + } +} + +// TestEmitFailureUsage_JSONEmitsStructuredRecord verifies the JSON form emits a +// parseable record to stderr with budget_exceeded=false (INV-4). +func TestEmitFailureUsage_JSONEmitsStructuredRecord(t *testing.T) { + ag := &mockResultProvider{ + filesReviewed: 2, + inputTokens: 200, + outputTokens: 80, + totalTokens: 280, + toolCalls: map[string]int64{"file_read": 1}, + } + got := captureStderr(t, func() { + emitFailureUsage(ag, 5*time.Second, "json") + }) + var out jsonOutput + if err := json.Unmarshal([]byte(got), &out); err != nil { + t.Fatalf("unmarshal stderr json: %v\ngot: %q", err, got) + } + if out.Status != "failed" { + t.Errorf("status = %q, want failed", out.Status) + } + if out.Summary == nil { + t.Fatal("expected summary on failure usage record") + } + if out.Summary.BudgetExceeded { + t.Error("budget_exceeded must be false on the non-budget failure path") + } + if out.Summary.TotalTokens != 280 { + t.Errorf("total_tokens = %d, want 280", out.Summary.TotalTokens) + } + if out.ToolCalls == nil || out.ToolCalls.Total != 1 { + t.Errorf("tool_calls.total = %v, want 1", out.ToolCalls) + } +} + +// TestEmitFailureUsage_BudgetExceededPropagated closes the residual edge the +// OCR review flagged: a budget gate can trip after dispatching N files, and if +// every dispatched file then fails, dispatchSubtasks returns (nil, error) — so +// the failure path is reached with BudgetExceeded()==true. The failure usage +// record must report the agent's ACTUAL budget state, never a hardcoded false, +// so it cannot contradict the agent's typed status. +func TestEmitFailureUsage_BudgetExceededPropagated(t *testing.T) { + // Text form. + ag := &mockResultProvider{ + filesReviewed: 1, + totalTokens: 100, + budgetExceeded: true, + } + got := captureStderr(t, func() { + emitFailureUsage(ag, 3*time.Second, "text") + }) + if !strings.Contains(got, "budget_exceeded=true") { + t.Errorf("text failure record must reflect budget_exceeded=true; got %q", got) + } + + // JSON form. + ag2 := &mockResultProvider{ + filesReviewed: 1, + totalTokens: 100, + budgetExceeded: true, + } + gotJSON := captureStderr(t, func() { + emitFailureUsage(ag2, 3*time.Second, "json") + }) + var out jsonOutput + if err := json.Unmarshal([]byte(gotJSON), &out); err != nil { + t.Fatalf("unmarshal stderr json: %v\ngot: %q", err, gotJSON) + } + if out.Summary == nil || !out.Summary.BudgetExceeded { + t.Errorf("JSON failure record must carry summary.budget_exceeded=true; got %+v", out.Summary) + } +} + +// TestEmitRunResult_BudgetExceededFalseOmittedFromJSON is a regression guard +// that the omitempty tag on BudgetExceeded keeps the success JSON free of the +// key (so old parsers see no diff on the common path). +func TestEmitRunResult_BudgetExceededFalseOmittedFromJSON(t *testing.T) { + ag := &mockResultProvider{ + filesReviewed: 1, + inputTokens: 10, + totalTokens: 10, + } + got := captureStdout(t, func() { + err := emitRunResult(context.Background(), ag, []model.LlmComment{}, time.Now(), "json", "developer", nil) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + }) + if strings.Contains(got, "budget_exceeded") { + t.Errorf("success JSON should omit budget_exceeded (omitempty), got %q", got) + } +} diff --git a/cmd/opencodereview/emit_run_result_test.go b/cmd/opencodereview/emit_run_result_test.go index 655f83b1..70eaa4d5 100644 --- a/cmd/opencodereview/emit_run_result_test.go +++ b/cmd/opencodereview/emit_run_result_test.go @@ -27,6 +27,7 @@ type mockResultProvider struct { toolCalls map[string]int64 resumeInfo *agent.ResumeInfo sessionID string + budgetExceeded bool } func (m *mockResultProvider) Diffs() []model.Diff { return m.diffs } @@ -41,6 +42,7 @@ func (m *mockResultProvider) ProjectSummary() string { return m.projectS func (m *mockResultProvider) ToolCalls() map[string]int64 { return m.toolCalls } func (m *mockResultProvider) ResumeInfo() *agent.ResumeInfo { return m.resumeInfo } func (m *mockResultProvider) SessionID() string { return m.sessionID } +func (m *mockResultProvider) BudgetExceeded() bool { return m.budgetExceeded } func TestEmitRunResult_JSONNoFiles(t *testing.T) { ag := &mockResultProvider{filesReviewed: 0} diff --git a/cmd/opencodereview/flags.go b/cmd/opencodereview/flags.go index a2a5ba6e..18250ed7 100644 --- a/cmd/opencodereview/flags.go +++ b/cmd/opencodereview/flags.go @@ -95,25 +95,26 @@ func expandShortFlags(args []string, shortMap map[string]string) []string { // --- review subcommand options --- type reviewOptions struct { - toolConfigPath string - rulePath string - repoDir string - from string - to string - commit string - resume string - excludes string // --exclude: comma-separated gitignore-style patterns - outputFormat string - audience string // --audience: "human" (default) or "agent" - background string // --background: optional requirement context - backgroundFile string // --background-file: path to a Markdown file used as background - model string // --model: override resolved LLM model for this review - concurrency int - perFileTimeout int - maxTools int - maxGitProcs int - preview bool - showHelp bool + toolConfigPath string + rulePath string + repoDir string + from string + to string + commit string + resume string + excludes string // --exclude: comma-separated gitignore-style patterns + outputFormat string + audience string // --audience: "human" (default) or "agent" + background string // --background: optional requirement context + backgroundFile string // --background-file: path to a Markdown file used as background + model string // --model: override resolved LLM model for this review + concurrency int + perFileTimeout int + maxTools int + maxGitProcs int + maxTokensBudget int // --max-tokens-budget: cap total token usage; 0 = unlimited + preview bool + showHelp bool } func parseReviewFlags(args []string) (reviewOptions, error) { @@ -138,6 +139,7 @@ func parseReviewFlags(args []string) (reviewOptions, error) { a.StringVar(&opts.model, "model", "", "override LLM model for this review (e.g., claude-opus-4-6)") a.IntVar(&opts.maxTools, "max-tools", 0, "max tool call rounds per file (0 = template default; min 10)") a.IntVar(&opts.maxGitProcs, "max-git-procs", 16, "max concurrent git subprocesses") + a.IntVar(&opts.maxTokensBudget, "max-tokens-budget", 0, "cap total token usage (input+output); dispatch stops once exceeded (0 = unlimited)") a.BoolVarP(&opts.preview, "preview", "p", false, "preview which files will be reviewed without running the LLM") if err := a.Parse(args); err != nil { @@ -188,6 +190,9 @@ func parseReviewFlags(args []string) (reviewOptions, error) { if opts.maxGitProcs < 0 { return opts, fmt.Errorf("--max-git-procs must be a non-negative integer (0 means use default 16)") } + if opts.maxTokensBudget < 0 { + return opts, fmt.Errorf("--max-tokens-budget must be a non-negative integer (0 means unlimited)") + } return opts, nil } @@ -241,6 +246,7 @@ Flags: --concurrency int max concurrent file reviews (default 8) --exclude string comma-separated gitignore-style patterns to exclude (merged with rule.json) --max-git-procs int max concurrent git subprocesses (default 16) + --max-tokens-budget int cap total token usage; dispatch stops once exceeded (0 = unlimited) --from string source ref to start diff from (e.g., 'main') --max-tools int max tool call rounds per file (0 = template default; min 10) --model string override LLM model for this review (e.g., claude-opus-4-6) diff --git a/cmd/opencodereview/flags_test.go b/cmd/opencodereview/flags_test.go index 8e259964..55a8b9df 100644 --- a/cmd/opencodereview/flags_test.go +++ b/cmd/opencodereview/flags_test.go @@ -84,6 +84,34 @@ func TestParseReviewFlags_NegativeMaxGitProcs(t *testing.T) { } } +func TestParseReviewFlags_NegativeMaxTokensBudget(t *testing.T) { + _, err := parseReviewFlags([]string{"--max-tokens-budget", "-1"}) + if err == nil { + t.Fatal("expected error for negative max-tokens-budget") + } +} + +func TestParseReviewFlags_BudgetFlagsDefaultZero(t *testing.T) { + // Unset budget flag defaults to 0 (unlimited) so existing behavior is unchanged. + opts, err := parseReviewFlags([]string{"--from", "main", "--to", "dev"}) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if opts.maxTokensBudget != 0 { + t.Errorf("maxTokensBudget = %d, want 0 (default unlimited)", opts.maxTokensBudget) + } +} + +func TestParseReviewFlags_BudgetFlagsParsed(t *testing.T) { + opts, err := parseReviewFlags([]string{"--max-tokens-budget", "120000"}) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if opts.maxTokensBudget != 120000 { + t.Errorf("maxTokensBudget = %d, want 120000", opts.maxTokensBudget) + } +} + func TestParseReviewFlags_ConflictingModes(t *testing.T) { _, err := parseReviewFlags([]string{"--from", "main", "--to", "dev", "--commit", "abc"}) if err == nil { diff --git a/cmd/opencodereview/output.go b/cmd/opencodereview/output.go index 58c14ba9..99c3f2ad 100644 --- a/cmd/opencodereview/output.go +++ b/cmd/opencodereview/output.go @@ -231,6 +231,7 @@ type jsonSummary struct { CacheReadTokens int64 `json:"cache_read_tokens,omitempty"` CacheWriteTokens int64 `json:"cache_write_tokens,omitempty"` Elapsed string `json:"elapsed"` + BudgetExceeded bool `json:"budget_exceeded,omitempty"` } type jsonToolCalls struct { @@ -266,7 +267,7 @@ func outputJSON(comments []model.LlmComment) error { func outputJSONWithWarnings(comments []model.LlmComment, warnings []agent.AgentWarning, filesReviewed, inputTokens, outputTokens, totalTokens, cacheReadTokens, cacheWriteTokens int64, - duration time.Duration, projectSummary string, toolCalls map[string]int64, traceID string, resumeInfo *agent.ResumeInfo, sessionID string) error { + duration time.Duration, projectSummary string, toolCalls map[string]int64, traceID string, resumeInfo *agent.ResumeInfo, sessionID string, budgetExceeded bool) error { out := jsonOutput{ Status: "success", TraceID: traceID, @@ -280,6 +281,7 @@ func outputJSONWithWarnings(comments []model.LlmComment, warnings []agent.AgentW CacheReadTokens: cacheReadTokens, CacheWriteTokens: cacheWriteTokens, Elapsed: duration.Round(time.Second).String(), + BudgetExceeded: budgetExceeded, }, ProjectSummary: projectSummary, Resume: resumeInfo, @@ -312,6 +314,13 @@ func outputJSONWithWarnings(comments []model.LlmComment, warnings []agent.AgentW out.Status = "completed_with_warnings" } } + // A tripped budget is a distinct typed terminal state (INV-3): it must + // read "budget_exceeded" regardless of any concurrent subtask warnings / + // errors, since the run stopped for the budget reason and partial results + // were still emitted. Takes precedence over the warning statuses above. + if budgetExceeded { + out.Status = "budget_exceeded" + } enc := json.NewEncoder(os.Stdout) enc.SetIndent("", " ") return enc.Encode(out) @@ -332,6 +341,62 @@ func outputJSONNoFiles(traceID string) error { return enc.Encode(out) } +// emitFailureUsage writes a best-effort structured usage record to stderr when +// a review fails (INV-4): the outer caller must still see the cost of the +// failed attempt instead of losing it. It carries only token/tool-call tallies +// and elapsed — never credentials or prompts (INV-4 no-secrets). +// +// The common failure path is a non-budget error (budget exhaustion returns +// partial comments with a nil error and reaches emitRunResult instead). But +// there is a residual edge: a budget gate can trip after dispatching N files, +// and if every dispatched file then fails, dispatchSubtasks returns +// (nil, error) — reaching here with ag.BudgetExceeded()==true. We therefore +// report the agent's actual BudgetExceeded() value rather than hardcoding +// false, so the record never contradicts the agent's state. +// +// In json format it emits a jsonOutput-shaped object to stderr (kept separate +// from stdout so it does not pollute the machine-readable result stream); +// otherwise a single human-readable [ocr] line. It must never return an error +// that masks the original failure — all writes are best-effort. +func emitFailureUsage(ag ResultProvider, duration time.Duration, outputFormat string) { + var toolTotal int64 + for _, v := range ag.ToolCalls() { + toolTotal += v + } + budgetExceeded := ag.BudgetExceeded() + if outputFormat == "json" { + out := jsonOutput{ + Status: "failed", + Summary: &jsonSummary{ + FilesReviewed: ag.FilesReviewed(), + TotalTokens: ag.TotalTokensUsed(), + InputTokens: ag.TotalInputTokens(), + OutputTokens: ag.TotalOutputTokens(), + CacheReadTokens: ag.TotalCacheReadTokens(), + CacheWriteTokens: ag.TotalCacheWriteTokens(), + Elapsed: duration.Round(time.Second).String(), + BudgetExceeded: budgetExceeded, + }, + ToolCalls: &jsonToolCalls{ + Total: toolTotal, + ByTool: ag.ToolCalls(), + }, + SessionID: ag.SessionID(), + } + enc := json.NewEncoder(os.Stderr) + enc.SetIndent("", " ") + _ = enc.Encode(out) + return + } + fmt.Fprintf(os.Stderr, "[ocr] usage on failure: %d file(s), %d input + %d output = %d total tokens, %d tool calls, elapsed %s, budget_exceeded=%v", + ag.FilesReviewed(), ag.TotalInputTokens(), ag.TotalOutputTokens(), ag.TotalTokensUsed(), + toolTotal, duration.Round(time.Second).String(), budgetExceeded) + if id := ag.SessionID(); id != "" { + fmt.Fprintf(os.Stderr, ", session %s", id) + } + fmt.Fprintln(os.Stderr) +} + func outputPreviewText(p *agent.DiffPreview) { if p.TotalFiles == 0 { fmt.Println("No files changed.") diff --git a/cmd/opencodereview/output_helpers_test.go b/cmd/opencodereview/output_helpers_test.go index dd4acd39..2d2f97e8 100644 --- a/cmd/opencodereview/output_helpers_test.go +++ b/cmd/opencodereview/output_helpers_test.go @@ -168,7 +168,7 @@ func TestOutputJSONWithWarnings_NoCommentsSubtaskError(t *testing.T) { os.Stdout = w warnings := []agent.AgentWarning{{Type: "subtask_error", File: "x.go", Message: "fail"}} - err := outputJSONWithWarnings(nil, warnings, 1, 10, 5, 15, 0, 0, time.Second, "", nil, "abc123trace", nil, "") + err := outputJSONWithWarnings(nil, warnings, 1, 10, 5, 15, 0, 0, time.Second, "", nil, "abc123trace", nil, "", false) _ = w.Close() os.Stdout = old @@ -281,7 +281,7 @@ func TestOutputJSONWithWarnings(t *testing.T) { comments := []model.LlmComment{{Path: "b.go", Content: "test"}} warnings := []agent.AgentWarning{{Type: "subtask_error", File: "c.go", Message: "failed"}} - err := outputJSONWithWarnings(comments, warnings, 5, 100, 50, 150, 10, 5, 3*time.Second, "summary", map[string]int64{"file_read": 3}, "trace-xyz-789", nil, "") + err := outputJSONWithWarnings(comments, warnings, 5, 100, 50, 150, 10, 5, 3*time.Second, "summary", map[string]int64{"file_read": 3}, "trace-xyz-789", nil, "", false) _ = w.Close() os.Stdout = old @@ -319,7 +319,7 @@ func TestOutputJSONWithWarnings_NoCommentsNoErrors(t *testing.T) { os.Stdout = w warnings := []agent.AgentWarning{{Type: "warning", Message: "something"}} - err := outputJSONWithWarnings(nil, warnings, 2, 50, 20, 70, 0, 0, time.Second, "", nil, "", nil, "") + err := outputJSONWithWarnings(nil, warnings, 2, 50, 20, 70, 0, 0, time.Second, "", nil, "", nil, "", false) _ = w.Close() os.Stdout = old @@ -387,6 +387,24 @@ func captureStdout(t *testing.T, fn func()) string { return buf.String() } +// captureStderr captures everything written to os.Stderr during fn. Mirrors +// captureStdout; used to assert structured usage emitted on the failure path. +func captureStderr(t *testing.T, fn func()) string { + t.Helper() + old := os.Stderr + r, w, err := os.Pipe() + if err != nil { + t.Fatalf("os.Pipe: %v", err) + } + os.Stderr = w + fn() + _ = w.Close() + os.Stderr = old + var buf bytes.Buffer + _, _ = buf.ReadFrom(r) + return buf.String() +} + func TestOutputText_NoComments(t *testing.T) { got := captureStdout(t, func() { outputText(nil) diff --git a/cmd/opencodereview/review_cmd.go b/cmd/opencodereview/review_cmd.go index 9ce703f5..ccf1b6cd 100644 --- a/cmd/opencodereview/review_cmd.go +++ b/cmd/opencodereview/review_cmd.go @@ -119,6 +119,7 @@ func runReview(args []string) error { Background: opts.background, GitRunner: cc.GitRunner, Resume: resumeState, + MaxTokensBudget: int64(opts.maxTokensBudget), }) // Silence progress output during execution; restored before the trace @@ -145,6 +146,17 @@ func runReview(args []string) error { if err != nil { span.SetStatus(codes.Error, err.Error()) span.RecordError(err) + // 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, so the record can never contradict it. + // Restore the quiet handle early so the stderr line is visible to + // agent-text audiences (Restore is idempotent; the deferred Restore is + // a no-op). + q.Restore() + emitFailureUsage(ag, time.Since(startTime), opts.outputFormat) if id := ag.SessionID(); id != "" { fmt.Fprintf(os.Stderr, "[ocr] Session: %s (retry with: --resume %s)\n", id, id) } diff --git a/cmd/opencodereview/shared.go b/cmd/opencodereview/shared.go index 42d46dd8..ce30296d 100644 --- a/cmd/opencodereview/shared.go +++ b/cmd/opencodereview/shared.go @@ -255,6 +255,11 @@ type ResultProvider interface { // in JSON output or failure diagnostics. Returns "" when no session was // created. SessionID() string + // BudgetExceeded reports whether a token/tool-call budget gate stopped the + // run before all files were reviewed. The run still returns partial + // comments; this lets the output layer set a typed "budget_exceeded" + // status distinct from success / warnings / errors. + BudgetExceeded() bool } type resumeInfoProvider interface { @@ -310,7 +315,7 @@ func emitRunResult( return outputJSONWithWarnings(comments, ag.Warnings(), ag.FilesReviewed(), ag.TotalInputTokens(), ag.TotalOutputTokens(), ag.TotalTokensUsed(), ag.TotalCacheReadTokens(), ag.TotalCacheWriteTokens(), duration, - ag.ProjectSummary(), ag.ToolCalls(), traceID, resumeInfo, ag.SessionID()) + ag.ProjectSummary(), ag.ToolCalls(), traceID, resumeInfo, ag.SessionID(), ag.BudgetExceeded()) } outputTextWithWarnings(comments, ag.Warnings()) if summary := ag.ProjectSummary(); summary != "" { diff --git a/internal/agent/agent.go b/internal/agent/agent.go index 6417f42f..8a5d4291 100644 --- a/internal/agent/agent.go +++ b/internal/agent/agent.go @@ -117,6 +117,11 @@ type Args struct { // Resume is an optional read-only checkpoint index from a previous review session. Resume *session.ResumeState + + // MaxTokensBudget caps the aggregate token usage (input+output) across the + // whole run; dispatch stops once the running total + a per-file look-ahead + // would exceed it. 0 = unlimited. Mirrors scan.Args.MaxTokensBudget. + MaxTokensBudget int64 } // Agent orchestrates the AI-powered code review. LLM tool-use loop / memory @@ -132,6 +137,7 @@ type Agent struct { subtaskFailed int64 // count of failed subtasks, accessed atomically runner *llmloop.Runner resumeInfo *ResumeInfo + budgetExceeded bool // set when a token/tool-call budget gate stopped dispatch } // ResumeInfo summarizes file-level reuse for a resumed review. @@ -226,6 +232,25 @@ func (a *Agent) Run(ctx context.Context) ([]model.LlmComment, error) { // Record file count metric. telemetry.RecordFilesReviewed(ctx, int64(reviewCount)) + // Pre-run cost projection so users aren't surprised by a large review + // (INV-5). Non-blocking warn-only: the estimate is order-of-magnitude and + // cannot account for agent tool-use inflation (≈300× in the #409 report), + // so it is a floor; real usage is reported from the API after the run. + // + // Gated behind MaxTokensBudget so users who never opt into a budget see no + // new output line (the estimate is only useful to budget-setters comparing + // projected cost against their cap). Keeps the prior text-mode output + // unchanged for the common unlimited path. + 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)) + } + } + // Step 2: Dispatch per-file subtasks concurrently comments, err := a.dispatchSubtasks(ctx) if len(comments) > 0 { @@ -294,6 +319,13 @@ func (a *Agent) Warnings() []AgentWarning { return a.runner.Warnings() } // ToolCalls returns per-tool call counts accumulated during review. func (a *Agent) ToolCalls() map[string]int64 { return a.runner.ToolCalls() } +// BudgetExceeded reports whether a token or tool-call budget gate stopped +// dispatch before all files were reviewed. The run still returns the partial +// comments collected up to that point (and a nil error) so the caller can +// emit a typed budget_exceeded status with the partial results instead of a +// bare failure. +func (a *Agent) BudgetExceeded() bool { return a.budgetExceeded } + // recordWarning adds a non-fatal warning to the agent's warning list. func (a *Agent) recordWarning(warningType, file, message string) { a.runner.RecordWarning(warningType, file, message) @@ -376,6 +408,29 @@ func (a *Agent) dispatchSubtasks(ctx context.Context) ([]model.LlmComment, error if toDispatch[i].IsDeleted { continue } + + // Per-file budget look-ahead, checked BEFORE acquiring the semaphore + // (mirrors scan/agent.go:472-486): if the tokens already spent PLUS a + // look-ahead estimate of this file's cost would exceed the budget, + // stop scheduling further files. Any worker already in flight is + // allowed to finish — its tokens flow into the atomic counter; we do + // NOT cancel it (matches scan, avoids half-written session records). + // Overrun is therefore bounded by the in-flight worker count + // (≤ concurrency, default 8), never a whole batch. + if a.args.MaxTokensBudget > 0 { + used := a.runner.TotalTokensUsed() + nextEst := estimateDiffFileTokens(toDispatch[i]) + projected := used + nextEst + if projected > a.args.MaxTokensBudget { + fmt.Fprintf(stdout.Writer(), "[ocr] token budget reached (used %s + next-file est %s = projected %s > budget %s) — skipping %s and remaining files\n", + humanTokens(used), humanTokens(nextEst), humanTokens(projected), humanTokens(a.args.MaxTokensBudget), toDispatch[i].NewPath) + a.recordWarning("token_budget_reached", toDispatch[i].NewPath, + fmt.Sprintf("stopped dispatch: used %d tokens + next-file estimate %d = projected %d exceeds budget %d", used, nextEst, projected, a.args.MaxTokensBudget)) + a.budgetExceeded = true + break + } + } + dispatched++ wg.Add(1) sem <- struct{}{} // acquire semaphore diff --git a/internal/agent/budget_test.go b/internal/agent/budget_test.go new file mode 100644 index 00000000..d90aeeaa --- /dev/null +++ b/internal/agent/budget_test.go @@ -0,0 +1,170 @@ +package agent + +import ( + "context" + "sync/atomic" + "testing" + + "github.com/alibaba/open-code-review/internal/config/template" + "github.com/alibaba/open-code-review/internal/llm" + "github.com/alibaba/open-code-review/internal/model" + "github.com/alibaba/open-code-review/internal/tool" +) + +// fakeBudgetAgentClient returns a task_done tool call on every request and +// reports a fixed token usage, so each file completes in exactly one round +// and consumes a predictable number of tokens. Used to drive the diff-path +// token-budget gate deterministically. Mirrors scan/budget_test.go's +// fakeBudgetClient. +type fakeBudgetAgentClient struct { + perCallTokens int64 + calls int64 // atomic +} + +func (f *fakeBudgetAgentClient) CompletionsWithCtx(_ context.Context, _ llm.ChatRequest) (*llm.ChatResponse, error) { + atomic.AddInt64(&f.calls, 1) + 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", + }}, + Model: "fake", + Usage: &llm.UsageInfo{ + PromptTokens: f.perCallTokens, + CompletionTokens: 0, + TotalTokens: f.perCallTokens, + }, + }, nil +} + +// budgetAgentTestTemplate returns a minimal template sufficient to drive +// dispatchSubtasks without plan/dedup/summary phases. +func budgetAgentTestTemplate() template.Template { + return template.Template{ + MaxTokens: 100000, + MaxToolRequestTimes: 5, + MainTask: template.LlmConversation{ + Messages: []template.ChatMessage{ + {Role: "system", Content: "review"}, + {Role: "user", Content: "review {{diff}} for {{current_file_path}}"}, + }, + }, + } +} + +// makeBudgetDiffs returns n small, non-deleted diffs that survive filterDiffs +// and filterLargeDiffs. +func makeBudgetDiffs(n int) []model.Diff { + diffs := make([]model.Diff, n) + for i := range diffs { + name := "f" + string(rune('0'+i)) + ".go" + diffs[i] = model.Diff{ + NewPath: name, + OldPath: name, + Diff: "+package x\n", + Insertions: 1, + } + } + return diffs +} + +// TestDispatchSubtasks_TokenBudgetStopsDispatch verifies the per-file gate +// stops dispatch once the running token total + next-file look-ahead would +// blow the budget (INV-1), that a token_budget_reached warning is recorded, +// and that BudgetExceeded()==true with partial results returned (INV-3). +// Overrun is bounded by at most (concurrency) in-flight files — here 1. +func TestDispatchSubtasks_TokenBudgetStopsDispatch(t *testing.T) { + const perCall = 50_000 + fake := &fakeBudgetAgentClient{perCallTokens: perCall} + a := New(Args{ + LLMClient: fake, + Model: "fake", + CommentCollector: tool.NewCommentCollector(), + Tools: tool.NewRegistry(), + MaxConcurrency: 1, // serialize so the gate is deterministic + MaxTokensBudget: 120_000, + Template: budgetAgentTestTemplate(), + MainToolDefs: []llm.ToolDef{ + {Type: "function", Function: llm.FunctionDef{Name: "task_done", Description: "done"}}, + }, + }) + a.diffs = makeBudgetDiffs(10) + a.currentDate = "2025-06-26 10:00" + a.args.Tools.Freeze() + + comments, err := a.dispatchSubtasks(context.Background()) + if err != nil { + t.Fatalf("dispatchSubtasks: %v", err) + } + + // Budget 120K, each file ~50K actual. The look-ahead adds a per-file + // estimate so the gate should stop well before all 10 files run. + calls := atomic.LoadInt64(&fake.calls) + if calls == 0 { + t.Fatal("expected at least one file to be dispatched") + } + if calls >= 10 { + t.Errorf("budget gate did not stop dispatch: all %d files ran (budget should have cut it short)", calls) + } + + // A token_budget_reached warning must be recorded. + var found bool + for _, w := range a.Warnings() { + if w.Type == "token_budget_reached" { + found = true + break + } + } + if !found { + t.Error("expected a token_budget_reached warning") + } + + // Budget exhaustion must signal out-of-band, not via an error (INV-3). + if !a.BudgetExceeded() { + t.Error("expected BudgetExceeded()==true after token budget trip") + } + + // Partial comments returned as a non-nil slice (INV-3 edge — partial + // findings). Even when empty it must not be nil-with-error. + if comments == nil { + t.Error("expected partial comments slice (non-nil), got nil") + } +} + +// TestDispatchSubtasks_UnlimitedBudget verifies MaxTokensBudget=0 runs every +// file (default behavior unchanged — regression guard). +func TestDispatchSubtasks_UnlimitedBudget(t *testing.T) { + fake := &fakeBudgetAgentClient{perCallTokens: 50_000} + a := New(Args{ + LLMClient: fake, + Model: "fake", + CommentCollector: tool.NewCommentCollector(), + Tools: tool.NewRegistry(), + MaxConcurrency: 1, + MaxTokensBudget: 0, // unlimited + Template: budgetAgentTestTemplate(), + MainToolDefs: []llm.ToolDef{ + {Type: "function", Function: llm.FunctionDef{Name: "task_done", Description: "done"}}, + }, + }) + a.diffs = makeBudgetDiffs(5) + a.currentDate = "2025-06-26 10:00" + a.args.Tools.Freeze() + + if _, err := a.dispatchSubtasks(context.Background()); err != nil { + t.Fatalf("dispatchSubtasks: %v", err) + } + if calls := atomic.LoadInt64(&fake.calls); calls != 5 { + t.Errorf("unlimited budget should run all 5 files, ran %d", calls) + } + if a.BudgetExceeded() { + t.Error("unlimited budget must not set BudgetExceeded") + } +} diff --git a/internal/agent/estimate.go b/internal/agent/estimate.go new file mode 100644 index 00000000..7caac6e7 --- /dev/null +++ b/internal/agent/estimate.go @@ -0,0 +1,107 @@ +package agent + +import ( + "fmt" + + "github.com/alibaba/open-code-review/internal/llm" + "github.com/alibaba/open-code-review/internal/model" +) + +// Cost-estimation heuristics for the diff-review path. These mirror +// internal/scan/estimate.go deliberately: scan and review are two paths over +// the same underlying review model, and their estimates should be comparable. +// The values are intentionally re-declared here (rather than imported from +// scan) to avoid a new agent→scan cross-package dependency; keep them in sync +// if scan's change. +// +// As in scan, these are rough — their job is to give the user an +// order-of-magnitude warning before a large review, not to be +// billing-accurate. Real usage is always reported from the API response after +// the run. The estimate also cannot account for agent tool-use inflation +// (≈300× in the #409 report: 308 files → ≈90.4M tokens), so it is a floor. +const ( + // promptOverheadTokens approximates the fixed prompt scaffolding per LLM + // call (system prompt + template wrappers + tool definitions). + // KEEP IN SYNC with internal/scan/estimate.go. + promptOverheadTokens = 2000 + // avgMainRoundsPerFile is the assumed number of MAIN_TASK tool-use rounds + // for a typical file. Observed ~6 on real repos; round up. + // KEEP IN SYNC with internal/scan/estimate.go. + avgMainRoundsPerFile = 7 + // avgOutputTokensPerRound approximates completion tokens per round. + // KEEP IN SYNC with internal/scan/estimate.go. + avgOutputTokensPerRound = 700 +) + +// Estimate is a pre-run, order-of-magnitude projection of review cost. It +// mirrors scan.Estimate so callers can treat the two paths uniformly. +type Estimate struct { + Files int + InputTokens int64 + OutputTokens int64 + TotalTokens int64 +} + +// estimateDiffFileTokens projects the input+output token cost of reviewing a +// single diff (PLAN + MAIN_TASK rounds). It mirrors scan.estimateFileTokens +// but counts tokens of the diff text (d.Diff) rather than whole-file content, +// since the diff path reviews patches, not full files. Returns 0 for deleted +// files (they are skipped before dispatch and must not trip the gate). Used +// both by the aggregate estimate (estimateDiffCost) and by the per-file budget +// look-ahead in dispatchSubtasks. +func estimateDiffFileTokens(d model.Diff) int64 { + if d.IsDeleted || d.Diff == "" { + return 0 + } + diffTokens := int64(llm.CountTokens(d.Diff)) + + // PLAN phase input/output (small JSON output). + total := diffTokens + promptOverheadTokens + total += 400 + // MAIN_TASK: diff content carried across rounds + per-round overhead. + total += (diffTokens + promptOverheadTokens) * avgMainRoundsPerFile + total += avgOutputTokensPerRound * avgMainRoundsPerFile + return total +} + +// estimateDiffCost projects token usage for reviewing the given diffs. It is +// the diff-path analogue of scan.estimateCost and is used for the pre-review +// scale warning. +func estimateDiffCost(diffs []model.Diff) Estimate { + var est Estimate + for _, d := range diffs { + if d.IsDeleted || d.Diff == "" { + continue + } + est.Files++ + diffTokens := int64(llm.CountTokens(d.Diff)) + est.InputTokens += diffTokens + promptOverheadTokens + est.OutputTokens += 400 + est.InputTokens += (diffTokens + promptOverheadTokens) * avgMainRoundsPerFile + est.OutputTokens += avgOutputTokensPerRound * avgMainRoundsPerFile + } + est.TotalTokens = est.InputTokens + est.OutputTokens + return est +} + +// String renders a one-line human-readable estimate, matching scan.Estimate's +// format so diff and scan warnings read consistently. Money is intentionally +// omitted — pricing varies per provider/model and we don't imply a precise +// dollar figure. +func (e Estimate) String() string { + return fmt.Sprintf("~%d file(s), est. %s input + %s output ≈ %s total tokens (rough; agent tool-use inflates this — actual reported after run)", + e.Files, humanTokens(e.InputTokens), humanTokens(e.OutputTokens), humanTokens(e.TotalTokens)) +} + +// humanTokens formats a token count as e.g. "1.2M" / "850K" / "420". +// Mirrors internal/scan/estimate.go; kept here so agent has no scan import. +func humanTokens(n int64) string { + switch { + case n >= 1_000_000: + return fmt.Sprintf("%.1fM", float64(n)/1_000_000) + case n >= 1_000: + return fmt.Sprintf("%.0fK", float64(n)/1_000) + default: + return fmt.Sprintf("%d", n) + } +} diff --git a/internal/agent/estimate_test.go b/internal/agent/estimate_test.go new file mode 100644 index 00000000..2ed823ff --- /dev/null +++ b/internal/agent/estimate_test.go @@ -0,0 +1,99 @@ +package agent + +import ( + "strings" + "testing" + + "github.com/alibaba/open-code-review/internal/model" +) + +// These tests mirror internal/scan/estimate_test.go deliberately: the agent +// estimate helpers are re-declared copies that must stay in sync with scan's +// (see internal/agent/estimate.go). The humanTokens table in particular is +// duplicated verbatim so a divergence between the two copies is caught here. + +// TestHumanTokens mirrors scan's TestHumanTokens so the diff- and scan-path +// formatters stay byte-identical. If this table and scan's ever disagree, one +// copy has drifted. +func TestHumanTokens(t *testing.T) { + cases := map[int64]string{ + 0: "0", + 420: "420", + 999: "999", + 1000: "1K", + 1500: "2K", // rounds + 850_000: "850K", + 1_000_000: "1.0M", + 2_400_000: "2.4M", + } + for in, want := range cases { + if got := humanTokens(in); got != want { + t.Errorf("humanTokens(%d) = %q, want %q", in, got, want) + } + } +} + +// TestEstimateDiffFileTokens_ZeroForSkipped verifies deleted and empty diffs +// project to zero tokens so they never trip the budget gate (they are skipped +// before dispatch). Also asserts a normal diff projects a sane positive value +// bounded below by the fixed per-file overhead. +func TestEstimateDiffFileTokens_ZeroForSkipped(t *testing.T) { + if got := estimateDiffFileTokens(model.Diff{IsDeleted: true, Diff: "+x"}); got != 0 { + t.Errorf("deleted diff projected %d tokens, want 0", got) + } + if got := estimateDiffFileTokens(model.Diff{NewPath: "a.go", Diff: ""}); got != 0 { + t.Errorf("empty diff projected %d tokens, want 0", got) + } + got := estimateDiffFileTokens(model.Diff{NewPath: "a.go", Diff: "+package main\nfunc f() {}\n"}) + if got <= 0 { + t.Errorf("normal diff projected %d tokens, want > 0", got) + } + // The look-ahead must be at least the fixed per-file overhead + // (promptOverhead*(1+rounds) + plan output + main output) so even a tiny + // diff carries meaningful cost in the projection. + const minExpected = int64(promptOverheadTokens*(1+avgMainRoundsPerFile) + 400 + avgOutputTokensPerRound*avgMainRoundsPerFile) + if got < minExpected { + t.Errorf("projected %d, want >= %d (fixed overhead floor)", got, minExpected) + } +} + +// TestEstimateDiffCost verifies the aggregate estimate sums per-file costs, +// skips deleted/empty diffs, and satisfies TotalTokens == Input + Output. Used +// by the pre-review scale warning (INV-5). +func TestEstimateDiffCost(t *testing.T) { + diffs := []model.Diff{ + {NewPath: "a.go", Diff: "+a\n"}, + {NewPath: "b.go", Diff: "+b\n"}, + {NewPath: "c.go", IsDeleted: true, Diff: "+c\n"}, // skipped + {NewPath: "d.go", Diff: ""}, // skipped + } + est := estimateDiffCost(diffs) + if est.Files != 2 { + t.Errorf("Files = %d, want 2 (deleted + empty skipped)", est.Files) + } + perFile := estimateDiffFileTokens(diffs[0]) + if est.TotalTokens != perFile*2 { + t.Errorf("TotalTokens = %d, want 2*%d = %d", est.TotalTokens, perFile, perFile*2) + } + if est.InputTokens <= 0 || est.OutputTokens <= 0 { + t.Errorf("expected non-zero input/output splits, got in=%d out=%d", est.InputTokens, est.OutputTokens) + } + // Split invariant the String() rendering relies on. + if est.TotalTokens != est.InputTokens+est.OutputTokens { + t.Errorf("TotalTokens = %d, want Input(%d)+Output(%d) = %d", + est.TotalTokens, est.InputTokens, est.OutputTokens, est.InputTokens+est.OutputTokens) + } + if s := est.String(); s == "" || !strings.Contains(s, "token") { + t.Errorf("expected non-empty estimate string mentioning tokens, got %q", s) + } +} + +// TestEstimateDiffCost_ScalesWithContent verifies that a larger diff projects +// strictly more tokens than a small one (the estimate isn't a flat constant). +func TestEstimateDiffCost_ScalesWithContent(t *testing.T) { + small := estimateDiffFileTokens(model.Diff{NewPath: "a.go", Diff: "+x\n"}) + large := estimateDiffFileTokens(model.Diff{NewPath: "a.go", Diff: strings.Repeat("line of code\n", 200)}) + if large <= small { + t.Errorf("expected large diff to project more tokens than small: large=%d small=%d", large, small) + } +} diff --git a/internal/scan/agent.go b/internal/scan/agent.go index 18b55aa1..240fac74 100644 --- a/internal/scan/agent.go +++ b/internal/scan/agent.go @@ -191,6 +191,14 @@ 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 } + func (a *Agent) recordWarning(warningType, file, message string) { a.runner.RecordWarning(warningType, file, message) }