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)
}

// 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
}

type adminQuotasData struct {
	Rows []quotaRow
	// 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:
	default:
		return "", errors.New("a quota belongs to an organization, a team or a project")
	}
	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{
		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))

	orgs, _ := s.AllOrgs(r.Context())
	teams, _ := s.AllTeams(r.Context())
	projects, _ := s.AllProjects(r.Context())

	// Depth first, so a team reads as belonging to the organization above it
	// rather than as another row in a flat list.
	teamsByOrg := map[string][]store.TeamPath{}
	for _, t := range teams {
		teamsByOrg[t.OrgID] = append(teamsByOrg[t.OrgID], t)
	}
	pathOf := map[string]string{}
	for _, pr := range projects {
		pathOf[pr.ID] = pr.Path
	}

	for _, o := range orgs {
		d.Rows = append(d.Rows, quotaRow2(store.ScopeOrg, o.ID, o.Name, o.Name, 0, limits, spend))
		for _, t := range teamsByOrg[o.ID] {
			d.Rows = append(d.Rows,
				quotaRow2(store.ScopeTeam, t.ID, t.Name, o.Name+" / "+t.Name, 1, limits, spend))
			for _, pr := range projects {
				if !strings.HasPrefix(pathOf[pr.ID], o.Name+" / "+t.Name+" / ") {
					continue
				}
				name := strings.TrimPrefix(pathOf[pr.ID], o.Name+" / "+t.Name+" / ")
				d.Rows = append(d.Rows,
					quotaRow2(store.ScopeProject, pr.ID, name, pathOf[pr.ID], 2, limits, spend))
			}
		}
	}
	markShadowed(d.Rows, limits, orgs, teamsByOrg, projects, pathOf)
	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 QuotaAdmin,
	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
}

func quotaRow2(scope store.Scope, id, name, crumb string, indent int,
	limits map[string]store.Limits, spend map[string]int64,
) quotaRow {
	key := store.LimitOwner{Scope: scope, ID: id}.CounterKey()
	l := limits[key]
	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]),
	}
	if l.Budget != nil && *l.Budget > 0 {
		row.Percent = int(spend[key] * 100 / *l.Budget)
		row.HasPct = true
	}
	return row
}

// markShadowed flags rows whose budget an ancestor already caps lower.
func markShadowed(rows []quotaRow, limits map[string]store.Limits,
	orgs []store.Org, teamsByOrg map[string][]store.TeamPath,
	projects []store.ProjectPath, pathOf map[string]string,
) {
	budget := func(scope store.Scope, id string) *int64 {
		return limits[store.LimitOwner{Scope: scope, ID: id}.CounterKey()].Budget
	}
	// Parent of each team and project, so a row can look upward.
	parent := map[string]struct {
		scope store.Scope
		id    string
		name  string
	}{}
	for _, o := range orgs {
		for _, t := range teamsByOrg[o.ID] {
			parent[t.ID] = struct {
				scope store.Scope
				id    string
				name  string
			}{store.ScopeOrg, o.ID, o.Name}
			for _, pr := range projects {
				if strings.HasPrefix(pathOf[pr.ID], o.Name+" / "+t.Name+" / ") {
					parent[pr.ID] = struct {
						scope store.Scope
						id    string
						name  string
					}{store.ScopeTeam, t.ID, t.Name}
				}
			}
		}
	}

	for i := range rows {
		own := budget(rows[i].Scope, rows[i].ID)
		if own == nil {
			continue
		}
		for cur := parent[rows[i].ID]; cur.id != ""; cur = parent[cur.id] {
			if up := budget(cur.scope, cur.id); up != nil && *up < *own {
				rows[i].Shadowed = fmt.Sprintf("%s caps this at %s", cur.name, units(*up))
				break
			}
		}
	}
}

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)
}