diff --git a/cmd/tinyagent/attempts.go b/cmd/tinyagent/attempts.go index a7e2e43..81e0083 100644 --- a/cmd/tinyagent/attempts.go +++ b/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) }, diff --git a/cmd/tinyagent/check.go b/cmd/tinyagent/check.go index 8e11fad..8aa1508 100644 --- a/cmd/tinyagent/check.go +++ b/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 { diff --git a/cmd/tinyagent/check_attempts.go b/cmd/tinyagent/check_attempts.go index b30d2eb..cdd3eb3 100644 --- a/cmd/tinyagent/check_attempts.go +++ b/cmd/tinyagent/check_attempts.go @@ -55,7 +55,7 @@ func checkAttempts(fs *flag.FlagSet, o *options, cfg config.AgentConfig) (string name = "" } 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; 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; 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(), diff --git a/cmd/tinyagent/check_attempts_test.go b/cmd/tinyagent/check_attempts_test.go index f96f549..e723e5a 100644 --- a/cmd/tinyagent/check_attempts_test.go +++ b/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; cost and time limits are shared by all attempts, which can cost up to 3x a single run" + filepath.Join(top, worktreeDir, "best") + "-a; 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, "") + "-a; cost and time limits are shared by all attempts, which can cost up to 2x a single run" + filepath.Join(top, worktreeDir, "") + "-a; 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) } diff --git a/cmd/tinyagent/check_test.go b/cmd/tinyagent/check_test.go index acea668..3b330f0 100644 --- a/cmd/tinyagent/check_test.go +++ b/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)") { diff --git a/cmd/tinyagent/check_token_limit_test.go b/cmd/tinyagent/check_token_limit_test.go new file mode 100644 index 0000000..7ae5bf6 --- /dev/null +++ b/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) + } +} diff --git a/cmd/tinyagent/init_config.go b/cmd/tinyagent/init_config.go index d560348..da10e52 100644 --- a/cmd/tinyagent/init_config.go +++ b/cmd/tinyagent/init_config.go @@ -12,6 +12,7 @@ const starterConfig = `# tinyagent user config. Flags override these values. Ful # model: openai/gpt-4o # same / 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 diff --git a/docs/eval-framework.md b/docs/eval-framework.md index 8c1f579..7393534 100644 --- a/docs/eval-framework.md +++ b/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 diff --git a/docs/project-context.md b/docs/project-context.md index 5440978..49e460e 100644 --- a/docs/project-context.md +++ b/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. diff --git a/docs/runs.md b/docs/runs.md index 3dff762..4f7c194 100644 --- a/docs/runs.md +++ b/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: `, a summary line per attempt (status, checks, cost, worktree), then `winner: attempt i (branch ...)` with a `git diff` hint, or `no attempt passed (); 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. diff --git a/internal/attempts/attempts.go b/internal/attempts/attempts.go index b838011..b7ee8ea 100644 --- a/internal/attempts/attempts.go +++ b/internal/attempts/attempts.go @@ -76,8 +76,10 @@ type Options struct { // and the report to .attempts. 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 +} diff --git a/internal/attempts/token_budget_test.go b/internal/attempts/token_budget_test.go new file mode 100644 index 0000000..c2f616a --- /dev/null +++ b/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) + } +} diff --git a/internal/eval/run.go b/internal/eval/run.go index 87d1361..9aa4c9d 100644 --- a/internal/eval/run.go +++ b/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)) } diff --git a/internal/eval/suite.go b/internal/eval/suite.go index 7b17cf8..f398185 100644 --- a/internal/eval/suite.go +++ b/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 == "" { diff --git a/internal/eval/token_limit_test.go b/internal/eval/token_limit_test.go new file mode 100644 index 0000000..24ee6f8 --- /dev/null +++ b/internal/eval/token_limit_test.go @@ -0,0 +1,40 @@ +package eval + +import ( + "context" + "io" + "path/filepath" + "slices" + "strings" + "testing" +) + +func TestTaskTokenLimitForwarded(t *testing.T) { + repo, _ := gitRepo(t) + body := "version: 1\nmodels: [openai/gpt-4o]\ntrials: 1\ntasks:\n" + + " - {id: capped, repo: " + repo + ", base: HEAD, task: status=Submitted, limits: {token_limit: 1234}, checks: [{argv: [git, --version]}]}\n" + + " - {id: open, repo: " + repo + ", base: HEAD, task: status=Submitted, checks: [{argv: [git, --version]}]}\n" + s, err := Load(writeSuite(t, body), nil) + if err != nil { + t.Fatal(err) + } + fa := &fakeAgent{} + Run(context.Background(), Options{Suite: s, OutDir: filepath.Join(t.TempDir(), "out"), MinPassRate: -1, Stderr: io.Discard, RunAgent: fa.run}) + if len(fa.calls) != 2 { + t.Fatalf("%d calls", len(fa.calls)) + } + if got := argAfter(fa.calls[0], "--token-limit"); got != "1234" { + t.Errorf("capped task args %v", fa.calls[0]) + } + if slices.Contains(fa.calls[1], "--token-limit") { + t.Errorf("uncapped task got a token limit: %v", fa.calls[1]) + } +} + +func TestTaskTokenLimitNegativeRejected(t *testing.T) { + repo, _ := gitRepo(t) + body := "version: 1\nmodels: [openai/gpt-4o]\ntrials: 1\ntasks:\n - {id: bad, repo: " + repo + ", base: HEAD, task: x, limits: {token_limit: -1}, checks: [{argv: [git, --version]}]}\n" + if _, err := Load(writeSuite(t, body), nil); err == nil || !strings.Contains(err.Error(), "limits must be >= 0") { + t.Fatalf("err = %v", err) + } +} diff --git a/internal/harness/loop.go b/internal/harness/loop.go index a69a2c2..d77439f 100644 --- a/internal/harness/loop.go +++ b/internal/harness/loop.go @@ -224,6 +224,9 @@ func runAttempt(ctx context.Context, o Options, st Stage, idx, n int, input stri if st.CostLimit > 0 { args = append(args, "--cost-limit", strconv.FormatFloat(st.CostLimit, 'g', -1, 64)) } + if st.TokenLimit > 0 { + args = append(args, "--token-limit", strconv.Itoa(st.TokenLimit)) + } if st.TimeLimit > 0 { args = append(args, "--time-limit", strconv.Itoa(st.TimeLimit)) } diff --git a/internal/harness/token_limit_test.go b/internal/harness/token_limit_test.go new file mode 100644 index 0000000..d513a47 --- /dev/null +++ b/internal/harness/token_limit_test.go @@ -0,0 +1,49 @@ +package harness + +import ( + "context" + "io" + "os" + "strings" + "testing" +) + +func TestStageTokenLimit(t *testing.T) { + w := loadTestWorkflow(t, "version: 1\ndefaults:\n token_limit: 5000\nworkflows:\n build:\n stages:\n - skill: plan-tasks\n inputs: {spec: \"s\"}\n - skill: plan-tasks\n token_limit: 700\n inputs: {spec: \"s\"}\n") + r, _ := w.Route("", "build") + var calls [][]string + agent := func(_ context.Context, args []string, _ io.Writer) error { + calls = append(calls, args) + for i := range args { + if args[i] == "-o" { + os.WriteFile(args[i+1], []byte(`{"exit_status":"Submitted","submission":"SUCCESS CHECK: PASS"}`), 0o600) + } + } + return nil + } + Run(context.Background(), Options{Route: r, Task: "t", Cwd: "/w", Dir: t.TempDir(), RunAgent: agent}) + if len(calls) != 2 { + t.Fatalf("%d stage runs, want 2", len(calls)) + } + if s := strings.Join(calls[0], " "); !strings.Contains(s, "--token-limit 5000 ") { + t.Errorf("stage 1 args %q missing the default token limit", s) + } + if s := strings.Join(calls[1], " "); !strings.Contains(s, "--token-limit 700 ") { + t.Errorf("stage 2 args %q missing its own token limit", s) + } +} + +func TestStageTokenLimitUnsetAndNegative(t *testing.T) { + w := loadTestWorkflow(t, "version: 1\nworkflows:\n build:\n stages:\n - skill: plan-tasks\n inputs: {spec: \"s\"}\n") + r, _ := w.Route("", "build") + if r.Pipeline.Stages[0].TokenLimit != 0 { + t.Fatalf("token limit defaults to %d, want 0 (off)", r.Pipeline.Stages[0].TokenLimit) + } + if err := checkSettings("defaults", Settings{TokenLimit: new(int)}); err != nil { + t.Fatalf("token_limit 0 rejected: %v", err) + } + neg := -1 + if err := checkSettings("defaults", Settings{TokenLimit: &neg}); err == nil || !strings.Contains(err.Error(), "token_limit must be >= 0") { + t.Fatalf("negative token_limit: %v", err) + } +} diff --git a/internal/harness/workflow.go b/internal/harness/workflow.go index ae87d88..fd702c6 100644 --- a/internal/harness/workflow.go +++ b/internal/harness/workflow.go @@ -34,6 +34,7 @@ type Settings struct { Profile string `yaml:"profile" json:"profile,omitempty"` StepLimit *int `yaml:"step_limit" json:"step_limit,omitempty"` CostLimit *float64 `yaml:"cost_limit" json:"cost_limit,omitempty"` + TokenLimit *int `yaml:"token_limit" json:"token_limit,omitempty"` TimeLimit *int `yaml:"time_limit" json:"time_limit,omitempty"` Retries *int `yaml:"retries" json:"retries,omitempty"` CheckTimeout *int `yaml:"check_timeout" json:"check_timeout,omitempty"` @@ -52,6 +53,7 @@ type Stage struct { Profile string `json:"profile,omitempty"` StepLimit int `json:"step_limit"` CostLimit float64 `json:"cost_limit,omitempty"` + TokenLimit int `json:"token_limit,omitempty"` TimeLimit int `json:"time_limit,omitempty"` Inputs map[string]string `json:"inputs,omitempty"` } @@ -220,6 +222,9 @@ func apply(st *Stage, s Settings) { if s.CostLimit != nil { st.CostLimit = *s.CostLimit } + if s.TokenLimit != nil { + st.TokenLimit = *s.TokenLimit + } if s.TimeLimit != nil { st.TimeLimit = *s.TimeLimit } @@ -237,6 +242,8 @@ func checkSettings(where string, s Settings) error { return fmt.Errorf("%s: step_limit must be >= 0", where) case s.CostLimit != nil && *s.CostLimit < 0: return fmt.Errorf("%s: cost_limit must be >= 0", where) + case s.TokenLimit != nil && *s.TokenLimit < 0: + return fmt.Errorf("%s: token_limit must be >= 0", where) case s.TimeLimit != nil && *s.TimeLimit < 0: return fmt.Errorf("%s: time_limit must be >= 0", where) case s.Retries != nil && (*s.Retries < 0 || *s.Retries > MaxRetries):