llmgw/internal/store/store.go

// Package store persists what the cluster cannot reconstruct: who may call the
// gateway, what they are allowed, and what they have used (plan §7.9).
//
// Everything else llmgw knows is derived from the fleet and is deliberately not
// stored — a database that describes the cluster can be wrong, and the cluster
// cannot (§6.5). This package holds only the facts that have no other source.
//
// The interface exists so the engine can change. SQLite is right for a single
// instance: it is embedded, so the store shares the gateway's fate instead of
// adding a second thing that can be down, and an inference gateway that stops
// serving because its accounting database is unreachable is worse than one that
// buffers some accounting. That argument stops holding the moment there are two
// gateway instances, because budgets and rate limits then have to be counted
// somewhere shared; Postgres arrives with that, not before (plan §7.9.6).
package store

import (
	"context"
	"encoding/json"
	"errors"
	"strings"
	"time"
)

// ErrNotFound is returned when a lookup finds nothing.
var ErrNotFound = errors.New("not found")

// Scope is a level of the tenancy tree. Limits hang off any of them.
type Scope string

// The tenancy levels, from widest to narrowest.
const (
	ScopeOrg     Scope = "org"
	ScopeTeam    Scope = "team"
	ScopeProject Scope = "project"
	ScopeKey     Scope = "key"

	// ScopeUser is a person, across every key they own and every project they
	// use: "Ayşe may spend at most this, wherever she spends it".
	ScopeUser Scope = "user"

	// The member scopes are a limit written on an org, team or project that
	// applies to each member separately: "each member of this project gets at
	// most this here". The limit is stored under the scope's own id; each
	// member's use is counted under that id and theirs, see MemberOwner.
	ScopeOrgMember     Scope = "org-member"
	ScopeTeamMember    Scope = "team-member"
	ScopeProjectMember Scope = "project-member"
)

// MemberScope is the per-member scope for a level of the tree, or "".
func MemberScope(level Scope) Scope {
	switch level {
	case ScopeOrg:
		return ScopeOrgMember
	case ScopeTeam:
		return ScopeTeamMember
	case ScopeProject:
		return ScopeProjectMember
	}
	return ""
}

// MemberOwner is the counter one member's use of a per-member limit goes to.
func MemberOwner(scope Scope, scopeID, userID string) LimitOwner {
	return LimitOwner{Scope: scope, ID: scopeID + "/" + userID}
}

// Org, Team and Project are the tenancy tree.
type Org struct {
	ID, Name string
}

// Team belongs to an organization.
type Team struct {
	ID, OrgID, Name string
}

// Project belongs to a team and owns keys.
type Project struct {
	ID, TeamID, Name string
}

// Key is an API key as stored: never the key itself.
type Key struct {
	ID string
	// UserID owns the key; empty for one an operator issued in a project.
	UserID string
	// Hash is sha256 of the key, and the only copy of anything derived from it.
	Hash []byte
	// Prefix is shown in listings so a key can be identified without revealing
	// it. It is never used for authentication.
	Prefix    string
	Name      string
	ProjectID string
	// Band and RoutingProfile are empty when the key inherits them.
	Band           string
	RoutingProfile string

	CreatedAt time.Time
	// ExpiresAt and RevokedAt are zero when unset. Revocation is a soft delete:
	// usage rows reference the key, and an audit trail that loses its subject
	// is not an audit trail.
	ExpiresAt  time.Time
	RevokedAt  time.Time
	LastUsedAt time.Time
}

// Usable reports whether a key may authenticate at this instant.
func (k Key) Usable(now time.Time) bool {
	if !k.RevokedAt.IsZero() {
		return false
	}
	return k.ExpiresAt.IsZero() || now.Before(k.ExpiresAt)
}

// Limits are the caps at one scope. A nil field means "no limit here", which is
// different from zero: zero forbids everything.
type Limits struct {
	Scope   Scope
	ScopeID string

	RPM *int
	TPM *int
	// Budget is spend per period, in micro-units — integers rather than floats
	// because a budget is accumulated by repeated addition, and repeated
	// addition of floats drifts. What a unit means is the operator's choice;
	// self-hosted models have no per-token price (plan §7.9.4).
	Budget       *int64
	BudgetPeriod Period
	// AllowedModels nil means every model.
	AllowedModels []string
}

// Period is how often a budget resets.
type Period string

// The supported budget periods.
const (
	PeriodNone  Period = ""
	PeriodDay   Period = "day"
	PeriodMonth Period = "month"
)

// Effective is a key's resolved limits, one value per field.
//
// Resolution takes the most restrictive value along key → project → team → org.
// A team cannot raise a limit its organization set, which is the only
// composition rule under which a budget means anything.
type Effective struct {
	KeyID string
	// UserID is the key's owner, empty for a key an operator issued in a
	// project. A personal key has an owner and no project.
	UserID    string
	ProjectID string
	TeamID    string
	OrgID     string

	Band           string
	RoutingProfile string

	RPM          *int
	TPM          *int
	Budget       *int64
	BudgetPeriod Period
	// AllowedModels is the intersection along the chain. nil means every model.
	AllowedModels []string

	// RPMOwner, TPMOwner and BudgetOwner say which scope set the value that
	// won, so consumption can be counted against that scope rather than
	// against each key under it.
	//
	// Without this the tree hands down values but never pools spending: a
	// project given a budget of 1000 and ten keys permits ten thousand, because
	// each key gets its own counter holding the project's number. That is "each
	// member may spend 1000", which is a useful thing to be able to say and is
	// not what somebody writing a budget on a project means.
	//
	// Each limit is owned separately because they resolve separately: a budget
	// can come from the team while the rate limit comes from the key.
	RPMOwner    LimitOwner
	TPMOwner    LimitOwner
	BudgetOwner LimitOwner

	// Chain is every level that set a limit, widest first.
	//
	// The collapsed values above are the ceiling a single caller faces, and
	// they are what a header or a profile page should show. They are not enough
	// to enforce with: once a quota is a pool, two levels are two pools and both
	// have to hold. An organization given 1000 and four teams given 500 each
	// would otherwise spend 2000 — the team's value wins as the more
	// restrictive, and the organization's is never looked at again.
	Chain []Limits
}

// LimitOwner is the scope a resolved limit came from.
type LimitOwner struct {
	Scope Scope
	ID    string
}

// CounterKey is what a usage counter for this limit should be keyed by.
//
// Empty when no scope set the limit, which means nothing is being counted and
// the caller should not be asking. Scope and id together, because ids are only
// unique within a level and a project could otherwise share a counter with a
// team that happened to be named the same.
func (o LimitOwner) CounterKey() string {
	if o.ID == "" {
		return ""
	}
	return string(o.Scope) + ":" + o.ID
}

// Allows reports whether a model may be served to this key.
func (e Effective) Allows(model string) bool {
	if e.AllowedModels == nil {
		return true
	}
	for _, m := range e.AllowedModels {
		if m == model {
			return true
		}
	}
	return false
}

// Usage is one request's accounting record.
type Usage struct {
	At    time.Time
	KeyID string
	// UserID is who spent it, empty for a key with no owner. Recorded with the
	// spend rather than looked up from the key later, so a key's past is not
	// moved onto whoever owns it next.
	UserID string
	// ProjectID is denormalized so reporting does not join on every row, and so
	// a row stays meaningful if the key is later revoked.
	ProjectID string
	Model     string
	// ServedModel differs from Model when a fallback served the request.
	ServedModel      string
	PromptTokens     int64
	CompletionTokens int64
	// CostMicros is what this request spent against the budget.
	CostMicros int64
	Status     int
	LatencyMs  int64

	// RequestID is the id the client sent in X-Request-Id or the gateway made.
	RequestID string
	// Replica answered the request; empty when none did.
	Replica string
	// QueueMs is time spent waiting in admission; TTFTMs is time to the first
	// token, 0 when none came.
	QueueMs int64
	TTFTMs  int64

	// GuardAction is what a guard did — flagged, redacted, blocked — empty
	// when it let the request through untouched or there is none;
	// GuardPolicy and GuardCategory what decided it, GuardMs how long
	// judging took. Never the text.
	GuardAction   string
	GuardPolicy   string
	GuardCategory string
	GuardMs       int64

	// EndUser is OpenAI's user field as the client sent it: its own name
	// for whoever it serves, not an account here. Tags are its metadata
	// and x-llmgw-tags. Both for attributing spend inside a project.
	EndUser string
	Tags    map[string]string
}

// tagsJSON is how tags are kept: a JSON object, keys sorted, or empty.
func tagsJSON(t map[string]string) string {
	if len(t) == 0 {
		return ""
	}
	b, _ := json.Marshal(t) // encoding/json sorts map keys
	return string(b)
}

// parseTags reads tagsJSON's form back.
func parseTags(s string) map[string]string {
	if s == "" {
		return nil
	}
	var t map[string]string
	if json.Unmarshal([]byte(s), &t) != nil {
		return nil
	}
	return t
}

// Store is the persistent half of tenancy and accounting.
//
// Nothing here is called on the request path. Authentication reads a cache that
// LookupKey fills on a miss, and usage is written in batches (plan §7.9.3).
type Store interface {
	// Tenancy. Put is an upsert so applying a configuration twice is a no-op.
	PutOrg(ctx context.Context, o Org) error
	PutTeam(ctx context.Context, t Team) error
	PutProject(ctx context.Context, p Project) error
	Orgs(ctx context.Context) ([]Org, error)
	Teams(ctx context.Context, orgID string) ([]Team, error)
	Projects(ctx context.Context, teamID string) ([]Project, error)

	// Keys.
	CreateKey(ctx context.Context, k Key) error
	KeyByHash(ctx context.Context, hash []byte) (Key, error)
	Keys(ctx context.Context, projectID string) ([]Key, error)
	RevokeKey(ctx context.Context, id string, at time.Time) error
	TouchKeys(ctx context.Context, seen map[string]time.Time) error

	// Limits.
	PutLimits(ctx context.Context, l Limits) error
	LimitsFor(ctx context.Context, scope Scope, id string) (Limits, error)
	// Resolve walks a key's chain and returns the most restrictive of each.
	Resolve(ctx context.Context, keyID string) (Effective, error)

	// Accounting.
	WriteUsage(ctx context.Context, records []Usage) error
	// SpendSince totals cost per key since a moment, so an in-memory budget
	// counter can be rebuilt after a restart. A restart must not reset a
	// budget, which it would if the counter lived only in memory.
	SpendSince(ctx context.Context, since time.Time) (map[string]int64, error)
	// SpendSinceByScope is the same totals with each key's place in the tree,
	// so a budget owned by a team can be rebuilt by summing every key under it.
	// Per key alone cannot do that: the counter is one number for the whole
	// team, and the rows that made it up are spread across its projects.
	SpendSinceByScope(ctx context.Context, since time.Time) ([]ScopeSpend, error)
	// UsageSince summarizes activity per key. A reporting query: it scans and
	// groups, so it belongs behind a dashboard refresh rather than a request.
	UsageSince(ctx context.Context, since time.Time) ([]KeyUsage, error)

	Close() error
}

// KeyUsage summarizes one key's activity in a period, for the dashboard and
// `ctl`. It is a reporting query and never runs on the request path.
type KeyUsage struct {
	KeyID     string
	KeyName   string
	KeyPrefix string
	ProjectID string
	// Who spent it and where, as names: the person's email, and the
	// project's path. Empty for a key with no owner, and for personal use.
	Email       string
	ProjectPath string
	Requests    int64
	Prompt      int64
	Completion  int64
	CostMicros  int64
	Errors      int64
	LastAt      time.Time
}

// Chat reports whether this is the chat page's traffic, which has no key.
func (k KeyUsage) Chat() bool { return strings.HasPrefix(k.KeyID, "session:") }

// Bucket is one time slice of usage, for the dashboard's charts.
type Bucket struct {
	Start      time.Time
	Requests   int64
	Errors     int64
	Prompt     int64
	Completion int64
	CostMicros int64
	// LatencyMs is the mean over the bucket. A mean rather than a percentile
	// because the raw rows are already gone by rollup time; this is a shape, not
	// an SLO measurement.
	LatencyMs int64
}

// ScopeSpend is one key's spend in a period, with where that key sits in the
// tenancy tree.
//
// Returned rather than pre-aggregated per scope because the roll-up depends on
// which scopes own budgets, which the store does not know and the accountant
// does. One query, and the folding happens where it can be tested.
type ScopeSpend struct {
	KeyID     string
	UserID    string
	ProjectID string
	TeamID    string
	OrgID     string
	Micros    int64
}