Part of #236: token budget and budget warning

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 {

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

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"}}},

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

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)

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

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.

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.

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: