feat(budget): cap what an agent may spend in money, tokens, or working time#3702
feat(budget): cap what an agent may spend in money, tokens, or working time#3702dwin-gharibi wants to merge 23 commits into
Conversation
|
@Sayt-0 Deps fixed. |
|
#3714 already cover a part of this pr, but I like the idea of budget limit |
|
👋 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 |
… to define amazing budgets
…tor for docker agent
…level config keys
…ent core into runtime of docker aagent
…re and wiring into agent loop and session
…er agent core components like agent loops and runtime and client
…events and exceptions
… into tui sidebar to announcing them realtime
…zing tests for budget lines in tui
…into docker agent cmd runner
…nnouncing new config
…ple into docker agent
…new configuration
…debar budget manager
…design tests for sidebar tui
75dfbfa to
3589dcd
Compare
|
@Sayt-0 Ready! |
aheritier
left a comment
There was a problem hiding this comment.
🤖 Automated implementer agent — this 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
newBudgetSetwiring 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 inbudgetsFor(). budgetStartedflag correctly prevents per-message resets without a special-case for unbudgeted runs.budget_exceededreusingnotificationrather than adding a new hook surface is the pragmatic call and the justification is solid.computeMessageCostreuse ensures the ceiling and the invoice can never disagree.- Schema
$ref, HCL block rule, andtopLevelConfigKeysadditions are all present. - Token-counting using a dedicated accumulator (not
session.Usage()) correctly survives compaction. - Race detector tests and the concurrent
budgetTrackerare clean.
The config docs example file is excellent — runnable, well-commented, and the shared-pot behavior is made concrete.
All 3 bugs from the review have been fixed. Three commits were pushed to Bug 1 — notification path (f6d26d2) Bug 2 — MaxTime comment (7d4f0e1) Bug 3 — docs first-impressions (7a5d0e7) Re: Please pull these commits from |
aheritier
left a comment
There was a problem hiding this comment.
🤖 Automated implementer agent — this 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.
f6d26d2 — handleBudgetExceeded 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. ✓
7d4f0e1 — MaxTime 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.
Thanks for the thorough review. All items addressed in commit Bug 1 — Bug 2 — Bug 3 — Docs Question — Minor — duplicate |
- 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
All five review findings from #3702 (review) have been addressed and landed in PR #3746 (#3746, branch Summary of changes:
|
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
left a comment
There was a problem hiding this comment.
🤖 Automated implementer agent — this 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. ✓
Adds budgets that stop an agent once it crosses a configured ceiling, and tracks every budget live in the TUI, broken down per agent.
Every limit is optional, unset means unlimited, and declaring nothing leaves runs unbudgeted. Existing configs are unaffected - the feature is inert unless opted into.
The named form follows the existing top-level
mcps/rag/toolsetsreference-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 toafter_llm_callhooks, which the docs suggest using for "a sidecar cost ledger".But a ledger can only record.
executeAfterLLMCallHooks(hooks.go) discards the hookResult- both callsites invoke it as a bare statement, andafter_llm_call_test.golocks that in.Continue: falseis only honored inexecutor.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.
rootandresearcherabove both referenceshell-work, so together they cannot exceed$0.03. Per-agent copies would let a run spendmax_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 sameLocalRuntime(loop.goalready computesrootStream := !sess.IsSubSession()). This intentionally differs frommax_iterations, where per-child caps (agent_delegation.go) are correct: money doesn't behave like iterations.A budget spans the session, not a message.
RunStreamruns once per user message, so building the trackers there would silently redefinemax_costas "per message" - a session could re-spend the whole ceiling on every turn.ensureBudgetbuilds once and leaves them alone. AbudgetStartedflag distinguishes "not built yet" from "built, nothing to budget", so unbudgeted runs don't rebuild every message.max_timemeasures 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 tripsmax_time: 2m. It accumulates turn durations instead. This also made the tracker simpler: no start timestamp, no clock.max_tokensuses its own accumulator, notsession.Usage().ApplyCompactionresets the session counters because they measure context length, not spend. Reading them would makemax_tokenssilently stop working on precisely the long runs it exists for.Unpriced models can't be enforced against, and say so.
computeMessageCostreturnsnilfor a model with no pricing table. There's no honest number to add, so it can't count towardsmax_cost. Failing closed would break local/custom-endpoint users; counting$0would 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-levelcost: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 keepsresumeChanand the headless consumers (chatserver,embeddedchat) untouched.Boundary-checked. Checked between turns beside
enforceMaxIterations, so a run overshoots by at most the turn in flight andmax_timewon'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).agent-schema.json:budgetvia$ref, abudgetsmap, and the agent-levelbudgetsarray. Plus an HCL block rule sobudgets "tight" { ... }maps correctly.max_timereuses the existinglatest.Durationtype, so10m/30s/1h30mall work.CloneThroughJSON), so additive fields arrive as zero values.Versionstays"12".Runtime
pkg/runtime/budget.go:budgetTracker(one per budget name) andbudgetSet(the run-wide budget plus each agent's named budgets). Nil-safe throughout, so the loop needs no nil checks.enforceMaxIterations; newturnEndReasonBudgetExceededso a budget-killed run is distinguishable from a completed one in telemetry.budget_usage(oneBudgetStatusper active budget, each with a per-agent breakdown) andbudget_exceeded(carryingbudget,limit,used,max,config_path), registered inclient.gofor wire decoding.notificationhook rather than addingon_budget_exceeded-budget_exceededalready carries the structured payload. Easy to add later if asked.TUI
budget_exceededsurfaces 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 budgetmax_tokens≠ modelmax_tokens.examples/budget.yaml- OpenRouter viabase_urlwith an explicitcost:block (somax_costis 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.5toolchain is unfetchable here (403 fromproxy.golang.org), so tests ran against cachedgo1.26.4via a local-only-modfile;go.modis untouched.golangci-lint run(v2.12.2, repo's own.golangci.yml) over every touched package → 0 issuesgo test ./pkg/runtime/ ./pkg/config/...→ passing, including the schema drift guards (TestSchemaMatchesGoTypes,TestHCLBlockRulesCoverSchemaMaps,TestJsonSchemaWorksForExamples,TestDocYAMLSnippetsAreValid)go test -raceon the shared trackers → clean (sub-sessions and background agents record concurrently)markdownlint-cli2@0.22.1fromdocs/→ 0 errors, 111 filesagents.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:
"budget"/"budgets"totopLevelConfigKeys(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.$reffor thebudgetproperty rather than inlining.TestHCLBlockRulesCoverSchemaMapsdoesn't resolve refs and typesAdditionalPropertiesasany, so an inlined"additionalProperties": falsewould demand a nonexistent HCL rule.budgetsmap 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'sTestListGatewayModels_*/TestAvailableModels_Gateway*fail identically on a cleanorigin/mainworktree with none of this code present (verified, not assumed).pkg/tui/components/sidebartests cannot build under the-modfileworkaround (it resolves acellbufincompatible with the pinnedansi); 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
Watch the sidebar track each budget until the run stops.