feat: warn the agent at 80% of a run budget before hard stop

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 {

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())