feat: cost controls — pre-flight estimation, token budget, file priority

README.md

@@ -134,6 +134,7 @@ Settings are applied in priority order: **CLI flags > env vars > `.code-reviewer
 | `--sarif` | Write SARIF 2.1.0 output to file | — |
 | `--no-color` | Disable ANSI color output | `false` |
 | `--no-context` | Disable repo-aware cross-file context injection | `false` |
+| `--max-tokens` | Maximum total tokens per review (0 = unlimited) | `0` |
 | `--incremental` | Only review files changed in latest push (CI mode) | `false` |
 | `--proxy-url` | Route model calls through an LLM proxy (e.g. Candela) | — |
 | `--version` | Print version and exit | — |
@@ -157,6 +158,7 @@ Settings are applied in priority order: **CLI flags > env vars > `.code-reviewer
 | `SARIF_OUTPUT` | Write SARIF output to this file path | — |
 | `INCREMENTAL` | Only review changed files in latest push (`true`/`false`) | `false` |
 | `EXCLUDED_PATTERNS` | Glob patterns to skip | `go.sum,*.lock,vendor/*` |
+| `REVIEW_MAX_TOKENS` | Maximum total tokens per review (0 = unlimited) | `0` |
 | `NO_COLOR` | Disable ANSI colors ([no-color.org](https://no-color.org)) | — |
 
 ### Per-Repo Config
@@ -175,6 +177,7 @@ excluded_patterns:
 extra_rules: |
   Always flag raw SQL string concatenation.
   Check that zerolog is used instead of log/fmt.
+max_tokens: 50000  # Optional: cap total tokens per review
 ```
 
 See [`.code-reviewer.example.yaml`](.code-reviewer.example.yaml) for all options.

internal/config/config.go

@@ -5,8 +5,10 @@ package config
 import (
 	"flag"
 	"fmt"
+	"log/slog"
 	"os"
 	"path/filepath"
+	"strconv"
 	"strings"
 
 	"gopkg.in/yaml.v3"
@@ -117,6 +119,9 @@ type Config struct {
 
 	// Context discovery.
 	DisableContext bool // Skip repo-aware context discovery (--no-context).
+
+	// Budget.
+	MaxTokens int // Maximum total tokens per review (0 = unlimited).
 }
 
 // repoConfig represents the .code-reviewer.yaml file.
@@ -131,6 +136,7 @@ type repoConfig struct {
 	OutputJSON       bool     `yaml:"output_json"`
 	CustomPrompt     string   `yaml:"custom_prompt"`
 	ProxyURL         string   `yaml:"proxy_url"`
+	MaxTokens        int      `yaml:"max_tokens"`
 }
 
 // DefaultExcludedPatterns are file patterns excluded by default.
@@ -250,6 +256,9 @@ func (c *Config) applyRepoConfig(data []byte) error {
 	if rc.ProxyURL != "" {
 		c.ProxyURL = rc.ProxyURL
 	}
+	if rc.MaxTokens > 0 {
+		c.MaxTokens = rc.MaxTokens
+	}
 	return nil
 }
 
@@ -314,6 +323,17 @@ func (c *Config) loadEnv() {
 	if v := os.Getenv("REVIEW_PROXY_URL"); v != "" {
 		c.ProxyURL = v
 	}
+	if v := os.Getenv("REVIEW_MAX_TOKENS"); v != "" {
+		n, err := strconv.Atoi(v)
+		if err != nil {
+			slog.Warn("ignoring invalid REVIEW_MAX_TOKENS", "value", v, "error", err)
+		} else if n < 0 {
+			slog.Warn("ignoring negative REVIEW_MAX_TOKENS", "value", n)
+		} else {
+			// 0 = unlimited (clears any YAML cap), >0 = token budget.
+			c.MaxTokens = n
+		}
+	}
 }
 
 func (c *Config) loadFlags() error {
@@ -340,6 +360,7 @@ func (c *Config) loadFlags() error {
 	incremental := fs.Bool("incremental", false, "Only review files changed in the latest push (CI mode)")
 	proxyURL := fs.String("proxy-url", "", "LLM proxy URL for observability (e.g., http://localhost:8181/proxy/google/)")
 	noContext := fs.Bool("no-context", false, "Disable repo-aware context discovery")
+	maxTokens := fs.Int("max-tokens", 0, "Maximum total tokens (input+output) per review (0 = unlimited)")
 
 	if err := fs.Parse(os.Args[1:]); err != nil {
 		return err
@@ -408,6 +429,16 @@ func (c *Config) loadFlags() error {
 	if *noContext {
 		c.DisableContext = true
 	}
+	// Detect if --max-tokens was explicitly set (including to 0 for unlimited).
+	fs.Visit(func(f *flag.Flag) {
+		if f.Name == "max-tokens" {
+			if *maxTokens < 0 {
+				slog.Warn("ignoring negative --max-tokens", "value", *maxTokens)
+				return
+			}
+			c.MaxTokens = *maxTokens
+		}
+	})
 
 	return nil
 }

internal/diff/priority.go

@@ -0,0 +1,113 @@
+package diff
+
+import (
+	"path/filepath"
+	"sort"
+	"strings"
+)
+
+// securityPatterns are path substrings that indicate security-sensitive files.
+var securityPatterns = []string{
+	"auth", "session", "token", "secret", "secrets", "password", "credential",
+	"crypto", "security", "permission", "rbac", "acl", "oauth",
+	"saml", "jwt", "cert", "tls", "ssl", "key",
+}
+
+// securityFiles are exact filenames or extensions that are security-sensitive.
+var securityFiles = []string{
+	".env", "Dockerfile", "docker-compose",
+	"config.yaml", "config.yml", "config.json",
+}
+
+// generatedPatterns are path substrings that indicate generated code.
+var generatedPatterns = []string{
+	".pb.go", ".pb.gw.go", "_generated", "_gen.go",
+	"generated.go", "mock_", "mocks/", "zz_generated",
+	"vendor/", "node_modules/", "dist/",
+}
+
+// FilePriority assigns a review priority score to a file diff.
+// Higher score = review first. Used to decide which files to skip
+// when a token budget is exceeded.
+func FilePriority(fd FileDiff) int {
+	score := 0
+	path := strings.ToLower(fd.NewPath)
+
+	// Security-sensitive files get highest priority.
+	if isSecuritySensitive(path) {
+		score += 100
+	}
+
+	// More changed lines = higher priority.
+	score += fd.LineCount()
+
+	// New files get a boost — no prior review coverage.
+	if fd.IsNew {
+		score += 30
+	}
+
+	// Penalize low-value files.
+	if isGenerated(path) {
+		score -= 50
+	}
+	if isTestFile(path) {
+		score -= 10 // Still review tests, but prioritize prod code.
+	}
+	if fd.IsRename && fd.LineCount() == 0 {
+		score -= 40 // Pure rename with no content changes.
+	}
+
+	return score
+}
+
+// SortByPriority sorts diffs highest-priority first.
+func SortByPriority(diffs []FileDiff) {
+	sort.SliceStable(diffs, func(i, j int) bool {
+		return FilePriority(diffs[i]) > FilePriority(diffs[j])
+	})
+}
+
+func isSecuritySensitive(path string) bool {
+	// Split into path components and check each against security patterns.
+	parts := strings.FieldsFunc(path, func(r rune) bool {
+		return r == '/' || r == '\\' || r == '_' || r == '-' || r == '.'
+	})
+	for _, part := range parts {
+		part = strings.ToLower(part)
+		for _, p := range securityPatterns {
+			if part == p {
+				return true
+			}
+		}
+	}
+	base := filepath.Base(path)
+	for _, f := range securityFiles {
+		if base == f || strings.HasPrefix(base, f) {
+			return true
+		}
+	}
+	return false
+}
+
+func isGenerated(path string) bool {
+	for _, p := range generatedPatterns {
+		if strings.Contains(path, p) {
+			return true
+		}
+	}
+	return false
+}
+
+func isTestFile(path string) bool {
+	base := filepath.Base(path)
+	return strings.HasSuffix(base, "_test.go") ||
+		strings.HasSuffix(base, ".test.ts") ||
+		strings.HasSuffix(base, ".test.js") ||
+		strings.HasSuffix(base, ".spec.ts") ||
+		strings.HasSuffix(base, ".spec.js") ||
+		strings.HasSuffix(base, "Test.java") ||
+		strings.HasSuffix(base, "Test.kt") ||
+		strings.HasPrefix(base, "test_") ||
+		strings.Contains(path, "/test/") ||
+		strings.Contains(path, "/tests/")
+}

internal/diff/priority_test.go

@@ -0,0 +1,122 @@
+package diff
+
+import "testing"
+
+func TestFilePriority_SecurityFiles(t *testing.T) {
+	tests := []struct {
+		path string
+		want bool // should score > 100
+	}{
+		{"internal/auth/handler.go", true},
+		{"pkg/crypto/encrypt.go", true},
+		{"config/secrets.yaml", true},
+		{".env", true},
+		{"internal/handler/list.go", false},
+		{"README.md", false},
+		// Negative cases — should NOT match despite containing security substrings.
+		{"pkg/monkey.go", false},        // contains "key" but is not security-related
+		{"docs/authors.md", false},      // contains "auth" but is not security-related
+		{"internal/tokenizer.go", false}, // contains "token" but is not security-related
+	}
+
+	for _, tt := range tests {
+		fd := FileDiff{NewPath: tt.path, Hunks: []Hunk{{Lines: []DiffLine{{Type: LineAdded}}}}}
+		score := FilePriority(fd)
+		if tt.want && score < 100 {
+			t.Errorf("FilePriority(%q) = %d, expected >= 100 (security-sensitive)", tt.path, score)
+		}
+		if !tt.want && score >= 100 {
+			t.Errorf("FilePriority(%q) = %d, expected < 100 (not security-sensitive)", tt.path, score)
+		}
+	}
+}
+
+func TestFilePriority_NewVsModified(t *testing.T) {
+	newFile := FileDiff{
+		NewPath: "pkg/handler.go",
+		IsNew:   true,
+		Hunks:   []Hunk{{Lines: []DiffLine{{Type: LineAdded}, {Type: LineAdded}}}},
+	}
+	modFile := FileDiff{
+		NewPath: "pkg/handler.go",
+		Hunks:   []Hunk{{Lines: []DiffLine{{Type: LineAdded}, {Type: LineAdded}}}},
+	}
+
+	newScore := FilePriority(newFile)
+	modScore := FilePriority(modFile)
+
+	if newScore <= modScore {
+		t.Errorf("new file (%d) should score higher than modified file (%d)", newScore, modScore)
+	}
+}
+
+func TestFilePriority_Generated(t *testing.T) {
+	gen := FileDiff{
+		NewPath: "api/service.pb.go",
+		Hunks:   []Hunk{{Lines: []DiffLine{{Type: LineAdded}}}},
+	}
+	normal := FileDiff{
+		NewPath: "api/service.go",
+		Hunks:   []Hunk{{Lines: []DiffLine{{Type: LineAdded}}}},
+	}
+
+	if FilePriority(gen) >= FilePriority(normal) {
+		t.Errorf("generated file should score lower than normal file")
+	}
+}
+
+func TestFilePriority_RenameOnly(t *testing.T) {
+	rename := FileDiff{
+		NewPath:  "pkg/new_name.go",
+		OldPath:  "pkg/old_name.go",
+		IsRename: true,
+		// No hunks = no changed lines.
+	}
+	modified := FileDiff{
+		NewPath: "pkg/handler.go",
+		Hunks:   []Hunk{{Lines: []DiffLine{{Type: LineAdded}, {Type: LineAdded}, {Type: LineAdded}}}},
+	}
+
+	if FilePriority(rename) >= FilePriority(modified) {
+		t.Errorf("pure rename (%d) should score lower than modified file (%d)",
+			FilePriority(rename), FilePriority(modified))
+	}
+}
+
+func TestFilePriority_TestFile(t *testing.T) {
+	test := FileDiff{
+		NewPath: "pkg/handler_test.go",
+		Hunks:   []Hunk{{Lines: []DiffLine{{Type: LineAdded}, {Type: LineAdded}}}},
+	}
+	prod := FileDiff{
+		NewPath: "pkg/handler.go",
+		Hunks:   []Hunk{{Lines: []DiffLine{{Type: LineAdded}, {Type: LineAdded}}}},
+	}
+
+	if FilePriority(test) >= FilePriority(prod) {
+		t.Errorf("test file (%d) should score lower than prod file (%d)",
+			FilePriority(test), FilePriority(prod))
+	}
+}
+
+func TestSortByPriority(t *testing.T) {
+	diffs := []FileDiff{
+		{NewPath: "README.md"},                               // low priority
+		{NewPath: "internal/auth/handler.go", IsNew: true,    // high priority
+			Hunks: []Hunk{{Lines: []DiffLine{{Type: LineAdded}, {Type: LineAdded}}}}},
+		{NewPath: "vendor/lib.go"},                            // generated = very low
+		{NewPath: "pkg/server.go",
+			Hunks: []Hunk{{Lines: []DiffLine{{Type: LineAdded}}}}}, // medium
+	}
+
+	SortByPriority(diffs)
+
+	// Auth file should be first.
+	if diffs[0].NewPath != "internal/auth/handler.go" {
+		t.Errorf("expected auth file first, got %q", diffs[0].NewPath)
+	}
+	// Vendor file should be last.
+	if diffs[len(diffs)-1].NewPath != "vendor/lib.go" {
+		t.Errorf("expected vendor file last, got %q", diffs[len(diffs)-1].NewPath)
+	}
+}