llmgw/internal/accounting/accounting.go

package accounting

import (
	"context"
	"errors"
	"fmt"
	"log/slog"
	"strings"
	"sync/atomic"
	"time"

	"gitea.mynodes.xyz/ilker/llmgw/internal/auth"
	"gitea.mynodes.xyz/ilker/llmgw/internal/notify"
	"gitea.mynodes.xyz/ilker/llmgw/internal/store"
)

// Refusals a key can hit. They are distinguished so the client is told which
// limit it hit, and so retrying is only suggested where retrying helps.
var (
	ErrRateLimited  = errors.New("request rate limit exceeded")
	ErrTokenLimited = errors.New("token rate limit exceeded")
	ErrBudgetSpent  = errors.New("budget exhausted")
)

// LimitError carries which limit was hit and, where it makes sense, when to
// come back.
type LimitError struct {
	Err error
	// RetryAfter is zero when waiting does not help, which is the case for a
	// spent budget: it will not refill until the period rolls over, and telling
	// a client to retry in a second would be a lie.
	RetryAfter time.Duration
}

func (e *LimitError) Error() string { return e.Err.Error() }
func (e *LimitError) Unwrap() error { return e.Err }

// Accountant enforces limits and records what was used.
//
// One object because the three halves are not independent: a budget check needs
// the price, the price needs the model, and the record closes the loop by
// feeding the counters that the next check reads.
type Accountant struct {
	pricing  *Pricing
	budgets  *Budgets
	requests *Limiter
	tokens   *Limiter
	recorder *Recorder
	log      *slog.Logger
	// notify is told when a budget runs out; nil tells no one.
	notify atomic.Pointer[notify.Notifier]
}

// SetNotifier gives the accountant somewhere to say a budget is spent.
func (a *Accountant) SetNotifier(n *notify.Notifier) { a.notify.Store(n) }

// New builds an accountant.
func New(pricing *Pricing, recorder *Recorder, log *slog.Logger) *Accountant {
	if log == nil {
		log = slog.Default()
	}
	return &Accountant{
		pricing:  pricing,
		budgets:  NewBudgets(),
		requests: NewLimiter(),
		tokens:   NewLimiter(),
		recorder: recorder,
		log:      log,
	}
}

// RestoreBudgets seeds the in-memory counters from the store.
//
// Called at startup. A restart must not reset a budget, which it would if the
// counters only ever lived in memory.
func (a *Accountant) RestoreBudgets(ctx context.Context, s store.Store, period store.Period) error {
	start := PeriodStart(period, time.Now())
	if start.IsZero() {
		return nil
	}
	rows, err := s.SpendSinceByScope(ctx, start)
	if err != nil {
		return fmt.Errorf("restoring budgets: %w", err)
	}
	// Every scope a key sits under, not just the key. A budget owned by a team
	// is one counter fed by every key beneath it, and restoring only per-key
	// totals would bring the control plane back up believing the team had spent
	// nothing — which is a budget that resets on every restart, for exactly the
	// scopes where it matters most.
	// And every pool a person is in: their own, and each per-member limit of
	// the project, team and organization they spent it in.
	spend := make(map[string]int64, len(rows)*8)
	for _, r := range rows {
		owners := []store.LimitOwner{
			{Scope: store.ScopeKey, ID: r.KeyID},
			{Scope: store.ScopeProject, ID: r.ProjectID},
			{Scope: store.ScopeTeam, ID: r.TeamID},
			{Scope: store.ScopeOrg, ID: r.OrgID},
			{Scope: store.ScopeUser, ID: r.UserID},
		}
		if r.UserID != "" && r.ProjectID != "" {
			owners = append(owners,
				store.MemberOwner(store.ScopeProjectMember, r.ProjectID, r.UserID),
				store.MemberOwner(store.ScopeTeamMember, r.TeamID, r.UserID),
				store.MemberOwner(store.ScopeOrgMember, r.OrgID, r.UserID))
		}
		for _, o := range owners {
			if o.ID == "" || strings.HasPrefix(o.ID, "/") {
				continue
			}
			if k := o.CounterKey(); k != "" {
				spend[k] += r.Micros
			}
		}
	}
	a.budgets.Restore(period, spend)
	a.log.Info("restored budgets from the store",
		"keys", len(rows), "counters", len(spend), "since", start)
	return nil
}

// Reservation is what Check returns: the estimate it charged, so Record can
// settle the difference against the real figure.
// pool is one level's counter and the period its budget runs over.
type pool struct {
	Counter string
	Period  store.Period
}

type Reservation struct {
	KeyID string
	// Budgets and Tokens are every counter the check went through, carried so
	// the settlement lands on the same ones. Every, not the tightest: an
	// organization's pool only fills up if each team's spend is added to it
	// as well as to the team's own.
	//
	// Derived again at settlement would mean re-resolving the limits, and a
	// budget checked against the team and settled against the key is a budget
	// that never fills up.
	Budgets         []pool
	Tokens          []string
	EstimatedTokens int64
	Price           Price
	// Period is the key's budget period, carried from the check so the record
	// counts against the same window the check read. Deriving it later would
	// mean guessing, and a monthly budget credited to a daily counter is a
	// budget that never binds.
	Period store.Period
}

// Check decides whether a request may proceed.
//
// It runs before scheduling and touches no database: every counter it reads is
// in memory. The token limit is checked against an estimate, because the real
// figure only exists once the request has run; Record settles the difference.
func (a *Accountant) Check(p auth.Principal, model string, estimatedTokens int64) (Reservation, error) {
	price := a.pricing.For(model)
	lim := p.Limits
	// Every level that set a limit, each against its own pool. A quota written
	// on a team is one pool the team shares, and the organization above it is a
	// second pool that the team's spending also fills — both have to hold, or
	// an organization given 1000 with four teams given 500 spends 2000.
	chain := levels(lim, p.KeyID)

	res := Reservation{
		KeyID: p.KeyID, EstimatedTokens: estimatedTokens,
		Price: price, Period: lim.BudgetPeriod,
	}
	for _, l := range chain {
		if l.Budget != nil {
			res.Budgets = append(res.Budgets, pool{l.Counter, l.Period})
		}
		if l.TPM != nil {
			res.Tokens = append(res.Tokens, l.Counter)
		}
	}
	if p.KeyID == "" {
		return res, nil
	}

	// Budget first. It is the limit whose refusal cannot be retried away, so
	// telling the client about it first is more useful than reporting a rate
	// limit it would hit again in a second.
	//
	// A model that costs nothing is exempt, which is what lets a key with no
	// budget left still be served a cheap fallback (plan §7.7).
	// Budgets first, and all of them. Exhausted only reads, so the order costs
	// nothing and a refusal here charges no counter.
	if !price.Zero() {
		for _, l := range chain {
			if l.Budget != nil && a.budgets.Exhausted(l.Counter, l.Period, *l.Budget) {
				a.notify.Load().Notify(notify.BudgetSpent, l.Counter,
					fmt.Sprintf("the %s budget of %s is spent: requests to priced models are refused until it renews",
						or(string(l.Period), "total"), l.Counter),
					map[string]any{"counter": l.Counter, "period": string(l.Period), "budget": *l.Budget})
				return res, &LimitError{Err: ErrBudgetSpent}
			}
		}
	}

	// Rate limits charge as they check — deliberately, so two concurrent
	// requests cannot both see room. With several levels that means a refusal
	// at one has already taken from the levels before it, so those are handed
	// back. Without the refund a team at its ceiling would quietly eat its
	// organization's rate limit with requests that were never served.
	if err := a.allowAll(chain, func(l level) (string, int, int64, bool) {
		return l.Counter, valInt(l.RPM), 1, l.RPM != nil
	}, a.requests); err != nil {
		return res, &LimitError{Err: ErrRateLimited, RetryAfter: time.Second}
	}
	if err := a.allowAll(chain, func(l level) (string, int, int64, bool) {
		return l.Counter, valInt(l.TPM), estimatedTokens, l.TPM != nil
	}, a.tokens); err != nil {
		// The request was not charged against the request-rate limit in vain:
		// it was, and that is correct — it did arrive.
		return res, &LimitError{Err: ErrTokenLimited, RetryAfter: time.Second}
	}
	return res, nil
}

// level is one scope's limits with the counter they are enforced against.
type level struct {
	Counter string
	Period  store.Period
	RPM     *int
	TPM     *int
	Budget  *int64
}

// levels turns a resolved chain into the pools a request must pass.
//
// Falls back to the collapsed values when the chain is empty, which is what a
// principal assembled by hand looks like — in tests, and anywhere that builds
// limits without walking the tree.
func levels(e store.Effective, keyID string) []level {
	if len(e.Chain) == 0 {
		var out []level
		if e.RPM != nil || e.TPM != nil || e.Budget != nil {
			out = append(out, level{
				Counter: counterFor(e.BudgetOwner, keyID),
				Period:  e.BudgetPeriod, Budget: e.Budget,
			})
			out[0].RPM, out[0].TPM = e.RPM, e.TPM
			// One entry, so a hand-built principal with different owners per
			// limit still counts each where it said to.
			if r, b := counterFor(e.RPMOwner, keyID), out[0].Counter; r != b && e.RPM != nil {
				out[0].RPM = nil
				out = append(out, level{Counter: r, RPM: e.RPM})
			}
			if tp, b := counterFor(e.TPMOwner, keyID), out[0].Counter; tp != b && e.TPM != nil {
				out[0].TPM = nil
				out = append(out, level{Counter: tp, TPM: e.TPM})
			}
		}
		return out
	}
	out := make([]level, 0, len(e.Chain))
	for _, l := range e.Chain {
		out = append(out, level{
			Counter: store.LimitOwner{Scope: l.Scope, ID: l.ScopeID}.CounterKey(),
			Period:  l.BudgetPeriod, RPM: l.RPM, TPM: l.TPM, Budget: l.Budget,
		})
	}
	return out
}

// allowAll charges every level that has this limit, refunding on a refusal.
func (a *Accountant) allowAll(chain []level,
	pick func(level) (key string, limit int, n int64, has bool), lim *Limiter,
) error {
	var charged []struct {
		key string
		n   int64
	}
	for _, l := range chain {
		key, limit, n, has := pick(l)
		if !has {
			continue
		}
		if !lim.Allow(key, limit, n) {
			for _, c := range charged {
				lim.Charge(c.key, -c.n)
			}
			return errRefused
		}
		charged = append(charged, struct {
			key string
			n   int64
		}{key, n})
	}
	return nil
}

var errRefused = errors.New("refused")

func valInt(v *int) int {
	if v == nil {
		return 0
	}
	return *v
}

// Record settles a finished request: it reconciles the token estimate, adds to
// the budget, and queues the usage row.
func (a *Accountant) Record(res Reservation, u store.Usage) {
	if res.KeyID == "" {
		return
	}
	actual := u.PromptTokens + u.CompletionTokens

	// The estimate was charged up front; charge the difference now. A caller
	// whose requests are consistently larger than estimated is still bounded,
	// one window later than with a perfect estimate.
	// Every level the check went through, so an organization's pool fills up
	// from its teams' spending rather than only from requests that named it.
	for _, c := range res.Tokens {
		a.tokens.Charge(c, actual-res.EstimatedTokens)
	}

	// Priced by the model that answered: a fallback served by another model
	// costs what that one costs, not what the one asked for would have.
	price := res.Price
	if u.ServedModel != "" && u.ServedModel != u.Model && a.pricing != nil {
		price = a.pricing.For(u.ServedModel)
	}
	u.CostMicros = price.Cost(u.PromptTokens, u.CompletionTokens, u.Images)
	for _, b := range res.Budgets {
		a.budgets.Add(b.Counter, b.Period, u.CostMicros)
	}

	if a.recorder != nil {
		a.recorder.Record(u)
	}
}

// RecordRefused queues the row of a request a limit refused: nothing is
// charged — it used nothing — but it is counted, so a team running into its
// limits shows as one rather than as a team that went quiet.
func (a *Accountant) RecordRefused(u store.Usage) {
	if u.KeyID == "" || a.recorder == nil {
		return
	}
	u.PromptTokens, u.CompletionTokens, u.CostMicros = 0, 0, 0
	a.recorder.Record(u)
}

// Status returns what a key has used, for headers and `ctl`.
func (a *Accountant) Status(p auth.Principal) (spent int64, rpm, tpm float64) {
	if p.KeyID == "" {
		return 0, 0, 0
	}
	// The pool that will refuse them, which is not always the one with the
	// smallest number on it. A team given 2000 inside an organization given the
	// same is bounded by whichever is nearer its ceiling, and that is the
	// organization as soon as a second team spends anything. A header that
	// reported the team's would disagree with the refusal that follows it.
	for _, l := range levels(p.Limits, p.KeyID) {
		if l.Budget != nil {
			if v := a.budgets.Spent(l.Counter, l.Period); v > spent {
				spent = v
			}
		}
		if l.RPM != nil {
			if v := a.requests.Rate(l.Counter); v > rpm {
				rpm = v
			}
		}
		if l.TPM != nil {
			if v := a.tokens.Rate(l.Counter); v > tpm {
				tpm = v
			}
		}
	}
	return spent, rpm, tpm
}

// counterFor is the counter a limit owned by this scope should use.
//
// Falls back to the key's own counter when no scope owns the limit, which
// happens when nothing in the chain set one.
//
// Namespaced even then. A bare key id here and a "key:"-prefixed one in the
// restore path would be two different counters for the same key, so a budget
// would check against an empty one and survive every restart — which is how
// this was wrong when it was first written.
func counterFor(o store.LimitOwner, keyID string) string {
	if k := o.CounterKey(); k != "" {
		return k
	}
	return keyCounter(keyID)
}

// keyCounter is a key's own counter, in the same namespace as every other.
func or(v, fallback string) string {
	if v == "" {
		return fallback
	}
	return v
}

func keyCounter(keyID string) string {
	if keyID == "" {
		return ""
	}
	return store.LimitOwner{Scope: store.ScopeKey, ID: keyID}.CounterKey()
}

// AnythingPriced reports whether a budget can ever be spent against.
//
// On the accountant rather than making callers reach for the price table: what
// they are asking is whether the budgets they are about to display mean
// anything, and the accountant is what would enforce them.
func (a *Accountant) AnythingPriced() bool { return a.pricing.AnythingPriced() }

// Headroom is the tightest of a principal's request and token limits: how
// much is left of each, and when the minute it is counted over rolls on.
type Headroom struct {
	RequestLimit, RequestsLeft int
	TokenLimit, TokensLeft     int
	Reset                      time.Duration
}

// Headroom reports what a principal has left, after whatever Check just
// charged. Of several levels the one with least left is the one that will
// refuse first, so that is the one reported. A limit of 0 means none is set.
func (a *Accountant) Headroom(p auth.Principal) Headroom {
	var h Headroom
	firstReq, firstTok := true, true
	for _, l := range levels(p.Limits, p.KeyID) {
		if l.RPM != nil && *l.RPM > 0 {
			left := max(0, *l.RPM-int(a.requests.Rate(l.Counter)))
			if firstReq || left < h.RequestsLeft {
				h.RequestLimit, h.RequestsLeft, firstReq = *l.RPM, left, false
			}
		}
		if l.TPM != nil && *l.TPM > 0 {
			left := max(0, *l.TPM-int(a.tokens.Rate(l.Counter)))
			if firstTok || left < h.TokensLeft {
				h.TokenLimit, h.TokensLeft, firstTok = *l.TPM, left, false
			}
		}
	}
	// The windows slide, so room comes back continuously; the next minute
	// boundary is when a whole window's worth has moved on, and is the figure
	// a client pacing itself can act on.
	now := time.Now()
	h.Reset = now.Truncate(windowLen).Add(windowLen).Sub(now)
	return h
}