Fix #236: show token limit in --check/init-config and forward it through eval, attempts and workflows

cmd/tinyagent/attempts.go

@@ -68,7 +68,7 @@ var attemptConflicts = []string{"resume", "rewind", "to", "transcript", "acp", "
 // attemptOwn are flags --attempts consumes itself; every other flag set on
 // the command line is forwarded to each attempt's child run.
 var attemptOwn = []string{"attempts", "attempt-model", "attempt-check", "attempt-check-timeout",
-	"m", "model", "o", "output", "worktree", "cost-limit", "time-limit"}
+	"m", "model", "o", "output", "worktree", "cost-limit", "token-limit", "time-limit"}
 
 func runAttempts(o *options, fs *flag.FlagSet, args []string, stdout, stderr io.Writer) int {
 	a := o.attempts
@@ -137,7 +137,7 @@ func runAttempts(o *options, fs *flag.FlagSet, args []string, stdout, stderr io.
 		AttemptsID: link.id, TraceParent: link.traceParent(),
 		CheckEnv: attemptCheckEnv(fs, *o, cfg, stderr),
 		Models:   models, Checks: checks, BaseName: base, Output: o.output,
-		CostLimit: o.costLimit, TimeLimit: time.Duration(o.timeLimit) * time.Second,
+		CostLimit: o.costLimit, TokenLimit: o.tokenLimit, TimeLimit: time.Duration(o.timeLimit) * time.Second,
 		ExtraArgs: forwardArgs(fs, args, attemptOwn), Stderr: stderr, Stdout: stdout,
 		Redact: red.Redact, RunAll: !cfg.Attempts.StopsOnPass(), Keep: cfg.Attempts.KeepMode(),
 		Remove: func(at attempts.Attempt) error { return removeAttemptWorktree(top, base, at) },

cmd/tinyagent/check.go

@@ -166,6 +166,7 @@ func resolveCheck(o *options, fs *flag.FlagSet) ([][2]string, error) {
 	add("limits", strings.Join([]string{
 		"steps " + limitText(fs, "step-limit", cfg.StepLimit > 0, fmt.Sprint(cfg.StepLimit)),
 		"cost " + limitText(fs, "cost-limit", cfg.CostLimit > 0, formatCost(cfg.CostLimit)),
+		"tokens " + limitText(fs, "token-limit", cfg.TokenLimit > 0, fmt.Sprint(cfg.TokenLimit)),
 		"time " + limitText(fs, "time-limit", cfg.WallTimeSec > 0, fmt.Sprintf("%ds", cfg.WallTimeSec)),
 	}, ", "))
 	for _, s := range cfg.MCPServers {

cmd/tinyagent/check_attempts.go

@@ -55,7 +55,7 @@ func checkAttempts(fs *flag.FlagSet, o *options, cfg config.AgentConfig) (string
 		name = "<generated>"
 	}
 	n := len(p.models)
-	return fmt.Sprintf("up to %d runs (%s; %s), %s (%s, timeout %s), stop_on_pass %t, keep %s, worktrees %s-a<N>; cost and time limits are shared by all attempts, which can cost up to %dx a single run",
+	return fmt.Sprintf("up to %d runs (%s; %s), %s (%s, timeout %s), stop_on_pass %t, keep %s, worktrees %s-a<N>; cost, token and time limits are shared by all attempts, which can cost up to %dx a single run",
 		n, strings.Join(shown, ", "), modelSrc,
 		countNoun(len(p.checks), "check"), checkSrc, strings.Join(timeouts, ", "),
 		cfg.Attempts.StopsOnPass(), cfg.Attempts.KeepMode(),

cmd/tinyagent/check_attempts_test.go

@@ -35,7 +35,7 @@ func TestCheckAttemptsPlan(t *testing.T) {
 	}
 	want := "up to 3 runs (anthropic/claude-sonnet-4-5, anthropic/claude-sonnet-4-5, anthropic/claude-sonnet-4-5; from -m), " +
 		"1 check (from --attempt-check, timeout 90s), stop_on_pass true, keep all, worktrees " +
-		filepath.Join(top, worktreeDir, "best") + "-a<N>; cost and time limits are shared by all attempts, which can cost up to 3x a single run"
+		filepath.Join(top, worktreeDir, "best") + "-a<N>; cost, token and time limits are shared by all attempts, which can cost up to 3x a single run"
 	if got := reportLine(t, stdout, "attempts"); got != want {
 		t.Errorf("attempts line\n got %q\nwant %q", got, want)
 	}
@@ -57,7 +57,7 @@ func TestCheckAttemptsPlan(t *testing.T) {
 	}
 	want = "up to 2 runs (anthropic/claude-sonnet-4-5, openai/gpt-4o; from --attempt-model), " +
 		"2 checks (from config, timeout 600s, 30s), stop_on_pass false, keep failed, worktrees " +
-		filepath.Join(top, worktreeDir, "<generated>") + "-a<N>; cost and time limits are shared by all attempts, which can cost up to 2x a single run"
+		filepath.Join(top, worktreeDir, "<generated>") + "-a<N>; cost, token and time limits are shared by all attempts, which can cost up to 2x a single run"
 	if got := reportLine(t, stdout, "attempts"); got != want {
 		t.Errorf("config attempts line\n got %q\nwant %q", got, want)
 	}

cmd/tinyagent/check_test.go

@@ -71,7 +71,7 @@ func TestCheckValidSetupExits0WithoutRequestOrSecret(t *testing.T) {
 	if got := reportLine(t, stdout, "cwd"); got != root+" (from --cwd)" {
 		t.Errorf("cwd line %q, want %q", got, root+" (from --cwd)")
 	}
-	if got := reportLine(t, stdout, "limits"); got != "steps unlimited, cost unlimited, time unlimited" {
+	if got := reportLine(t, stdout, "limits"); got != "steps unlimited, cost unlimited, tokens unlimited, time unlimited" {
 		t.Errorf("limits line %q", got)
 	}
 	if !strings.HasSuffix(stdout, "\nok\n") {
@@ -98,7 +98,7 @@ func TestCheckSourceLabels(t *testing.T) {
 	if got := reportLine(t, stdout, "config"); got != cfg+" (from -c)" {
 		t.Errorf("config line %q", got)
 	}
-	if got := reportLine(t, stdout, "limits"); got != "steps 40 (from config), cost $2.50 (from config), time 90s (from --time-limit)" {
+	if got := reportLine(t, stdout, "limits"); got != "steps 40 (from config), cost $2.50 (from config), tokens unlimited, time 90s (from --time-limit)" {
 		t.Errorf("limits line %q", got)
 	}
 	if got := reportLine(t, stdout, "cwd"); !strings.HasSuffix(got, " (default)") {

cmd/tinyagent/check_token_limit_test.go

@@ -0,0 +1,26 @@
+package main
+
+import (
+	"strings"
+	"testing"
+)
+
+func TestCheckShowsTokenLimit(t *testing.T) {
+	checkEnv(t)
+	t.Setenv("ANTHROPIC_API_KEY", "k")
+	cfg := writeConfig(t, "model: anthropic/claude-haiku-4-5\ntoken_limit: 9000\n")
+	_, stdout, stderr := runCLI(t, "--check", "-c", cfg)
+	if got := reportLine(t, stdout, "limits"); got != "steps unlimited, cost unlimited, tokens 9000 (from config), time unlimited" {
+		t.Errorf("limits line %q; stderr %s", got, stderr)
+	}
+	_, stdout, _ = runCLI(t, "--check", "-c", cfg, "--token-limit", "500")
+	if got := reportLine(t, stdout, "limits"); !strings.Contains(got, "tokens 500 (from --token-limit)") {
+		t.Errorf("limits line %q", got)
+	}
+}
+
+func TestStarterConfigMentionsTokenLimit(t *testing.T) {
+	if !strings.Contains(starterConfig, "\n# token_limit: 0 ") {
+		t.Fatalf("starter config has no token_limit line:\n%s", starterConfig)
+	}
+}

cmd/tinyagent/init_config.go

@@ -12,6 +12,7 @@ const starterConfig = `# tinyagent user config. Flags override these values. Ful
 # model: openai/gpt-4o        # same <provider>/<model> syntax as -m
 # step_limit: 50
 # cost_limit: 2.0
+# token_limit: 0              # max input+output tokens, delegates and workers included (0 = unlimited)
 # wall_time_seconds: 1800
 # read_only: false
 # redact_secrets:             # mask known secret values in command, MCP, web_fetch and hook output

docs/eval-framework.md

@@ -203,7 +203,7 @@ Strict keys: any unknown key, at any level, is an error. The suite is fully vali
 | `checks_on_limit` | bool | no | `true` | Run checks even when the trial ended on a limit. |
 | `output_schema` | path | no | — | JSON Schema for the final result (as `--output-schema`). |
 | `expected_json` | object | no | — | Expected-JSON assertions on the structured result. Key name and matching semantics are not fixed by the draft (**Recommendation:** subset match). |
-| `limits` | map | no | profile/config | Per-task limits. **Recommendation:** keys mirroring existing flags: `step_limit`, `cost_limit`, `time_limit`. |
+| `limits` | map | no | profile/config | Per-task limits. **Recommendation:** keys mirroring existing flags: `step_limit`, `cost_limit`, `token_limit`, `time_limit`. |
 
 ### Check
 

docs/project-context.md

@@ -81,7 +81,7 @@ workflows:
         check_timeout: 600
 ```
 
-Settings in `defaults` or on a stage: `model`, `profile`, `step_limit`, `cost_limit`, `time_limit`, `retries`, and `check_timeout` (seconds). The defaults are 1 retry (at most 5), 50 steps, a 300 second check timeout, and otherwise your usual config. Flags you pass with `--workflow` (such as `--sandbox` or `-m`) are forwarded to every stage, and stage settings override them. Workflows are sequential and have at most 20 stages. A complete example is in [docs/workflows](workflows).
+Settings in `defaults` or on a stage: `model`, `profile`, `step_limit`, `cost_limit`, `token_limit`, `time_limit`, `retries`, and `check_timeout` (seconds). The defaults are 1 retry (at most 5), 50 steps, a 300 second check timeout, and otherwise your usual config. Flags you pass with `--workflow` (such as `--sandbox` or `-m`) are forwarded to every stage, and stage settings override them. Workflows are sequential and have at most 20 stages. A complete example is in [docs/workflows](workflows).
 
 **Routing.** `--workflow-name name` runs that workflow, or a single skill, directly. Otherwise the router matches the task against each workflow's intent and each skill's intent line. A clear best match runs. A tie or no match falls back to `default`, or exits 2 with an error suggesting `--workflow-name`, before any model call.
 

docs/runs.md

@@ -59,7 +59,7 @@ tinyagent --attempts 2 --attempt-model openai/gpt-5 --attempt-model anthropic/cl
 - **What counts as a pass.** An attempt passes when its run ends `Submitted` or `Completed` and every check exits 0 before its timeout. Checks run in order with `bash -c` in the attempt's working directory (the worktree path that matches `--cwd`); when the run didn't finish successfully its checks are skipped. The winner is the lowest-numbered passing attempt. By default tinyagent stops at the first pass; `attempts.stop_on_pass: false` runs all N attempts and compares them.
 - **Checks.** `--attempt-check` (repeatable) with `--attempt-check-timeout` replaces the config's `attempts.checks` list, whose entries each take a `command` and an optional `timeout_seconds` (default 600, up to 3600). At least one check is required. Checks run where the attempts' commands run: with `--env container`, in a fresh container from `container.image` (same `run_args`) with the attempt's worktree and any additional directories mounted, started for that attempt's checks and removed after them; with `--sandbox`, under a sandbox rooted at the attempt's worktree with the same writable paths and network settings; with a `command_env` policy, with that environment. Otherwise they run on the host with your environment. stderr names the place (`attempts: checks run in the container, like the attempts' commands`). If the container or sandbox can't be set up, that attempt's checks fail with the reason and it can't win; they never fall back to the host. The report keeps the last 8 KiB of each check's output, with known secrets redacted.
 - **Models.** Without `--attempt-model`, every attempt uses the model a single run would use (`-m`, then `$AGENT_MODEL`, then the config's `model:`). With it, give one `--attempt-model` per attempt. Every attempt's model, provider key and client are checked before the first attempt runs, and a problem exits 2 naming the attempt.
-- **Budgets.** `--cost-limit` and `--time-limit` cover the whole invocation: each attempt gets what is left, and no new attempt starts once either is used up (stop reason `budget exhausted`). Without limits, N attempts can cost up to N times a single run, and the setup line says so.
+- **Budgets.** `--cost-limit`, `--token-limit` and `--time-limit` cover the whole invocation: each attempt gets what is left, and no new attempt starts once any of them is used up (stop reason `budget exhausted`). Without limits, N attempts can cost up to N times a single run, and the setup line says so.
 - **Uncommitted changes.** Like `--worktree`, attempts start from the last commit. The dirty-checkout warning prints once for the invocation.
 - **Output.** Progress goes to stderr: `attempt i/N: <model>`, a summary line per attempt (status, checks, cost, worktree), then `winner: attempt i (branch ...)` with a `git diff` hint, or `no attempt passed (<reason>); worktrees kept for inspection`. The winner's stdout (its submission) is printed to stdout at the end.
 - **JSON events.** With `--json`, stdout carries JSONL instead: every event from each attempt's run, with an added `attempt` field (otherwise unchanged), plus `attempt.started` (`attempt`, `attempts_requested`, `model`, `worktree_name`), one `attempt.check` per check (`command`, `exit_code`, `duration_ms`, `timed_out`, `skipped`, `output`, with known secrets redacted from the command and output), `attempt.result` (the attempt's report fields without `checks`) and one final `attempts.result` (`schema_version`, `attempts_requested`, `attempts_run`, `winner`, `winner_branch`, `stop_reason`, `stop_on_pass`, `keep`, `removed`, totals, and `report` when `-o` wrote one). The winner's output isn't printed again at the end; its `result` event already carries the submission.

internal/attempts/attempts.go

@@ -76,8 +76,10 @@ type Options struct {
 	// and the report to <stem>.attempts<ext>. Empty means trajectories go to
 	// a temporary directory and no report file is written.
 	Output    string
-	CostLimit float64       // invocation-wide USD budget; 0 = none
-	TimeLimit time.Duration // invocation-wide wall-clock budget; 0 = none
+	CostLimit float64 // invocation-wide USD budget; 0 = none
+	// TokenLimit is the invocation-wide input+output token budget; 0 = none.
+	TokenLimit int
+	TimeLimit  time.Duration // invocation-wide wall-clock budget; 0 = none
 	// ExtraArgs are forwarded to every child run (the task among them).
 	ExtraArgs []string
 	Stderr    io.Writer
@@ -256,6 +258,13 @@ func Run(ctx context.Context, o Options) (Report, int) {
 		n := i + 1
 		var remCost float64
 		var remTime time.Duration
+		remTokens := 0
+		if o.TokenLimit > 0 {
+			if remTokens = o.TokenLimit - usedTokens(rep.Attempts); remTokens <= 0 {
+				rep.StopReason = StopBudget
+				break
+			}
+		}
 		if o.CostLimit > 0 {
 			if remCost = o.CostLimit - rep.TotalCostUSD; remCost <= 0 {
 				rep.StopReason = StopBudget
@@ -278,7 +287,7 @@ func Run(ctx context.Context, o Options) (Report, int) {
 		} else {
 			traj, _ = filepath.Abs(TrajectoryPath(o.Output, n))
 		}
-		a := runAttempt(ctx, o, ev, n, m, traj, remCost, remTime, now)
+		a := runAttempt(ctx, o, ev, n, m, traj, remCost, remTokens, remTime, now)
 		if trajDir != "" {
 			a.Trajectory = ""
 		}
@@ -297,7 +306,8 @@ func Run(ctx context.Context, o Options) (Report, int) {
 			}
 		case a.Status == "Interrupted" || ctx.Err() != nil:
 			rep.StopReason = StopInterrupt
-		case a.Status == "CostLimitExceeded" && o.CostLimit > 0, a.Status == "TimeLimitExceeded" && o.TimeLimit > 0:
+		case a.Status == "CostLimitExceeded" && o.CostLimit > 0, a.Status == "TimeLimitExceeded" && o.TimeLimit > 0,
+			a.Status == "TokenLimitExceeded" && o.TokenLimit > 0:
 			rep.StopReason = StopBudget
 		case a.Error != "" && a.Status == "":
 			rep.StopReason = StopError
@@ -377,7 +387,7 @@ func attemptLine(a Attempt) string {
 	return fmt.Sprintf("attempt %d  %s  %s  %s  $%.4f  %s", a.Attempt, a.Model, status, checks, a.CostUSD, wt)
 }
 
-func runAttempt(ctx context.Context, o Options, ev *eventWriter, n int, m, traj string, remCost float64, remTime time.Duration, now func() time.Time) (a Attempt) {
+func runAttempt(ctx context.Context, o Options, ev *eventWriter, n int, m, traj string, remCost float64, remTokens int, remTime time.Duration, now func() time.Time) (a Attempt) {
 	a = Attempt{Attempt: n, Model: m, Trajectory: traj, Checks: []CheckResult{}}
 	t0 := now()
 	if o.AttemptsID != "" {
@@ -390,6 +400,9 @@ func runAttempt(ctx context.Context, o Options, ev *eventWriter, n int, m, traj
 	if remCost > 0 {
 		args = append(args, "--cost-limit", strconv.FormatFloat(remCost, 'g', -1, 64))
 	}
+	if remTokens > 0 {
+		args = append(args, "--token-limit", strconv.Itoa(remTokens))
+	}
 	if remTime > 0 {
 		args = append(args, "--time-limit", strconv.Itoa(int(math.Floor(remTime.Seconds()))))
 	}
@@ -622,3 +635,13 @@ func RemoveWorktree(repoRoot, path, branch string) error {
 	}
 	return run("branch", "-D", branch)
 }
+
+// usedTokens sums the provider-reported input and output tokens of the
+// attempts run so far.
+func usedTokens(as []Attempt) int {
+	n := 0
+	for _, a := range as {
+		n += a.InputTokens + a.OutputTokens
+	}
+	return n
+}

internal/attempts/token_budget_test.go

@@ -0,0 +1,53 @@
+package attempts
+
+import (
+	"context"
+	"io"
+	"strings"
+	"testing"
+)
+
+// Each fake attempt reports 10 input + 5 output tokens.
+func TestTokenBudgetPassedOnAndExhausted(t *testing.T) {
+	skipNoBash(t)
+	f := &fakeAgent{t: t, statuses: []string{"Submitted", "Submitted", "Submitted"}}
+	rep, code := Run(context.Background(), Options{
+		Models: []string{"a", "b", "c"}, BaseName: "x", TokenLimit: 25,
+		Checks: []Check{{Command: "false"}},
+		Stderr: io.Discard, RunAgent: f.run,
+	})
+	if code != 1 || rep.StopReason != StopBudget || len(f.calls) != 2 {
+		t.Fatalf("code=%d calls=%d rep=%+v", code, len(f.calls), rep)
+	}
+	if got := strings.Join(f.calls[0], " "); !strings.Contains(got, "--token-limit 25") {
+		t.Fatalf("first attempt should get the whole budget: %s", got)
+	}
+	if got := strings.Join(f.calls[1], " "); !strings.Contains(got, "--token-limit 10") {
+		t.Fatalf("second attempt should get the remaining budget: %s", got)
+	}
+}
+
+func TestTokenLimitExceededStopsAttempts(t *testing.T) {
+	skipNoBash(t)
+	f := &fakeAgent{t: t, statuses: []string{"TokenLimitExceeded", "Submitted"}}
+	rep, _ := Run(context.Background(), Options{
+		Models: []string{"a", "b"}, BaseName: "x", TokenLimit: 1000,
+		Checks: []Check{{Command: "true"}},
+		Stderr: io.Discard, RunAgent: f.run,
+	})
+	if rep.StopReason != StopBudget || len(f.calls) != 1 {
+		t.Fatalf("calls=%d rep=%+v", len(f.calls), rep)
+	}
+}
+
+func TestNoTokenLimitNotForwarded(t *testing.T) {
+	skipNoBash(t)
+	f := &fakeAgent{t: t, statuses: []string{"Submitted"}}
+	Run(context.Background(), Options{
+		Models: []string{"a"}, BaseName: "x",
+		Checks: []Check{{Command: "true"}}, Stderr: io.Discard, RunAgent: f.run,
+	})
+	if got := strings.Join(f.calls[0], " "); strings.Contains(got, "--token-limit") {
+		t.Fatalf("unset token limit forwarded: %s", got)
+	}
+}

internal/eval/run.go

@@ -228,6 +228,9 @@ func runTrial(ctx context.Context, o Options, tr Trial) (rec Record) {
 	if l.CostLimit > 0 {
 		args = append(args, "--cost-limit", strconv.FormatFloat(l.CostLimit, 'g', -1, 64))
 	}
+	if l.TokenLimit > 0 {
+		args = append(args, "--token-limit", strconv.Itoa(l.TokenLimit))
+	}
 	if l.TimeLimit > 0 {
 		args = append(args, "--time-limit", strconv.Itoa(l.TimeLimit))
 	}

internal/eval/suite.go

@@ -59,9 +59,10 @@ type Check struct {
 
 // Limits are per-trial run limits passed to the child tinyagent.
 type Limits struct {
-	StepLimit int     `yaml:"step_limit"`
-	CostLimit float64 `yaml:"cost_limit"`
-	TimeLimit int     `yaml:"time_limit"` // seconds
+	StepLimit  int     `yaml:"step_limit"`
+	CostLimit  float64 `yaml:"cost_limit"`
+	TokenLimit int     `yaml:"token_limit"` // input+output tokens
+	TimeLimit  int     `yaml:"time_limit"`  // seconds
 }
 
 var idRe = regexp.MustCompile(`^[A-Za-z0-9][A-Za-z0-9._-]*$`)
@@ -175,7 +176,7 @@ func (s *Suite) validate(knownProfile func(string) bool) error {
 				c.TimeoutSeconds = DefaultCheckTimeout
 			}
 		}
-		if t.Limits.StepLimit < 0 || t.Limits.CostLimit < 0 || t.Limits.TimeLimit < 0 {
+		if t.Limits.StepLimit < 0 || t.Limits.CostLimit < 0 || t.Limits.TokenLimit < 0 || t.Limits.TimeLimit < 0 {
 			return fmt.Errorf("%s: limits must be >= 0", where)
 		}
 		if t.Repo == "" {