From 27b56866dfb2d6bdb3e95c5a8f76169a0ae49f7f Mon Sep 17 00:00:00 2001 From: Mawen Salignat-Moandal Date: Sat, 5 Sep 2026 12:27:56 +0200 Subject: [PATCH] feat: warn the agent at 80% of a run budget before hard stop Give the model one chance to cheapen or finish before the existing kill-switch fires, without changing the budget_exceeded contract. --- pkg/runtime/budget.go | 156 +++++++++++++++++++++++++++++- pkg/runtime/budget_test.go | 104 ++++++++++++++++++++ pkg/runtime/budget_wiring_test.go | 61 ++++++++++++ 3 files changed, 320 insertions(+), 1 deletion(-) diff --git a/pkg/runtime/budget.go b/pkg/runtime/budget.go index ac585acd14..56fb346ea4 100644 --- a/pkg/runtime/budget.go +++ b/pkg/runtime/budget.go @@ -20,6 +20,12 @@ const ( budgetLimitCost budgetLimit = "max_cost" budgetLimitTokens budgetLimit = "max_tokens" budgetLimitTime budgetLimit = "max_time" + + // budgetWarnFraction is the share of a ceiling at which the runtime + // warns the agent once, before the hard stop at 100%. Internal, not + // a YAML knob: a run that still crosses the ceiling must stop the + // same way it does today. + budgetWarnFraction = 0.8 ) type budgetTracker struct { @@ -32,6 +38,9 @@ type budgetTracker struct { active time.Duration unpriced bool perAgent map[string]*agentSpend + // warned records which limits have already emitted the 80% warning + // so each tracker warns at most once per limit. + warned map[budgetLimit]bool } type agentSpend struct { @@ -164,6 +173,13 @@ func (br budgetBreach) Message() string { ) } +func (br budgetBreach) WarnMessage() string { + return fmt.Sprintf( + "You are approaching the configured budget (used %s of %s %s). Prefer cheaper tools, avoid redundant calls, summarize, and finish soon. The run will stop if the limit is reached.", + br.Used, br.Max, br.configPath(), + ) +} + func (br budgetBreach) configPath() string { if br.Budget == "" || br.Budget == runBudgetName { return "budget." + string(br.Limit) @@ -177,7 +193,10 @@ func (b *budgetTracker) exceeded() *budgetBreach { } b.mu.Lock() defer b.mu.Unlock() + return b.exceededLocked() +} +func (b *budgetTracker) exceededLocked() *budgetBreach { if b.maxCost > 0 && b.cost >= b.maxCost { return &budgetBreach{ Limit: budgetLimitCost, @@ -202,6 +221,94 @@ func (b *budgetTracker) exceeded() *budgetBreach { return nil } +// approaching reports the first limit that has crossed budgetWarnFraction +// but is not yet at its ceiling. Cost, then tokens, then time — the same +// order as [exceeded]. Does not consult or mutate the warned set; use +// [consumeApproaching] when emitting a one-shot warning. +func (b *budgetTracker) approaching() *budgetBreach { + if b == nil { + return nil + } + b.mu.Lock() + defer b.mu.Unlock() + return b.approachingLocked() +} + +// consumeApproaching returns the next unwarned approaching limit and +// marks it warned. Returns nil when nothing is approaching, when the +// ceiling is already exceeded, or when every approaching limit has +// already been warned. +func (b *budgetTracker) consumeApproaching() *budgetBreach { + if b == nil { + return nil + } + b.mu.Lock() + defer b.mu.Unlock() + if b.exceededLocked() != nil { + return nil + } + for _, limit := range []budgetLimit{budgetLimitCost, budgetLimitTokens, budgetLimitTime} { + if b.warned[limit] { + continue + } + br := b.breachIfApproachingLocked(limit) + if br == nil { + continue + } + if b.warned == nil { + b.warned = make(map[budgetLimit]bool) + } + b.warned[limit] = true + return br + } + return nil +} + +func (b *budgetTracker) approachingLocked() *budgetBreach { + for _, limit := range []budgetLimit{budgetLimitCost, budgetLimitTokens, budgetLimitTime} { + if br := b.breachIfApproachingLocked(limit); br != nil { + return br + } + } + return nil +} + +func (b *budgetTracker) breachIfApproachingLocked(limit budgetLimit) *budgetBreach { + switch limit { + case budgetLimitCost: + if b.maxCost > 0 && b.cost >= b.maxCost*budgetWarnFraction && b.cost < b.maxCost { + return &budgetBreach{ + Limit: budgetLimitCost, + Used: formatUSD(b.cost), + Max: formatUSD(b.maxCost), + } + } + case budgetLimitTokens: + if b.maxTokens > 0 { + warnAt := int64(float64(b.maxTokens) * budgetWarnFraction) + if b.tokens >= warnAt && b.tokens < b.maxTokens { + return &budgetBreach{ + Limit: budgetLimitTokens, + Used: fmt.Sprintf("%d tokens", b.tokens), + Max: fmt.Sprintf("%d tokens", b.maxTokens), + } + } + } + case budgetLimitTime: + if b.maxTime > 0 { + warnAt := time.Duration(float64(b.maxTime) * budgetWarnFraction) + if b.active >= warnAt && 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 @@ -341,8 +448,10 @@ func (r *LocalRuntime) enforceBudget( a *agent.Agent, events EventSink, ) iterationDecision { - breach := r.currentBudget().exceededFor(a.Name()) + budgets := r.currentBudget() + breach := budgets.exceededFor(a.Name()) if breach == nil { + r.warnBudgetIfApproaching(ctx, sess, a, events, budgets) return iterationContinue } @@ -372,7 +481,39 @@ func (r *LocalRuntime) enforceBudget( return iterationStop } +func (r *LocalRuntime) warnBudgetIfApproaching( + ctx context.Context, + sess *session.Session, + a *agent.Agent, + events EventSink, + budgets *budgetSet, +) { + warn := budgets.consumeApproachingFor(a.Name()) + if warn == nil { + return + } + + msg := warn.WarnMessage() + slog.InfoContext(ctx, "Run budget approaching", + "agent", a.Name(), + "session_id", sess.ID, + "budget", warn.Budget, + "limit", string(warn.Limit), + "used", warn.Used, + "max", warn.Max, + ) + events.Emit(Warning(msg, a.Name())) + addAgentMessage(sess, a, &chat.Message{ + Role: chat.MessageRoleSystem, + Content: msg, + CreatedAt: r.now().Format(time.RFC3339), + }, events) +} + func (s *budgetSet) exceededFor(agentName string) *budgetBreach { + if s == nil { + return nil + } for _, nt := range s.budgetsFor(agentName) { if br := nt.Tracker.exceeded(); br != nil { br.Budget = nt.Name @@ -382,6 +523,19 @@ func (s *budgetSet) exceededFor(agentName string) *budgetBreach { return nil } +func (s *budgetSet) consumeApproachingFor(agentName string) *budgetBreach { + if s == nil { + return nil + } + for _, nt := range s.budgetsFor(agentName) { + if br := nt.Tracker.consumeApproaching(); 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 { diff --git a/pkg/runtime/budget_test.go b/pkg/runtime/budget_test.go index 643c6b4bd5..8d9e426072 100644 --- a/pkg/runtime/budget_test.go +++ b/pkg/runtime/budget_test.go @@ -35,6 +35,8 @@ func TestNilBudgetTrackerIsInert(t *testing.T) { assert.NotPanics(t, func() { b.record("root", &chat.Usage{InputTokens: 10}, new(1.0), time.Second) assert.Nil(t, b.exceeded()) + assert.Nil(t, b.approaching()) + assert.Nil(t, b.consumeApproaching()) assert.Equal(t, budgetSnapshot{}, b.snapshot()) assert.False(t, b.unpricedSpend()) }) @@ -187,6 +189,8 @@ func TestBudgetTrackerIsConcurrencySafe(t *testing.T) { for range 50 { b.record("root", &chat.Usage{InputTokens: 1, OutputTokens: 1}, new(0.01), time.Second) b.exceeded() + b.approaching() + b.consumeApproaching() b.snapshot() } }() @@ -360,6 +364,106 @@ func TestBudgetSetSnapshotPerBudget(t *testing.T) { assert.InDelta(t, 0.10, snaps[2].Snapshot.MaxCost, 1e-9) } +func TestBudgetApproachingCostAtEightyPercent(t *testing.T) { + b := newBudgetTracker(&latest.BudgetConfig{MaxCost: 0.50}) + require.NotNil(t, b) + + b.record("root", &chat.Usage{InputTokens: 100}, new(0.39), time.Second) + assert.Nil(t, b.approaching(), "$0.39 of $0.50 is under 80%") + assert.Nil(t, b.exceeded()) + assert.Nil(t, b.consumeApproaching()) + + b.record("root", &chat.Usage{InputTokens: 100}, new(0.01), time.Second) + warn := b.approaching() + require.NotNil(t, warn, "$0.40 of $0.50 must warn") + assert.Equal(t, budgetLimitCost, warn.Limit) + assert.Equal(t, "$0.40", warn.Used) + assert.Equal(t, "$0.50", warn.Max) + assert.Nil(t, b.exceeded(), "80% must not hard-stop") + assert.Contains(t, warn.WarnMessage(), "used $0.40 of $0.50 budget.max_cost") +} + +func TestBudgetApproachingDoesNotFireAtCeiling(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()) + assert.Nil(t, b.approaching(), "at the ceiling exceeded wins; approaching is a pre-stop signal") + assert.Nil(t, b.consumeApproaching(), "a hard-stopped tracker must not emit a warning") +} + +func TestBudgetApproachingWarnsOncePerLimit(t *testing.T) { + b := newBudgetTracker(&latest.BudgetConfig{MaxCost: 0.50}) + b.record("root", &chat.Usage{}, new(0.40), time.Second) + + first := b.consumeApproaching() + require.NotNil(t, first) + assert.Equal(t, budgetLimitCost, first.Limit) + + b.record("root", &chat.Usage{}, new(0.05), time.Second) + assert.Nil(t, b.consumeApproaching(), "second consume after more spend must not re-warn the same limit") + assert.Nil(t, b.exceeded()) +} + +func TestBudgetApproachingTokensWhenCostUnset(t *testing.T) { + b := newBudgetTracker(&latest.BudgetConfig{MaxTokens: 1000}) + b.record("root", &chat.Usage{InputTokens: 700, OutputTokens: 100}, nil, time.Second) + warn := b.approaching() + require.NotNil(t, warn, "800 of 1000 tokens is 80%") + assert.Equal(t, budgetLimitTokens, warn.Limit) + assert.Equal(t, "800 tokens", warn.Used) + assert.Equal(t, "1000 tokens", warn.Max) + assert.Nil(t, b.exceeded()) +} + +func TestBudgetApproachingTime(t *testing.T) { + b := newBudgetTracker(&latest.BudgetConfig{MaxTime: latest.Duration{Duration: 10 * time.Minute}}) + b.record("root", &chat.Usage{}, nil, 8*time.Minute) + warn := b.approaching() + require.NotNil(t, warn, "8m of 10m is 80%") + assert.Equal(t, budgetLimitTime, warn.Limit) + assert.Equal(t, "8m0s", warn.Used) + assert.Equal(t, "10m0s", warn.Max) +} + +func TestBudgetApproachingCostPreferredOverTokens(t *testing.T) { + b := newBudgetTracker(&latest.BudgetConfig{MaxCost: 1, MaxTokens: 100}) + b.record("root", &chat.Usage{InputTokens: 80}, new(0.80), time.Second) + warn := b.approaching() + require.NotNil(t, warn) + assert.Equal(t, budgetLimitCost, warn.Limit, "cost has the same priority as exceeded()") +} + +func TestBudgetApproachingSkipsUnpricedCost(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()) + assert.Nil(t, b.approaching(), "unpriced spend must not invent an approaching-cost warning") + assert.Nil(t, b.consumeApproaching()) +} + +func TestBudgetApproachingTokensDespiteUnpricedCost(t *testing.T) { + b := newBudgetTracker(&latest.BudgetConfig{MaxCost: 0.50, MaxTokens: 1000}) + b.record("root", &chat.Usage{InputTokens: 800}, nil, time.Second) + warn := b.approaching() + require.NotNil(t, warn, "token ceiling is honest even when cost is unpriced") + assert.Equal(t, budgetLimitTokens, warn.Limit) +} + +func TestBudgetConsumeApproachingThenNextLimit(t *testing.T) { + b := newBudgetTracker(&latest.BudgetConfig{MaxCost: 1, MaxTokens: 100}) + b.record("root", &chat.Usage{InputTokens: 80}, new(0.80), time.Second) + + costWarn := b.consumeApproaching() + require.NotNil(t, costWarn) + assert.Equal(t, budgetLimitCost, costWarn.Limit) + + tokenWarn := b.consumeApproaching() + require.NotNil(t, tokenWarn, "after cost is warned, tokens at 80% must still warn once") + assert.Equal(t, budgetLimitTokens, tokenWarn.Limit) + + assert.Nil(t, b.consumeApproaching()) +} + func TestBudgetConfigIsZero(t *testing.T) { assert.True(t, (*latest.BudgetConfig)(nil).IsZero()) assert.True(t, (&latest.BudgetConfig{}).IsZero()) diff --git a/pkg/runtime/budget_wiring_test.go b/pkg/runtime/budget_wiring_test.go index cf9262ab18..9864c40dd9 100644 --- a/pkg/runtime/budget_wiring_test.go +++ b/pkg/runtime/budget_wiring_test.go @@ -1,6 +1,7 @@ package runtime import ( + "strings" "testing" "time" @@ -27,6 +28,16 @@ func (s *collectSink) budgetUsages() []*BudgetUsageEvent { return out } +func (s *collectSink) warnings() []*WarningEvent { + var out []*WarningEvent + for _, e := range s.events { + if w, ok := e.(*WarningEvent); ok { + out = append(out, w) + } + } + return out +} + func budgetRuntime(t *testing.T, clock func() time.Time) *LocalRuntime { t.Helper() r := &LocalRuntime{now: clock} @@ -198,3 +209,53 @@ func TestEnforceBudgetEmitsCanonicalStopMessage(t *testing.T) { require.NotNil(t, added.Message) assert.Equal(t, recorded, *added.Message) } + +func TestEnforceBudgetWarnsOnceThenStillHardStops(t *testing.T) { + now := budgetEpoch + r := &LocalRuntime{now: func() time.Time { return now }} + WithBudget(&latest.BudgetConfig{MaxCost: 0.50})(r) + r.ensureBudget() + + sess := session.New() + a := agent.New("root", "test") + sink := &collectSink{} + + r.recordBudget(sess, a, &chat.Usage{InputTokens: 100, OutputTokens: 100}, new(0.40), time.Second, sink) + require.Equal(t, iterationContinue, r.enforceBudget(t.Context(), sess, a, sink), + "80% of max_cost must not stop the run") + + warns := sink.warnings() + require.Len(t, warns, 1, "enforceBudget must emit exactly one Warning at 80%") + assert.Contains(t, warns[0].Message, "used $0.40 of $0.50 budget.max_cost") + assert.Contains(t, warns[0].Message, "Prefer cheaper tools") + + prompt := sess.GetMessages(a) + var sawSystem bool + for _, msg := range prompt { + if msg.Role == chat.MessageRoleSystem && strings.Contains(msg.Content, "approaching the configured budget") { + sawSystem = true + assert.Equal(t, warns[0].Message, msg.Content, "the model-visible message must match the Warning event") + } + } + assert.True(t, sawSystem, "the approaching warning must be in the next-turn prompt") + + eventCount := len(sink.events) + require.Equal(t, iterationContinue, r.enforceBudget(t.Context(), sess, a, sink)) + assert.Len(t, sink.warnings(), 1, "a second enforceBudget must not re-warn") + assert.Equal(t, eventCount, len(sink.events), "no extra events on the second approaching check") + + r.recordBudget(sess, a, &chat.Usage{InputTokens: 10, OutputTokens: 10}, new(0.15), time.Second, sink) + require.Equal(t, iterationStop, r.enforceBudget(t.Context(), sess, a, sink), + "crossing the ceiling after a warning must still hard-stop") + + var exceeded *BudgetExceededEvent + for _, e := range sink.events { + if ev, ok := e.(*BudgetExceededEvent); ok { + exceeded = ev + } + } + require.NotNil(t, exceeded, "hard-stop contract is unchanged") + assert.Equal(t, "max_cost", exceeded.Limit) + assert.Equal(t, "budget.max_cost", exceeded.ConfigPath) + assert.Contains(t, exceeded.Message, "Execution stopped") +}