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"
	"errors"
	"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"
)

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

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

// 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)
	// 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
	Requests   int64
	Prompt     int64
	Completion int64
	CostMicros int64
	Errors     int64
	LastAt     time.Time
}

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