1package subagents
2
3import (
4 "errors"
5 "fmt"
6 "math"
7 "sync"
8
9 "github.com/bnema/zut/packages/provider"
10)
11
12const (
13 DefaultBudgetRatio = 0.75
14 cachedInputBudgetWeight = 0.25
15)
16
17// BudgetState is the progressive constraint applied to a resident child.
18type BudgetState string
19
20const (
21 BudgetNormal BudgetState = "normal"
22 BudgetNotice BudgetState = "notice"
23 BudgetFocused BudgetState = "focused"
24 BudgetVerifying BudgetState = "verifying"
25 BudgetFinalizing BudgetState = "finalizing"
26 BudgetExceeded BudgetState = "exceeded"
27)
28
29var ErrBudgetExceeded = errors.New("resident subagent rollout budget exceeded")
30
31// BudgetSnapshot is the bounded public projection of a child's cumulative
32// rollout budget. Used is weighted usage, not the active context size.
33type BudgetSnapshot struct {
34 Used int64 `json:"used_tokens"`
35 Limit int64 `json:"limit_tokens"`
36 Percent int `json:"percent"`
37 State BudgetState `json:"state"`
38}
39
40func BudgetLimit(contextWindow int, ratio float64, tokens int64) (int64, float64, error) {
41 if tokens < 0 {
42 return 0, 0, errors.New("budget_tokens must be positive")
43 }
44 if tokens > 0 {
45 return tokens, 0, nil
46 }
47 if ratio == 0 {
48 ratio = DefaultBudgetRatio
49 }
50 if math.IsNaN(ratio) || math.IsInf(ratio, 0) || ratio <= 0 || ratio > 1 {
51 return 0, 0, errors.New("budget_ratio must be greater than 0 and at most 1")
52 }
53 if contextWindow <= 0 {
54 return 0, ratio, errors.New("cannot derive a subagent budget without a parent context window")
55 }
56 limit := int64(math.Floor(float64(contextWindow) * ratio))
57 if limit <= 0 {
58 return 0, ratio, errors.New("derived subagent budget is empty")
59 }
60 return limit, ratio, nil
61}
62
63// WeightedBudgetUsage charges uncached and cache-write input fully, cache-read
64// input at a reduced weight, and output once (provider OutputTokens already
65// includes reasoning tokens where the provider reports them).
66func WeightedBudgetUsage(usage provider.Usage) int64 {
67 weighted := float64(max(usage.InputTokens, 0)) +
68 float64(max(usage.CacheWriteTokens, 0)) +
69 float64(max(usage.OutputTokens, 0)) +
70 float64(max(usage.CacheReadTokens, 0))*cachedInputBudgetWeight
71 return int64(math.Ceil(weighted))
72}
73
74// EffectiveBudgetLimit preserves budgets accepted by current versions and
75// derives the default for resident specs written before budgets were stored.
76func EffectiveBudgetLimit(limit int64, contextWindow int) int64 {
77 if limit > 0 {
78 return limit
79 }
80 derived, _, err := BudgetLimit(contextWindow, DefaultBudgetRatio, 0)
81 if err != nil {
82 return 0
83 }
84 return derived
85}
86
87func BudgetSystemPrompt(limit int64) string {
88 if limit <= 0 {
89 return ""
90 }
91 return fmt.Sprintf("Your cumulative rollout budget is %d weighted tokens. At 50%% prioritize concrete evidence; at 70%% stop broad exploration; at 85%% verify existing candidates only; at 90%% stop using tools and return the best verified final result. The host enforces the finalization reserve and terminates work at 100%%.", limit)
92}
93
94func BudgetSnapshotFor(usage provider.Usage, limit int64) BudgetSnapshot {
95 if limit <= 0 {
96 return BudgetSnapshot{}
97 }
98 used := WeightedBudgetUsage(usage)
99 percent := int(math.Floor(float64(used) * 100 / float64(limit)))
100 return BudgetSnapshot{Used: used, Limit: limit, Percent: percent, State: budgetState(used, limit)}
101}
102
103func budgetState(used, limit int64) BudgetState {
104 if limit <= 0 {
105 return ""
106 }
107 ratio := float64(used) / float64(limit)
108 switch {
109 case ratio >= 1:
110 return BudgetExceeded
111 case ratio >= 0.90:
112 return BudgetFinalizing
113 case ratio >= 0.85:
114 return BudgetVerifying
115 case ratio >= 0.70:
116 return BudgetFocused
117 case ratio >= 0.50:
118 return BudgetNotice
119 default:
120 return BudgetNormal
121 }
122}
123
124// RolloutBudget accounts one resident child's cumulative provider usage and
125// produces progressive model instructions at threshold transitions.
126type RolloutBudget struct {
127 mu sync.Mutex
128 limit int64
129 snapshot BudgetSnapshot
130 delivered BudgetState
131}
132
133func NewRolloutBudget(limit int64, usage provider.Usage) *RolloutBudget {
134 budget := &RolloutBudget{limit: limit}
135 budget.snapshot = BudgetSnapshotFor(usage, limit)
136 return budget
137}
138
139func (b *RolloutBudget) Observe(usage provider.Usage) BudgetSnapshot {
140 if b == nil {
141 return BudgetSnapshot{}
142 }
143 b.mu.Lock()
144 b.snapshot = BudgetSnapshotFor(usage, b.limit)
145 snapshot := b.snapshot
146 b.mu.Unlock()
147 return snapshot
148}
149
150func (b *RolloutBudget) Snapshot() BudgetSnapshot {
151 if b == nil {
152 return BudgetSnapshot{}
153 }
154 b.mu.Lock()
155 snapshot := b.snapshot
156 b.mu.Unlock()
157 return snapshot
158}
159
160// TurnContext returns a one-shot model-visible transition alert. Core stores
161// the resulting developer context in the transcript, including across compaction.
162func (b *RolloutBudget) TurnContext() (string, bool) {
163 if b == nil {
164 return "", false
165 }
166 b.mu.Lock()
167 defer b.mu.Unlock()
168 state := b.snapshot.State
169 if state == BudgetExceeded {
170 return "", true
171 }
172 if budgetStateRank(state) <= budgetStateRank(b.delivered) {
173 return "", false
174 }
175 b.delivered = state
176 return budgetInstruction(b.snapshot), false
177}
178
179func budgetStateRank(state BudgetState) int {
180 switch state {
181 case BudgetNotice:
182 return 1
183 case BudgetFocused:
184 return 2
185 case BudgetVerifying:
186 return 3
187 case BudgetFinalizing:
188 return 4
189 case BudgetExceeded:
190 return 5
191 default:
192 return 0
193 }
194}
195
196func budgetInstruction(snapshot BudgetSnapshot) string {
197 remaining := max(snapshot.Limit-snapshot.Used, 0)
198 prefix := fmt.Sprintf("[subagent budget: %d%% used, %d weighted tokens remaining] ", snapshot.Percent, remaining)
199 switch snapshot.State {
200 case BudgetNotice:
201 return prefix + "Keep the remaining investigation bounded and prioritize concrete evidence."
202 case BudgetFocused:
203 return prefix + "Stop broad exploration. Investigate only concrete unresolved hypotheses."
204 case BudgetVerifying:
205 return prefix + "Do not identify new candidates. Verify or reject current candidates and prepare the final result."
206 case BudgetFinalizing:
207 return prefix + "Finalization is required now. Return the best verified result without using tools or starting new investigations."
208 default:
209 return ""
210 }
211}