From 27b56866dfb2d6bdb3e95c5a8f76169a0ae49f7f Mon Sep 17 00:00:00 2001 From: Mawen Salignat-Moandal Date: Sat, 5 Sep 2026 12:27:56 +0200 Subject: [PATCH 1/2] 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") +} From 1b54bb4fa0a1af9255dfc8bba1f874d2d8c23953 Mon Sep 17 00:00:00 2001 From: Mawen Salignat-Moandal Date: Wed, 16 Sep 2026 15:09:44 +0200 Subject: [PATCH 2/2] fix(runtime): inject 80% budget warning as a cache-stable extra Keep the model-visible approaching signal out of session history so it does not bust the prompt-cache prefix, survive resume as a stale system message, or vanish after compaction. --- pkg/runtime/budget.go | 49 ++++++++++++++++++++++---- pkg/runtime/budget_test.go | 5 +++ pkg/runtime/budget_wiring_test.go | 58 +++++++++++++++++++++++++------ pkg/runtime/loop.go | 9 +++-- 4 files changed, 101 insertions(+), 20 deletions(-) diff --git a/pkg/runtime/budget.go b/pkg/runtime/budget.go index 56fb346ea4..5a2caac819 100644 --- a/pkg/runtime/budget.go +++ b/pkg/runtime/budget.go @@ -26,6 +26,14 @@ const ( // a YAML knob: a run that still crosses the ceiling must stop the // same way it does today. budgetWarnFraction = 0.8 + + // budgetApproachingPrompt is the model-visible extra injected after + // any 80% warning. Wording is free of used/max amounts so prompt-cache + // checkpoints stay reusable on later turns. The TUI/JSON Warning still + // uses WarnMessage, which includes the numbers. + budgetApproachingPrompt = "You are approaching the configured run budget. Prefer cheaper tools, avoid redundant calls, summarize, and finish soon. The run will stop if a limit is reached." + + budgetWarningSourceKey = "runtime/budget-warning" ) type budgetTracker struct { @@ -38,9 +46,12 @@ 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 records which limits have already emitted the 80% TUI/JSON + // warning so each tracker warns at most once per limit. warned map[budgetLimit]bool + // softPrompt stays set after the first approaching warning so the + // stable extra is re-injected every remaining turn until hard stop. + softPrompt bool } type agentSpend struct { @@ -259,6 +270,7 @@ func (b *budgetTracker) consumeApproaching() *budgetBreach { b.warned = make(map[budgetLimit]bool) } b.warned[limit] = true + b.softPrompt = true return br } return nil @@ -309,6 +321,15 @@ func (b *budgetTracker) breachIfApproachingLocked(limit budgetLimit) *budgetBrea return nil } +func (b *budgetTracker) hasSoftPrompt() bool { + if b == nil { + return false + } + b.mu.Lock() + defer b.mu.Unlock() + return b.softPrompt +} + func (b *budgetTracker) unpricedSpend() bool { if b == nil { return false @@ -503,11 +524,25 @@ func (r *LocalRuntime) warnBudgetIfApproaching( "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) +} + +// budgetPromptMessages returns the sticky, cache-stable extra for this +// agent's next model call, or nil. Never persisted: callers thread it +// through extraSystemMessages / instruction sources at assembly time. +func (r *LocalRuntime) budgetPromptMessages(agentName string) []chat.Message { + return r.currentBudget().promptMessagesFor(agentName) +} + +func (s *budgetSet) promptMessagesFor(agentName string) []chat.Message { + if s == nil { + return nil + } + for _, nt := range s.budgetsFor(agentName) { + if nt.Tracker.hasSoftPrompt() { + return []chat.Message{{Role: chat.MessageRoleSystem, Content: budgetApproachingPrompt}} + } + } + return nil } func (s *budgetSet) exceededFor(agentName string) *budgetBreach { diff --git a/pkg/runtime/budget_test.go b/pkg/runtime/budget_test.go index 8d9e426072..f99d9c1908 100644 --- a/pkg/runtime/budget_test.go +++ b/pkg/runtime/budget_test.go @@ -37,6 +37,7 @@ func TestNilBudgetTrackerIsInert(t *testing.T) { assert.Nil(t, b.exceeded()) assert.Nil(t, b.approaching()) assert.Nil(t, b.consumeApproaching()) + assert.False(t, b.hasSoftPrompt()) assert.Equal(t, budgetSnapshot{}, b.snapshot()) assert.False(t, b.unpricedSpend()) }) @@ -395,12 +396,16 @@ func TestBudgetApproachingWarnsOncePerLimit(t *testing.T) { b := newBudgetTracker(&latest.BudgetConfig{MaxCost: 0.50}) b.record("root", &chat.Usage{}, new(0.40), time.Second) + require.NotNil(t, b.approaching()) + assert.False(t, b.hasSoftPrompt(), "approaching() must not arm the prompt extra") first := b.consumeApproaching() require.NotNil(t, first) assert.Equal(t, budgetLimitCost, first.Limit) + assert.True(t, b.hasSoftPrompt(), "consumeApproaching must arm the sticky prompt extra") 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.True(t, b.hasSoftPrompt(), "the prompt extra stays armed until hard stop") assert.Nil(t, b.exceeded()) } diff --git a/pkg/runtime/budget_wiring_test.go b/pkg/runtime/budget_wiring_test.go index 9864c40dd9..17f81304c7 100644 --- a/pkg/runtime/budget_wiring_test.go +++ b/pkg/runtime/budget_wiring_test.go @@ -229,20 +229,22 @@ func TestEnforceBudgetWarnsOnceThenStillHardStops(t *testing.T) { 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") + assert.Empty(t, sess.GetAllMessages(), "the approaching extra must not be persisted") + assert.False(t, promptContains(sess.GetMessages(a), budgetApproachingPrompt), + "GetMessages without extras must not carry the approaching warning") + + extras := r.budgetPromptMessages(a.Name()) + require.Len(t, extras, 1) + assert.Equal(t, chat.MessageRoleSystem, extras[0].Role) + assert.Equal(t, budgetApproachingPrompt, extras[0].Content) + assert.True(t, promptContains(sess.GetMessagesWithoutInstructionContext(a, extras...), budgetApproachingPrompt), + "the next-turn prompt must carry the stable extra") 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") + assert.Len(t, sink.events, eventCount, "no extra events on the second approaching check") + assert.Equal(t, extras, r.budgetPromptMessages(a.Name()), "the extra stays sticky after the one-shot Warning") 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), @@ -259,3 +261,39 @@ func TestEnforceBudgetWarnsOnceThenStillHardStops(t *testing.T) { assert.Equal(t, "budget.max_cost", exceeded.ConfigPath) assert.Contains(t, exceeded.Message, "Execution stopped") } + +func TestBudgetWarningInstructionSourceDoesNotRewritePrefix(t *testing.T) { + sess := session.New() + a := agent.New("root", "base prompt") + sess.AddMessage(session.UserMessage("hello")) + + env := []session.InstructionSource{{ + Key: "hooks/session-start", Label: "environment context", Content: "cwd: /tmp", Available: true, + }} + require.True(t, sess.PrepareInstructionContext(env)) + first := sess.GetMessages(a) + + budgetSrc := instructionSource(budgetWarningSourceKey, "run budget", []chat.Message{ + {Role: chat.MessageRoleSystem, Content: budgetApproachingPrompt}, + }) + sources := append(append([]session.InstructionSource{}, env...), budgetSrc) + require.True(t, sess.PrepareInstructionContext(sources), "first appearance must be a chronological update") + + updated := sess.GetMessages(a) + assert.Equal(t, first[0].Content, updated[0].Content, "invariant system prompt must stay byte-stable") + assert.Equal(t, first[1].Content, updated[1].Content, "frozen instruction prefix must stay byte-stable") + require.GreaterOrEqual(t, len(updated), 4) + assert.Equal(t, chat.MessageRoleUser, updated[len(updated)-1].Role) + assert.Contains(t, updated[len(updated)-1].Content, budgetApproachingPrompt) + + assert.False(t, sess.PrepareInstructionContext(sources), "stable wording must not rotate the cache on later turns") +} + +func promptContains(messages []chat.Message, text string) bool { + for _, msg := range messages { + if strings.Contains(msg.Content, text) { + return true + } + } + return false +} diff --git a/pkg/runtime/loop.go b/pkg/runtime/loop.go index f6d917bc13..7d695ce662 100644 --- a/pkg/runtime/loop.go +++ b/pkg/runtime/loop.go @@ -767,12 +767,15 @@ func (r *LocalRuntime) runTurn( // against what the model already knows. Changes extend the conversation; // they never rewrite the frozen instruction prefix. turnStartMsgs := r.executeTurnStartHooks(ctx, sess, a, events) - // Pending tool-mode structured-output reminder rides with the transient - // system extras: threaded per call, never persisted as a user message. + // Pending tool-mode structured-output reminder and the sticky 80% + // budget extra ride with the transient system extras: threaded per + // call, never persisted as a user message. reminderMsgs := ls.structuredOutputReminderMessages() - legacyExtras := slices.Concat(ls.sessionStartLegacyMsgs, ls.userPromptMsgs, turnStartMsgs.legacyMessages(), reminderMsgs) + budgetMsgs := r.budgetPromptMessages(a.Name()) + legacyExtras := slices.Concat(ls.sessionStartLegacyMsgs, ls.userPromptMsgs, turnStartMsgs.legacyMessages(), reminderMsgs, budgetMsgs) sources := instructionSources(ls.sessionStartMsgs, ls.userPromptMsgs, turnStartMsgs, ls.sessionStartSources...) sources = append(sources, instructionSource("runtime/structured-output", "structured-output reminder", reminderMsgs)) + sources = append(sources, instructionSource(budgetWarningSourceKey, "run budget", budgetMsgs)) messages := r.messagesWithDynamicContext(ctx, sess, a, sources, legacyExtras) slog.DebugContext(ctx, "Retrieved messages for processing", "agent", a.Name(), "message_count", len(messages))