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

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)
+	}
+}

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}

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)

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 {