| 1 | // The usage meter (#857, docs/ai-provisioning.md phase 1). The sidecar is the |
| 2 | // only place model spend is visible, so it counts here and reports upward: |
| 3 | // token counters per (day, model) plus activity counters, flushed to the |
| 4 | // *instance* (POST /api/bot-usage, bot-token-gated like charter-next-due) |
| 5 | // rather than straight to the cloud. Via-instance keeps the sidecar's egress |
| 6 | // surface at its two destinations (the instance + the model API — docs/ |
| 7 | // bots.md) and keeps the meter token in exactly one env file, the instance's; |
| 8 | // the instance relays to the cloud ledger and relays the budget state back. |
| 9 | // |
| 10 | // Counters are quantities only — token counts, request counts, run counts — |
| 11 | // never content, per the meter-only-what-ciphertext-exposes rule (#668). They |
| 12 | // live in memory: a sidecar restart loses at most one flush interval, which |
| 13 | // is acceptable for metering (billing-grade reconciliation is cloud-side). |
| 14 | package bot |
| 15 | |
| 16 | import ( |
| 17 | "context" |
| 18 | "crypto/rand" |
| 19 | "encoding/hex" |
| 20 | "sync" |
| 21 | "time" |
| 22 | ) |
| 23 | |
| 24 | // UsageRow is one meter cell: a (day, model, dimension) quantity. Model is |
| 25 | // empty for activity dimensions (charter_runs). Day is UTC — the ledger's |
| 26 | // grain, deliberately coarse. |
| 27 | type UsageRow struct { |
| 28 | Day string `json:"day"` |
| 29 | Model string `json:"model,omitempty"` |
| 30 | Dimension string `json:"dimension"` |
| 31 | Quantity int64 `json:"quantity"` |
| 32 | } |
| 33 | |
| 34 | type meterKey struct{ day, model, dim string } |
| 35 | |
| 36 | type usageMeter struct { |
| 37 | mu sync.Mutex |
| 38 | rows map[meterKey]int64 |
| 39 | } |
| 40 | |
| 41 | func newUsageMeter() *usageMeter { return &usageMeter{rows: map[meterKey]int64{}} } |
| 42 | |
| 43 | func meterDay() string { return time.Now().UTC().Format("2006-01-02") } |
| 44 | |
| 45 | // addLLM records one Messages call's usage block, all four token counters |
| 46 | // kept distinct (cache reads/writes price differently from fresh input), plus |
| 47 | // the request count. |
| 48 | func (m *usageMeter) addLLM(model string, u anthUsage) { |
| 49 | m.mu.Lock() |
| 50 | defer m.mu.Unlock() |
| 51 | day := meterDay() |
| 52 | for dim, n := range map[string]int64{ |
| 53 | "llm_requests": 1, |
| 54 | "llm_input_tokens": u.InputTokens, |
| 55 | "llm_output_tokens": u.OutputTokens, |
| 56 | "llm_cache_write_tokens": u.CacheCreationInputTokens, |
| 57 | "llm_cache_read_tokens": u.CacheReadInputTokens, |
| 58 | } { |
| 59 | if n > 0 { |
| 60 | m.rows[meterKey{day, model, dim}] += n |
| 61 | } |
| 62 | } |
| 63 | } |
| 64 | |
| 65 | // inc records a model-less activity dimension (e.g. charter_runs). |
| 66 | func (m *usageMeter) inc(dim string, n int64) { |
| 67 | if n <= 0 { |
| 68 | return |
| 69 | } |
| 70 | m.mu.Lock() |
| 71 | defer m.mu.Unlock() |
| 72 | m.rows[meterKey{meterDay(), "", dim}] += n |
| 73 | } |
| 74 | |
| 75 | // drain snapshots and clears the counters; the caller restores the snapshot |
| 76 | // if the flush fails, so a transient error loses nothing. |
| 77 | func (m *usageMeter) drain() []UsageRow { |
| 78 | m.mu.Lock() |
| 79 | defer m.mu.Unlock() |
| 80 | if len(m.rows) == 0 { |
| 81 | return nil |
| 82 | } |
| 83 | out := make([]UsageRow, 0, len(m.rows)) |
| 84 | for k, q := range m.rows { |
| 85 | out = append(out, UsageRow{Day: k.day, Model: k.model, Dimension: k.dim, Quantity: q}) |
| 86 | } |
| 87 | m.rows = map[meterKey]int64{} |
| 88 | return out |
| 89 | } |
| 90 | |
| 91 | // restore folds a failed flush's rows back into the live counters. |
| 92 | func (m *usageMeter) restore(rows []UsageRow) { |
| 93 | m.mu.Lock() |
| 94 | defer m.mu.Unlock() |
| 95 | for _, r := range rows { |
| 96 | m.rows[meterKey{r.Day, r.Model, r.Dimension}] += r.Quantity |
| 97 | } |
| 98 | } |
| 99 | |
| 100 | // recordUsage is the model client's accounting tap (its record hook). |
| 101 | func (r *Runtime) recordUsage(model string, u anthUsage) { |
| 102 | if r.meter != nil { |
| 103 | r.meter.addLLM(model, u) |
| 104 | } |
| 105 | } |
| 106 | |
| 107 | // meterInc bumps an activity dimension from anywhere in the runtime. |
| 108 | func (r *Runtime) meterInc(dim string, n int64) { |
| 109 | if r.meter != nil { |
| 110 | r.meter.inc(dim, n) |
| 111 | } |
| 112 | } |
| 113 | |
| 114 | // meterFlushEvery paces reporting. Coarse on purpose: the ledger's grain is a |
| 115 | // day, so nothing downstream needs freshness — this only bounds how much a |
| 116 | // crash can lose. |
| 117 | const meterFlushEvery = 5 * time.Minute |
| 118 | |
| 119 | // usageReport is one flush in flight: drained rows plus the id the ledger |
| 120 | // dedupes replays on. A failed flush keeps the report VERBATIM and resends |
| 121 | // it — never re-drains into it — because "failed" includes "the cloud |
| 122 | // committed but the ack got lost", and only a byte-identical replay under |
| 123 | // the same id lets the ledger recognize and drop the duplicate. |
| 124 | type usageReport struct { |
| 125 | ID string `json:"report_id"` |
| 126 | Rows []UsageRow `json:"rows"` |
| 127 | } |
| 128 | |
| 129 | func newReportID() string { |
| 130 | raw := make([]byte, 16) |
| 131 | if _, err := rand.Read(raw); err != nil { |
| 132 | // Out of entropy is unrecoverable weirdness; a constant would break |
| 133 | // dedupe silently, so fail the report instead (retried next tick). |
| 134 | return "" |
| 135 | } |
| 136 | return hex.EncodeToString(raw) |
| 137 | } |
| 138 | |
| 139 | func (r *Runtime) meterLoop(ctx context.Context) { |
| 140 | t := time.NewTicker(meterFlushEvery) |
| 141 | defer t.Stop() |
| 142 | for { |
| 143 | select { |
| 144 | case <-t.C: |
| 145 | r.flushUsage() |
| 146 | case <-ctx.Done(): |
| 147 | // Best-effort final flush — the graceful paths (sleep co-spawn |
| 148 | // SIGTERM, deploy restart) shouldn't drop their last interval. |
| 149 | r.flushUsage() |
| 150 | return |
| 151 | } |
| 152 | } |
| 153 | } |
| 154 | |
| 155 | // flushUsage reports counters to the instance and remembers the budget state |
| 156 | // the response carries (inert today; phase 3 gates the model path on it). |
| 157 | // Runs only on the meterLoop goroutine, so pendingUsage needs no lock: a |
| 158 | // pending report is retried as-is before anything new is drained, and fresh |
| 159 | // counts keep accumulating in the meter for the report after it. |
| 160 | func (r *Runtime) flushUsage() { |
| 161 | if r.meter == nil { |
| 162 | return |
| 163 | } |
| 164 | if r.pendingUsage == nil { |
| 165 | rows := r.meter.drain() |
| 166 | if len(rows) == 0 { |
| 167 | return |
| 168 | } |
| 169 | id := newReportID() |
| 170 | if id == "" { |
| 171 | r.meter.restore(rows) |
| 172 | return |
| 173 | } |
| 174 | r.pendingUsage = &usageReport{ID: id, Rows: rows} |
| 175 | } |
| 176 | var resp struct { |
| 177 | OK bool `json:"ok"` |
| 178 | Budget struct { |
| 179 | State string `json:"state"` |
| 180 | } `json:"budget"` |
| 181 | } |
| 182 | if err := r.post("/api/bot-usage", map[string]any{ |
| 183 | "token": r.Session.Token, "report_id": r.pendingUsage.ID, "rows": r.pendingUsage.Rows, |
| 184 | }, &resp); err != nil { |
| 185 | r.logf("usage: flush failed (%v) — report retried next tick", err) |
| 186 | return |
| 187 | } |
| 188 | r.pendingUsage = nil |
| 189 | if resp.Budget.State != "" { |
| 190 | r.budgetState.Store(resp.Budget.State) |
| 191 | if resp.Budget.State != "exhausted" { |
| 192 | // Back within budget (a new month, a raised allowance): re-arm |
| 193 | // the one-time notice for the next exhaustion. |
| 194 | r.budgetNotified.Store(false) |
| 195 | } |
| 196 | } |
| 197 | } |
| 198 | |
| 199 | // budgetExhausted is phase 3's enforcement gate (docs/ai-provisioning.md): |
| 200 | // when the cloud's last flush response said this space's AI budget is spent, |
| 201 | // every model path stands down and the deterministic command floor keeps |
| 202 | // answering — the same degrade shape as ai_enabled off, plus one honest |
| 203 | // notice. Unmetered, "ok", or never-flushed all read as within budget. |
| 204 | func (r *Runtime) budgetExhausted() bool { |
| 205 | v, _ := r.budgetState.Load().(string) |
| 206 | return v == "exhausted" |
| 207 | } |
| 208 | |
| 209 | // noteBudgetExhausted returns true exactly once per exhaustion (re-armed when |
| 210 | // the budget recovers), so the bot explains itself without repeating. Process |
| 211 | // memory only: a restart repeats the notice at most once more. |
| 212 | func (r *Runtime) noteBudgetExhausted() bool { |
| 213 | return r.budgetNotified.CompareAndSwap(false, true) |
| 214 | } |
| 215 | |
amadan