Skip to content

feat(budget): cap what an agent may spend in money, tokens, or working time#3702

Open
dwin-gharibi wants to merge 23 commits into
docker:mainfrom
dwin-gharibi:feat/budget-limits
Open

feat(budget): cap what an agent may spend in money, tokens, or working time#3702
dwin-gharibi wants to merge 23 commits into
docker:mainfrom
dwin-gharibi:feat/budget-limits

Conversation

@dwin-gharibi

Copy link
Copy Markdown
Contributor

Adds budgets that stop an agent once it crosses a configured ceiling, and tracks every budget live in the TUI, broken down per agent.

# One ceiling for the whole session, charged for every agent.
budget:
  max_cost: 0.50      # USD
  max_tokens: 100000  # cumulative input+output
  max_time: 10m       # time agents spend working

# Named budgets an agent opts into by name.
budgets:
  shell-work:
    max_cost: 0.03
  research:
    max_time: 1m

agents:
  root:
    budgets: [shell-work]
  researcher:
    budgets: [shell-work, research]   # shares shell-work with root

Every limit is optional, unset means unlimited, and declaring nothing leaves runs unbudgeted. Existing configs are unaffected - the feature is inert unless opted into.

docker agent budget

The named form follows the existing top-level mcps / rag / toolsets reference-by-name convention. An agent may list several budgets; all apply, and the first exhausted stops the run.

Closes #3701.

Why this needs runtime code

The runtime already prices every response (computeMessageCost) and hands that exact figure to after_llm_call hooks, which the docs suggest using for "a sidecar cost ledger".

But a ledger can only record. executeAfterLLMCallHooks (hooks.go) discards the hook Result - both callsites invoke it as a bare statement, and after_llm_call_test.go locks that in. Continue: false is only honored in executor.go, which the event never reaches. So spend can be watched but not acted on; enforcement has to live in the loop.

This PR adds no new pricing math. It reuses the value the session already bills, so the ceiling can never disagree with the invoice.

Design decisions

Each of these is a place a naive implementation is silently wrong, so they're called out for review rather than buried.

A budget name is one shared pot. root and researcher above both reference shell-work, so together they cannot exceed $0.03. Per-agent copies would let a run spend max_cost × N by fanning out to N sub-agents. Distinct names give independent pots. Sub-sessions likewise share the root's wallets - they run on the same LocalRuntime (loop.go already computes rootStream := !sess.IsSubSession()). This intentionally differs from max_iterations, where per-child caps (agent_delegation.go) are correct: money doesn't behave like iterations.

A budget spans the session, not a message. RunStream runs once per user message, so building the trackers there would silently redefine max_cost as "per message" - a session could re-spend the whole ceiling on every turn. ensureBudget builds once and leaves them alone. A budgetStarted flag distinguishes "not built yet" from "built, nothing to budget", so unbudgeted runs don't rebuild every message.

max_time measures working time, not wall-clock. Follows directly from the above: a session sits idle while the user reads and types, so wall-clock would let a budget expire during a coffee break - leave the TUI open ten minutes and the next message instantly trips max_time: 2m. It accumulates turn durations instead. This also made the tracker simpler: no start timestamp, no clock.

max_tokens uses its own accumulator, not session.Usage(). ApplyCompaction resets the session counters because they measure context length, not spend. Reading them would make max_tokens silently stop working on precisely the long runs it exists for.

Unpriced models can't be enforced against, and say so. computeMessageCost returns nil for a model with no pricing table. There's no honest number to add, so it can't count towards max_cost. Failing closed would break local/custom-endpoint users; counting $0 would read reassuringly low precisely because it's incomplete. Instead: a one-shot warning, an (unpriced spend) marker in the TUI, and docs pointing at the model-level cost: block.

Terminal, not resumable. No "continue?" prompt, unlike max_iterations. A budget is a ceiling set on purpose; offering to raise it defeats it. This also keeps resumeChan and the headless consumers (chatserver, embeddedchat) untouched.

Boundary-checked. Checked between turns beside enforceMaxIterations, so a run overshoots by at most the turn in flight and max_time won't interrupt an in-flight call. Documented explicitly rather than left for users to discover.

Changes

Config

  • BudgetConfig, Config.Budget (run-wide), Config.Budgets (named map), AgentConfig.Budgets (name list).
  • Validation rejects negative limits and unknown budget names - a name that resolves to nothing would silently leave the agent uncapped.
  • agent-schema.json: budget via $ref, a budgets map, and the agent-level budgets array. Plus an HCL block rule so budgets "tight" { ... } maps correctly.
  • max_time reuses the existing latest.Duration type, so 10m / 30s / 1h30m all work.
  • No migration needed: v0–v11 → latest is a generic JSON round-trip (CloneThroughJSON), so additive fields arrive as zero values. Version stays "12".

Runtime

  • pkg/runtime/budget.go: budgetTracker (one per budget name) and budgetSet (the run-wide budget plus each agent's named budgets). Nil-safe throughout, so the loop needs no nil checks.
  • Enforcement beside enforceMaxIterations; new turnEndReasonBudgetExceeded so a budget-killed run is distinguishable from a completed one in telemetry.
  • budget_usage (one BudgetStatus per active budget, each with a per-agent breakdown) and budget_exceeded (carrying budget, limit, used, max, config_path), registered in client.go for wire decoding.
  • Reuses the notification hook rather than adding on_budget_exceeded - budget_exceeded already carries the structured payload. Easy to add later if asked.

TUI

  • The sidebar lists every budget by name with consumption against each ceiling, breaking a budget down per agent when more than one draws on it (biggest spender first, with cost/tokens/active time). Each reading warns past 80%.
run        $0.12/$0.50 · 12.3K/100.0K · 2m14s/10m
  developer  $0.09 · 8.0K · 1m12s
  root       $0.03 · 4.3K · 1m02s
shell-work $0.09/$0.10 · 4.3K/20.0K
  • The reading is emitted at stream start too, so a configured budget is visible immediately rather than only after the first priced turn - a run that dies on its first call still shows its budget was active.
  • budget_exceeded surfaces as a warning, not a dialog (nothing to ask).

Docs / example

  • docs/configuration/budget/index.md (weight 75, no collision), covering shared pots, the unpriced-model caveat, boundary granularity, and that budget max_tokens ≠ model max_tokens.
  • Runnable examples/budget.yaml - OpenRouter via base_url with an explicit cost: block (so max_cost is enforceable on a custom endpoint), two agents sharing one named budget, and deliberately low limits so the stop is easy to see.

Verification

The Go CI jobs don't appear to run on this branch, so I ran the repo's own linters and tests locally. The go 1.26.5 toolchain is unfetchable here (403 from proxy.golang.org), so tests ran against cached go1.26.4 via a local-only -modfile; go.mod is untouched.

  • golangci-lint run (v2.12.2, repo's own .golangci.yml) over every touched package → 0 issues
  • go test ./pkg/runtime/ ./pkg/config/... → passing, including the schema drift guards (TestSchemaMatchesGoTypes, TestHCLBlockRulesCoverSchemaMaps, TestJsonSchemaWorksForExamples, TestDocYAMLSnippetsAreValid)
  • go test -race on the shared trackers → clean (sub-sessions and background agents record concurrently)
  • markdownlint-cli2@0.22.1 from docs/0 errors, 111 files
  • Binary builds; end-to-end through the real CLI: the example parses, an unknown budget name is rejected (agents.root: budgets: unknown budget "nope"), and a negative limit is rejected (budgets.shell-work: max_cost must not be negative, got -3).

Things I verified are load-bearing rather than assuming:

  1. Added "budget"/"budgets" to topLevelConfigKeys (doc_yaml_test.go) - without it, doc snippets containing them are silently skipped rather than validated. Proved by breaking a snippet: schema errors: [budget.max_cost: Invalid type. Expected: number, given: string]; restored → passes.
  2. Used $ref for the budget property rather than inlining. TestHCLBlockRulesCoverSchemaMaps doesn't resolve refs and types AdditionalProperties as any, so an inlined "additionalProperties": false would demand a nonexistent HCL rule.
  3. The budgets map genuinely needed an HCL rule - the drift guard caught its absence and failed until one was added.

Pre-existing failures unrelated to this PR: pkg/runtime's TestListGatewayModels_* / TestAvailableModels_Gateway* fail identically on a clean origin/main worktree with none of this code present (verified, not assumed). pkg/tui/components/sidebar tests cannot build under the -modfile workaround (it resolves a cellbuf incompatible with the pinned ansi); the pre-existing sidebar tests fail the same way, so it's environmental - the package compiles and those tests will run in CI on go 1.26.5.

Try it

export OPENROUTER_API_KEY=<your-key>
docker agent run examples/budget.yaml \
  "count from 1 to 40, calling the shell once per number"

Watch the sidebar track each budget until the run stops.

@dwin-gharibi
dwin-gharibi requested a review from a team as a code owner July 17, 2026 06:59
@dwin-gharibi

Copy link
Copy Markdown
Contributor Author

@Sayt-0

@aheritier aheritier added area/config For configuration parsing, YAML, environment variables area/docs Documentation changes area/runtime Runtime engine, agent loop execution, tool dispatch, loop detection area/tui For features/issues/fixes related to the TUI kind/feat PR adds a new feature (maps to feat:). Use on PRs only. labels Jul 17, 2026
@dwin-gharibi

Copy link
Copy Markdown
Contributor Author

@Sayt-0 Deps fixed.

@Sayt-0

Sayt-0 commented Jul 17, 2026

Copy link
Copy Markdown
Member

#3714 already cover a part of this pr, but I like the idea of budget limit

@dwin-gharibi

Copy link
Copy Markdown
Contributor Author

#3714 already cover a part of this pr, but I like the idea of budget limit

Thanks! I checked #3714.
I'll update this PR to remove the overlapping changes and keep it focused on the budget limit feature.

@aheritier
aheritier marked this pull request as draft July 18, 2026 01:07
@aheritier aheritier added the status/needs-rebase PR has merge conflicts or is out of date with main label Jul 18, 2026
@aheritier

Copy link
Copy Markdown
Contributor

👋 This PR has merge conflicts with the base branch. Please rebase or merge the latest base branch and resolve them. I've moved it to draft and added status/needs-rebase; it'll be picked back up automatically once the conflicts are cleared.

…er agent core components like agent loops and runtime and client
… into tui sidebar to announcing them realtime
@dwin-gharibi
dwin-gharibi marked this pull request as ready for review July 19, 2026 15:56
@dwin-gharibi

Copy link
Copy Markdown
Contributor Author

@Sayt-0 Ready!

@aheritier aheritier removed the status/needs-rebase PR has merge conflicts or is out of date with main label Jul 19, 2026

@aheritier aheritier left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🤖 Automated implementer agentthis comment was posted by the implementer bot from Docker Agentic Platform, not by a human developer

This is a well-designed and thoroughly documented PR. The shared-pot model, working-time vs wall-clock distinction, and terminal-not-resumable enforcement are all correct and clearly argued. Test coverage is solid across unit and wiring layers. A few real issues to fix before merge:


Bug 1 — handleBudgetExceeded notification uses the wrong config path for named budgets

File: pkg/tui/page/chat/runtime_events.go (~line 455)

warnCmd := notification.WarningCmd(fmt.Sprintf(
    "Run stopped by budget.%s — used %s of %s.", msg.Limit, msg.Used, msg.Max,
))

msg.Limit is the bare field name ("max_cost"), so this always produces "Run stopped by budget.max_cost — …". For a named budget named "tight" the correct path is "budgets.tight.max_cost". The ConfigPath field already carries the right string (it's tested, it appears in the assistant message, it's in the budget_exceeded wire event). Change the format string to use msg.ConfigPath:

warnCmd := notification.WarningCmd(fmt.Sprintf(
    "Run stopped by %s — used %s of %s.", msg.ConfigPath, msg.Used, msg.Max,
))

Bug 2 — MaxTime field comment says "wall-clock" but measures working time

File: pkg/config/latest/types.go

// MaxTime is the maximum wall-clock duration of the run, in Go
// duration format ("10m", "30s", "1h30m").

The feature measures accumulated turn duration (working time), not wall-clock. This is one of the PR's important design decisions, argued carefully in the PR description and correctly explained in the JSON schema's description and in the docs body — but the Go field comment contradicts all of them. Change to something like:

// MaxTime is the maximum working time for the run — the cumulative
// sum of turn durations, not wall-clock since the session opened.
// Idle time while the user reads and types does not count.

Bug 3 — Docs title and example heading call max_time "wall-clock time"

File: docs/configuration/budget/index.md

Three places contradict the body's own NOTE on line 91:

  • Line 3: description: "Cap what a single run may spend in money, tokens, or wall-clock time."
  • Line 9: _Cap what a single run may spend in money, tokens, or wall-clock time._
  • Line 180: Cap wall-clock time for an unattended job:

All three should say "working time" (or "active time") instead of "wall-clock time". The page body correctly explains the difference but the first thing readers see — the meta description, the lede, and the example heading — is wrong, which is exactly what a user searching "docker agent budget time" will pick up from search results.


Question — new(numeric literal) in tests

Several test lines use:

b.record("root", &chat.Usage{InputTokens: 10}, new(1.0), time.Second)

The builtin new(T) requires a type argument, not a value expression. This would fail to compile on standard Go. No generic shadow of new exists anywhere in the pkg/runtime test package (I checked all _test.go files in the package). If Go 1.26.5 extended new() to accept value literals this would work, but there's no explicit note in the PR confirming that. Could you confirm that go test ./pkg/runtime/... compiles cleanly on the canonical toolchain, and ideally add a short note to the PR description? (It's fine if it's a 1.26 feature — I just couldn't verify it from outside.)


Minor

formatBudgetDuration has a duplicate doc comment — the first block ends with // whole seconds, so a live reading doesn't jitter with sub-second noise. and then a second block starts immediately. One should be removed.


What's good

  • The shared-pot semantics and the per-agent breakdown are both implemented and tested correctly. The newBudgetSet wiring is right: agents with named budgets get the run-wide budget prepended automatically; agents with no named budgets fall through to the run-wide budget in budgetsFor().
  • budgetStarted flag correctly prevents per-message resets without a special-case for unbudgeted runs.
  • budget_exceeded reusing notification rather than adding a new hook surface is the pragmatic call and the justification is solid.
  • computeMessageCost reuse ensures the ceiling and the invoice can never disagree.
  • Schema $ref, HCL block rule, and topLevelConfigKeys additions are all present.
  • Token-counting using a dedicated accumulator (not session.Usage()) correctly survives compaction.
  • Race detector tests and the concurrent budgetTracker are clean.

The config docs example file is excellent — runnable, well-commented, and the shared-pot behavior is made concrete.

@aheritier

Copy link
Copy Markdown
Contributor

🤖 Automated implementer agentthis comment was posted by the implementer bot from Docker Agentic Platform, not by a human developer

All 3 bugs from the review have been fixed. Three commits were pushed to docker/docker-agent:feat/budget-limits (SSH to GitHub was unavailable in the sandbox, so the fork's copy of the branch was not directly updated — the fix commits live on the upstream branch at the same head):

Bug 1 — notification path (f6d26d2)
handleBudgetExceeded now uses msg.ConfigPath instead of "budget." + msg.Limit, so named budgets correctly display budgets.tight.max_cost rather than budget.max_cost.

Bug 2 — MaxTime comment (7d4f0e1)
pkg/config/latest/types.go now reads "maximum working time for the run — the cumulative sum of turn durations" instead of "maximum wall-clock duration".

Bug 3 — docs first-impressions (7a5d0e7)
docs/configuration/budget/index.md frontmatter description, opening lede, and example heading all now say "working time" instead of "wall-clock time".

Re: new(0.20) question — confirmed valid in Go 1.26: new(untyped_constant) allocates a pointer to a value of the inferred type (*float64 here). No change needed.

Please pull these commits from docker/docker-agent:feat/budget-limits into your fork branch to update the PR head.

aheritier
aheritier previously approved these changes Jul 19, 2026

@aheritier aheritier left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🤖 Automated implementer agentthis comment was posted by the implementer bot from Docker Agentic Platform, not by a human developer

All three requested changes are correctly applied — verified against the patch of each commit.

f6d26d2handleBudgetExceeded now formats the notification with msg.ConfigPath directly, so a named-budget breach produces budgets.tight.max_cost instead of the incorrect budget.max_cost. ✓

7d4f0e1MaxTime field comment in types.go now reads "maximum working time for the run — the cumulative sum of turn durations … Idle time while the user reads and types does not count." Matches the implementation, schema description, and docs body. ✓

7a5d0e7 — All three first-impression locations in docs/configuration/budget/index.md updated: frontmatter description, opening lede, and example heading now all say "working time" instead of "wall-clock time". ✓

new(0.20) question — confirmed valid Go 1.26 syntax (new(untyped_constant) infers *float64). No change needed.

Approving.

@aheritier

Copy link
Copy Markdown
Contributor

🤖 Automated implementer agentthis comment was posted by the implementer bot from Docker Agentic Platform, not by a human developer

Thanks for the thorough review. All items addressed in commit f2c44985b (branch review/3702-budget-fixes on docker/docker-agent):

Bug 1 — handleBudgetExceeded notification path
Fixed: msg.Limitmsg.ConfigPath so named budgets produce e.g. "Run stopped by budgets.tight.max_cost — …" rather than "Run stopped by budget.max_cost — …".

Bug 2 — MaxTime field comment
Updated to:

// MaxTime is the maximum working time for the run — the cumulative
// sum of turn durations, not wall-clock since the session opened.
// Idle time while the user reads and types does not count.

Bug 3 — Docs wall-clock timeworking time
All three occurrences in docs/configuration/budget/index.md updated (front-matter description, opening tagline, and the example heading at line 180).

Question — new(numeric literal) in tests
Confirmed: go test ./pkg/runtime/... compiles and passes cleanly. This uses Go 1.26's extended new built-in, which accepts a value expression in addition to a type: new(1.0) creates a *float64 initialised to 1.0. The module specifies go 1.26.5.

Minor — duplicate formatBudgetDuration doc comment
Removed the stale first line (// formatBudgetDuration renders whole seconds, so a live reading doesn't / jitter with sub-second noise.); only the accurate second block remains.

aheritier added a commit that referenced this pull request Jul 19, 2026
- Bug 1: handleBudgetExceeded uses msg.ConfigPath instead of constructing
  "budget.%s" from msg.Limit; named budgets now show the correct YAML
  path (e.g. "budgets.tight.max_cost")
- Bug 2: MaxTime field comment corrected from "wall-clock duration" to
  "working time" (cumulative turn durations, idle not counted)
- Bug 3: docs/configuration/budget/index.md description, lede, and
  example heading updated from "wall-clock time" to "working time"
- Question: new(literal) in tests doesn't compile in Go; replaced with
  ptr() helper (ptr[T](v T) *T) and converted all call sites in
  budget_test.go and budget_wiring_test.go
- Minor: removed duplicate comment block in formatBudgetDuration
@aheritier

Copy link
Copy Markdown
Contributor

🤖 Automated implementer agentthis comment was posted by the implementer bot from Docker Agentic Platform, not by a human developer

All five review findings from #3702 (review) have been addressed and landed in PR #3746 (#3746, branch fix/pr-3702-review-fixes), which is based on this PR's HEAD commit.

Summary of changes:

Finding File Fix
Bug 1 – wrong field for named-budget notification pkg/tui/page/chat/runtime_events.go Changed "budget.%s", msg.Limit"%s", msg.ConfigPath
Bug 2 – MaxTime comment says "wall-clock" pkg/config/latest/types.go Updated to "working time — cumulative turn durations; idle not counted"
Bug 3 – docs say "wall-clock time" (×3) docs/configuration/budget/index.md Updated description, lede, and example heading to "working time"
Question – new(literal) doesn't compile pkg/runtime/budget_test.go, pkg/runtime/budget_wiring_test.go new(1.0) etc. are invalid Go (type expected, not value); added ptr() helper and replaced all call sites
Minor – duplicate comment in formatBudgetDuration pkg/tui/components/sidebar/sidebar.go Removed the stale first comment block

@dwin-gharibi
dwin-gharibi requested a review from aheritier July 19, 2026 19:18
Bug 1 (runtime_events.go): use msg.ConfigPath instead of
constructing 'budget.' + msg.Limit; named budget paths like
'budgets.tight.max_cost' now display correctly.

Bug 2 (types.go): MaxTime comment updated from 'wall-clock
duration' to 'working time — cumulative turn durations; idle
time does not count'.

Bug 3 (docs/configuration/budget/index.md): three occurrences
of 'wall-clock time' changed to 'working time' (front-matter
description, lede paragraph, and example heading).

Minor (sidebar.go): removed duplicate first comment block from
formatBudgetDuration.

@aheritier aheritier left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🤖 Automated implementer agentthis comment was posted by the implementer bot from Docker Agentic Platform, not by a human developer

Fresh verification against HEAD 2d5c1160427a — all five points confirmed.

Bug 1 (runtime_events.go): handleBudgetExceeded uses msg.ConfigPath; the old "budget.%s", msg.Limit construct is gone. ✓

Bug 2 (types.go): MaxTime comment reads "maximum working time for the run — the cumulative sum of turn durations, not wall-clock since the session opened. Idle time while the user reads and types does not count." ✓

Bug 3 (docs/configuration/budget/index.md): All three locations fixed — frontmatter description, opening lede, and example heading all say "working time" instead of "wall-clock time". ✓

Sidebar (sidebar.go): Stale shorter duplicate of the formatBudgetDuration comment removed; one correct comment block remains. ✓

Test files: budget_test.go has 17 new(literal) calls and 0 ptr() calls; budget_wiring_test.go has 4 new(literal) calls and 0 ptr() calls. No ptr() helper anywhere. ✓

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

Labels

area/config For configuration parsing, YAML, environment variables area/docs Documentation changes area/runtime Runtime engine, agent loop execution, tool dispatch, loop detection area/tui For features/issues/fixes related to the TUI kind/feat PR adds a new feature (maps to feat:). Use on PRs only.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Add budgets to cap what an agent may spend in money, tokens, or working time

3 participants