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	used := WeightedBudgetUsage(usage)
58	percent := int(math.Floor(float64(used) * 100 / float64(limit)))
59	return BudgetSnapshot{Used: used, Limit: limit, Percent: percent, State: budgetState(used, limit)}
60}
61
62func budgetState(used, limit int64) BudgetState {
63	if limit <= 0 {
64		return ""
65	}
66	if used >= limit {
67		return BudgetExceeded
68	}
69	return BudgetNormal
70}
71
72// RolloutBudget accounts one resident child's cumulative provider usage.
73type RolloutBudget struct {
74	mu       sync.Mutex
75	limit    int64
76	snapshot BudgetSnapshot
77}
78
79func NewRolloutBudget(limit int64, usage provider.Usage) *RolloutBudget {
80	budget := &RolloutBudget{limit: limit}
81	budget.snapshot = BudgetSnapshotFor(usage, limit)
82	return budget
83}
84
85func (b *RolloutBudget) Observe(usage provider.Usage) BudgetSnapshot {
86	if b == nil {
87		return BudgetSnapshot{}
88	}
89	b.mu.Lock()
90	b.snapshot = BudgetSnapshotFor(usage, b.limit)
91	snapshot := b.snapshot
92	b.mu.Unlock()
93	return snapshot
94}
95
96func (b *RolloutBudget) Snapshot() BudgetSnapshot {
97	if b == nil {
98		return BudgetSnapshot{}
99	}
100	b.mu.Lock()
101	snapshot := b.snapshot
102	b.mu.Unlock()
103	return snapshot
104}