diff --git a/agent-schema.json b/agent-schema.json index f9183793b4..de9f2ec5b0 100644 --- a/agent-schema.json +++ b/agent-schema.json @@ -107,6 +107,17 @@ "runtime": { "$ref": "#/definitions/RuntimeDefaults", "description": "Execution-time defaults the agent author wants applied. Values act as defaults only — explicit CLI flags or user-config settings always win." + }, + "budget": { + "$ref": "#/definitions/BudgetConfig", + "description": "Ceilings on what a single run may consume before the agent is stopped, regardless of which agent spends. Scoped to a run and shared by every sub-session spawned inside it, so delegated work spends against the same budget as its parent." + }, + "budgets": { + "type": "object", + "description": "Reusable, named budget definitions. An agent opts into one by listing its name in its 'budgets' field. A named budget is a single shared pot: when several agents reference the same name they draw from the SAME ceiling rather than each getting a copy, so a fan-out to N sub-agents cannot spend N times the limit. Give agents distinct budget names when you want independent pots.", + "additionalProperties": { + "$ref": "#/definitions/BudgetConfig" + } } }, "additionalProperties": false, @@ -637,6 +648,22 @@ "description": "Maximum number of iterations", "minimum": 0 }, + "budgets": { + "type": "array", + "description": "Names of top-level 'budgets' entries this agent spends against. Every listed budget must be defined under the top-level 'budgets' key. An agent may list several, in which case all apply and the first to be exhausted stops the run. Agents sharing a budget name share one pot, so a fan-out cannot multiply the ceiling.", + "items": { + "type": "string" + }, + "examples": [ + [ + "tight" + ], + [ + "tight", + "daily" + ] + ] + }, "max_consecutive_tool_calls": { "type": "integer", "description": "Maximum consecutive identical tool calls before the agent is terminated. Prevents degenerate loops. 0 uses the default of 5.", @@ -1784,6 +1811,42 @@ }, "additionalProperties": false }, + "BudgetConfig": { + "type": "object", + "description": "Ceilings on what a single run may consume. Every limit is optional; an unset or zero limit is unlimited. The budget is scoped to a run and shared by every sub-session spawned inside it, so a fan-out to sub-agents spends against one budget rather than one per child. Exceeding any limit stops the run with an assistant message and a budget_exceeded event; raising a limit means editing the config, not answering a prompt.", + "properties": { + "max_cost": { + "type": "number", + "description": "Maximum spend for the run, in USD. Only responses the runtime can price count towards it: a model with no pricing data (unknown ID, or a custom endpoint without a model-level 'cost' block) contributes nothing, and such a run emits a warning rather than behaving as if the spend were free.", + "minimum": 0, + "examples": [ + 0.5, + 5, + 25 + ] + }, + "max_tokens": { + "type": "integer", + "description": "Maximum cumulative input+output tokens for the run, counted over its whole lifetime. This is not the session's context length: compaction resets that, while this keeps counting.", + "minimum": 0, + "examples": [ + 100000, + 1000000 + ] + }, + "max_time": { + "type": "string", + "description": "Maximum time the agents spend working, in Go duration format (e.g., '10m', '30s', '1h30m'). This is the sum of turn durations, NOT wall-clock since the session opened, so idle time while the user reads and types does not count. Checked at turn boundaries, so it will not interrupt an in-flight model call or tool.", + "pattern": "^([0-9]+(ns|us|µs|ms|s|m|h))+$", + "examples": [ + "30s", + "10m", + "1h30m" + ] + } + }, + "additionalProperties": false + }, "PermissionsConfig": { "type": "object", "description": "Tool permission configuration. Controls tool call approval behavior with optional argument matching.", diff --git a/cmd/root/run.go b/cmd/root/run.go index 7ef73f0c70..7c067c66f0 100644 --- a/cmd/root/run.go +++ b/cmd/root/run.go @@ -812,6 +812,8 @@ func (f *runExecFlags) runtimeOpts(loadResult *teamloader.LoadResult, runConfig runtime.WithWorkingDir(runConfig.WorkingDir), runtime.WithTracer(otel.Tracer(AppName)), runtime.WithModelSwitcherConfig(modelSwitcherCfg), + runtime.WithBudget(loadResult.Budget), + runtime.WithNamedBudgets(loadResult.Budgets, loadResult.AgentBudgets), } return opts } diff --git a/docs/configuration/budget/index.md b/docs/configuration/budget/index.md new file mode 100644 index 0000000000..d668ce5c6b --- /dev/null +++ b/docs/configuration/budget/index.md @@ -0,0 +1,194 @@ +--- +title: "Budget" +description: "Cap what a single run may spend in money, tokens, or working time." +keywords: docker agent, ai agents, configuration, yaml, budget, cost, limits +weight: 75 +canonical: https://docs.docker.com/ai/docker-agent/configuration/budget/ +--- + +_Cap what a single run may spend in money, tokens, or working time._ + +## Overview + +A budget sets ceilings on a run. When a ceiling is crossed the run stops with a message naming the exact limit that tripped, and the TUI tracks consumption live while the run is still going. + +Every limit is optional and an unset limit is unlimited, so you can cap money, tokens, time, or any combination. Declare no budget at all and runs are unbudgeted, which is the default. + +There are two ways to declare one, and they compose: + +| Block | Scope | +| --- | --- | +| `budget` | One run-wide ceiling, charged for every agent. | +| `budgets` | Named budgets an agent opts into by name. | + +### A run-wide budget + +```yaml +agents: + root: + model: openai/gpt-4o-mini + description: An agent on a short leash. + instruction: You are a helpful assistant. + +budget: + max_cost: 0.50 + max_tokens: 100000 + max_time: 10m +``` + +| Field | Type | Description | +| --- | --- | --- | +| `max_cost` | number | Maximum spend, in USD. | +| `max_tokens` | integer | Maximum cumulative input+output tokens. | +| `max_time` | string | Maximum time the agents spend working, in Go duration format (`10m`, `30s`, `1h30m`). | + +### Named budgets + +Define budgets by name under the top-level `budgets` key, then have each agent opt in by listing names in its own `budgets` field. The fields are the same three. + +```yaml +budgets: + shell-work: + max_cost: 0.03 + max_tokens: 8000 + research: + max_time: 1m + +agents: + root: + model: openai/gpt-4o-mini + description: Does shell work. + instruction: You are a helpful assistant. + budgets: [shell-work] + researcher: + model: openai/gpt-4o-mini + description: Answers questions. + instruction: You answer questions concisely. + budgets: [shell-work, research] +``` + +An agent may list several budgets; all of them apply, and the first to be exhausted stops the run. A run-wide `budget` applies on top of any named budget, so the ceiling that binds is whichever runs out first. + +Referencing a budget name that isn't defined is a config error, caught at parse time rather than silently leaving the agent uncapped. + +## A name is one shared pot + +When several agents reference the same budget name they draw from the **same** ceiling — not a copy each. Above, `root` and `researcher` share `shell-work`: together they cannot spend more than `$0.03`. + +This is deliberate, and it is the whole reason budgets are worth having. If each agent received its own allowance, a run could spend `max_cost` × N simply by fanning out to N sub-agents, and the ceiling would mean nothing for exactly the workloads that most need one. + +Give agents **distinct budget names** when you want independent pots. + +The same applies to the run-wide `budget`: every sub-session inside a run (transferred tasks, sub-agents, skills) spends from that one wallet. + +## Scope: a budget spans the session + +Spend accumulates for the life of the **session**, across every message you send — it does not reset each time you hit enter. A `max_cost: 0.50` you could re-spend on every message would not be a ceiling at all. + +Starting a new session starts a fresh budget. It is a per-session ceiling, not a lifetime quota across sessions. + +> [!NOTE] +> `max_time` measures the time the agents actually spend **working** — the sum of their turn durations — not wall-clock since the session opened. Because a budget spans a session, and a session sits idle while you read and type, wall-clock would let a budget expire during a coffee break: leave the TUI open for ten minutes and your next message would instantly trip a `max_time` of `2m` without the agent having done anything. + +## What stops a run + +Crossing any limit — run-wide or named — stops the run and produces: + +- an assistant message in the transcript naming the limit and the amounts, +- a `budget_exceeded` event carrying `budget`, `limit`, `used`, `max` and `config_path`, +- a `notification` hook at `warning` level, +- a stream end reason of `budget_exceeded`, so a stopped run is distinguishable from a completed one in telemetry. + +The message names the exact YAML path to raise, so there is no ambiguity about which of several budgets tripped: + +```text +Execution stopped after reaching the configured budgets.shell-work.max_cost +limit (used $0.0312 of $0.0300). +``` + +Unlike [`max_iterations`](../agents/index.md), a budget stop is **terminal** — there is no prompt offering to continue. A budget is a ceiling you set deliberately, so raising it means editing the config rather than answering a dialog. + +## Tracking spend in the TUI + +The sidebar's Token Usage panel lists every active budget by name, with consumption against each ceiling it declares: + +```text +run $0.12/$0.50 · 12.3K/100.0K · 2m14s/10m +shell-work $0.09/$0.10 · 4.3K/20.0K +``` + +Only the ceilings you configured appear. Each reading is colored by the sidebar's shared gauge bands — the same ones a context gauge uses as it nears compaction — so a budget turns amber well before its ceiling and red just short of it, and a run about to be stopped is visible before it stops. + +For a per-agent view of who spent what, set **Sidebar info mode** to `Detailed` in `/settings` → Appearance: the Agents section then reports each agent's cost alongside its effort and context. The budget line deliberately does not repeat that breakdown — the same numbers twice would crowd the sidebar's narrowest column. The per-agent split is still carried on the `budget_usage` event for programmatic consumers, and `/cost` has a **By Agent** section. + +## Limits and caveats + +### Unpriced models do not count towards `max_cost` + +Only responses the runtime can price count towards `max_cost`. A model with no pricing data — an unknown model ID, or a custom endpoint such as a local or private deployment — contributes nothing, because there is no honest number to add. + +Such a run emits a warning and the TUI marks the reading `(unpriced spend)`, rather than silently reading low because the spend is invisible. To make a custom endpoint count, price it explicitly with a model-level [`cost`](../models/index.md#custom-token-pricing) block: + +```yaml +models: + local: + provider: openai + model: my-model + base_url: http://localhost:8000/v1 + cost: + input: 0.15 + output: 0.60 + +agents: + root: + model: local + description: A locally-served agent with real cost accounting. + instruction: You are a helpful assistant. + +budget: + max_cost: 0.50 +``` + +### Limits are checked at turn boundaries + +A run is checked between turns, so it can overshoot by at most the turn already in flight. `max_time` in particular will not interrupt a model call or tool that has already started; the run stops at the first boundary after the limit is reached. + +This is the same granularity [`max_iterations`](../agents/index.md) has, and it keeps the ceiling out of the streaming hot path. Set limits with a little headroom rather than at the exact number you cannot exceed. + +### `max_tokens` here is not the model's `max_tokens` + +The `max_tokens` in a budget is a **cumulative** count of input+output tokens across the whole run. It is unrelated to the provider- or model-level [`max_tokens`](../models/index.md), which caps the output of a single response. + +It is also not the session's context length: compaction resets that, while the budget keeps counting. + +## Examples + +Cap money only, and let the run take as long as it needs: + +```yaml +agents: + root: + model: openai/gpt-4o + description: A cost-capped agent. + instruction: You are a helpful assistant. + +budget: + max_cost: 5.00 +``` + +Cap working time for an unattended job: + +```yaml +agents: + root: + model: openai/gpt-4o-mini + description: A time-boxed agent. + instruction: You are a helpful assistant. + toolsets: + - type: shell + +budget: + max_time: 15m +``` + +See [`examples/budget.yaml`](https://github.com/docker/docker-agent/blob/main/examples/budget.yaml) for a runnable configuration. diff --git a/examples/budget.yaml b/examples/budget.yaml new file mode 100644 index 0000000000..dbe62003d2 --- /dev/null +++ b/examples/budget.yaml @@ -0,0 +1,97 @@ +# yaml-language-server: $schema=../agent-schema.json +# +# Run budgets — stop an agent before it spends more than you meant +# ========================================================================== +# +# A budget caps what a session may consume in money, tokens, or working +# time. Every limit is optional and an unset limit is unlimited. +# +# Spend accumulates for the life of the session, across every message — +# it does not reset each time you hit enter. And max_time counts the time +# the agents actually spend working, so reading and typing is free. +# +# There are two ways to declare one: +# +# budget: a single run-wide ceiling, charged for every agent +# budgets: named budgets an agent opts into by name +# +# This example runs against OpenRouter (an OpenAI-compatible endpoint) so +# you can try it with an OpenRouter key instead of a first-party provider: +# +# export OPENROUTER_API_KEY= +# docker agent run examples/budget.yaml \ +# "count from 1 to 40, calling the shell once per number" +# +# The TUI's Token Usage panel lists every active budget by name and tracks +# spend against it live, breaking each down per agent once more than one +# agent draws on it. When a ceiling is crossed the run stops with a message +# naming the exact YAML path to raise. The token cap below is deliberately +# low so the counting demo trips it partway through. + +providers: + # OpenRouter exposes an OpenAI-compatible API, so it is configured as a + # custom provider with an explicit base_url. Its key is read from + # OPENROUTER_API_KEY. + openrouter: + base_url: https://openrouter.ai/api/v1 + token_key: OPENROUTER_API_KEY + +models: + budget-model: + provider: openrouter + model: openai/gpt-4o-mini + # A custom base_url endpoint is not in the models.dev catalogue, so the + # runtime cannot price it automatically — and an unpriced model does not + # count against a max_cost limit. Declaring cost here (USD per 1M + # tokens) makes max_cost enforceable. These match gpt-4o-mini's list + # price; edit to match whatever model you point at. + cost: + input: 0.15 + output: 0.60 + +# Named budgets. An agent opts in by listing a name under its `budgets`. +# +# A name is ONE shared pot: `shell-work` below is referenced by both agents, +# so together they cannot spend more than it allows. That is what keeps a +# budget meaningful when an agent fans out — ten sub-agents on one budget +# still cannot spend ten times the limit. Use distinct names when you want +# independent pots. +budgets: + shell-work: + max_cost: 0.03 + max_tokens: 8000 + research: + max_time: 1m + +# The run-wide budget applies to every agent on top of any named budget it +# references. The first ceiling to be exhausted — run-wide or named — stops +# the run. +budget: + max_cost: 0.05 + max_tokens: 15000 + # Time the agents spend working — the sum of their turn durations, not + # wall-clock since the session opened. + max_time: 2m + +agents: + root: + model: budget-model + description: An agent on a short leash. + instruction: | + You are a helpful assistant with shell access. Be concise. + Delegate research questions to the researcher sub-agent. + toolsets: + - type: shell + sub_agents: + - researcher + # root draws on the run-wide budget plus the shared shell-work pot. + budgets: [shell-work] + + researcher: + model: budget-model + description: Answers questions without touching the shell. + instruction: | + You answer questions concisely. You have no tools. + # researcher shares shell-work with root — the same pot, not a copy — + # and additionally has its own working-time ceiling. + budgets: [shell-work, research] diff --git a/pkg/config/doc_yaml_test.go b/pkg/config/doc_yaml_test.go index 7882c1c29c..39c413051a 100644 --- a/pkg/config/doc_yaml_test.go +++ b/pkg/config/doc_yaml_test.go @@ -30,6 +30,8 @@ var topLevelConfigKeys = map[string]bool{ "rag": true, "metadata": true, "permissions": true, + "budget": true, + "budgets": true, } // TestDocYAMLSnippetsAreValid extracts every ```yaml code block from docs/ diff --git a/pkg/config/hcl/hcl.go b/pkg/config/hcl/hcl.go index 6874256fd6..2a5de336ad 100644 --- a/pkg/config/hcl/hcl.go +++ b/pkg/config/hcl/hcl.go @@ -148,6 +148,10 @@ var blockRules = map[string]blockRule{ // toolsets: { name: { ... } }. Distinct from the agent-level `toolset` // (singular) block, which aggregates into a list under the same key. "toolsets": {mode: modeMapByLabel, outKey: "toolsets"}, + // Top-level named budget definitions: `budgets "tight" { ... }` becomes + // budgets: { tight: { ... } }. The run-wide `budget` block is a + // singleton and needs no rule — the 0-label default already covers it. + "budgets": {mode: modeMapByLabel, outKey: "budgets"}, // `shell "name" { ... }` is used inside script toolsets as a map of // scripted shell commands. "shell": {mode: modeMapByLabel, outKey: "shell"}, diff --git a/pkg/config/latest/budget_test.go b/pkg/config/latest/budget_test.go new file mode 100644 index 0000000000..58e36e0254 --- /dev/null +++ b/pkg/config/latest/budget_test.go @@ -0,0 +1,194 @@ +package latest + +import ( + "testing" + "time" + + "github.com/goccy/go-yaml" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestBudgetConfigValidate(t *testing.T) { + tests := []struct { + name string + budget *BudgetConfig + wantErr string + }{ + {name: "nil is valid (unbudgeted)", budget: nil}, + {name: "empty is valid (no ceilings)", budget: &BudgetConfig{}}, + { + name: "all limits set", + budget: &BudgetConfig{MaxCost: 0.5, MaxTokens: 100000, MaxTime: Duration{Duration: 10 * time.Minute}}, + }, + { + name: "negative max_cost", + budget: &BudgetConfig{MaxCost: -1}, + wantErr: "max_cost must not be negative", + }, + { + name: "negative max_tokens", + budget: &BudgetConfig{MaxTokens: -5}, + wantErr: "max_tokens must not be negative", + }, + { + name: "negative max_time", + budget: &BudgetConfig{MaxTime: Duration{Duration: -time.Second}}, + wantErr: "max_time must not be negative", + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + err := tt.budget.validate() + if tt.wantErr == "" { + assert.NoError(t, err) + return + } + require.Error(t, err) + assert.Contains(t, err.Error(), tt.wantErr) + }) + } +} + +func TestConfigValidateRejectsNegativeBudget(t *testing.T) { + cfg := Config{ + Agents: Agents{{Name: "root", Model: "openai/gpt-4o", Description: "d", Instruction: "i"}}, + Budget: &BudgetConfig{MaxCost: -0.5}, + } + err := cfg.Validate() + require.Error(t, err) + assert.Contains(t, err.Error(), "budget:") + assert.Contains(t, err.Error(), "max_cost must not be negative") +} + +func TestBudgetParsesDurationFromYAML(t *testing.T) { + var cfg Config + require.NoError(t, yaml.Unmarshal([]byte(` +agents: + root: + model: openai/gpt-4o + description: d + instruction: i +budget: + max_cost: 0.50 + max_tokens: 100000 + max_time: 10m +`), &cfg)) + + require.NotNil(t, cfg.Budget) + assert.InDelta(t, 0.50, cfg.Budget.MaxCost, 1e-9) + assert.Equal(t, int64(100000), cfg.Budget.MaxTokens) + assert.Equal(t, 10*time.Minute, cfg.Budget.MaxTime.Duration) + assert.False(t, cfg.Budget.IsZero()) +} + +func TestBudgetAbsentIsUnbudgeted(t *testing.T) { + var cfg Config + require.NoError(t, yaml.Unmarshal([]byte(` +agents: + root: + model: openai/gpt-4o + description: d + instruction: i +`), &cfg)) + + assert.Nil(t, cfg.Budget) + assert.True(t, cfg.Budget.IsZero(), "a nil budget must read as inert, not as a zero ceiling") +} + +func TestNamedBudgetsParseAndBind(t *testing.T) { + var cfg Config + require.NoError(t, yaml.Unmarshal([]byte(` +budgets: + tight: + max_cost: 0.10 + roomy: + max_tokens: 1000000 + max_time: 30m +agents: + root: + model: openai/gpt-4o + description: d + instruction: i + budgets: [tight] + developer: + model: openai/gpt-4o + description: d + instruction: i + budgets: [tight, roomy] +`), &cfg)) + + require.Len(t, cfg.Budgets, 2) + assert.InDelta(t, 0.10, cfg.Budgets["tight"].MaxCost, 1e-9) + assert.Equal(t, int64(1000000), cfg.Budgets["roomy"].MaxTokens) + assert.Equal(t, 30*time.Minute, cfg.Budgets["roomy"].MaxTime.Duration) + + byName := map[string][]string{} + for _, a := range cfg.Agents { + byName[a.Name] = a.Budgets + } + assert.Equal(t, []string{"tight"}, byName["root"]) + assert.Equal(t, []string{"tight", "roomy"}, byName["developer"]) +} + +func TestAgentBudgetsRejectsUnknownName(t *testing.T) { + var cfg Config + err := yaml.Unmarshal([]byte(` +budgets: + tight: + max_cost: 0.10 +agents: + root: + model: openai/gpt-4o + description: d + instruction: i + budgets: [tihgt] +`), &cfg) + require.Error(t, err) + assert.Contains(t, err.Error(), `unknown budget "tihgt"`) +} + +func TestNamedBudgetValidationNamesTheBudget(t *testing.T) { + cfg := Config{ + Agents: Agents{{Name: "root", Model: "openai/gpt-4o", Description: "d", Instruction: "i"}}, + Budgets: map[string]BudgetConfig{"tight": {MaxCost: -1}}, + } + err := cfg.Validate() + require.Error(t, err) + assert.Contains(t, err.Error(), "budgets.tight:") + assert.Contains(t, err.Error(), "max_cost must not be negative") +} + +func TestAgentWithoutBudgetsIsValid(t *testing.T) { + var cfg Config + require.NoError(t, yaml.Unmarshal([]byte(` +budgets: + tight: + max_cost: 0.10 +agents: + root: + model: openai/gpt-4o + description: d + instruction: i +`), &cfg)) + assert.Empty(t, cfg.Agents[0].Budgets) +} + +func TestBudgetPartialLeavesOthersUnlimited(t *testing.T) { + var cfg Config + require.NoError(t, yaml.Unmarshal([]byte(` +agents: + root: + model: openai/gpt-4o + description: d + instruction: i +budget: + max_time: 30s +`), &cfg)) + + require.NotNil(t, cfg.Budget) + assert.Equal(t, 30*time.Second, cfg.Budget.MaxTime.Duration) + assert.Zero(t, cfg.Budget.MaxCost) + assert.Zero(t, cfg.Budget.MaxTokens) + assert.False(t, cfg.Budget.IsZero()) +} diff --git a/pkg/config/latest/types.go b/pkg/config/latest/types.go index dc688903e0..696abb65b6 100644 --- a/pkg/config/latest/types.go +++ b/pkg/config/latest/types.go @@ -44,6 +44,82 @@ type Config struct { Metadata Metadata `json:"metadata"` Permissions *PermissionsConfig `json:"permissions,omitempty"` Runtime *RuntimeDefaults `json:"runtime,omitempty"` + // Budget caps the whole run, regardless of which agent spends. It is + // the simple case: one ceiling for everything. + Budget *BudgetConfig `json:"budget,omitempty"` + // Budgets are reusable, named budget definitions. An agent opts into + // one by listing its name in AgentConfig.Budgets, mirroring the + // top-level MCPs/RAG/Toolsets reference-by-name convention. + // + // A named budget is a single shared pot: when several agents reference + // the same name they draw from the *same* ceiling rather than each + // getting a copy of it. That is what keeps a budget meaningful when an + // agent fans out — ten sub-agents on one "tight" budget still cannot + // spend more than "tight" allows. Give agents distinct budget names + // when you want independent pots. + Budgets map[string]BudgetConfig `json:"budgets,omitempty"` +} + +// BudgetConfig caps what a single run may consume before the agent is +// stopped. Every limit is optional; a zero/unset limit is unlimited, and +// an absent budget block disables the whole mechanism. +// +// The budget is scoped to a run — one root RunStream — and is shared by +// every sub-session spawned inside it (task transfers, sub-agents, +// skills). Delegated work therefore spends against the same wallet as +// its parent, rather than each child receiving a fresh allowance. Per-run +// (not per-session) is the only accounting that makes a cost ceiling +// meaningful when an agent can fan out. +type BudgetConfig struct { + // MaxCost is the maximum spend for the run, in USD. + // + // Only responses the runtime can price count towards it: a model with + // no pricing data (an unknown ID, or a custom endpoint without a + // model-level `cost:` block) contributes nothing, because there is no + // honest number to add. Such a run emits a warning rather than + // silently behaving as if it were free — see [CostConfig] to price a + // custom endpoint explicitly. + MaxCost float64 `json:"max_cost,omitempty" yaml:"max_cost,omitempty"` + + // MaxTokens is the maximum cumulative input+output tokens for the run. + // + // This counts every token the run sends and receives over its whole + // lifetime. It is deliberately not the session's token counters, which + // compaction resets — those measure current context length, not spend. + MaxTokens int64 `json:"max_tokens,omitempty" yaml:"max_tokens,omitempty"` + + // 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. In Go + // duration format ("10m", "30s", "1h30m"). + // + // It is checked at turn boundaries, so a run stops at the first + // boundary after the deadline rather than interrupting an in-flight + // model call or tool. Expect overshoot up to the length of one turn. + MaxTime Duration `json:"max_time,omitzero" yaml:"max_time,omitempty"` +} + +// IsZero reports whether no limit is set, i.e. the block is inert. +func (b *BudgetConfig) IsZero() bool { + return b == nil || (b.MaxCost == 0 && b.MaxTokens == 0 && b.MaxTime.Duration == 0) +} + +// validate rejects negative limits. All-zero is valid and simply means +// "no limit", which keeps `budget: {}` from being an error. +func (b *BudgetConfig) validate() error { + if b == nil { + return nil + } + if b.MaxCost < 0 { + return fmt.Errorf("max_cost must not be negative, got %v", b.MaxCost) + } + if b.MaxTokens < 0 { + return fmt.Errorf("max_tokens must not be negative, got %d", b.MaxTokens) + } + if b.MaxTime.Duration < 0 { + return fmt.Errorf("max_time must not be negative, got %s", b.MaxTime.Duration) + } + return nil } // RuntimeDefaults captures execution-time defaults the agent author @@ -550,8 +626,13 @@ type AgentConfig struct { CodeModeTools bool `json:"code_mode_tools,omitempty"` AddDescriptionParameter bool `json:"add_description_parameter,omitempty"` MaxIterations int `json:"max_iterations,omitempty"` - MaxConsecutiveToolCalls int `json:"max_consecutive_tool_calls,omitempty"` - MaxOldToolCallTokens int `json:"max_old_tool_call_tokens,omitempty"` + // Budgets names the top-level `budgets` entries this agent spends + // against. Every listed budget must exist. An agent may list several, + // in which case all of them apply and the first to be exhausted stops + // the run. Agents sharing a name share one pot — see [Config.Budgets]. + Budgets []string `json:"budgets,omitempty" yaml:"budgets,omitempty"` + MaxConsecutiveToolCalls int `json:"max_consecutive_tool_calls,omitempty"` + MaxOldToolCallTokens int `json:"max_old_tool_call_tokens,omitempty"` // MaxToolResultTokens caps each textual tool result when it enters the // session. Oversized results are middle-out truncated (head and tail kept, // middle replaced with a marker); textual documents attached to the result diff --git a/pkg/config/latest/validate.go b/pkg/config/latest/validate.go index 6b5231f29d..c6627f6793 100644 --- a/pkg/config/latest/validate.go +++ b/pkg/config/latest/validate.go @@ -19,6 +19,15 @@ func (t *Config) UnmarshalYAML(unmarshal func(any) error) error { } func (t *Config) Validate() error { + if err := t.Budget.validate(); err != nil { + return fmt.Errorf("budget: %w", err) + } + for name := range t.Budgets { + b := t.Budgets[name] + if err := b.validate(); err != nil { + return fmt.Errorf("budgets.%s: %w", name, err) + } + } for name, p := range t.Providers { if err := p.Auth.Validate(p.Provider); err != nil { return fmt.Errorf("providers.%s: %w", name, err) @@ -64,6 +73,11 @@ func (t *Config) Validate() error { if err := validateCompactionThreshold(agent.CompactionThreshold); err != nil { return fmt.Errorf("agents.%s: %w", agent.Name, err) } + for _, name := range agent.Budgets { + if _, ok := t.Budgets[name]; !ok { + return fmt.Errorf("agents.%s: budgets: unknown budget %q; define it under the top-level 'budgets'", agent.Name, name) + } + } for j := range agent.Toolsets { if err := agent.Toolsets[j].validate(); err != nil { diff --git a/pkg/config/schema_test.go b/pkg/config/schema_test.go index d663348ee9..e7827e665f 100644 --- a/pkg/config/schema_test.go +++ b/pkg/config/schema_test.go @@ -143,6 +143,7 @@ func TestSchemaMatchesGoTypes(t *testing.T) { "ScriptShellToolConfig": reflect.TypeFor[latest.ScriptShellToolConfig](), "PostEditConfig": reflect.TypeFor[latest.PostEditConfig](), "PermissionsConfig": reflect.TypeFor[latest.PermissionsConfig](), + "BudgetConfig": reflect.TypeFor[latest.BudgetConfig](), "HooksConfig": reflect.TypeFor[latest.HooksConfig](), "HookMatcherConfig": reflect.TypeFor[latest.HookMatcherConfig](), "HookDefinition": reflect.TypeFor[latest.HookDefinition](), diff --git a/pkg/runtime/budget.go b/pkg/runtime/budget.go new file mode 100644 index 0000000000..668e2fd565 --- /dev/null +++ b/pkg/runtime/budget.go @@ -0,0 +1,431 @@ +package runtime + +import ( + "context" + "fmt" + "log/slog" + "sort" + "sync" + "time" + + "github.com/docker/docker-agent/pkg/agent" + "github.com/docker/docker-agent/pkg/chat" + "github.com/docker/docker-agent/pkg/config/latest" + "github.com/docker/docker-agent/pkg/session" +) + +type budgetLimit string + +const ( + budgetLimitCost budgetLimit = "max_cost" + budgetLimitTokens budgetLimit = "max_tokens" + budgetLimitTime budgetLimit = "max_time" +) + +type budgetTracker struct { + maxCost float64 + maxTokens int64 + maxTime time.Duration + mu sync.Mutex + cost float64 + tokens int64 + active time.Duration + unpriced bool + perAgent map[string]*agentSpend +} + +type agentSpend struct { + cost float64 + tokens int64 + active time.Duration +} + +func newBudgetTracker(cfg *latest.BudgetConfig) *budgetTracker { + if cfg.IsZero() { + return nil + } + return &budgetTracker{ + maxCost: cfg.MaxCost, + maxTokens: cfg.MaxTokens, + maxTime: cfg.MaxTime.Duration, + perAgent: make(map[string]*agentSpend), + } +} + +func (b *budgetTracker) record(agentName string, usage *chat.Usage, cost *float64, active time.Duration) { + if b == nil { + return + } + b.mu.Lock() + defer b.mu.Unlock() + + var addTokens int64 + if usage != nil { + addTokens = usage.InputTokens + usage.OutputTokens + b.tokens += addTokens + } + var addCost float64 + switch { + case cost != nil: + addCost = *cost + b.cost += addCost + case usage != nil: + b.unpriced = true + } + + if active > 0 { + b.active += active + } + + spend := b.perAgent[agentName] + if spend == nil { + spend = &agentSpend{} + b.perAgent[agentName] = spend + } + spend.cost += addCost + spend.tokens += addTokens + if active > 0 { + spend.active += active + } +} + +type budgetSnapshot struct { + Cost float64 + MaxCost float64 + Tokens int64 + MaxTokens int64 + Elapsed time.Duration + MaxTime time.Duration + Unpriced bool + PerAgent []agentBudgetSpend +} + +type agentBudgetSpend struct { + AgentName string + Cost float64 + Tokens int64 + Active time.Duration +} + +func (b *budgetTracker) snapshot() budgetSnapshot { + if b == nil { + return budgetSnapshot{} + } + b.mu.Lock() + defer b.mu.Unlock() + return budgetSnapshot{ + Cost: b.cost, + MaxCost: b.maxCost, + Tokens: b.tokens, + MaxTokens: b.maxTokens, + Elapsed: b.active, + MaxTime: b.maxTime, + Unpriced: b.unpriced, + PerAgent: b.perAgentSpendLocked(), + } +} + +func (b *budgetTracker) perAgentSpendLocked() []agentBudgetSpend { + if len(b.perAgent) == 0 { + return nil + } + out := make([]agentBudgetSpend, 0, len(b.perAgent)) + for name, s := range b.perAgent { + out = append(out, agentBudgetSpend{ + AgentName: name, + Cost: s.cost, + Tokens: s.tokens, + Active: s.active, + }) + } + sort.Slice(out, func(i, j int) bool { + if out[i].Cost != out[j].Cost { + return out[i].Cost > out[j].Cost + } + if out[i].Tokens != out[j].Tokens { + return out[i].Tokens > out[j].Tokens + } + return out[i].AgentName < out[j].AgentName + }) + return out +} + +type budgetBreach struct { + Budget string + Limit budgetLimit + Used string + Max string +} + +func (br budgetBreach) Message() string { + return fmt.Sprintf( + "Execution stopped after reaching the configured %s limit (used %s of %s).", + br.configPath(), br.Used, br.Max, + ) +} + +func (br budgetBreach) configPath() string { + if br.Budget == "" || br.Budget == runBudgetName { + return "budget." + string(br.Limit) + } + return "budgets." + br.Budget + "." + string(br.Limit) +} + +func (b *budgetTracker) exceeded() *budgetBreach { + if b == nil { + return nil + } + b.mu.Lock() + defer b.mu.Unlock() + + if b.maxCost > 0 && b.cost >= b.maxCost { + return &budgetBreach{ + Limit: budgetLimitCost, + Used: formatUSD(b.cost), + Max: formatUSD(b.maxCost), + } + } + if b.maxTokens > 0 && b.tokens >= b.maxTokens { + return &budgetBreach{ + Limit: budgetLimitTokens, + Used: fmt.Sprintf("%d tokens", b.tokens), + Max: fmt.Sprintf("%d tokens", b.maxTokens), + } + } + if b.maxTime > 0 && b.active >= b.maxTime { + return &budgetBreach{ + Limit: budgetLimitTime, + Used: b.active.Round(time.Second).String(), + Max: b.maxTime.String(), + } + } + return nil +} + +func (b *budgetTracker) unpricedSpend() bool { + if b == nil { + return false + } + b.mu.Lock() + defer b.mu.Unlock() + return b.unpriced && b.maxCost > 0 +} + +func formatUSD(v float64) string { + if v < 0 { + return "-" + formatUSD(-v) + } + if v < 0.0001 { + return "$0.00" + } + if v < 0.01 { + return fmt.Sprintf("$%.4f", v) + } + return fmt.Sprintf("$%.2f", v) +} + +const runBudgetName = "run" + +type budgetSet struct { + trackers map[string]*budgetTracker + agentBudgets map[string][]string + order []string +} + +func (s *budgetSet) budgetsFor(agentName string) []namedTracker { + if s == nil { + return nil + } + names := s.agentBudgets[agentName] + if len(names) == 0 { + if t := s.trackers[runBudgetName]; t != nil { + return []namedTracker{{Name: runBudgetName, Tracker: t}} + } + return nil + } + out := make([]namedTracker, 0, len(names)) + for _, n := range names { + if t := s.trackers[n]; t != nil { + out = append(out, namedTracker{Name: n, Tracker: t}) + } + } + return out +} + +func (s *budgetSet) all() []namedTracker { + if s == nil { + return nil + } + out := make([]namedTracker, 0, len(s.order)) + for _, n := range s.order { + if t := s.trackers[n]; t != nil { + out = append(out, namedTracker{Name: n, Tracker: t}) + } + } + return out +} + +type namedTracker struct { + Name string + Tracker *budgetTracker +} + +func newBudgetSet(runBudget *latest.BudgetConfig, named map[string]latest.BudgetConfig, agentBudgets map[string][]string) *budgetSet { + s := &budgetSet{ + trackers: make(map[string]*budgetTracker), + agentBudgets: make(map[string][]string, len(agentBudgets)), + } + if t := newBudgetTracker(runBudget); t != nil { + s.trackers[runBudgetName] = t + s.order = append(s.order, runBudgetName) + } + + referenced := make(map[string]bool) + for _, names := range agentBudgets { + for _, n := range names { + referenced[n] = true + } + } + namedOrder := make([]string, 0, len(referenced)) + for n := range referenced { + cfg, ok := named[n] + if !ok { + continue + } + if t := newBudgetTracker(&cfg); t != nil { + s.trackers[n] = t + namedOrder = append(namedOrder, n) + } + } + sort.Strings(namedOrder) + s.order = append(s.order, namedOrder...) + + for agentName, names := range agentBudgets { + list := make([]string, 0, len(names)+1) + if _, ok := s.trackers[runBudgetName]; ok { + list = append(list, runBudgetName) + } + for _, n := range names { + if _, ok := s.trackers[n]; ok { + list = append(list, n) + } + } + s.agentBudgets[agentName] = list + } + + if len(s.trackers) == 0 { + return nil + } + return s +} + +func (r *LocalRuntime) ensureBudget() { + r.budgetMu.Lock() + defer r.budgetMu.Unlock() + if r.budgetStarted { + return + } + r.budgetStarted = true + r.budget = newBudgetSet(r.budgetCfg, r.budgetsCfg, r.agentBudgets) +} + +func (r *LocalRuntime) currentBudget() *budgetSet { + r.budgetMu.Lock() + defer r.budgetMu.Unlock() + return r.budget +} + +func (r *LocalRuntime) enforceBudget( + ctx context.Context, + sess *session.Session, + a *agent.Agent, + events EventSink, +) iterationDecision { + breach := r.currentBudget().exceededFor(a.Name()) + if breach == nil { + return iterationContinue + } + + slog.InfoContext(ctx, "Run budget exceeded", + "agent", a.Name(), + "session_id", sess.ID, + "budget", breach.Budget, + "limit", string(breach.Limit), + "used", breach.Used, + "max", breach.Max, + ) + + events.Emit(BudgetExceeded(sess.ID, a.Name(), *breach)) + r.notifyBudgetExceeded(ctx, a, sess.ID, breach.Message()) + + addAgentMessage(sess, a, &chat.Message{ + Role: chat.MessageRoleAssistant, + Content: breach.Message(), + CreatedAt: r.now().Format(time.RFC3339), + }, events) + + return iterationStop +} + +func (s *budgetSet) exceededFor(agentName string) *budgetBreach { + for _, nt := range s.budgetsFor(agentName) { + if br := nt.Tracker.exceeded(); br != nil { + br.Budget = nt.Name + return br + } + } + return nil +} + +func (r *LocalRuntime) recordBudget(sess *session.Session, a *agent.Agent, usage *chat.Usage, cost *float64, active time.Duration, events EventSink) { + s := r.currentBudget() + if s == nil { + return + } + targets := s.budgetsFor(a.Name()) + if len(targets) == 0 { + return + } + + warnUnpriced := !s.unpricedSpend() + for _, nt := range targets { + nt.Tracker.record(a.Name(), usage, cost, active) + } + if warnUnpriced && s.unpricedSpend() { + events.Emit(Warning( + "This run has a max_cost limit, but the model reported usage the runtime cannot price, "+ + "so that spend does not count against the limit. Set a model-level `cost:` block to price it.", + a.Name(), + )) + } + events.Emit(BudgetUsage(sess.ID, a.Name(), s.snapshot())) +} + +func (s *budgetSet) unpricedSpend() bool { + for _, nt := range s.all() { + if nt.Tracker.unpricedSpend() { + return true + } + } + return false +} + +func (s *budgetSet) snapshot() []namedBudgetSnapshot { + all := s.all() + if len(all) == 0 { + return nil + } + out := make([]namedBudgetSnapshot, 0, len(all)) + for _, nt := range all { + out = append(out, namedBudgetSnapshot{ + Name: nt.Name, + Snapshot: nt.Tracker.snapshot(), + }) + } + return out +} + +type namedBudgetSnapshot struct { + Name string + Snapshot budgetSnapshot +} diff --git a/pkg/runtime/budget_test.go b/pkg/runtime/budget_test.go new file mode 100644 index 0000000000..643c6b4bd5 --- /dev/null +++ b/pkg/runtime/budget_test.go @@ -0,0 +1,369 @@ +package runtime + +import ( + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/docker/docker-agent/pkg/chat" + "github.com/docker/docker-agent/pkg/config/latest" +) + +var budgetEpoch = time.Date(2026, 1, 1, 0, 0, 0, 0, time.UTC) + +func TestNewBudgetTrackerNilWhenNoLimitSet(t *testing.T) { + tests := []struct { + name string + cfg *latest.BudgetConfig + }{ + {"nil config", nil}, + {"empty config", &latest.BudgetConfig{}}, + {"all zero", &latest.BudgetConfig{MaxCost: 0, MaxTokens: 0}}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + assert.Nil(t, newBudgetTracker(tt.cfg), + "a config with no ceilings must leave the run unbudgeted") + }) + } +} + +func TestNilBudgetTrackerIsInert(t *testing.T) { + var b *budgetTracker + assert.NotPanics(t, func() { + b.record("root", &chat.Usage{InputTokens: 10}, new(1.0), time.Second) + assert.Nil(t, b.exceeded()) + assert.Equal(t, budgetSnapshot{}, b.snapshot()) + assert.False(t, b.unpricedSpend()) + }) +} + +func TestBudgetMaxCost(t *testing.T) { + b := newBudgetTracker(&latest.BudgetConfig{MaxCost: 0.50}) + require.NotNil(t, b) + + b.record("root", &chat.Usage{InputTokens: 100}, new(0.20), time.Second) + assert.Nil(t, b.exceeded(), "under budget must not trip") + + b.record("root", &chat.Usage{InputTokens: 100}, new(0.20), time.Second) + assert.Nil(t, b.exceeded(), "$0.40 of $0.50 must not trip") + + b.record("root", &chat.Usage{InputTokens: 100}, new(0.15), time.Second) + breach := b.exceeded() + require.NotNil(t, breach, "$0.55 of $0.50 must trip") + assert.Equal(t, budgetLimitCost, breach.Limit) + assert.Equal(t, "$0.55", breach.Used) + assert.Equal(t, "$0.50", breach.Max) + assert.Contains(t, breach.Message(), "budget.max_cost") +} + +func TestBudgetTripsOnExactLimit(t *testing.T) { + b := newBudgetTracker(&latest.BudgetConfig{MaxCost: 0.50}) + b.record("root", &chat.Usage{}, new(0.50), time.Second) + require.NotNil(t, b.exceeded()) +} + +func TestBudgetMaxTokens(t *testing.T) { + b := newBudgetTracker(&latest.BudgetConfig{MaxTokens: 1000}) + + b.record("root", &chat.Usage{InputTokens: 400, OutputTokens: 100}, nil, time.Second) + assert.Nil(t, b.exceeded(), "500 of 1000 must not trip") + + b.record("root", &chat.Usage{InputTokens: 400, OutputTokens: 200}, nil, time.Second) + breach := b.exceeded() + require.NotNil(t, breach, "1100 of 1000 must trip") + assert.Equal(t, budgetLimitTokens, breach.Limit) + assert.Equal(t, "1100 tokens", breach.Used) + assert.Equal(t, "1000 tokens", breach.Max) +} + +func TestBudgetTokensAreMonotonicAndCountBothDirections(t *testing.T) { + b := newBudgetTracker(&latest.BudgetConfig{MaxTokens: 100}) + for range 10 { + b.record("root", &chat.Usage{InputTokens: 3, OutputTokens: 2}, nil, time.Second) + } + assert.Equal(t, int64(50), b.snapshot().Tokens, + "10 turns x (3 in + 2 out) must accumulate to 50") + assert.Nil(t, b.exceeded()) + + for range 11 { + b.record("root", &chat.Usage{InputTokens: 3, OutputTokens: 2}, nil, time.Second) + } + assert.Equal(t, int64(105), b.snapshot().Tokens) + require.NotNil(t, b.exceeded()) +} + +func TestBudgetMaxTime(t *testing.T) { + b := newBudgetTracker(&latest.BudgetConfig{MaxTime: latest.Duration{Duration: 10 * time.Minute}}) + + b.record("root", &chat.Usage{}, nil, 9*time.Minute) + assert.Nil(t, b.exceeded(), "9m of active work in a 10m budget must not trip") + + b.record("root", &chat.Usage{}, nil, 2*time.Minute) + breach := b.exceeded() + require.NotNil(t, breach, "11m of active work must trip a 10m budget") + assert.Equal(t, budgetLimitTime, breach.Limit) + assert.Equal(t, "11m0s", breach.Used) + assert.Equal(t, "10m0s", breach.Max) +} + +func TestBudgetMaxTimeIgnoresIdleTime(t *testing.T) { + b := newBudgetTracker(&latest.BudgetConfig{MaxTime: latest.Duration{Duration: time.Minute}}) + b.record("root", &chat.Usage{}, nil, 2*time.Second) + b.record("root", &chat.Usage{}, nil, 3*time.Second) + assert.Equal(t, 5*time.Second, b.snapshot().Elapsed, + "only active turn time counts, not wall-clock since the session began") + assert.Nil(t, b.exceeded()) +} + +func TestBudgetUnsetLimitsNeverTrip(t *testing.T) { + b := newBudgetTracker(&latest.BudgetConfig{MaxTime: latest.Duration{Duration: time.Hour}}) + b.record("root", &chat.Usage{InputTokens: 1_000_000, OutputTokens: 1_000_000}, new(999.0), time.Second) + assert.Nil(t, b.exceeded(), "only max_time is set; cost and tokens must be ignored") +} + +func TestBudgetUnpricedSpendIsFlaggedNotCountedAsFree(t *testing.T) { + b := newBudgetTracker(&latest.BudgetConfig{MaxCost: 0.50}) + + b.record("root", &chat.Usage{InputTokens: 5000, OutputTokens: 5000}, nil, time.Second) + assert.True(t, b.unpricedSpend(), "usage with no price must flag the run as unpriced") + assert.Zero(t, b.snapshot().Cost, "unpriceable spend must not invent a number") + assert.Nil(t, b.exceeded(), "unpriced spend cannot trip a cost ceiling") + assert.True(t, b.snapshot().Unpriced) +} + +func TestBudgetPricedFreeCallIsNotUnpriced(t *testing.T) { + b := newBudgetTracker(&latest.BudgetConfig{MaxCost: 0.50}) + b.record("root", &chat.Usage{InputTokens: 10}, new(0.0), time.Second) + assert.False(t, b.unpricedSpend(), "a priced free call is not unpriced") + assert.False(t, b.snapshot().Unpriced) +} + +func TestBudgetUnpricedIrrelevantWithoutCostLimit(t *testing.T) { + b := newBudgetTracker(&latest.BudgetConfig{MaxTokens: 1000}) + b.record("root", &chat.Usage{InputTokens: 10}, nil, time.Second) + assert.False(t, b.unpricedSpend(), "no max_cost means unpriced spend is not a problem") +} + +func TestBudgetReportsCostFirstWhenSeveralTrip(t *testing.T) { + b := newBudgetTracker(&latest.BudgetConfig{ + MaxCost: 0.10, + MaxTokens: 10, + MaxTime: latest.Duration{Duration: time.Minute}, + }) + b.record("root", &chat.Usage{InputTokens: 100, OutputTokens: 100}, new(5.0), time.Second) + + breach := b.exceeded() + require.NotNil(t, breach) + assert.Equal(t, budgetLimitCost, breach.Limit) +} + +func TestBudgetSnapshotReportsLimitsAndTotals(t *testing.T) { + b := newBudgetTracker(&latest.BudgetConfig{ + MaxCost: 2, + MaxTokens: 1000, + MaxTime: latest.Duration{Duration: 10 * time.Minute}, + }) + b.record("root", &chat.Usage{InputTokens: 30, OutputTokens: 20}, new(0.25), time.Second) + + s := b.snapshot() + assert.InDelta(t, 0.25, s.Cost, 1e-9) + assert.InDelta(t, 2.0, s.MaxCost, 1e-9) + assert.Equal(t, int64(50), s.Tokens) + assert.Equal(t, int64(1000), s.MaxTokens) + assert.Equal(t, time.Second, s.Elapsed, "Elapsed is accumulated active time") + assert.Equal(t, 10*time.Minute, s.MaxTime) +} + +func TestBudgetTrackerIsConcurrencySafe(t *testing.T) { + b := newBudgetTracker(&latest.BudgetConfig{MaxCost: 1000}) + + done := make(chan struct{}) + for range 8 { + go func() { + defer func() { done <- struct{}{} }() + for range 50 { + b.record("root", &chat.Usage{InputTokens: 1, OutputTokens: 1}, new(0.01), time.Second) + b.exceeded() + b.snapshot() + } + }() + } + for range 8 { + <-done + } + + s := b.snapshot() + assert.Equal(t, int64(800), s.Tokens, "8 goroutines x 50 turns x 2 tokens") + assert.InDelta(t, 4.0, s.Cost, 1e-6, "8 x 50 x $0.01") +} + +func TestBudgetRecordToleratesNilUsage(t *testing.T) { + b := newBudgetTracker(&latest.BudgetConfig{MaxTokens: 10}) + assert.NotPanics(t, func() { b.record("root", nil, nil, time.Second) }) + assert.Zero(t, b.snapshot().Tokens) +} + +func TestBudgetPerAgentBreakdown(t *testing.T) { + b := newBudgetTracker(&latest.BudgetConfig{MaxCost: 10}) + + b.record("root", &chat.Usage{InputTokens: 100, OutputTokens: 100}, new(0.02), 2*time.Second) + b.record("developer", &chat.Usage{InputTokens: 300, OutputTokens: 300}, new(0.09), 5*time.Second) + b.record("root", &chat.Usage{InputTokens: 50, OutputTokens: 50}, new(0.01), time.Second) + + per := b.snapshot().PerAgent + require.Len(t, per, 2) + + assert.Equal(t, "developer", per[0].AgentName) + assert.InDelta(t, 0.09, per[0].Cost, 1e-9) + assert.Equal(t, int64(600), per[0].Tokens) + assert.Equal(t, 5*time.Second, per[0].Active) + + assert.Equal(t, "root", per[1].AgentName) + assert.InDelta(t, 0.03, per[1].Cost, 1e-9) + assert.Equal(t, int64(300), per[1].Tokens) + assert.Equal(t, 3*time.Second, per[1].Active) + + s := b.snapshot() + assert.InDelta(t, per[0].Cost+per[1].Cost, s.Cost, 1e-9) + assert.Equal(t, per[0].Tokens+per[1].Tokens, s.Tokens) +} + +func TestBudgetPerAgentEmptyBeforeSpend(t *testing.T) { + b := newBudgetTracker(&latest.BudgetConfig{MaxCost: 1}) + assert.Empty(t, b.snapshot().PerAgent) +} + +func budgetSetFixture() *budgetSet { + return newBudgetSet( + &latest.BudgetConfig{MaxCost: 1.00}, + map[string]latest.BudgetConfig{ + "tight": {MaxCost: 0.10}, + "roomy": {MaxTokens: 1_000_000}, + "unused": {MaxCost: 99}, + }, + map[string][]string{ + "root": {"tight"}, + "developer": {"tight", "roomy"}, + }, + ) +} + +func TestBudgetSetNamedBudgetIsSharedNotCopied(t *testing.T) { + s := budgetSetFixture() + require.NotNil(t, s) + + for _, agent := range []string{"root", "developer"} { + for _, nt := range s.budgetsFor(agent) { + nt.Tracker.record(agent, &chat.Usage{OutputTokens: 10}, new(0.04), time.Second) + } + } + + tight := s.trackers["tight"] + assert.InDelta(t, 0.08, tight.snapshot().Cost, 1e-9, + "a named budget must be one shared pot across referencing agents") + assert.Nil(t, s.exceededFor("root"), "$0.08 of $0.10 must not trip") + + for _, nt := range s.budgetsFor("root") { + nt.Tracker.record("root", &chat.Usage{OutputTokens: 10}, new(0.04), time.Second) + } + for _, agent := range []string{"root", "developer"} { + br := s.exceededFor(agent) + require.NotNil(t, br, "shared pot exhausted must trip for %s", agent) + assert.Equal(t, "tight", br.Budget) + assert.Equal(t, budgetLimitCost, br.Limit) + } +} + +func TestBudgetSetChargesEveryReferencedBudget(t *testing.T) { + s := budgetSetFixture() + + names := make([]string, 0) + for _, nt := range s.budgetsFor("developer") { + names = append(names, nt.Name) + } + assert.Equal(t, []string{runBudgetName, "tight", "roomy"}, names, + "developer draws on the run-wide budget and both of its named budgets") + + names = names[:0] + for _, nt := range s.budgetsFor("root") { + names = append(names, nt.Name) + } + assert.Equal(t, []string{runBudgetName, "tight"}, names) +} + +func TestBudgetSetUndeclaredAgentUsesRunBudget(t *testing.T) { + s := budgetSetFixture() + nts := s.budgetsFor("some-other-agent") + require.Len(t, nts, 1) + assert.Equal(t, runBudgetName, nts[0].Name) +} + +func TestBudgetSetSkipsUnreferencedBudgets(t *testing.T) { + s := budgetSetFixture() + assert.NotContains(t, s.trackers, "unused") + assert.NotContains(t, s.order, "unused") +} + +func TestBudgetSetOrderIsStable(t *testing.T) { + s := budgetSetFixture() + assert.Equal(t, []string{runBudgetName, "roomy", "tight"}, s.order) +} + +func TestBudgetSetRunWideBudgetTrips(t *testing.T) { + s := newBudgetSet( + &latest.BudgetConfig{MaxCost: 0.05}, + map[string]latest.BudgetConfig{"roomy": {MaxCost: 100}}, + map[string][]string{"root": {"roomy"}}, + ) + for _, nt := range s.budgetsFor("root") { + nt.Tracker.record("root", &chat.Usage{OutputTokens: 10}, new(0.06), time.Second) + } + br := s.exceededFor("root") + require.NotNil(t, br) + assert.Equal(t, runBudgetName, br.Budget) + assert.Contains(t, br.Message(), "budget.max_cost") +} + +func TestBudgetBreachConfigPath(t *testing.T) { + assert.Equal(t, "budget.max_cost", + budgetBreach{Budget: runBudgetName, Limit: budgetLimitCost}.configPath()) + assert.Equal(t, "budgets.tight.max_tokens", + budgetBreach{Budget: "tight", Limit: budgetLimitTokens}.configPath()) + assert.Equal(t, "budget.max_cost", + budgetBreach{Limit: budgetLimitCost}.configPath()) +} + +func TestBudgetSetNilWhenNothingConfigured(t *testing.T) { + assert.Nil(t, newBudgetSet(nil, nil, nil)) + assert.Nil(t, newBudgetSet(&latest.BudgetConfig{}, map[string]latest.BudgetConfig{}, map[string][]string{})) + assert.Nil(t, newBudgetSet(nil, map[string]latest.BudgetConfig{"x": {MaxCost: 1}}, nil)) +} + +func TestBudgetSetSnapshotPerBudget(t *testing.T) { + s := budgetSetFixture() + for _, nt := range s.budgetsFor("developer") { + nt.Tracker.record("developer", &chat.Usage{InputTokens: 5, OutputTokens: 5}, new(0.02), time.Second) + } + + snaps := s.snapshot() + require.Len(t, snaps, 3) + assert.Equal(t, runBudgetName, snaps[0].Name) + assert.Equal(t, "roomy", snaps[1].Name) + assert.Equal(t, "tight", snaps[2].Name) + + assert.InDelta(t, 0.02, snaps[0].Snapshot.Cost, 1e-9) + assert.InDelta(t, 1.00, snaps[0].Snapshot.MaxCost, 1e-9) + assert.Equal(t, int64(10), snaps[1].Snapshot.Tokens) + assert.InDelta(t, 0.10, snaps[2].Snapshot.MaxCost, 1e-9) +} + +func TestBudgetConfigIsZero(t *testing.T) { + assert.True(t, (*latest.BudgetConfig)(nil).IsZero()) + assert.True(t, (&latest.BudgetConfig{}).IsZero()) + assert.False(t, (&latest.BudgetConfig{MaxCost: 0.5}).IsZero()) + assert.False(t, (&latest.BudgetConfig{MaxTokens: 1}).IsZero()) + assert.False(t, (&latest.BudgetConfig{MaxTime: latest.Duration{Duration: time.Second}}).IsZero()) +} diff --git a/pkg/runtime/budget_wiring_test.go b/pkg/runtime/budget_wiring_test.go new file mode 100644 index 0000000000..d621d8e8f6 --- /dev/null +++ b/pkg/runtime/budget_wiring_test.go @@ -0,0 +1,148 @@ +package runtime + +import ( + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/docker/docker-agent/pkg/agent" + "github.com/docker/docker-agent/pkg/chat" + "github.com/docker/docker-agent/pkg/config/latest" + "github.com/docker/docker-agent/pkg/session" +) + +type collectSink struct{ events []Event } + +func (s *collectSink) Emit(e Event) { s.events = append(s.events, e) } + +func (s *collectSink) budgetUsages() []*BudgetUsageEvent { + var out []*BudgetUsageEvent + for _, e := range s.events { + if b, ok := e.(*BudgetUsageEvent); ok { + out = append(out, b) + } + } + return out +} + +func budgetRuntime(t *testing.T, clock func() time.Time) *LocalRuntime { + t.Helper() + r := &LocalRuntime{now: clock} + WithBudget(&latest.BudgetConfig{MaxCost: 0.05, MaxTokens: 15000, MaxTime: latest.Duration{Duration: 2 * time.Minute}})(r) + WithNamedBudgets( + map[string]latest.BudgetConfig{ + "shell-work": {MaxCost: 0.03, MaxTokens: 8000}, + }, + map[string][]string{"root": {"shell-work"}}, + )(r) + return r +} + +func TestRecordBudgetEmitsNonZeroReading(t *testing.T) { + now := budgetEpoch + r := budgetRuntime(t, func() time.Time { return now }) + r.ensureBudget() + + sess := session.New() + a := agent.New("root", "test") + sink := &collectSink{} + + now = budgetEpoch.Add(30 * time.Second) + r.recordBudget(sess, a, &chat.Usage{InputTokens: 1000, OutputTokens: 200}, new(0.01), 30*time.Second, sink) + + usages := sink.budgetUsages() + require.Len(t, usages, 1, "recordBudget must emit exactly one budget_usage event") + + byName := map[string]BudgetStatus{} + for _, b := range usages[0].Budgets { + byName[b.Name] = b + } + require.Contains(t, byName, runBudgetName) + require.Contains(t, byName, "shell-work") + + run := byName[runBudgetName] + assert.InDelta(t, 0.01, run.Cost, 1e-9, "run budget must see the turn's cost") + assert.Equal(t, int64(1200), run.Tokens, "run budget must see the turn's tokens") + assert.InDelta(t, 30, run.ElapsedSeconds, 0.001, "elapsed must advance with the clock") + + sw := byName["shell-work"] + assert.InDelta(t, 0.01, sw.Cost, 1e-9) + assert.Equal(t, int64(1200), sw.Tokens) + + require.Len(t, run.PerAgent, 1) + assert.Equal(t, "root", run.PerAgent[0].AgentName) + assert.Equal(t, int64(1200), run.PerAgent[0].Tokens) +} + +func TestRecordBudgetAccumulatesAcrossTurns(t *testing.T) { + now := budgetEpoch + r := budgetRuntime(t, func() time.Time { return now }) + r.ensureBudget() + + sess := session.New() + a := agent.New("root", "test") + sink := &collectSink{} + + for i := range 3 { + now = budgetEpoch.Add(time.Duration(i+1) * 10 * time.Second) + r.recordBudget(sess, a, &chat.Usage{InputTokens: 100, OutputTokens: 100}, new(0.001), 10*time.Second, sink) + } + + usages := sink.budgetUsages() + require.Len(t, usages, 3) + last := usages[2] + for _, b := range last.Budgets { + if b.Name != runBudgetName { + continue + } + assert.Equal(t, int64(600), b.Tokens, "3 turns x 200 tokens must accumulate") + assert.InDelta(t, 0.003, b.Cost, 1e-9, "3 turns x $0.001 must accumulate") + assert.InDelta(t, 30, b.ElapsedSeconds, 0.001) + } +} + +func TestRecordBudgetNilUsageIsNotSilentZero(t *testing.T) { + now := budgetEpoch + r := budgetRuntime(t, func() time.Time { return now }) + r.ensureBudget() + + sink := &collectSink{} + r.recordBudget(session.New(), agent.New("root", "test"), nil, nil, time.Second, sink) + + usages := sink.budgetUsages() + require.Len(t, usages, 1) + for _, b := range usages[0].Budgets { + assert.Zero(t, b.Tokens) + } +} + +func TestBudgetSurvivesAcrossMessages(t *testing.T) { + now := budgetEpoch + r := budgetRuntime(t, func() time.Time { return now }) + + sess := session.New() + a := agent.New("root", "test") + sink := &collectSink{} + + r.ensureBudget() + now = budgetEpoch.Add(10 * time.Second) + r.recordBudget(sess, a, &chat.Usage{InputTokens: 500, OutputTokens: 500}, new(0.02), 10*time.Second, sink) + + r.ensureBudget() + now = budgetEpoch.Add(20 * time.Second) + r.recordBudget(sess, a, &chat.Usage{InputTokens: 500, OutputTokens: 500}, new(0.02), 10*time.Second, sink) + + usages := sink.budgetUsages() + last := usages[len(usages)-1] + for _, b := range last.Budgets { + if b.Name != runBudgetName { + continue + } + assert.Equal(t, int64(2000), b.Tokens, + "spend from earlier messages in the same session must still count") + assert.InDelta(t, 0.04, b.Cost, 1e-9, + "a budget that resets per message lets a session spend the ceiling on every turn") + } +} diff --git a/pkg/runtime/client.go b/pkg/runtime/client.go index 8c0819943c..2d674b9e20 100644 --- a/pkg/runtime/client.go +++ b/pkg/runtime/client.go @@ -93,6 +93,8 @@ func NewClient(baseURL string, opts ...ClientOption) (*Client, error) { "session_compaction": func() Event { return &SessionCompactionEvent{} }, "partial_tool_call": func() Event { return &PartialToolCallEvent{} }, "max_iterations_reached": func() Event { return &MaxIterationsReachedEvent{} }, + "budget_usage": func() Event { return &BudgetUsageEvent{} }, + "budget_exceeded": func() Event { return &BudgetExceededEvent{} }, "error": func() Event { return &ErrorEvent{} }, "elicitation_request": func() Event { return &ElicitationRequestEvent{} }, "authorization_event": func() Event { return &AuthorizationEvent{} }, diff --git a/pkg/runtime/event.go b/pkg/runtime/event.go index e38b0201d5..5223d39fa8 100644 --- a/pkg/runtime/event.go +++ b/pkg/runtime/event.go @@ -672,6 +672,132 @@ func MaxIterationsReached(maxIterations int) Event { } } +// BudgetUsageEvent reports a run's spend against its configured budget. +// It is emitted after every priced turn so a UI can track consumption +// while the run is still going, rather than only learning about the +// budget at the moment it stops the agent. +// +// Fields are absent when the corresponding limit is unset, so a consumer +// can render only the ceilings the operator actually configured. +type BudgetUsageEvent struct { + AgentContext + + Type string `json:"type"` + SessionID string `json:"session_id"` + // Budgets carries one reading per active budget, run-wide first then + // named budgets alphabetically, so a consumer can render them in a + // stable order. + Budgets []BudgetStatus `json:"budgets,omitempty"` +} + +// BudgetStatus is one budget's reading: its ceilings, what has been spent +// against them, and which agents did the spending. +type BudgetStatus struct { + // Name is "run" for the top-level `budget:`, otherwise the key under + // `budgets:`. + Name string `json:"name"` + Cost float64 `json:"cost"` + MaxCost float64 `json:"max_cost,omitempty"` + Tokens int64 `json:"tokens"` + MaxTokens int64 `json:"max_tokens,omitempty"` + // ElapsedSeconds and MaxTimeSeconds are seconds rather than a + // duration string so consumers can compute a ratio without parsing. + ElapsedSeconds float64 `json:"elapsed_seconds"` + MaxTimeSeconds float64 `json:"max_time_seconds,omitempty"` + // Unpriced reports that at least one response could not be priced, so + // Cost understates what was really spent. + Unpriced bool `json:"unpriced,omitempty"` + // PerAgent attributes this budget's totals to each agent that spent + // from it, biggest spender first. Empty until a turn completes. + PerAgent []AgentBudgetUsage `json:"per_agent,omitempty"` +} + +// AgentBudgetUsage is one agent's attributed slice of a budget. +type AgentBudgetUsage struct { + AgentName string `json:"agent_name"` + Cost float64 `json:"cost"` + Tokens int64 `json:"tokens"` + ActiveSeconds float64 `json:"active_seconds"` +} + +// GetSessionID implements SessionScoped so sub-session budget events can +// be attributed to the sub-session that produced them, even though they +// count against the root run's shared budget. +func (e *BudgetUsageEvent) GetSessionID() string { return e.SessionID } + +func BudgetUsage(sessionID, agentName string, snaps []namedBudgetSnapshot) Event { + budgets := make([]BudgetStatus, 0, len(snaps)) + for _, ns := range snaps { + s := ns.Snapshot + var perAgent []AgentBudgetUsage + if len(s.PerAgent) > 0 { + perAgent = make([]AgentBudgetUsage, len(s.PerAgent)) + for i, a := range s.PerAgent { + perAgent[i] = AgentBudgetUsage{ + AgentName: a.AgentName, + Cost: a.Cost, + Tokens: a.Tokens, + ActiveSeconds: a.Active.Seconds(), + } + } + } + budgets = append(budgets, BudgetStatus{ + Name: ns.Name, + Cost: s.Cost, + MaxCost: s.MaxCost, + Tokens: s.Tokens, + MaxTokens: s.MaxTokens, + ElapsedSeconds: s.Elapsed.Seconds(), + MaxTimeSeconds: s.MaxTime.Seconds(), + Unpriced: s.Unpriced, + PerAgent: perAgent, + }) + } + return &BudgetUsageEvent{ + Type: "budget_usage", + SessionID: sessionID, + AgentContext: newAgentContext(agentName), + Budgets: budgets, + } +} + +// BudgetExceededEvent is emitted once, when a run crosses a configured +// budget ceiling and is stopped. Limit names the YAML key that tripped +// ("max_cost", "max_tokens", "max_time") so an operator knows which one +// to raise. +type BudgetExceededEvent struct { + AgentContext + + Type string `json:"type"` + SessionID string `json:"session_id"` + // Budget names which budget tripped: "run" for the top-level + // `budget:`, otherwise the key under `budgets:`. + Budget string `json:"budget"` + Limit string `json:"limit"` + Used string `json:"used"` + Max string `json:"max"` + // ConfigPath is the YAML path of the limit that tripped, e.g. + // "budgets.tight.max_cost", so an operator knows exactly what to edit. + ConfigPath string `json:"config_path"` + Message string `json:"message"` +} + +func (e *BudgetExceededEvent) GetSessionID() string { return e.SessionID } + +func BudgetExceeded(sessionID, agentName string, br budgetBreach) Event { + return &BudgetExceededEvent{ + Type: "budget_exceeded", + SessionID: sessionID, + AgentContext: newAgentContext(agentName), + Budget: br.Budget, + Limit: string(br.Limit), + Used: br.Used, + Max: br.Max, + ConfigPath: br.configPath(), + Message: br.Message(), + } +} + // MCPInitStartedEvent is for MCP initialization lifecycle events type MCPInitStartedEvent struct { AgentContext diff --git a/pkg/runtime/hooks.go b/pkg/runtime/hooks.go index 990f9da1f9..8e3c776ead 100644 --- a/pkg/runtime/hooks.go +++ b/pkg/runtime/hooks.go @@ -182,6 +182,11 @@ const ( // runtime routed the conversation to the agent's configured // force_handoff target, prompting a new iteration. turnEndReasonForceHandoff = "force_handoff" + // turnEndReasonBudgetExceeded — the run crossed a configured + // budget ceiling (max_cost / max_tokens / max_time) and was + // stopped. Distinct from "normal" so a budget-killed run is + // distinguishable from a completed one in downstream analytics. + turnEndReasonBudgetExceeded = "budget_exceeded" ) // executeTurnEndHooks fires turn_end once per turn — symmetric to @@ -337,6 +342,17 @@ func (r *LocalRuntime) notifyMaxIterations(ctx context.Context, a *agent.Agent, r.notify(ctx, a, hooks.EventOnMaxIterations, sessionID, "warning", message) } +// notifyBudgetExceeded fires notification(level=warning) when a run is +// stopped by its budget. It deliberately reuses the generic notification +// event rather than introducing a dedicated on_budget_exceeded: the +// budget_exceeded runtime event already carries the structured payload +// (limit, used, max) that an audit consumer wants, so a second hook +// surface would duplicate it for no gain. A dedicated event is easy to +// add later if hook authors ask to match on it specifically. +func (r *LocalRuntime) notifyBudgetExceeded(ctx context.Context, a *agent.Agent, sessionID, message string) { + r.notify(ctx, a, hooks.EventNotification, sessionID, "warning", message) +} + // notify is the shared dispatch path for the (level, message)-shaped // hook events: notification, on_error, on_max_iterations. They all // take the same Input fields and are observational (no Result is diff --git a/pkg/runtime/loop.go b/pkg/runtime/loop.go index 509d2dd0af..1760fc256f 100644 --- a/pkg/runtime/loop.go +++ b/pkg/runtime/loop.go @@ -230,6 +230,10 @@ func (r *LocalRuntime) RunStream(ctx context.Context, sess *session.Session) <-c rootStream := !sess.IsSubSession() if rootStream { r.activeRootStreams.Add(1) + // Install a fresh budget for this run. Sub-sessions deliberately + // skip this: they share the root's tracker so a fan-out spends + // against one wallet rather than one allowance per child. + r.ensureBudget() } // Register before the run goroutine starts so the session is listed in @@ -349,6 +353,14 @@ func (r *LocalRuntime) runStreamLoop(ctx context.Context, sess *session.Session, // Emit team information sink.Emit(TeamInfo(r.agentDetailsFromTeam(ctx), a.Name())) + // Surface the run budget from the first frame, before any model call, + // so the configured ceilings are visible immediately rather than only + // after the first priced turn. A run that fails on its very first call + // (e.g. a bad API key) still shows its budget was active. + if b := r.currentBudget(); b != nil { + sink.Emit(BudgetUsage(sess.ID, a.Name(), b.snapshot())) + } + r.emitAgentWarnings(a, sink) r.configureToolsetHandlers(a, sink) @@ -479,6 +491,14 @@ func (r *LocalRuntime) runStreamLoop(ctx context.Context, sess *session.Session, } ls.maxIterations = newMax + // Check the run budget. Placed next to the iteration cap because + // both are "should this run keep going" questions answered at the + // turn boundary, before a model call is paid for. + if r.enforceBudget(ctx, sess, a, sink) == iterationStop { + streamReason = turnEndReasonBudgetExceeded + return + } + ls.iteration++ // Exit immediately if the stream context has been cancelled (e.g., Ctrl+C) @@ -677,6 +697,9 @@ func (r *LocalRuntime) runTurn( ls *loopState, events EventSink, ) turnControl { + // turnStart bounds this turn's active time, attributed to the agent + // below so the run budget can show how long each agent spent working. + turnStart := r.now() streamAttrs := []attribute.KeyValue{ attribute.String(genai.AttrConversationID, sess.ID), attribute.String(genai.AttrAgentNameRuntime, a.Name()), @@ -789,6 +812,13 @@ func (r *LocalRuntime) runTurn( // table); see computeMessageCost. msgCost := computeMessageCost(res.Usage, m) + // Fold this turn into the run budget from the same computed value, so + // what the ceiling counts can never disagree with what the session + // bills or what after_llm_call reports. Sub-sessions reach the root + // run's tracker here, which is how delegated spend lands on the + // parent's budget. The turn's duration is attributed to this agent. + r.recordBudget(sess, a, res.Usage, msgCost, r.now().Sub(turnStart), events) + // after_llm_call hooks fire on success only; failed calls // fire on_error above. The assistant text content is passed // via stop_response, matching the stop event's payload, so diff --git a/pkg/runtime/runtime.go b/pkg/runtime/runtime.go index dcfe935212..b57c3116c8 100644 --- a/pkg/runtime/runtime.go +++ b/pkg/runtime/runtime.go @@ -309,6 +309,26 @@ type LocalRuntime struct { // exactly ONE message after the model stops and stop-hooks have run. followUpQueue MessageQueue + // budgetCfg is the run-wide budget (nil when unset), budgetsCfg the + // named budget definitions, and agentBudgets the names each agent + // declared. All three are immutable templates; budget below is the + // live accumulator built from them. + budgetCfg *latest.BudgetConfig + budgetsCfg map[string]latest.BudgetConfig + agentBudgets map[string][]string + + // budgetMu guards budget and budgetStarted. The first root stream + // installs the trackers and every later one reuses them, so a budget + // spans the session rather than resetting on each message; sub-sessions + // read the same set, so delegated work spends against the root run's + // wallets rather than getting a fresh allowance each. + budgetMu sync.Mutex + budget *budgetSet + // budgetStarted distinguishes "not built yet" from "built, and there + // was nothing to budget" — without it a nil budget would be rebuilt on + // every message, and the reset bug would come back for unbudgeted runs. + budgetStarted bool + // activeRootStreams tracks top-level RunStream loops. Recalls use this to // preserve mid-turn steering while still waking the embedder once the parent // stream has gone idle. @@ -535,6 +555,31 @@ func WithMaxOverflowCompactions(n int) Opt { } } +// WithBudget sets the run-wide budget: the cost, token and wall-clock +// ceilings that stop a run once crossed, regardless of which agent spends. +// A nil or all-zero config leaves runs unbudgeted, which is the default. +func WithBudget(cfg *latest.BudgetConfig) Opt { + return func(r *LocalRuntime) { + if cfg.IsZero() { + return + } + r.budgetCfg = cfg + } +} + +// WithNamedBudgets registers the manifest's named budget definitions and +// the budget names each agent declared. A named budget referenced by +// several agents is one shared pot, not a copy per agent. +func WithNamedBudgets(defs map[string]latest.BudgetConfig, agentBudgets map[string][]string) Opt { + return func(r *LocalRuntime) { + if len(defs) == 0 || len(agentBudgets) == 0 { + return + } + r.budgetsCfg = defs + r.agentBudgets = agentBudgets + } +} + // WithToolListTimeout overrides how long EmitStartupInfo waits for a single // toolset to enumerate its tools before skipping it. Defaults to // defaultToolListTimeout. A non-positive value is ignored so the default diff --git a/pkg/teamloader/teamloader.go b/pkg/teamloader/teamloader.go index df8dfd78a8..9d5be6439a 100644 --- a/pkg/teamloader/teamloader.go +++ b/pkg/teamloader/teamloader.go @@ -121,6 +121,16 @@ type LoadResult struct { ProviderRegistry *provider.Registry // AgentDefaultModels maps agent names to their configured default model references AgentDefaultModels map[string]string + // Budget is the manifest's run-wide budget, or nil when the manifest + // sets no run-wide ceiling. It is per-run rather than per-agent, so it + // lives on the load result next to the team rather than on any + // individual agent. + Budget *latest.BudgetConfig + // Budgets are the manifest's named budget definitions, and + // AgentBudgets maps each agent to the budget names it declared. A name + // referenced by several agents is one shared pot. + Budgets map[string]latest.BudgetConfig + AgentBudgets map[string][]string } // Load loads an agent team from the given source @@ -456,8 +466,12 @@ func LoadWithConfig(ctx context.Context, agentSource config.Source, runConfig *c // Retain the resolved per-agent configs so inspection surfaces (the agent // inspector modal) can show declared toolset allow-lists, limits and flags. agentConfigs := make(map[string]latest.AgentConfig, len(cfg.Agents)) + agentBudgets := make(map[string][]string, len(cfg.Agents)) for i := range cfg.Agents { agentConfigs[cfg.Agents[i].Name] = cfg.Agents[i] + if len(cfg.Agents[i].Budgets) > 0 { + agentBudgets[cfg.Agents[i].Name] = cfg.Agents[i].Budgets + } } return &LoadResult{ @@ -470,6 +484,9 @@ func LoadWithConfig(ctx context.Context, agentSource config.Source, runConfig *c Providers: cfg.Providers, ProviderRegistry: loadOpts.providerRegistry, AgentDefaultModels: agentDefaultModels, + Budget: cfg.Budget, + Budgets: cfg.Budgets, + AgentBudgets: agentBudgets, }, nil } diff --git a/pkg/tui/components/sidebar/budget_line_test.go b/pkg/tui/components/sidebar/budget_line_test.go new file mode 100644 index 0000000000..2dee3a313e --- /dev/null +++ b/pkg/tui/components/sidebar/budget_line_test.go @@ -0,0 +1,183 @@ +package sidebar + +import ( + "strings" + "testing" + + "github.com/charmbracelet/x/ansi" + "github.com/stretchr/testify/assert" + + "github.com/docker/docker-agent/pkg/runtime" +) + +func (s *testSidebar) feedBudget(ev *runtime.BudgetUsageEvent) { + updated, _ := s.Update(ev) + s.model = updated.(*model) +} + +func TestBudgetLine_AbsentWhenUnbudgeted(t *testing.T) { + t.Parallel() + + m := newTestSidebar(t) + m.startStream("s1", "root") + m.recordUsageTokens("s1", "root", 5000, 3000) + + assert.NotContains(t, ansi.Strip(m.tokenUsage(60)), "/$") +} + +func TestBudgetLine_ShowsCeilingsAtZero(t *testing.T) { + t.Parallel() + + m := newTestSidebar(t) + m.startStream("s1", "root") + m.feedBudget(&runtime.BudgetUsageEvent{ + SessionID: "s1", + AgentContext: runtime.AgentContext{AgentName: "root"}, + Budgets: []runtime.BudgetStatus{{ + Name: "run", MaxCost: 0.50, MaxTokens: 100000, MaxTimeSeconds: 120, + }}, + }) + + out := ansi.Strip(m.tokenUsage(60)) + assert.Contains(t, out, "run") + assert.Contains(t, out, "$0.00/$0.50") + assert.Contains(t, out, "/100.0K") + assert.Contains(t, out, "/2m", "a round ceiling reads 2m, not 2m0s") +} + +func TestFormatBudgetDuration(t *testing.T) { + t.Parallel() + + for _, tc := range []struct { + seconds float64 + want string + }{ + {0, "0s"}, + {45, "45s"}, + {60, "1m"}, + {134, "2m14s"}, + {600, "10m"}, + {3600, "1h"}, + {5400, "1h30m"}, + } { + assert.Equal(t, tc.want, formatBudgetDuration(tc.seconds)) + } +} + +func TestBudgetLine_ShowsEachBudgetByName(t *testing.T) { + t.Parallel() + + m := newTestSidebar(t) + m.startStream("s1", "root") + m.feedBudget(&runtime.BudgetUsageEvent{ + SessionID: "s1", + AgentContext: runtime.AgentContext{AgentName: "developer"}, + Budgets: []runtime.BudgetStatus{ + {Name: "run", Cost: 0.12, MaxCost: 0.50}, + {Name: "shell-work", Cost: 0.09, MaxCost: 0.10}, + }, + }) + + out := ansi.Strip(m.tokenUsage(80)) + assert.Contains(t, out, "run") + assert.Contains(t, out, "$0.12/$0.50") + assert.Contains(t, out, "shell-work") + assert.Contains(t, out, "$0.09/$0.10") + assert.Less(t, strings.Index(out, "run"), strings.Index(out, "shell-work"), + "run-wide budget leads, then named budgets") +} + +func TestBudgetLine_OmitsPerAgentBreakdown(t *testing.T) { + t.Parallel() + + m := newTestSidebar(t) + m.startStream("s1", "root") + m.feedBudget(&runtime.BudgetUsageEvent{ + SessionID: "s1", + AgentContext: runtime.AgentContext{AgentName: "developer"}, + Budgets: []runtime.BudgetStatus{{ + Name: "run", Cost: 0.12, MaxCost: 0.50, + PerAgent: []runtime.AgentBudgetUsage{ + {AgentName: "developer", Cost: 0.09, Tokens: 8000, ActiveSeconds: 72}, + {AgentName: "solo-agent", Cost: 0.03, Tokens: 4300, ActiveSeconds: 62}, + }, + }}, + }) + + out := ansi.Strip(m.tokenUsage(80)) + assert.Contains(t, out, "$0.12/$0.50", "the budget total still shows") + assert.NotContains(t, out, "solo-agent", + "per-agent rows are the Agents section's job, not the budget line's") +} + +func TestBudgetLine_SubCentUsesPreciseFormat(t *testing.T) { + t.Parallel() + + m := newTestSidebar(t) + m.startStream("s1", "root") + m.feedBudget(&runtime.BudgetUsageEvent{ + SessionID: "s1", + AgentContext: runtime.AgentContext{AgentName: "root"}, + Budgets: []runtime.BudgetStatus{{Name: "run", Cost: 0.0012, MaxCost: 0.05}}, + }) + + out := ansi.Strip(m.tokenUsage(60)) + assert.Contains(t, out, "$0.0012", "sub-cent spend must not read as $0.00") +} + +func TestBudgetLine_TokensOnlyBudgetOmitsCost(t *testing.T) { + t.Parallel() + + m := newTestSidebar(t) + m.startStream("s1", "root") + m.feedBudget(&runtime.BudgetUsageEvent{ + SessionID: "s1", + AgentContext: runtime.AgentContext{AgentName: "root"}, + Budgets: []runtime.BudgetStatus{{Name: "roomy", Tokens: 500, MaxTokens: 1000}}, + }) + + var row string + for line := range strings.SplitSeq(ansi.Strip(m.tokenUsage(80)), "\n") { + if strings.Contains(line, "roomy") { + row = line + break + } + } + assert.NotEmpty(t, row, "the tokens-only budget must be rendered") + assert.Contains(t, row, "500/1.0K") + assert.NotContains(t, row, "$", "a tokens-only budget must not render costs") +} + +func TestBudgetLine_MarksUnpricedSpend(t *testing.T) { + t.Parallel() + + m := newTestSidebar(t) + m.startStream("s1", "root") + m.feedBudget(&runtime.BudgetUsageEvent{ + SessionID: "s1", + AgentContext: runtime.AgentContext{AgentName: "root"}, + Budgets: []runtime.BudgetStatus{{Name: "run", Cost: 0, MaxCost: 0.50, Unpriced: true}}, + }) + + assert.Contains(t, ansi.Strip(m.tokenUsage(80)), "unpriced spend") +} + +func TestBudgetLine_UpdatesLive(t *testing.T) { + t.Parallel() + + m := newTestSidebar(t) + m.startStream("s1", "root") + m.feedBudget(&runtime.BudgetUsageEvent{ + SessionID: "s1", AgentContext: runtime.AgentContext{AgentName: "root"}, + Budgets: []runtime.BudgetStatus{{Name: "run", Cost: 0.01, MaxCost: 0.50}}, + }) + assert.Contains(t, ansi.Strip(m.tokenUsage(60)), "$0.01/$0.50") + + m.feedBudget(&runtime.BudgetUsageEvent{ + SessionID: "s1", AgentContext: runtime.AgentContext{AgentName: "root"}, + Budgets: []runtime.BudgetStatus{{Name: "run", Cost: 0.44, MaxCost: 0.50}}, + }) + out := ansi.Strip(m.tokenUsage(60)) + assert.Contains(t, out, "$0.44/$0.50") + assert.NotContains(t, out, "$0.01/$0.50", "stale reading must be replaced") +} diff --git a/pkg/tui/components/sidebar/sidebar.go b/pkg/tui/components/sidebar/sidebar.go index 98622bceb1..6ed853c71a 100644 --- a/pkg/tui/components/sidebar/sidebar.go +++ b/pkg/tui/components/sidebar/sidebar.go @@ -271,12 +271,16 @@ func transferTimer(d time.Duration, gen int64, kind transferTimerKind) TransferT // model implements Model type model struct { - width int - height int - xPos int // absolute x position on screen - yPos int // absolute y position on screen - layoutCfg LayoutConfig // layout configuration for spacing - sessionUsage map[string]*runtime.Usage // sessionID -> latest usage snapshot + width int + height int + xPos int // absolute x position on screen + yPos int // absolute y position on screen + layoutCfg LayoutConfig // layout configuration for spacing + sessionUsage map[string]*runtime.Usage // sessionID -> latest usage snapshot + // budgetUsage is the newest run-budget snapshot, or nil on an + // unbudgeted run. It is a single value rather than a per-session map + // because one budget covers the whole run, sub-sessions included. + budgetUsage *runtime.BudgetUsageEvent todoComp *todotool.SidebarComponent ragIndexing map[string]*ragIndexingState // strategy name -> indexing state spinner spinner.Spinner @@ -1081,6 +1085,13 @@ func (m *model) Update(msg tea.Msg) (layout.Model, tea.Cmd) { case *runtime.TokenUsageEvent: m.SetTokenUsage(msg) return m, nil + case *runtime.BudgetUsageEvent: + // The budget is per-run and shared by every sub-session, so the + // latest event always describes the whole run regardless of which + // session emitted it. Keep one snapshot rather than a per-session + // map — summing across sessions would double-count the wallet. + m.budgetUsage = msg + return m, nil case *runtime.SessionCompactionEvent: // The "compacting…" gauge describes the displayed session only: a // targeted compaction of a session that is not currently shown @@ -1738,7 +1749,11 @@ func (m *model) computeUsageStats() usageStats { } func (m *model) tokenUsage(contentWidth int) string { - return m.renderTab("Token Usage", m.tokenUsageLine(), contentWidth) + body := m.tokenUsageLine() + if b := m.budgetLine(contentWidth); b != "" { + body += "\n" + b + } + return m.renderTab("Token Usage", body, contentWidth) } // tokenUsageLine renders the usage line shared by the vertical Token Usage @@ -1772,6 +1787,110 @@ func (m *model) tokenUsageSummary() string { return m.tokenUsageLine() } +// budgetLine renders each active budget by name against the ceilings it +// declares, e.g. +// +// run $0.12/$0.50 · 12.3K/100.0K · 2m14s/10m +// shell-work $0.09/$0.10 · 4.3K/20.0K +// +// Only the ceilings the operator configured appear, and an unbudgeted run +// renders nothing. Each reading is colored by the shared gauge bands, so a +// budget nearing its ceiling reads the same way a context gauge nearing +// compaction does. +// +// Per-agent attribution is deliberately not repeated here: the Agents +// section already reports each agent's cost in AgentInfoDetailed mode, so a +// second per-agent breakdown under every budget would be duplicate reading +// in the sidebar's tightest column. The structured per-agent split is still +// carried on the budget_usage event for programmatic consumers. +func (m *model) budgetLine(contentWidth int) string { + b := m.budgetUsage + if b == nil || len(b.Budgets) == 0 { + return "" + } + + nameWidth := 0 + for _, s := range b.Budgets { + nameWidth = max(nameWidth, lipgloss.Width(s.Name)) + } + if limit := contentWidth / 3; limit > 0 && nameWidth > limit { + nameWidth = limit + } + + var lines []string + for _, s := range b.Budgets { + if line := m.oneBudgetLine(s, nameWidth); line != "" { + lines = append(lines, line) + } + } + return strings.Join(lines, "\n") +} + +// oneBudgetLine renders a single named budget's totals against its ceilings. +// Returns "" when the budget declares no ceilings at all. +func (m *model) oneBudgetLine(s runtime.BudgetStatus, nameWidth int) string { + var parts []string + if s.MaxCost > 0 { + parts = append(parts, budgetPartStyle(s.Cost, s.MaxCost).Render( + toolcommon.FormatCostPrecise(s.Cost)+"/"+toolcommon.FormatCostPrecise(s.MaxCost))) + } + if s.MaxTokens > 0 { + parts = append(parts, budgetPartStyle(float64(s.Tokens), float64(s.MaxTokens)).Render( + toolcommon.FormatTokenCount(s.Tokens)+"/"+toolcommon.FormatTokenCount(s.MaxTokens))) + } + if s.MaxTimeSeconds > 0 { + parts = append(parts, budgetPartStyle(s.ElapsedSeconds, s.MaxTimeSeconds).Render( + formatBudgetDuration(s.ElapsedSeconds)+"/"+formatBudgetDuration(s.MaxTimeSeconds))) + } + if len(parts) == 0 { + return "" + } + + name := toolcommon.TruncateText(s.Name, nameWidth) + line := styles.MutedStyle.Render(padRight(name, nameWidth)) + " " + + strings.Join(parts, metricSeparator()) + if s.Unpriced { + // A cost ceiling that cannot see some of the spend is a silent + // failure otherwise: the reading would look reassuringly low + // precisely because it is incomplete. + line += " " + styles.WarningStyle.Render("(unpriced spend)") + } + return line +} + +func budgetPartStyle(used, limit float64) lipgloss.Style { + if limit <= 0 { + return styles.TabAccentStyle + } + level := styles.ContextGaugeLevelFor(used/limit, 1) + return contextGaugeStyle(level, styles.TabAccentStyle) +} + +// formatBudgetDuration renders a whole-second reading compactly, dropping +// zero components so a round ceiling reads "10m" and "1h" rather than +// Duration.String's "10m0s" and "1h0m0s" — the budget column is the +// narrowest place in the sidebar to spend characters on trailing zeros. +func formatBudgetDuration(seconds float64) string { + d := time.Duration(seconds) * time.Second + if d <= 0 { + return "0s" + } + switch { + case d < time.Minute: + return fmt.Sprintf("%ds", int(d.Seconds())) + case d < time.Hour: + if s := int(d.Seconds()) % 60; s != 0 { + return fmt.Sprintf("%dm%ds", int(d.Minutes()), s) + } + return fmt.Sprintf("%dm", int(d.Minutes())) + default: + if mn := int(d.Minutes()) % 60; mn != 0 { + return fmt.Sprintf("%dh%dm", int(d.Hours()), mn) + } + return fmt.Sprintf("%dh", int(d.Hours())) + } +} + func (m *model) sessionInfo(contentWidth int) string { star := m.starIndicator() diff --git a/pkg/tui/page/chat/runtime_events.go b/pkg/tui/page/chat/runtime_events.go index 4502807e52..4edb9942a0 100644 --- a/pkg/tui/page/chat/runtime_events.go +++ b/pkg/tui/page/chat/runtime_events.go @@ -42,6 +42,10 @@ import ( // // Sidebar Updates (forwarded): // - TokenUsageEvent, AgentInfoEvent, TeamInfoEvent, etc. +// - BudgetUsageEvent → Live run-budget reading +// +// Warnings: +// - BudgetExceededEvent → Run stopped by its budget // // Dialogs: // - MaxIterationsReachedEvent → Show max iterations dialog @@ -172,6 +176,17 @@ func (p *chatPage) handleRuntimeEvent(msg tea.Msg) (bool, tea.Cmd) { *runtime.RAGIndexingCompletedEvent: return true, p.forwardToSidebar(msg) + // ===== Budget Events ===== + // budget_usage feeds the sidebar's live spend reading. budget_exceeded + // is rendered as a warning rather than a dialog: the run is already + // stopping and there is nothing to ask the user, unlike the iteration + // cap, which offers to continue. + case *runtime.BudgetUsageEvent: + return true, p.forwardToSidebar(msg) + + case *runtime.BudgetExceededEvent: + return true, p.handleBudgetExceeded(msg) + // ===== Dialog Events ===== case *runtime.MaxIterationsReachedEvent: return true, p.handleMaxIterationsReached(msg) @@ -427,6 +442,21 @@ func (p *chatPage) handleToolCallResponse(msg *runtime.ToolCallResponseEvent) te return tea.Batch(toolCmd, p.messages.ScrollToBottom(), spinnerCmd, sidebarCmd) } +// handleBudgetExceeded reports a run stopped by its budget. It stops the +// spinner and raises a warning naming the limit that tripped, alongside +// the assistant stop message the runtime already appended. +// +// Deliberately not a dialog: unlike the iteration cap there is nothing to +// ask. A budget is a ceiling the operator set on purpose, so raising it +// means editing the config rather than answering a prompt. +func (p *chatPage) handleBudgetExceeded(msg *runtime.BudgetExceededEvent) tea.Cmd { + spinnerCmd := p.setWorking(false) + warnCmd := notification.WarningCmd(fmt.Sprintf( + "Run stopped by %s — used %s of %s.", msg.ConfigPath, msg.Used, msg.Max, + )) + return tea.Batch(spinnerCmd, warnCmd) +} + func (p *chatPage) handleMaxIterationsReached(msg *runtime.MaxIterationsReachedEvent) tea.Cmd { spinnerCmd := p.setWorking(false) dialogCmd := core.CmdHandler(dialog.OpenDialogMsg{