budget.go RawHistoryFolder

1package subagents
2
3import (
4	"errors"
5	"math"
6	"sync"
7
8	"github.com/bnema/zut/packages/provider"
9)
10
11const cachedInputBudgetWeight = 0.25
12
13// BudgetState is the host-visible state of a resident child's cumulative
14// rollout budget.
15type BudgetState string
16
17const (
18	BudgetNormal   BudgetState = "normal"
19	BudgetExceeded BudgetState = "exceeded"
20)
21
22var ErrBudgetExceeded = errors.New("resident subagent rollout budget exceeded")
23
24// BudgetSnapshot is the bounded public projection of a child's cumulative
25// rollout budget. Used is weighted usage, not the active context size.
26type BudgetSnapshot struct {
27	Used    int64       `json:"used_tokens"`
28	Limit   int64       `json:"limit_tokens"`
29	Percent int         `json:"percent"`
30	State   BudgetState `json:"state"`
31}
32
33// ContextBudgetLimit derives the deterministic rollout limit for a child from
34// its resolved model metadata.
35func ContextBudgetLimit(contextWindow int) (int64, error) {
36	if contextWindow <= 0 {
37		return 0, errors.New("cannot derive a subagent budget without a model context window")
38	}
39	return int64(contextWindow), nil
40}
41
42// WeightedBudgetUsage charges uncached and cache-write input fully, cache-read
43// input at a reduced weight, and output once (provider OutputTokens already
44// includes reasoning tokens where the provider reports them).
45func WeightedBudgetUsage(usage provider.Usage) int64 {
46	weighted := float64(max(usage.InputTokens, 0)) +
47		float64(max(usage.CacheWriteTokens, 0)) +
48		float64(max(usage.OutputTokens, 0)) +
49		float64(max(usage.CacheReadTokens, 0))*cachedInputBudgetWeight
50	return int64(math.Ceil(weighted))
51}
52
53func BudgetSnapshotFor(usage provider.Usage, limit int64) BudgetSnapshot {
54	if limit <= 0 {
55		return BudgetSnapshot{}
56	}
57	return budgetSnapshotSince(usage, limit, 0)
58}
59
60func budgetSnapshotSince(usage provider.Usage, limit, baseline int64) BudgetSnapshot {
61	if limit <= 0 {
62		return BudgetSnapshot{}
63	}
64	used := max(WeightedBudgetUsage(usage)-baseline, 0)
65	percent := int(math.Floor(float64(used) * 100 / float64(limit)))
66	return BudgetSnapshot{Used: used, Limit: limit, Percent: percent, State: budgetState(used, limit)}
67}
68
69func budgetState(used, limit int64) BudgetState {
70	if limit <= 0 {
71		return ""
72	}
73	if used >= limit {
74		return BudgetExceeded
75	}
76	return BudgetNormal
77}
78
79// RolloutBudget accounts one resident child's cumulative provider usage.
80type RolloutBudget struct {
81	mu       sync.Mutex
82	limit    int64
83	baseline int64
84	usage    provider.Usage
85	snapshot BudgetSnapshot
86}
87
88func NewRolloutBudget(limit int64, usage provider.Usage) *RolloutBudget {
89	budget := &RolloutBudget{limit: limit, usage: usage}
90	budget.snapshot = BudgetSnapshotFor(usage, limit)
91	return budget
92}
93
94func (b *RolloutBudget) Observe(usage provider.Usage) BudgetSnapshot {
95	if b == nil {
96		return BudgetSnapshot{}
97	}
98	b.mu.Lock()
99	b.usage = usage
100	b.snapshot = budgetSnapshotSince(usage, b.limit, b.baseline)
101	snapshot := b.snapshot
102	b.mu.Unlock()
103	return snapshot
104}
105
106// SetBaseline applies an explicitly accepted recovery allowance without
107// discarding cumulative cost or conversation history.
108func (b *RolloutBudget) SetBaseline(baseline int64) {
109	if b == nil {
110		return
111	}
112	b.mu.Lock()
113	b.baseline = baseline
114	b.snapshot = budgetSnapshotSince(b.usage, b.limit, baseline)
115	b.mu.Unlock()
116}
117
118func (b *RolloutBudget) Snapshot() BudgetSnapshot {
119	if b == nil {
120		return BudgetSnapshot{}
121	}
122	b.mu.Lock()
123	snapshot := b.snapshot
124	b.mu.Unlock()
125	return snapshot
126}