llmgw/internal/web/admin_quotas.go

package web

import (
	"context"
	"errors"
	"fmt"
	"net/http"
	"strconv"
	"strings"
	"time"

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

// Priced says whether anything a request could name costs money. The
// accountant implements it.
type Priced interface{ AnythingPriced() bool }

// QuotaAdmin is the store half of the quota page.
type QuotaAdmin interface {
	AllOrgs(ctx context.Context) ([]store.Org, error)
	AllTeams(ctx context.Context) ([]store.TeamPath, error)
	AllProjects(ctx context.Context) ([]store.ProjectPath, error)
	AllLimits(ctx context.Context) (map[string]store.Limits, error)
	PutLimits(ctx context.Context, l store.Limits) error
	ClearLimits(ctx context.Context, scope store.Scope, id string) error
	UsageByScope(ctx context.Context, scope store.Scope, since time.Time) ([]store.ScopeUsage, error)
	// Tree is the same tree with ids all the way down, so a project is placed
	// under its team by id rather than by matching the start of its path —
	// which put projects under the wrong team when two organizations shared a
	// name.
	Tree(ctx context.Context) ([]store.TreeOrg, error)
	// People have limits of their own, across every key they hold.
	Users(ctx context.Context, state store.UserState) ([]store.User, error)
	UsageByUser(ctx context.Context, since time.Time) ([]store.UserUsage, error)
}

// scopeUsage is what spend per branch is read from.
type scopeUsage interface {
	UsageByScope(ctx context.Context, scope store.Scope, since time.Time) ([]store.ScopeUsage, error)
}

// personRow is one person's own limits: across all their keys and projects.
type personRow struct {
	ID, Email      string
	RPM, TPM       string
	Budget, Period string
	Models         string
	HasAny         bool
	Spend          string
	Percent        int
	HasPct         bool
	Editing        bool
}

// quotaRow is one place a quota can be written, and what is written there.
type quotaRow struct {
	Scope store.Scope
	ID    string
	// Path is how a person recognises it: "Scania / platform / gateway".
	Path string
	// Indent is how deep in the tree it sits, so the table reads as one.
	Indent int
	// Crumb is the full path, for a chooser where indentation cannot carry it.
	Crumb string

	RPM    string
	TPM    string
	Budget string
	Period string
	Models string
	HasAny bool

	Spend   string
	Percent int
	HasPct  bool

	// Editing is true for the one row whose form is open.
	//
	// One at a time, chosen by a link rather than opened by script. A
	// disclosure on every row would put twenty collapsed forms in the markup of
	// a page whose job is to be read as a tree.
	Editing bool

	// Shadowed names an ancestor whose limit is stricter, which means the value
	// written here does nothing.
	//
	// Resolution takes the most restrictive along the chain, so a project given
	// a budget larger than its team's is not a bigger budget — it is a number
	// somebody typed that will never be reached. Without this the page would
	// show it as though it applied, and the first sign would be a refusal that
	// disagrees with what is on screen.
	Shadowed string

	// RPMFrom, TPMFrom and BudgetFrom name the ancestor a limit comes from when
	// this row sets none of its own, with the value in RPMUp, TPMUp and
	// BudgetUp. A project with "—" in every column read as unlimited when its
	// organization held it to sixty requests a minute.
	RPMUp, TPMUp, BudgetUp       string
	RPMFrom, TPMFrom, BudgetFrom string
	// ModelsFrom names the nearest ancestor with a model list, for a row with
	// none of its own.
	ModelsFrom string

	// Member is the per-member limit written here, summarised; MemberRaw holds
	// its fields for the form, and MemberEditing opens that form.
	Member        string
	MemberRaw     quotaFields
	MemberEditing bool
}

// quotaFields is a limit as the form shows it.
type quotaFields struct {
	RPM, TPM, Budget, Period, Models string
	HasAny                           bool
}

type adminQuotasData struct {
	Nav    Nav
	Rows   []quotaRow
	People []personRow
	// Personal summarises the defaults a person's own use falls under when
	// nobody set a limit for them.
	Personal string
	// BudgetsBind is false when no model has a price, which makes every budget
	// on this page a number that cannot move. Shown, because a limit that binds
	// on nothing and reads as "0% used" is worse than no limit at all.
	BudgetsBind bool
	// Editing is the "scope:id" whose form is open, so the template can keep
	// the link on every other row.
	Editing string
	Since   string
	Message string
	Error   string
	CSRF    string
	Version string
	Now     string
}

// RegisterAdminQuotas mounts quota administration under /admin/quotas.
//
// Reading and writing both, which is a change of position. Setting a quota was
// kept to the shell on the grounds that raising a budget is a financial act and
// a stolen admin session should not manage one. That argument holds for exactly
// one row on this page: resolution takes the most restrictive value along the
// chain, so raising a team or project limit cannot exceed what the organization
// allows. Only the organization's own limit is bounded by nothing.
//
// So the trade is narrower than it looked, and against it stands a quota nobody
// could see or change without a shell on the control plane — which in practice
// means quotas nobody sets. Every change is logged with who made it and what it
// was before.
func (p *AccountPages) RegisterAdminQuotas(mux *http.ServeMux, s QuotaAdmin, priced Priced) {
	p.hasQuotas = true
	p.priced = priced

	mux.HandleFunc("GET /admin/quotas", p.adminOnly(
		func(w http.ResponseWriter, r *http.Request, _ store.User) {
			p.renderQuotas(w, r, s, "", "")
		}))

	mux.HandleFunc("POST /admin/quotas", p.adminOnly(
		func(w http.ResponseWriter, r *http.Request, admin store.User) {
			sess, _, _ := p.session(r)
			if !p.accounts.CheckCSRF(sess, r.FormValue("csrf")) {
				http.Error(w, "bad request", http.StatusBadRequest)
				return
			}
			msg, err := p.applyQuota(r, s, admin)
			if err != nil {
				p.log.Warn("a quota change was refused",
					"admin", admin.Email, "scope", r.FormValue("scope"),
					"id", r.FormValue("id"), "err", err)
				p.renderQuotas(w, r, s, "", err.Error())
				return
			}
			p.renderQuotas(w, r, s, msg, "")
		}))
}

// applyQuota writes one scope's limits, and says what it did.
func (p *AccountPages) applyQuota(r *http.Request, s QuotaAdmin, admin store.User) (string, error) {
	// Carried by the form the row opened, so the branch a quota lands on is the
	// one that was clicked rather than one somebody named again in a box.
	scope, id := store.Scope(r.FormValue("scope")), r.FormValue("id")
	switch scope {
	case store.ScopeOrg, store.ScopeTeam, store.ScopeProject, store.ScopeUser,
		store.ScopeOrgMember, store.ScopeTeamMember, store.ScopeProjectMember:
	default:
		return "", errors.New("a quota belongs to an organization, team, project or person")
	}
	if id == "" {
		return "", errors.New("which one?")
	}

	// Before the write, so the audit line can say what it was. A record that a
	// budget changed without the number it changed from answers nothing later.
	before := "unknown"
	if all, err := s.AllLimits(r.Context()); err == nil {
		before = describeLimits(all[store.LimitOwner{Scope: scope, ID: id}.CounterKey()])
	}

	if r.FormValue("clear") != "" {
		if err := s.ClearLimits(r.Context(), scope, id); err != nil {
			return "", err
		}
		p.log.Info("a quota was cleared", "admin", admin.Email,
			"scope", scope, "id", id, "was", before)
		return string(scope) + " " + id + ": every limit here removed", nil
	}

	l := store.Limits{Scope: scope, ScopeID: id}
	var err error
	if l.RPM, err = optInt(r.FormValue("rpm"), "requests per minute"); err != nil {
		return "", err
	}
	if l.TPM, err = optInt(r.FormValue("tpm"), "tokens per minute"); err != nil {
		return "", err
	}
	// Whole units in, micro-units stored: the operator set a budget in the unit
	// they think in, and the store keeps integers so repeated addition of a
	// cost cannot drift.
	if l.Budget, err = optUnits(r.FormValue("budget")); err != nil {
		return "", err
	}
	if l.Budget != nil {
		l.BudgetPeriod = store.Period(r.FormValue("period"))
		if l.BudgetPeriod != store.PeriodDay && l.BudgetPeriod != store.PeriodMonth {
			return "", errors.New("a budget needs a period: day or month")
		}
	}
	if m := strings.TrimSpace(r.FormValue("models")); m != "" {
		for _, name := range strings.Split(m, ",") {
			if name = strings.TrimSpace(name); name != "" {
				l.AllowedModels = append(l.AllowedModels, name)
			}
		}
	}

	if err := s.PutLimits(r.Context(), l); err != nil {
		return "", err
	}
	now := describeLimits(l)
	p.log.Info("a quota was changed", "admin", admin.Email,
		"scope", scope, "id", id, "was", before, "now", now)
	return string(scope) + " " + id + ": " + now, nil
}

func (p *AccountPages) renderQuotas(w http.ResponseWriter, r *http.Request,
	s QuotaAdmin, msg, errMsg string,
) {
	const window = 30 * 24 * time.Hour
	d := adminQuotasData{
		Nav:   navFrom(r, "quotas"),
		Since: "30 days", Message: msg, Error: errMsg,
		CSRF: csrfFor(r), Version: p.version,
		Now: time.Now().UTC().Format(time.RFC3339),
	}

	d.BudgetsBind = p.priced != nil && p.priced.AnythingPriced()
	d.Editing = r.URL.Query().Get("edit")

	limits, err := s.AllLimits(r.Context())
	if err != nil {
		p.log.Error("could not read the limits", "err", err)
		d.Error = "could not read the limits"
		render(w, "admin_quotas", d)
		return
	}
	spend := p.spendByScope(r.Context(), s, time.Now().Add(-window))

	tree, err := s.Tree(r.Context())
	if err != nil {
		p.log.Error("could not read the tree", "err", err)
		d.Error = "could not read the organizations"
	}
	d.Rows = quotaRows(tree, limits, spend)
	for i := range d.Rows {
		ms := store.MemberScope(d.Rows[i].Scope)
		if l, ok := limits[store.LimitOwner{Scope: ms, ID: d.Rows[i].ID}.CounterKey()]; ok {
			d.Rows[i].MemberRaw = fieldsOf(l)
			d.Rows[i].Member = describeLimits(l)
		}
		d.Rows[i].MemberEditing = d.Editing == string(ms)+":"+d.Rows[i].ID
	}
	d.People = p.personRows(r.Context(), s, limits, time.Now().Add(-window), d.Editing)
	d.Personal = p.personalSummary()
	for i := range d.Rows {
		d.Rows[i].Editing = d.Editing ==
			string(d.Rows[i].Scope)+":"+d.Rows[i].ID
	}
	render(w, "admin_quotas", d)
}

// spendByScope is what each group has used in the window, keyed like a limit.
func (p *AccountPages) spendByScope(ctx context.Context, s scopeUsage,
	since time.Time,
) map[string]int64 {
	out := map[string]int64{}
	for _, sc := range []store.Scope{store.ScopeOrg, store.ScopeTeam, store.ScopeProject} {
		rows, err := s.UsageByScope(ctx, sc, since)
		if err != nil {
			// One level failing must not blank the other two: a page showing
			// two thirds of the answer is more use than an error.
			p.log.Warn("could not read usage for the quota page", "scope", sc, "err", err)
			continue
		}
		for _, u := range rows {
			out[store.LimitOwner{Scope: sc, ID: u.ID}.CounterKey()] = u.CostMicros
		}
	}
	return out
}

// quotaRows lays the tree out depth first, each row with what it sets and
// what it is held to from above.
func quotaRows(tree []store.TreeOrg, limits map[string]store.Limits, spend map[string]int64) []quotaRow {
	var rows []quotaRow
	for _, o := range tree {
		org := quotaRow2(store.ScopeOrg, o.ID, o.Name, o.Name, 0, limits, spend, nil)
		rows = append(rows, org)
		for _, t := range o.Teams {
			crumb := o.Name + " / " + t.Name
			team := quotaRow2(store.ScopeTeam, t.ID, t.Name, crumb, 1, limits, spend,
				[]ancestor{{o.Name, limits[key(store.ScopeOrg, o.ID)]}})
			rows = append(rows, team)
			for _, pr := range t.Projects {
				rows = append(rows, quotaRow2(store.ScopeProject, pr.ID, pr.Name, crumb+" / "+pr.Name, 2,
					limits, spend, []ancestor{
						{t.Name, limits[key(store.ScopeTeam, t.ID)]},
						{o.Name, limits[key(store.ScopeOrg, o.ID)]},
					}))
			}
		}
	}
	return rows
}

// ancestor is one level above a row, nearest first.
type ancestor struct {
	name string
	l    store.Limits
}

func key(scope store.Scope, id string) string {
	return store.LimitOwner{Scope: scope, ID: id}.CounterKey()
}

func quotaRow2(scope store.Scope, id, name, crumb string, indent int,
	limits map[string]store.Limits, spend map[string]int64, up []ancestor,
) quotaRow {
	l := limits[key(scope, id)]
	row := quotaRow{
		Scope: scope, ID: id, Path: name, Crumb: crumb, Indent: indent,
		RPM: dash(l.RPM), TPM: dash(l.TPM),
		Budget: dashUnits(l.Budget), Period: string(l.BudgetPeriod),
		Models: strings.Join(l.AllowedModels, ", "),
		HasAny: l.RPM != nil || l.TPM != nil || l.Budget != nil || l.AllowedModels != nil,
		Spend:  units(spend[key(scope, id)]),
	}
	if l.Budget != nil && *l.Budget > 0 {
		row.Percent = int(spend[key(scope, id)] * 100 / *l.Budget)
		row.HasPct = true
	}

	// The strictest value above, and who set it. Resolution takes the most
	// restrictive along the chain, so this is what applies when the row sets
	// nothing, and what makes a larger value set here unreachable.
	var rpm, tpm *int
	var budget *int64
	var period store.Period
	for _, a := range up {
		if a.l.RPM != nil && (rpm == nil || *a.l.RPM < *rpm) {
			rpm, row.RPMFrom = a.l.RPM, a.name
		}
		if a.l.TPM != nil && (tpm == nil || *a.l.TPM < *tpm) {
			tpm, row.TPMFrom = a.l.TPM, a.name
		}
		if a.l.Budget != nil && (budget == nil || *a.l.Budget < *budget) {
			budget, period, row.BudgetFrom = a.l.Budget, a.l.BudgetPeriod, a.name
		}
		if row.ModelsFrom == "" && a.l.AllowedModels != nil {
			row.ModelsFrom = a.name
		}
	}
	if rpm != nil {
		row.RPMUp = strconv.Itoa(*rpm)
	}
	if tpm != nil {
		row.TPMUp = strconv.Itoa(*tpm)
	}
	if budget != nil {
		row.BudgetUp = units(*budget)
		if period != "" {
			row.BudgetUp += "/" + string(period)
		}
	}

	var capped []string
	if l.RPM != nil && rpm != nil && *rpm < *l.RPM {
		capped = append(capped, fmt.Sprintf("%s requests/min", row.RPMUp))
	}
	if l.TPM != nil && tpm != nil && *tpm < *l.TPM {
		capped = append(capped, fmt.Sprintf("%s tokens/min", row.TPMUp))
	}
	if l.Budget != nil && budget != nil && *budget < *l.Budget {
		capped = append(capped, "a budget of "+units(*budget))
	}
	if len(capped) > 0 {
		from := row.RPMFrom
		if from == "" {
			from = row.TPMFrom
		}
		if from == "" {
			from = row.BudgetFrom
		}
		row.Shadowed = "above it, " + strings.Join(capped, " and ") + " still apply"
		if from != "" && len(capped) == 1 {
			row.Shadowed = fmt.Sprintf("%s caps this at %s", from, capped[0])
		}
	}
	return row
}

func optInt(v, what string) (*int, error) {
	if strings.TrimSpace(v) == "" {
		return nil, nil
	}
	n, err := strconv.Atoi(strings.TrimSpace(v))
	if err != nil || n < 0 {
		return nil, fmt.Errorf("%s must be a whole number, or empty for no limit", what)
	}
	return &n, nil
}

// optUnits reads a budget written in whole units and stores micro-units.
func optUnits(v string) (*int64, error) {
	if strings.TrimSpace(v) == "" {
		return nil, nil
	}
	f, err := strconv.ParseFloat(strings.TrimSpace(v), 64)
	if err != nil || f < 0 {
		return nil, errors.New("the budget must be a number, or empty for no budget")
	}
	m := int64(f * 1e6)
	return &m, nil
}

func describeLimits(l store.Limits) string {
	var parts []string
	if l.RPM != nil {
		parts = append(parts, "rpm "+strconv.Itoa(*l.RPM))
	}
	if l.TPM != nil {
		parts = append(parts, "tpm "+strconv.Itoa(*l.TPM))
	}
	if l.Budget != nil {
		parts = append(parts, "budget "+units(*l.Budget)+"/"+string(l.BudgetPeriod))
	}
	if l.AllowedModels != nil {
		parts = append(parts, "models "+strings.Join(l.AllowedModels, "|"))
	}
	if len(parts) == 0 {
		return "no limits"
	}
	return strings.Join(parts, ", ")
}

func dash(v *int) string {
	if v == nil {
		return ""
	}
	return strconv.Itoa(*v)
}

func dashUnits(v *int64) string {
	if v == nil {
		return ""
	}
	return units(*v)
}

func fieldsOf(l store.Limits) quotaFields {
	return quotaFields{
		RPM: dash(l.RPM), TPM: dash(l.TPM), Budget: dashUnits(l.Budget),
		Period: string(l.BudgetPeriod), Models: strings.Join(l.AllowedModels, ", "),
		HasAny: l.RPM != nil || l.TPM != nil || l.Budget != nil || l.AllowedModels != nil,
	}
}

// personRows lists every approved person with the limits written for them and
// what they spent, wherever they spent it.
func (p *AccountPages) personRows(ctx context.Context, s QuotaAdmin,
	limits map[string]store.Limits, since time.Time, editing string,
) []personRow {
	users, err := s.Users(ctx, "")
	if err != nil {
		p.log.Warn("could not list people for the quota page", "err", err)
		return nil
	}
	spent := map[string]int64{}
	if usage, err := s.UsageByUser(ctx, since); err == nil {
		for _, u := range usage {
			spent[u.UserID] = u.CostMicros
		}
	}
	var out []personRow
	for _, u := range users {
		if u.State == store.UserPending {
			continue
		}
		l := limits[store.LimitOwner{Scope: store.ScopeUser, ID: u.ID}.CounterKey()]
		f := fieldsOf(l)
		row := personRow{ID: u.ID, Email: u.Email, RPM: f.RPM, TPM: f.TPM, Budget: f.Budget,
			Period: f.Period, Models: f.Models, HasAny: f.HasAny, Spend: units(spent[u.ID]),
			Editing: editing == "user:"+u.ID}
		if l.Budget != nil && *l.Budget > 0 {
			row.Percent, row.HasPct = int(spent[u.ID]*100 / *l.Budget), true
		}
		out = append(out, row)
	}
	return out
}

// SetPersonalDefaults tells the quota page what a person's own use falls under
// when nobody set a limit for them.
func (p *AccountPages) SetPersonalDefaults(l *store.Limits) { p.personalDefaults = l }

func (p *AccountPages) personalSummary() string {
	if p.personalOff {
		return "Personal use is off: people use the gateway only through their projects."
	}
	if p.personalDefaults == nil {
		return "Personal use has no default limit: someone with no limit of their own is bounded only by the cluster-wide defaults."
	}
	return "With no limit of their own, a person's personal use is held to " + describeLimits(*p.personalDefaults) + "."
}