diff --git a/cmd/tinyagent/acp.go b/cmd/tinyagent/acp.go index 4bd6994..03c4eba 100644 --- a/cmd/tinyagent/acp.go +++ b/cmd/tinyagent/acp.go @@ -453,6 +453,9 @@ func (s *acpSession) Prompt(ctx context.Context, p acp.Prompt, u *acp.Updates) ( if cfg.CostLimit > 0 { cfg.CostLimit += prev.Cost } + if cfg.TokenLimit > 0 { + cfg.TokenLimit += prev.InputTokens + prev.OutputTokens + } } if s.turn == 1 { for _, n := range s.notes { diff --git a/cmd/tinyagent/budget_test.go b/cmd/tinyagent/budget_test.go new file mode 100644 index 0000000..a5a015b --- /dev/null +++ b/cmd/tinyagent/budget_test.go @@ -0,0 +1,50 @@ +package main + +import ( + "path/filepath" + "testing" + + "github.com/vicgupta/tinyagent/internal/agent" + "github.com/vicgupta/tinyagent/internal/config" + "github.com/vicgupta/tinyagent/internal/trajectory" +) + +func TestTokenLimitFlagOverridesConfig(t *testing.T) { + var o options + fs := newFlagSet(&o) + if err := fs.Parse([]string{"--token-limit", "5000"}); err != nil { + t.Fatal(err) + } + cfg := config.AgentConfig{TokenLimit: 100, StepLimit: 7} + applyFlagOverrides(fs, &cfg, o.stepLimit, o.costLimit, o.timeLimit) + if cfg.TokenLimit != 5000 || cfg.StepLimit != 7 { + t.Errorf("token limit %d, step limit %d", cfg.TokenLimit, cfg.StepLimit) + } + + fs = newFlagSet(&o) + cfg = config.AgentConfig{TokenLimit: 100} + applyFlagOverrides(fs, &cfg, 0, 0, 0) + if cfg.TokenLimit != 100 { + t.Errorf("unset flag changed the config's token_limit to %d", cfg.TokenLimit) + } +} + +func TestTokenLimitMustNotBeNegative(t *testing.T) { + if msg := validateLimits(options{tokenLimit: -1, maxFormatErrors: 1}); msg != "--token-limit must be >= 0 (0 = unlimited)" { + t.Errorf("validateLimits = %q", msg) + } + if msg := validateLimits(options{tokenLimit: 0, maxFormatErrors: 1}); msg != "" { + t.Errorf("validateLimits(0) = %q", msg) + } +} + +func TestResumeCommandDoublesTokenLimit(t *testing.T) { + dir := t.TempDir() + run := filepath.Join(dir, "run.json") + tr := trajectory.Trajectory{ExitStatus: agent.StatusTokenLimit} + got := resumeCommand(&tr, run, "/src", config.AgentConfig{TokenLimit: 1000}) + want := "tinyagent --resume " + run + " --cwd /src --token-limit 2000 -o " + filepath.Join(dir, "run-2.json") + if got != want { + t.Errorf("resumeCommand = %q, want %q", got, want) + } +} diff --git a/cmd/tinyagent/help.go b/cmd/tinyagent/help.go index 16dacb5..e083687 100644 --- a/cmd/tinyagent/help.go +++ b/cmd/tinyagent/help.go @@ -65,7 +65,7 @@ var helpGroups = []flagGroup{ {"", [][3]string{{"h", "help", ""}}}, {"Task", [][3]string{{"t", "task", "string"}, {"", "cwd", "dir"}, {"", "worktree", "[=name]"}, {"", "resume", "path"}, {"", "ask-user", ""}, {"", "acp", ""}, {"", "image", "path"}, {"", "images", ""}, {"", "untrusted-task", ""}}}, {"Model and provider", [][3]string{{"m", "model", "string"}, {"", "reasoning-effort", "level"}, {"", "fallback-model", "spec"}, {"", "no-fallback", ""}, {"", "openai-api", "api"}, {"c", "config", "path"}, {"", "profile", "name"}}}, - {"Run control", [][3]string{{"", "step-limit", "n"}, {"", "cost-limit", "usd"}, {"", "time-limit", "sec"}, {"", "max-format-errors", "n"}, {"", "max-output-tokens", "n"}, {"", "no-compact", ""}, {"", "subagents", ""}, {"", "workers", ""}}}, + {"Run control", [][3]string{{"", "step-limit", "n"}, {"", "cost-limit", "usd"}, {"", "token-limit", "n"}, {"", "time-limit", "sec"}, {"", "max-format-errors", "n"}, {"", "max-output-tokens", "n"}, {"", "no-compact", ""}, {"", "subagents", ""}, {"", "workers", ""}}}, {"Verify before stop", verifyGroup.flags}, {"Attempts (best of N)", attemptsGroup.flags}, {"Review", [][3]string{{"", "review", ""}, {"", "review-base", "ref"}, {"", "review-commit", "rev"}}}, diff --git a/cmd/tinyagent/main.go b/cmd/tinyagent/main.go index 44b1193..e008300 100644 --- a/cmd/tinyagent/main.go +++ b/cmd/tinyagent/main.go @@ -72,6 +72,7 @@ type options struct { redactSecrets bool stepLimit, timeLimit, maxFormatErrors int maxOutputTokens int + tokenLimit int costLimit float64 verbose, noStream, noAgentsMD, noCompact bool readOnly, sandbox, sandboxNetwork bool @@ -128,6 +129,7 @@ func newFlagSet(o *options) *flag.FlagSet { fs.Var(&o.worktree, "worktree", "run in a new Git worktree on branch tinyagent/ from HEAD of the repo containing --cwd, kept afterwards (default name generated)") fs.IntVar(&o.stepLimit, "step-limit", 0, "max model calls, 0 = unlimited") fs.Float64Var(&o.costLimit, "cost-limit", 0, "max USD, 0 = unlimited") + fs.IntVar(&o.tokenLimit, "token-limit", 0, "max input+output tokens, delegates and workers included, 0 = unlimited") fs.IntVar(&o.timeLimit, "time-limit", 0, "max wall-clock seconds, 0 = unlimited") fs.IntVar(&o.maxFormatErrors, "max-format-errors", config.DefaultMaxFormatErrors, "stop after this many consecutive steps of malformed tool calls (>= 1)") fs.IntVar(&o.maxOutputTokens, "max-output-tokens", 0, "per-request output-token budget sent to the provider, 0 = provider default (8192 for anthropic; config max_output_tokens)") @@ -1756,7 +1758,7 @@ var usageGroups = []struct { {"", [][3]string{{"h", "help", ""}}}, {"Task", [][3]string{{"t", "task", "string"}, {"m", "model", "string"}, {"", "reasoning-effort", "level"}, {"", "fallback-model", "spec"}, {"", "no-fallback", ""}, {"", "openai-api", "api"}, {"c", "config", "path"}, {"", "profile", "name"}, {"", "cwd", "dir"}, {"", "worktree", "[=name]"}, {"", "resume", "path"}, {"", "ask-user", ""}, {"", "acp", ""}}}, {"Review", [][3]string{{"", "review", ""}, {"", "review-base", "ref"}, {"", "review-commit", "rev"}}}, - {"Limits", [][3]string{{"", "step-limit", "n"}, {"", "cost-limit", "usd"}, {"", "time-limit", "sec"}, {"", "max-format-errors", "n"}, {"", "max-output-tokens", "n"}, {"", "no-compact", ""}}}, + {"Limits", [][3]string{{"", "step-limit", "n"}, {"", "cost-limit", "usd"}, {"", "token-limit", "n"}, {"", "time-limit", "sec"}, {"", "max-format-errors", "n"}, {"", "max-output-tokens", "n"}, {"", "no-compact", ""}}}, {"Output", [][3]string{{"o", "output", "path"}, {"v", "verbose", ""}, {"", "no-stream", ""}, {"", "json", ""}, {"", "output-schema", "path"}, {"", "log-db", "[=path]"}, {"", "otel", ""}, {"", "no-agents-md", ""}, {"", "scoped-agents-md", ""}, {"", "read-only", ""}, {"", "add-dir-read", "path"}, {"", "add-dir-write", "path"}, {"", "sandbox", ""}, {"", "sandbox-network", ""}, {"", "sandbox-allow-domain", "host"}, {"", "env", "name"}, {"", "command-env", "mode"}, {"", "snapshots", ""}}}, {"Rewind", [][3]string{{"", "rewind", "path"}, {"", "to", "target"}, {"", "dry-run", ""}}}, {"Other", [][3]string{{"V", "version", ""}, {"", "update", ""}, {"", "check", ""}, {"", "list-profiles", ""}, {"", "list-plugins", ""}, {"", "plugins-prune", ""}, {"", "plugins-keep", "n"}, {"", "plugins-rollback", "module@sha256"}, {"", "tinygo-install", ""}, {"", "tinygo-update", ""}, {"", "tinygo-status", ""}, {"", "transcript", "path"}, {"", "init-config", ""}, {"", "completion", "shell"}}}, @@ -1853,6 +1855,10 @@ func resumeCommand(t *trajectory.Trajectory, output, root string, cfg config.Age if cfg.CostLimit > 0 { limit = " --cost-limit " + strconv.FormatFloat(2*cfg.CostLimit, 'f', -1, 64) } + case agent.StatusTokenLimit: + if cfg.TokenLimit > 0 { + limit = " --token-limit " + strconv.Itoa(2*cfg.TokenLimit) + } case agent.StatusTimeLimit: if cfg.WallTimeSec > 0 { limit = " --time-limit " + strconv.Itoa(2*cfg.WallTimeSec) @@ -1996,6 +2002,8 @@ func validateLimits(o options) string { return "--step-limit must be >= 0 (0 = unlimited)" case !(o.costLimit >= 0) || math.IsInf(o.costLimit, 1): return "--cost-limit must be finite and >= 0 (0 = unlimited)" + case o.tokenLimit < 0: + return "--token-limit must be >= 0 (0 = unlimited)" case o.timeLimit < 0: return "--time-limit must be >= 0" case o.maxFormatErrors < 1: @@ -2017,6 +2025,12 @@ func applyFlagOverrides(fs *flag.FlagSet, cfg *config.AgentConfig, stepLimit int cfg.CostLimit = costLimit case "time-limit": cfg.WallTimeSec = timeLimit + case "token-limit": + if g, ok := f.Value.(flag.Getter); ok { + if n, ok := g.Get().(int); ok { + cfg.TokenLimit = n + } + } } }) } diff --git a/cmd/tinyagent/main_test.go b/cmd/tinyagent/main_test.go index 3439b55..523f398 100644 --- a/cmd/tinyagent/main_test.go +++ b/cmd/tinyagent/main_test.go @@ -383,6 +383,7 @@ Review: Limits: --step-limit n max model calls, 0 = unlimited --cost-limit usd max USD, 0 = unlimited + --token-limit n max input+output tokens, delegates and workers included, 0 = unlimited --time-limit sec max wall-clock seconds, 0 = unlimited --max-format-errors n stop after this many consecutive steps of malformed tool calls (>= 1) (default 3) --max-output-tokens n per-request output-token budget sent to the provider, 0 = provider default (8192 for anthropic; config max_output_tokens) diff --git a/docs/config.md b/docs/config.md index 33fa450..130b5ae 100644 --- a/docs/config.md +++ b/docs/config.md @@ -52,6 +52,7 @@ additional_dirs: # omitted by default (only --cwd); see Additional dir - {path: /home/me/src/shared-lib, mode: read} step_limit: 50 cost_limit: 2.0 +token_limit: 0 # max input+output tokens per run, delegates and workers included; 0 = unlimited wall_time_seconds: 1800 max_format_errors: 3 # consecutive malformed-call steps before RepeatedFormatError (>= 1) bash_timeout_seconds: 120 diff --git a/docs/long-runs.md b/docs/long-runs.md index 3ebd02d..91f21ed 100644 --- a/docs/long-runs.md +++ b/docs/long-runs.md @@ -101,7 +101,7 @@ A checkpoint records whether `apply_patch` was enabled (`apply_patch_enabled`), `--subagents` (or `subagents.enabled: true`) adds a `delegate` tool. A call such as `{"task": "find every caller of LoadConfig and say which ones ignore its error", "label": "callers"}` starts a child agent with a fresh context: the same system prompt and `AGENTS.md`, plus a notice that it is a subagent that must answer first and then give evidence with `file:line` references. The child runs to completion, and the parent gets back only its final report followed by a footer such as `[delegate d1: Submitted, 4 steps, 9200 in / 410 out tokens, $0.0123]`. It is off by default; with it off, requests, tools, and checkpoints are unchanged. - **Read-only.** The child has only `read_file`, `list_files`, and `grep`, whatever the parent has. It gets no `write_file`, `edit_file`, `apply_patch`, `bash`, process tools, MCP tools, `skill`, `view_image`, `update_plan`, or `delegate`, and calls to them are refused. Paths follow the parent's rules (`--cwd` and `file_access`). The tool is also offered with `--read-only`. -- **Limits.** Each child stops after `subagents.max_steps` model steps (default 30); hitting the cap, or any other unfinished status, returns a tool error such as `error: delegate d1 did not finish: StepLimit` rather than a report. Child steps don't count toward `--step-limit`; they're reported separately as `subagent_steps`. Child tokens and cost do count toward the run's totals and `--cost-limit`, and the child only gets what's left of the budget, so a delegated task can end the whole run as `CostLimit`. `--time-limit` covers children too. Reports longer than `max_output_chars` are truncated like any tool output. Several `delegate` calls in one turn run one after another. +- **Limits.** Each child stops after `subagents.max_steps` model steps (default 30); hitting the cap, or any other unfinished status, returns a tool error such as `error: delegate d1 did not finish: StepLimit` rather than a report. Child steps don't count toward `--step-limit`; they're reported separately as `subagent_steps`. Child tokens and cost do count toward the run's totals, `--cost-limit` and `--token-limit`, and the child only gets what's left of the budget, so a delegated task can end the whole run as `CostLimit` or `TokenLimitExceeded`. `--time-limit` covers children too. Reports longer than `max_output_chars` are truncated like any tool output. Several `delegate` calls in one turn run one after another. - **Model.** `subagents.model` picks a different (typically cheaper) model; empty uses the parent's. `reasoning_effort` and `max_output_tokens` carry over only when the child uses the same provider as the parent. An unknown provider, or a missing key for a different provider, exits 2 before any request. - **Tokens.** Delegation saves the parent's context, not tokens overall: the child resends its own history each step, so a task the parent could answer with one `grep` is cheaper done directly. Use it for searches that would otherwise pull many large files into the parent's history. - **Observability.** `--json` and `--log-db` add `subagent.started` (`id`, `label`, `model`, `tool_call_id`, `parent_step`) and `subagent.finished` (`id`, `status`, `steps`, tokens, `cost_usd`) events, and the child's own events carry `"subagent": "d1"`. `-v` prints a line when each child starts and finishes; the child's model output is not streamed. Hooks run for child tool calls with `subagent` set in their input, and `pre_tool`/`post_tool` hooks also see the `delegate` call itself; stop hooks run only for the parent. With `-o`, the trajectory records `subagents_enabled`, `subagent_steps`, and a `subagents` array with each child's task, status, usage, report, and full message history. diff --git a/docs/runs.md b/docs/runs.md index b1efc7e..3dff762 100644 --- a/docs/runs.md +++ b/docs/runs.md @@ -151,7 +151,7 @@ cat next.md | tinyagent --resume fix.json -t - --cwd . -o next.json - A finished checkpoint requires a non-empty follow-up; `--resume finished.json` without `-t` exits 2 and says so. An unfinished checkpoint still refuses `-t` (exit 2): continue it with plain `--resume`. - Every other resume rule applies unchanged: the same model, root, system prompt, tools and loading mode, a distinct new `-o`, the unchanged source file, and every refusal before provider setup. Runs with `--output-schema` or MCP tools, runs continued by `stop` hooks, and `VerificationFailed` runs are not eligible. - The request is the saved history (or the rebuilt compacted view) followed by one user message: a notice that the previous invocation finished and the workspace may have changed, then the follow-up text. No historical tool call is re-executed. The trajectory's top-level `task` stays the original task; the follow-up text is recorded in that message. -- Steps, cost, and tokens stay cumulative, so `--step-limit` and `--cost-limit` count the earlier runs too. If the ceiling is already used up, no request is made and the run ends `StepLimitExceeded` or `CostLimitExceeded`; raise the ceiling for the follow-up. +- Steps, cost, and tokens stay cumulative, so `--step-limit`, `--cost-limit` and `--token-limit` count the earlier runs too. If the ceiling is already used up, no request is made and the run ends `StepLimitExceeded`, `CostLimitExceeded` or `TokenLimitExceeded`; raise the ceiling for the follow-up. - Follow-ups chain: a follow-up's finished result can take another follow-up, and a follow-up stopped by a limit continues with plain `--resume`. Each boundary is recorded in `resume` with `follow_up: true`, and in `--json` `run.started` carries `resumed: true` and `follow_up: true`. - Checkpoints with a follow-up boundary use the existing schema version; older tinyagent binaries refuse them with an unknown-field error rather than misreading them. - A `NeedsInput` run is answered the same way, but the answer is delivered differently; see [Asking for input](#asking-for-input). @@ -177,3 +177,9 @@ tinyagent --resume run.json --cwd . -t "per API key" -o run-2.json - **Where it applies.** It is offered with `--read-only` and in review mode, never to [subagents](long-runs.md#subagents). `--ask-user` without `-o`, or with `--output-schema`, exits 2 before any request, since the question could never be answered. With it off, tools, prompts, trajectories, events and exit codes are unchanged; older tinyagent binaries refuse `NeedsInput` checkpoints rather than misreading them. The question is model-written text that can quote workspace files or fetched pages; treat it as untrusted if you forward it to other systems. + +## Token and cost budgets + +`--token-limit n` (config `token_limit`) caps the run's provider-reported input plus output tokens; prompt-cache tokens are part of input. `--cost-limit usd` (config `cost_limit`) caps its cost, computed from the same usage and the model's price (the built-in table, or `input_price_per_mtok` and `output_price_per_mtok`). Both are off (0) by default, and delegates and workers count against the parent's budget. A model with no known price can't be held to a cost limit: tinyagent warns once at start, and only the token limit is enforced. + +When a token or cost budget reaches 80%, the model gets one user message saying how much is used and asking it to wrap up; the JSON stream, log database and OpenTelemetry trace get a `budget.warning` event. At 100% the run stops cleanly with `TokenLimitExceeded` or `CostLimitExceeded` (exit 1), and the printed `--resume` command doubles the limit that was hit. diff --git a/docs/subcommands.md b/docs/subcommands.md index 8e12b9b..bd33e01 100644 --- a/docs/subcommands.md +++ b/docs/subcommands.md @@ -120,7 +120,7 @@ Five subcommands are shortcuts for existing flags. They rewrite the command line ## Exit status -The run ends when the model writes `COMPLETE_TASK_AND_SUBMIT_FINAL_OUTPUT` (`Submitted`), stops calling tools (`Completed`), or hits a limit (`StepLimitExceeded`, `CostLimitExceeded`, `TimeLimitExceeded`), Ctrl-C (`Interrupted`), a non-retryable API error (`Error`), too many consecutive malformed tool-call steps (`RepeatedFormatError`, see below), or, with `--output-schema`, 3 consecutive invalid results (`InvalidOutput`), or stop hooks still rejecting the result after `max_stop_continuations` (`VerificationFailed`, see [Stop hooks](permissions.md#stop-hooks)), or too many consecutive responses cut off at the output-token limit (`OutputTruncated`, see [Output truncation](troubleshooting.md#output-truncation)), or, with `--ask-user`, a question for the user (`NeedsInput`, see [Asking for input](runs.md#asking-for-input)). The process exits 0 only for `Submitted` and `Completed`, 3 for `OutputTruncated`, 4 for `NeedsInput`, and 1 otherwise (2 is reserved for usage errors). +The run ends when the model writes `COMPLETE_TASK_AND_SUBMIT_FINAL_OUTPUT` (`Submitted`), stops calling tools (`Completed`), or hits a limit (`StepLimitExceeded`, `CostLimitExceeded`, `TokenLimitExceeded`, `TimeLimitExceeded`), Ctrl-C (`Interrupted`), a non-retryable API error (`Error`), too many consecutive malformed tool-call steps (`RepeatedFormatError`, see below), or, with `--output-schema`, 3 consecutive invalid results (`InvalidOutput`), or stop hooks still rejecting the result after `max_stop_continuations` (`VerificationFailed`, see [Stop hooks](permissions.md#stop-hooks)), or too many consecutive responses cut off at the output-token limit (`OutputTruncated`, see [Output truncation](troubleshooting.md#output-truncation)), or, with `--ask-user`, a question for the user (`NeedsInput`, see [Asking for input](runs.md#asking-for-input)). The process exits 0 only for `Submitted` and `Completed`, 3 for `OutputTruncated`, 4 for `NeedsInput`, and 1 otherwise (2 is reserved for usage errors). Every run ends with one summary line on stderr, so stdout keeps only the model's text: diff --git a/internal/agent/budget.go b/internal/agent/budget.go new file mode 100644 index 0000000..ae2cc6d --- /dev/null +++ b/internal/agent/budget.go @@ -0,0 +1,74 @@ +package agent + +import ( + "fmt" + + "github.com/vicgupta/tinyagent/internal/model" + "github.com/vicgupta/tinyagent/internal/trajectory" +) + +// StatusTokenLimit ends a run whose provider-reported tokens (input plus +// output, delegates and workers included) reached the token limit. +const StatusTokenLimit = "TokenLimitExceeded" + +// BudgetWarnFraction is the share of a token or cost budget at which the +// model is told, once, to wrap up. +const BudgetWarnFraction = 0.8 + +const budgetNotice = "Budget notice: this run has used %s. The run stops when the budget is used up, so wrap up now: finish the most important remaining work, check it, and submit your result." + +// usedTokens is the run's total provider-reported token usage counted +// against the token limit. Cache tokens are part of InputTokens. +func usedTokens(t *trajectory.Trajectory) int { return t.InputTokens + t.OutputTokens } + +// tokensLeft is what remains of the token limit, or 0 with no limit. +func (a *Agent) tokensLeft(t *trajectory.Trajectory) (left int, limited bool) { + if a.Config.TokenLimit <= 0 { + return 0, false + } + return a.Config.TokenLimit - usedTokens(t), true +} + +// budgetUse describes each token or cost budget at or past the warning +// threshold, or "" when none is. +func (a *Agent) budgetUse(t *trajectory.Trajectory) string { + s := "" + add := func(part string) { + if s != "" { + s += " and " + } + s += part + } + if n := a.Config.TokenLimit; n > 0 && float64(usedTokens(t)) >= BudgetWarnFraction*float64(n) { + add(fmt.Sprintf("%d%% of its token budget (%d of %d tokens)", percent(float64(usedTokens(t)), float64(n)), usedTokens(t), n)) + } + // An unpriced model reports no cost, so only its token budget warns. + if c := a.Config.CostLimit; c > 0 && t.Cost >= BudgetWarnFraction*c { + add(fmt.Sprintf("%d%% of its cost budget ($%.4f of $%.4f)", percent(t.Cost, c), t.Cost, c)) + } + return s +} + +func percent(used, limit float64) int { return int(100 * used / limit) } + +// warnBudget tells the model once, with a user message, that a token or +// cost budget is nearly used up. It re-arms when usage falls back under the +// threshold, as after a resume with a raised limit. +func (a *Agent) warnBudget(t *trajectory.Trajectory, v *view) { + use := a.budgetUse(t) + if use == "" { + t.BudgetWarned = false + return + } + if t.BudgetWarned { + return + } + t.BudgetWarned = true + t.Messages = append(t.Messages, model.Message{Role: "user", Content: fmt.Sprintf(budgetNotice, use)}) + v.add(t, t.Messages[len(t.Messages)-1]) + a.event("budget.warning", map[string]any{"step": t.Steps, "usage": use, + "tokens": usedTokens(t), "token_limit": a.Config.TokenLimit, "cost_usd": t.Cost, "cost_limit": a.Config.CostLimit}) + if a.Verbose { + fmt.Fprintf(a.Err, "budget warning: %s\n", use) + } +} diff --git a/internal/agent/budget_test.go b/internal/agent/budget_test.go new file mode 100644 index 0000000..919fbd6 --- /dev/null +++ b/internal/agent/budget_test.go @@ -0,0 +1,206 @@ +package agent + +import ( + "context" + "errors" + "strings" + "testing" + "time" + + "github.com/vicgupta/tinyagent/internal/model" +) + +// loopReplies returns n responses that each call list_files. +func loopReplies(n int) []model.Message { + out := make([]model.Message, n) + for i := range out { + out[i] = assistantCalls(call("c", "list_files", `{}`)) + } + return out +} + +func budgetNotices(msgs []model.Message) int { + n := 0 + for _, m := range msgs { + if m.Role == "user" && strings.HasPrefix(m.Content, "Budget notice:") { + n++ + } + } + return n +} + +func TestTurnCapStopsAtN(t *testing.T) { + m := &scripted{replies: loopReplies(10)} + a := newAgent(t, m) + a.Config.StepLimit = 3 + tr := a.Run(context.Background(), "x") + if tr.ExitStatus != StatusStepLimit || tr.Steps != 3 || len(m.seen) != 3 { + t.Errorf("status %s, steps %d, requests %d", tr.ExitStatus, tr.Steps, len(m.seen)) + } +} + +func TestTokenLimitStopsRun(t *testing.T) { + m := &scripted{replies: loopReplies(10)} // 120 tokens per response + a := newAgent(t, m) + a.Config.TokenLimit = 300 + tr := a.Run(context.Background(), "x") + if tr.ExitStatus != StatusTokenLimit || tr.Steps != 3 || len(m.seen) != 3 { + t.Errorf("status %s, steps %d, requests %d", tr.ExitStatus, tr.Steps, len(m.seen)) + } + if got := tr.InputTokens + tr.OutputTokens; got != 360 { + t.Errorf("tokens = %d, want 360", got) + } +} + +func TestTokenBudgetWarnsOnceAt80Percent(t *testing.T) { + m := &scripted{replies: loopReplies(10)} + a := newAgent(t, m) + a.Config.TokenLimit = 600 // warn at 480 tokens, after the 4th response + var warnings int + a.OnEvent = func(kind string, _ map[string]any) { + if kind == "budget.warning" { + warnings++ + } + } + tr := a.Run(context.Background(), "x") + if tr.ExitStatus != StatusTokenLimit || tr.Steps != 5 { + t.Fatalf("status %s, steps %d", tr.ExitStatus, tr.Steps) + } + if n := budgetNotices(tr.Messages); n != 1 || warnings != 1 || !tr.BudgetWarned { + t.Errorf("notices %d, events %d, warned %v; want one", n, warnings, tr.BudgetWarned) + } + if budgetNotices(m.seen[3]) != 0 { + t.Error("warned before reaching 80%") + } + last := m.seen[4][len(m.seen[4])-1] + if last.Role != "user" || !strings.Contains(last.Content, "80% of its token budget (480 of 600 tokens)") { + t.Errorf("5th request ends with %+v", last) + } +} + +func TestCostBudgetWarnsOnce(t *testing.T) { + m := &scripted{replies: loopReplies(10)} // $0.01 per response + a := newAgent(t, m) + a.Config.CostLimit = 0.045 + tr := a.Run(context.Background(), "x") + if tr.ExitStatus != StatusCostLimit || tr.Steps != 5 { + t.Fatalf("status %s, steps %d", tr.ExitStatus, tr.Steps) + } + if n := budgetNotices(tr.Messages); n != 1 || !strings.Contains(m.seen[4][len(m.seen[4])-1].Content, "cost budget") { + t.Errorf("notices %d", n) + } +} + +func TestNoBudgetNoWarning(t *testing.T) { + m := &scripted{replies: loopReplies(5)} + a := newAgent(t, m) + a.Config.StepLimit = 5 + tr := a.Run(context.Background(), "x") + if budgetNotices(tr.Messages) != 0 || tr.BudgetWarned { + t.Error("warned without a token or cost budget") + } +} + +func TestBudgetWarningRearmsWhenLimitRaised(t *testing.T) { + m := &scripted{replies: loopReplies(10)} + a := newAgent(t, m) + a.Config.TokenLimit = 600 + tr := a.Run(context.Background(), "x") // warned, then stopped at 600 + a.Config.TokenLimit = 6000 + a.warnBudget(tr, viewFromTrajectory(tr)) + if tr.BudgetWarned { + t.Error("raised limit did not re-arm the warning") + } +} + +// unpricedModel reports tokens but no cost, like a model with no price. +type unpricedModel struct{ scripted } + +func (u *unpricedModel) Query(ctx context.Context, msgs []model.Message, specs []model.ToolSpec) (model.Message, model.Cost, error) { + msg, c, err := u.scripted.Query(ctx, msgs, specs) + c.USD = 0 + return msg, c, err +} + +func TestUnpricedModelFallsBackToTokenBudget(t *testing.T) { + m := &unpricedModel{scripted{replies: loopReplies(10)}} + a := newAgent(t, m) + a.Config.CostLimit = 0.01 // can't be enforced without a price + a.Config.TokenLimit = 300 + tr := a.Run(context.Background(), "x") + if tr.ExitStatus != StatusTokenLimit || tr.Cost != 0 || tr.Steps != 3 { + t.Errorf("status %s, cost %v, steps %d", tr.ExitStatus, tr.Cost, tr.Steps) + } + for _, msg := range tr.Messages { + if strings.Contains(msg.Content, "cost budget") { + t.Error("cost warning for an unpriced model") + } + } +} + +func TestDelegateTokensCountTowardParentLimit(t *testing.T) { + a, pm, cm := subagentAgent(t, + []model.Message{assistantCalls(delegateCall("p1", "spend")), {Role: "assistant", Content: "done"}}, + loopReplies(4)) + a.Config.TokenLimit = 300 + tr := a.Run(context.Background(), "x") + if tr.ExitStatus != StatusTokenLimit { + t.Fatalf("status = %s", tr.ExitStatus) + } + // Parent 120; the child gets the 180 left: 120, then 240 >= 180. + if len(cm.seen) != 2 || len(pm.seen) != 1 || tr.Subagents[0].Status != StatusTokenLimit { + t.Errorf("child requests %d, parent requests %d, record %+v", len(cm.seen), len(pm.seen), tr.Subagents[0]) + } + if got := tr.InputTokens + tr.OutputTokens; got != 360 { + t.Errorf("tokens = %d, want 360", got) + } +} + +func TestDelegateNotStartedWhenTokensSpent(t *testing.T) { + a, _, cm := subagentAgent(t, + []model.Message{assistantCalls(delegateCall("p1", "spend")), {Role: "assistant", Content: "done"}}, nil) + a.Config.TokenLimit = 120 + tr := a.Run(context.Background(), "x") + if tr.ExitStatus != StatusTokenLimit || len(cm.seen) != 0 { + t.Errorf("status %s, child requests %d", tr.ExitStatus, len(cm.seen)) + } +} + +func TestWorkersSharedTokenLimitCancelsAll(t *testing.T) { + loop := assistantCalls(call("c", "write_file", `{"path":"spin.txt","content":"x\n"}`)) + tm := &taskModel{scripts: map[string][]model.Message{"task-spend": {loop}, "task-wait": {loop}}} + tm.before = func(ctx context.Context, key string) error { + if key == "task-wait" { + select { + case <-ctx.Done(): + return ctx.Err() + case <-time.After(20 * time.Second): + return errors.New("never cancelled") + } + } + time.Sleep(time.Millisecond) + return nil + } + a := batchAgent(t, []model.Message{ + assistantCalls(workerCallMsg("p1", "task-spend"), workerCallMsg("p2", "task-wait")), + {Role: "assistant", Content: "done" + submitMarker}, + }, tm, 2) + a.Config.TokenLimit = 200 // the parent's first call uses 120, each worker call 15 + a.Workers.MaxSteps = 1000 + start := time.Now() + tr := a.Run(context.Background(), "go") + if time.Since(start) > 15*time.Second { + t.Fatal("workers were not cancelled") + } + if tr.ExitStatus != StatusTokenLimit { + t.Fatalf("exit = %s %s", tr.ExitStatus, tr.Submission) + } + for _, s := range tr.Subagents { + if s.Status != StatusTokenLimit { + t.Errorf("worker %s status = %s", s.ID, s.Status) + } + } + if got := tr.InputTokens + tr.OutputTokens; got < 200 { + t.Errorf("parent tokens = %d; worker usage not counted", got) + } +} diff --git a/internal/agent/loop.go b/internal/agent/loop.go index 2165bfe..ca04bed 100644 --- a/internal/agent/loop.go +++ b/internal/agent/loop.go @@ -521,6 +521,7 @@ func (a *Agent) Run(ctx context.Context, task string) *trajectory.Trajectory { if status, ok := a.checkLimits(ctx, t); !ok { return a.finish(t, status, "") } + a.warnBudget(t, v) if a.needsCompaction(t, v) { if err := a.compact(ctx, t, v, "proactive"); err != nil { return a.compactFailed(ctx, t, err) @@ -648,7 +649,7 @@ func (a *Agent) Run(ctx context.Context, task string) *trajectory.Trajectory { return a.saveFailed(t, err) } malformed := 0 - costStop := false + limitStop := "" batch := a.runWorkerBatch(ctx, t, msg.ToolCalls) for _, call := range msg.ToolCalls { a.event("tool.call", map[string]any{"id": call.ID, "name": call.Function.Name, "arguments": eventArguments(call.Function.Arguments)}) @@ -697,8 +698,8 @@ func (a *Agent) Run(ctx context.Context, task string) *trajectory.Trajectory { // The delegate call stays pending; a resume closes it // as interrupted, like any call cut off mid-run. return a.finish(t, stop, "") - case StatusCostLimit: - costStop = true + case StatusCostLimit, StatusTokenLimit: + limitStop = stop } } fields := map[string]any{"id": call.ID, "name": call.Function.Name, "output": out, "is_error": isError} @@ -720,8 +721,8 @@ func (a *Agent) Run(ctx context.Context, task string) *trajectory.Trajectory { if err := a.step(t); err != nil { return a.saveFailed(t, err) } - if costStop { - return a.finish(t, StatusCostLimit, "") + if limitStop != "" { + return a.finish(t, limitStop, "") } } if malformed == len(msg.ToolCalls) { @@ -1327,6 +1328,9 @@ func (a *Agent) checkLimits(ctx context.Context, t *trajectory.Trajectory) (stri if a.Config.CostLimit > 0 && t.Cost >= a.Config.CostLimit { return StatusCostLimit, false } + if left, ok := a.tokensLeft(t); ok && left <= 0 { + return StatusTokenLimit, false + } return "", true } diff --git a/internal/agent/subagent.go b/internal/agent/subagent.go index 72e15c5..8f86280 100644 --- a/internal/agent/subagent.go +++ b/internal/agent/subagent.go @@ -84,6 +84,13 @@ func (a *Agent) delegate(ctx context.Context, t *trajectory.Trajectory, call mod } cfg.CostLimit = left } + if left, ok := a.tokensLeft(t); ok { + if left <= 0 { + a.delegateStop = StatusTokenLimit + return tools.Result{Text: fmt.Sprintf("error: delegate %s not started: the token limit is reached", id), IsError: true} + } + cfg.TokenLimit = left + } out := io.Discard if a.Verbose { out = a.Out @@ -172,7 +179,7 @@ func (a *Agent) delegate(ctx context.Context, t *trajectory.Trajectory, call mod return res } switch ct.ExitStatus { - case StatusInterrupted, StatusTimeLimit, StatusCostLimit: + case StatusInterrupted, StatusTimeLimit, StatusCostLimit, StatusTokenLimit: a.delegateStop = ct.ExitStatus } msg := fmt.Sprintf("error: delegate %s did not finish: %s", id, ct.ExitStatus) diff --git a/internal/agent/worker.go b/internal/agent/worker.go index 921975b..0c04257 100644 --- a/internal/agent/worker.go +++ b/internal/agent/worker.go @@ -210,6 +210,13 @@ func (a *Agent) prepareWorker(t *trajectory.Trajectory, call model.ToolCall, id } cfg.CostLimit = left } + if left, ok := a.tokensLeft(t); ok { + if left <= 0 { + a.delegateStop = StatusTokenLimit + return nil, toolErr("worker %s not started: the token limit is reached", id), false + } + cfg.TokenLimit = left + } only := a.workerTools() if len(only) == 0 { return nil, toolErr("worker %s not started: none of the configured worker tools are available in this run", id), false @@ -396,7 +403,7 @@ func (a *Agent) finishWorker(t *trajectory.Trajectory, j *workerJob) tools.Resul return res } switch ct.ExitStatus { - case StatusInterrupted, StatusTimeLimit, StatusCostLimit: + case StatusInterrupted, StatusTimeLimit, StatusCostLimit, StatusTokenLimit: a.delegateStop = ct.ExitStatus } msg := fmt.Sprintf("error: worker %s did not finish: %s", id, ct.ExitStatus) diff --git a/internal/agent/worker_batch.go b/internal/agent/worker_batch.go index 5620105..fc93b18 100644 --- a/internal/agent/worker_batch.go +++ b/internal/agent/worker_batch.go @@ -27,6 +27,10 @@ func (w *Workers) maxConcurrent() int { // run's cost limit. var errWorkersCostLimit = errors.New("the workers together reached the cost limit") +// errWorkersTokenLimit cancels a batch whose workers together reached the +// run's token limit. +var errWorkersTokenLimit = errors.New("the workers together reached the token limit") + // workerBatch holds the results of a response's worker calls that ran // concurrently, for the tool loop to record in call order. type workerBatch struct { @@ -86,7 +90,7 @@ func (a *Agent) runWorkerBatch(ctx context.Context, t *trajectory.Trajectory, ca bctx, cancel := context.WithCancelCause(ctx) defer cancel(nil) - gate := &costGate{spent: t.Cost, limit: a.Config.CostLimit, cancel: cancel} + gate := &costGate{spent: t.Cost, limit: a.Config.CostLimit, tokens: usedTokens(t), tokenLimit: a.Config.TokenLimit, cancel: cancel} var jobs []*workerJob for _, c := range batch { @@ -136,7 +140,7 @@ func (a *Agent) runWorkerBatch(ctx context.Context, t *trajectory.Trajectory, ca } j.ct = j.child.Run(bctx, j.task) switch j.ct.ExitStatus { - case StatusInterrupted, StatusTimeLimit, StatusCostLimit: + case StatusInterrupted, StatusTimeLimit, StatusCostLimit, StatusTokenLimit: // A shared limit: stop the other workers too. cancel(fmt.Errorf("worker %s stopped: %s", j.id, j.ct.ExitStatus)) } @@ -145,10 +149,14 @@ func (a *Agent) runWorkerBatch(ctx context.Context, t *trajectory.Trajectory, ca wg.Wait() costHit := errors.Is(context.Cause(bctx), errWorkersCostLimit) + tokenHit := errors.Is(context.Cause(bctx), errWorkersTokenLimit) for _, j := range jobs { if costHit && j.ct.ExitStatus == StatusInterrupted { j.ct.ExitStatus = StatusCostLimit } + if tokenHit && j.ct.ExitStatus == StatusInterrupted { + j.ct.ExitStatus = StatusTokenLimit + } res := a.finishWorker(t, j) b.results[j.call.ID] = a.postWorkerHooks(ctx, t, j.call, res) if s := a.takeDelegateStop(); s != "" && b.stop == "" { @@ -161,6 +169,8 @@ func (a *Agent) runWorkerBatch(ctx context.Context, t *trajectory.Trajectory, ca b.stop = s } else if costHit { b.stop = StatusCostLimit + } else if tokenHit { + b.stop = StatusTokenLimit } a.event("workers.batch", map[string]any{"workers": len(jobs), "max_concurrent": a.Workers.maxConcurrent(), "stop": b.stop}) return b @@ -216,15 +226,22 @@ type costGate struct { spent float64 limit float64 cancel context.CancelCauseFunc + // tokens and tokenLimit track input plus output tokens the same way. + tokens int + tokenLimit int } -func (g *costGate) add(usd float64) { +func (g *costGate) add(c model.Cost) { g.mu.Lock() defer g.mu.Unlock() - g.spent += usd + g.spent += c.USD + g.tokens += c.InputTokens + c.OutputTokens if g.limit > 0 && g.spent >= g.limit { g.cancel(errWorkersCostLimit) } + if g.tokenLimit > 0 && g.tokens >= g.tokenLimit { + g.cancel(errWorkersTokenLimit) + } } // gatedModel adds each response's cost to a batch's costGate. @@ -235,7 +252,7 @@ type gatedModel struct { func (m *gatedModel) Query(ctx context.Context, msgs []model.Message, specs []model.ToolSpec) (model.Message, model.Cost, error) { msg, cost, err := m.Model.Query(ctx, msgs, specs) - m.gate.add(cost.USD) + m.gate.add(cost) return msg, cost, err } diff --git a/internal/config/config.go b/internal/config/config.go index efbad0f..0efd5e6 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -67,6 +67,9 @@ type AgentConfig struct { StepLimit int `yaml:"step_limit"` CostLimit float64 `yaml:"cost_limit"` WallTimeSec int `yaml:"wall_time_seconds"` + // TokenLimit caps the run's provider-reported input plus output + // tokens, delegates and workers included; 0 (the default) = unlimited. + TokenLimit int `yaml:"token_limit"` // MaxFormatErrors ends the run after this many consecutive steps whose // tool calls are all malformed. Must be >= 1; default 3. MaxFormatErrors int `yaml:"max_format_errors"` diff --git a/internal/eval/run.go b/internal/eval/run.go index 447959e..87d1361 100644 --- a/internal/eval/run.go +++ b/internal/eval/run.go @@ -267,7 +267,7 @@ func runTrial(ctx context.Context, o Options, tr Trial) (rec Record) { switch ts.ExitStatus { case "Submitted", "Completed": rec.Outcome = OutcomePassed - case "StepLimitExceeded", "CostLimitExceeded", "TimeLimitExceeded": + case "StepLimitExceeded", "CostLimitExceeded", "TokenLimitExceeded", "TimeLimitExceeded": rec.Outcome = OutcomeLimit if !tr.Task.checksOnLimit() { return rec diff --git a/internal/trajectory/resume.go b/internal/trajectory/resume.go index 28d5f1f..06cd91a 100644 --- a/internal/trajectory/resume.go +++ b/internal/trajectory/resume.go @@ -170,7 +170,7 @@ func Read(path string) (*Trajectory, error) { func unfinished(status string) bool { switch status { - case "", "StepLimitExceeded", "CostLimitExceeded", "TimeLimitExceeded", "Interrupted", "Error", "RepeatedFormatError": + case "", "StepLimitExceeded", "CostLimitExceeded", "TokenLimitExceeded", "TimeLimitExceeded", "Interrupted", "Error", "RepeatedFormatError": return true } return false diff --git a/internal/trajectory/trajectory.go b/internal/trajectory/trajectory.go index 1b063fb..d129029 100644 --- a/internal/trajectory/trajectory.go +++ b/internal/trajectory/trajectory.go @@ -305,6 +305,9 @@ type Trajectory struct { // QuestionsAsked counts the ask_user questions that ended a run, over // the whole resume chain. QuestionsAsked int `json:"questions_asked,omitempty"` + // BudgetWarned is true once the model was told a token or cost budget + // is nearly used up; it re-arms when usage drops under the threshold. + BudgetWarned bool `json:"budget_warned,omitempty"` // PendingQuestion is the unanswered question of a NeedsInput run; its // call is the only pending call. A resume delivers the answer as that // call's result.