Skip to content

Part of #236: token budget and budget warning - #246

Merged
vicgupta merged 1 commit into
mainfrom
amf/issue-236
Oct 30, 2026
Merged

vicgupta merged 1 commit into
mainfrom
amf/issue-236

Conversation

@vicgupta

Copy link
Copy Markdown
Owner

Part of #236

Summary

Most of #236 was already on main: --step-limit (turn cap), --cost-limit with a per-model price table and price overrides, a warn-once message for unpriced models, and delegate cost counted against the parent. This slice adds what was missing:

  • Token budget: --token-limit n / config token_limit caps provider-reported input+output tokens for the whole run. At 100% the run stops cleanly with the new status TokenLimitExceeded (exit 1, resumable; the printed --resume command doubles the limit).
  • 80% warning: when a token or cost budget reaches 80%, one user message is injected before the next request ("Budget notice: … wrap up"), and a budget.warning event (step, usage, tokens, token_limit, cost_usd, cost_limit) goes to the JSON stream, log DB and OTel trace.
  • Delegates and workers: subagents and workers get only the tokens left in the parent's budget, their usage counts toward it, and a worker batch shares one token gate (like the shared cost gate) so all workers stop when it's spent.
  • Docs: docs/runs.md (new "Token and cost budgets" section), docs/config.md, docs/long-runs.md, docs/subcommands.md; help text gets the new flag.

How it was tested

New internal/agent/budget_test.go and cmd/tinyagent/budget_test.go, all with fake models (no network or keys):

  • turn cap stops at N; token limit stops the run with TokenLimitExceeded
  • 80% warning injected exactly once (token and cost), not before 80%, no warning without a budget, re-arms when the limit is raised
  • unpriced model (zero cost) is still stopped by the token budget and gets no cost warning
  • delegate tokens count toward the parent and end it TokenLimitExceeded; a delegate isn't started when tokens are spent
  • worker batch shares the token limit and cancels all workers
  • --token-limit flag override, negative-value validation, resume command doubling

Locally: gofmt -l . clean, go build ./..., go vet ./..., go test -count=1 ./... pass, plus cross-builds for windows/amd64, darwin/arm64, linux/arm64.

Decisions on open questions

  • Tokens counted are input + output as reported by the provider; cache read/write tokens are already part of input, so they count.
  • The token budget is off by default (0 = unlimited), same as the existing limits, so no current run changes behaviour.
  • One 80% warning per run, covering whichever budget crosses first; it re-arms if usage drops back under 80% (for example after resuming with a higher limit).
  • The warning is a user message added before the next model request, not a system-prompt change, so cached prefixes are untouched.
  • Unpriced models keep the existing warn-once message; the cost limit can't be enforced for them, the token limit is.
  • No turn-cap default was added: turning one on by default changes behaviour for existing users and configs, so it's left for a follow-up.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant