Help us shape the future of Go! Take the 2026 Go Developer Survey

store

package
v0.15.2 Latest Latest
Warning

This package is not in the latest version of its module.

 Go to latest Published: Oct 26, 2026 License: BSD-3-Clause Imports: 19 Imported by: 0

Documentation

Overview

Package store holds the gateway's persistent state: the tenant hierarchy (organisation → team → project), virtual keys (hashed), budgets and per-request usage records. Two implementations share one contract test: Memory (tests, trials; state is lost on restart) and Postgres (pgx).

The hierarchy mirrors the Slurm account tree (sacctmgr: accounts nested under a parent account, users associated with accounts): an organisation is an institution, a team a lab, a project a piece of work inside a team. Keys belong to a team and optionally to one of its projects.

Index

Constants

View Source 
const AdminTokenPrefix = "llma-"

AdminTokenPrefix starts every user admin token.

View Source 
const GPUHours = "gpu_hours"

GPUHours is the unit of the costs of vLLM models, as the API and the reports spell it; it is stored as "".

View Source 
const KeyPrefix = "llmf-"

KeyPrefix starts every virtual key, so a leaked key is recognisable by secret scanners.

View Source 
const MaxReportSpan = 366 * 24 * time.Hour

MaxReportSpan bounds the period of one report: a report reads every record of its period.

Variables

View Source 
var (
	ErrNotFound = errors.New("not found")
	ErrExists   = errors.New("already exists")
	// ErrConflict reports an existing object that differs from the
	// request (e.g. a team that already belongs to another organisation).
	ErrConflict = errors.New("conflicts with an existing object")
)

Errors.

View Source 
var ErrAuditUpdate = errors.New("the audit log is append-only")

ErrAuditUpdate is what the PostgreSQL trigger raises on an UPDATE.

View Source 
var ErrFenced = errors.New("store: fenced: the lease is held by a newer owner or has expired")

ErrFenced: the write came from an owner whose lease is no longer the current one (expired, or taken over with a higher epoch).

View Source 
var ErrFull = errors.New("store: too many entries")

ErrFull: a bound on stored logins or sessions was reached.

View Source 
var ReportDims = []string{"org", "lab", "team", "project", "user", "key", "model", "pool", "tag", "day"}

ReportDims lists the accepted dimensions, in canonical order: the hierarchy org → lab → team → project, then user, key, model, pool, tag, and day.

Functions

func GenerateAdminToken added in v0.2.0 

func GenerateAdminToken() (secret, id string, hash [32]byte)

GenerateAdminToken returns a new admin token secret (256 random bits), its ID and its SHA-256.

func GenerateKey 

func GenerateKey() (secret, id string, hash [32]byte)

GenerateKey returns a new secret (256 random bits), its public ID and its SHA-256. The secret is shown once to the caller and never stored. An unsalted hash is enough for a uniformly random 256-bit secret: there is no dictionary to precompute.

func HashEqual 

func HashEqual(a, b [32]byte) bool

HashEqual compares two hashes in constant time.

func HashSecret 

func HashSecret(secret string) [32]byte

HashSecret hashes a presented key.

func NewKeySecret added in v0.4.0 

func NewKeySecret() (secret string, hash [32]byte)

NewKeySecret makes a new secret for an existing key (rotation).

func NormalizeUnit added in v0.12.0 

func NormalizeUnit(u string) (string, error)

NormalizeUnit maps the spelled-out GPU-hours unit to "" and checks a currency unit: 1 to 16 letters, digits or underscores (USD, EUR, usd_cents…). Units are compared exactly.

func UnitName added in v0.12.0 

func UnitName(u string) string

UnitName is the unit as the API and the reports show it.

func ValidateBudget 

func ValidateBudget(b Budget) error

ValidateBudget checks a budget before it is stored.

func ValidateDims added in v0.2.0 

func ValidateDims(dims []string) ([]string, error)

ValidateDims checks and orders report dimensions.

Types

type AllowKey added in v0.4.0 

type AllowKey struct {
	Kind ScopeKind `json:"kind"` // ScopeOrg, ScopeTeam or ScopeKey
	Name string    `json:"name"`
}

AllowKey names what an IP allowlist applies to: an organisation, a team or a key (by ID).

type Allowlist added in v0.4.0 

type Allowlist struct {
	AllowKey
	CIDRs []string `json:"cidrs"`
}

Allowlist is the set of client networks allowed for one scope.

func ValidateAllowlist added in v0.4.0 

func ValidateAllowlist(a Allowlist) (Allowlist, error)

ValidateAllowlist checks the kind and canonicalises the prefixes.

type AuditCheckpoint added in v0.5.0 

type AuditCheckpoint struct {
	Seq       int64
	Hash      []byte
	Time      time.Time
	KeyID     string
	Signature []byte
}

AuditCheckpoint signs the chain up to Seq, whose entry hash is Hash.

type AuditEntry added in v0.5.0 

type AuditEntry struct {
	// V is the canonical form the hash was computed with (package audit).
	V         int
	Seq       int64
	Time      time.Time // microsecond precision
	Kind      string    // admin, auth, scim, system, guardrail, tool
	Actor     string
	Roles     string // JSON list of the actor's role bindings
	Action    string
	Target    string
	Before    string
	After     string
	Status    int
	SourceIP  string
	RequestID string
	PrevHash  []byte
	Hash      []byte
}

AuditEntry is one recorded event. Before and After are JSON, kept as the exact bytes that were hashed.

type AuditHasher added in v0.5.0 

type AuditHasher func(e *AuditEntry) []byte

AuditHasher computes an entry's hash once its Seq and PrevHash are known.

type Binding added in v0.2.0 

type Binding struct {
	User string `json:"user"`
	Role Role   `json:"role"`
	Org  string `json:"org,omitempty"`
	Team string `json:"team,omitempty"`
}

Binding grants a role on a scope. Exactly the fields the role needs are set: none for proxy roles, Org for org_admin (and viewer of an org), Team for team_admin, member (and viewer of a team).

func (Binding) Validate added in v0.2.0 

func (b Binding) Validate() error

Validate checks that a binding names the scope its role needs.

type Budget 

type Budget struct {
	BudgetKey
	Period       Period     `json:"period"`
	MaxCost      float64    `json:"max_cost"`
	SoftLimit    float64    `json:"soft_limit,omitempty"`
	TempIncrease float64    `json:"temp_increase,omitempty"`
	TempExpiry   *time.Time `json:"temp_expiry,omitempty"`
}

Budget caps the cost of a scope per period, in its unit. A temporary increase adds to MaxCost until it expires; SoftLimit (0 = none) triggers one alert per period when the spend reaches it.

func (Budget) Limit added in v0.2.0 

func (b Budget) Limit(now time.Time) float64

Limit returns the effective maximum at now.

type BudgetKey added in v0.2.0 

type BudgetKey struct {
	Kind  ScopeKind `json:"kind"`
	Name  string    `json:"name"`
	Model string    `json:"model,omitempty"`
	Unit  string    `json:"unit,omitempty"`
}

BudgetKey identifies a budget: a scope, optionally restricted to one model group (Model "" covers every model), in one cost unit (Unit "" is GPU-hours; "USD" or another currency code caps what an external provider charges). A budget counts only the costs of its own unit: GPU-hours and dollars are never added together.

func (BudgetKey) String added in v0.2.0 

func (k BudgetKey) String() string

type Claim added in v0.15.2 

type Claim struct {
	Pool, Owner, Addr string
	// Held is the epoch the claimant owns (0: none).
	Held int64
	TTL  time.Duration
	// MayTake: the claimant may take an expired lease (it knows the
	// pool's membership); MayTakeYielded: it may take a live lease whose
	// owner yields (its own view of the membership is fresh).
	MayTake, MayTakeYielded bool
	// Yield, Capacity and Reserved are written with the claimant's lease.
	Yield    bool
	Capacity int
	Reserved string
}

Claim is one claim of a pool's lease.

type Duration added in v0.4.0 

type Duration time.Duration

Duration is a time.Duration written as a Go duration string in JSON.

func (Duration) MarshalJSON added in v0.4.0 

func (d Duration) MarshalJSON() ([]byte, error)

MarshalJSON writes "720h0m0s".

func (*Duration) UnmarshalJSON added in v0.4.0 

func (d *Duration) UnmarshalJSON(b []byte) error

UnmarshalJSON reads a Go duration string.

type HAStateStore added in v0.13.1 

type HAStateStore interface {
	// PutSealed stores sealed bytes under (kind, id) for ttl (by the
	// store's clock), unless a limit is reached (ErrFull).
	PutSealed(ctx context.Context, kind string, id [32]byte, net string, sealed []byte, ttl time.Duration, lim SealedLimits) error
	// GetSealed reads a live entry.
	GetSealed(ctx context.Context, kind string, id [32]byte) ([]byte, bool, error)
	// TakeSealed reads and deletes a live entry (one use).
	TakeSealed(ctx context.Context, kind string, id [32]byte) ([]byte, bool, error)
	DropSealed(ctx context.Context, kind string, id [32]byte) error
	// PurgeSealed deletes expired entries and old finished alerts.
	PurgeSealed(ctx context.Context) error

	// EnqueueAlert adds one alert per channel, once per key: a second
	// enqueue of the same (key, channel) changes nothing.
	EnqueueAlert(ctx context.Context, key string, channels []string, payload []byte) error
	// ClaimAlerts takes up to n due alerts and hides them for lease (a
	// replica that dies while sending leaves them to the next one).
	ClaimAlerts(ctx context.Context, n int, lease time.Duration) ([]OutboxAlert, error)
	// FinishAlert records a delivery attempt: sent, given up (failed), or
	// retried after the given delay.
	FinishAlert(ctx context.Context, key, channel, state string, retryIn time.Duration, lastErr string) error

	// NotifyKeys broadcasts a revocation to every replica.
	NotifyKeys(ctx context.Context, payload string) error
	// ListenKeys calls on for every revocation until ctx ends, and
	// connected each time it (re)connects (notifications sent while it was
	// not connected are lost, PostgreSQL does not keep them).
	ListenKeys(ctx context.Context, connected func(), on func(payload string))
}

HAStateStore is the shared state besides pool ownership.

type HAStore added in v0.13.1 

type HAStore interface {
	// ClaimLease renews the lease when the claimant holds it at epoch Held
	// and it has not expired; takes it (epoch + 1) when it has expired (or
	// does not exist) and MayTake, or when it is live but its owner yields
	// and MayTakeYielded (Quiet); otherwise changes nothing. It returns the
	// lease as it stands afterwards.
	ClaimLease(ctx context.Context, c Claim) (Lease, error)
	// ReleaseLease ends owner's lease at epoch now (graceful shutdown).
	ReleaseLease(ctx context.Context, pool, owner string, epoch int64) error
	// SaveUsage writes the owner's state, fenced: ErrFenced unless owner
	// holds the live lease at epoch, checked under a row lock that a
	// takeover must wait for.
	SaveUsage(ctx context.Context, pool, owner string, epoch int64, rows []PoolUsage) error
	// LoadUsage reads what the last owner persisted.
	LoadUsage(ctx context.Context, pool string) ([]PoolUsage, error)
}

HAStore is what package ha needs from the store. Only PostgreSQL has it.

type Key 

type Key struct {
	ID   string   `json:"id"`
	Hash [32]byte `json:"-"`
	// Org is resolved from the team when the key is read.
	Org     string `json:"org"`
	Team    string `json:"team"`
	Project string `json:"project,omitempty"`
	// Owner is the user who created the key as a team member; empty for
	// team keys created by an administrator.
	Owner string `json:"owner,omitempty"`
	// Tags label every request made with the key (budgets and reports).
	Tags        []string  `json:"tags,omitempty"`
	Name        string    `json:"name"`
	RPM         int       `json:"rpm"`
	TPM         int       `json:"tpm"`
	MaxInFlight int       `json:"max_inflight"`
	CreatedAt   time.Time `json:"created_at"`
	Revoked     bool      `json:"revoked"`
	// PrevHash is the secret replaced by the last rotation, valid until
	// PrevExpires (zero: no previous secret).
	PrevHash    [32]byte  `json:"-"`
	PrevExpires time.Time `json:"prev_expires_at,omitzero"`
	RotatedAt   time.Time `json:"rotated_at,omitzero"`
	// Rotation is the key's rotation schedule, if any.
	Rotation Rotation `json:"rotation,omitzero"`
	// Policy names the request policy of the key (empty: its team's).
	Policy string `json:"policy,omitempty"`
	// Guardrails names guardrails the key adds to those of its org and
	// team (it can add, never remove).
	Guardrails []string `json:"guardrails,omitempty"`
}

Key is a virtual key. The secret itself is never stored, only its SHA-256.

func (Key) Matches added in v0.4.0 

func (k Key) Matches(h [32]byte, now time.Time) bool

Matches reports whether h is the key's secret, or its previous secret still inside the rotation grace period at now.

type Lease added in v0.13.1 

type Lease struct {
	Pool, Owner, Addr string
	// Epoch is the fencing token: it grows by one at every change of
	// owner (or when an owner takes back a lease it let expire).
	Epoch int64
	// Live: the lease has not expired (by PostgreSQL's clock).
	Live bool
	// Yield: the owner asks to be replaced (it cannot read the pool's
	// membership any more, #75); Quiet: this claim took such a lease, so
	// the new owner must not decide for one lease duration.
	Yield, Quiet bool
	// Capacity and Reserved are what the owner published: the pool's C
	// and the integer reservations (JSON) it decides with at this epoch
	// (-1 and "": not published).
	Capacity int
	Reserved string
}

Lease is a pool's lease as PostgreSQL holds it.

type Memory 

type Memory struct {
	MaxRecords int
	// Fail, when set, is returned by every call (fault injection). Set it
	// with SetFail.
	Fail error
	// contains filtered or unexported fields
}

Memory is an in-process Store. Usage records are bounded by MaxRecords (oldest dropped first).

func NewMemory 

func NewMemory() *Memory

NewMemory returns an empty in-memory store.

func (*Memory) Allowlists added in v0.4.0 

func (m *Memory) Allowlists(context.Context) ([]Allowlist, error)

func (*Memory) AuditAddCheckpoint added in v0.5.0 

func (m *Memory) AuditAddCheckpoint(_ context.Context, c AuditCheckpoint) error

func (*Memory) AuditAppend added in v0.5.0 

func (m *Memory) AuditAppend(_ context.Context, e AuditEntry, hash AuditHasher) (AuditEntry, error)

func (*Memory) AuditCheckpoints added in v0.5.0 

func (m *Memory) AuditCheckpoints(context.Context) ([]AuditCheckpoint, error)

func (*Memory) AuditPrune added in v0.5.0 

func (m *Memory) AuditPrune(_ context.Context, through int64) (int64, error)

func (*Memory) AuditRange added in v0.5.0 

func (m *Memory) AuditRange(_ context.Context, from int64, limit int) ([]AuditEntry, error)

func (*Memory) Bindings added in v0.2.0 

func (m *Memory) Bindings(_ context.Context, user string) ([]Binding, error)

func (*Memory) Budgets 

func (m *Memory) Budgets(context.Context) ([]Budget, error)

func (*Memory) Close 

func (m *Memory) Close()

func (*Memory) CreateKey 

func (m *Memory) CreateKey(_ context.Context, k Key) error

func (*Memory) CreateOrg added in v0.2.0 

func (m *Memory) CreateOrg(_ context.Context, name string) error

func (*Memory) CreateProject added in v0.2.0 

func (m *Memory) CreateProject(_ context.Context, team, name string) error

func (*Memory) CreateTeam 

func (m *Memory) CreateTeam(_ context.Context, name, org string) error

func (*Memory) CreateUser added in v0.2.0 

func (m *Memory) CreateUser(_ context.Context, name string) error

func (*Memory) CreateUserToken added in v0.2.0 

func (m *Memory) CreateUserToken(_ context.Context, t UserToken) error

func (*Memory) DeleteAudit added in v0.5.0 

func (m *Memory) DeleteAudit(seq int64)

DeleteAudit removes one stored entry, bypassing everything: for tests that must show the verifier catches it.

func (*Memory) DeleteBudget 

func (m *Memory) DeleteBudget(_ context.Context, k BudgetKey) error

func (*Memory) DisableOrg added in v0.2.0 

func (m *Memory) DisableOrg(_ context.Context, name string) error

func (*Memory) DisableProject added in v0.2.0 

func (m *Memory) DisableProject(_ context.Context, team, name string) error

func (*Memory) DisableTeam 

func (m *Memory) DisableTeam(_ context.Context, name string) error

func (*Memory) DisableUser added in v0.2.0 

func (m *Memory) DisableUser(_ context.Context, name string) error

func (*Memory) DueKeys added in v0.4.0 

func (m *Memory) DueKeys(_ context.Context, now time.Time) ([]Key, error)

func (*Memory) GetKey added in v0.2.0 

func (m *Memory) GetKey(_ context.Context, id string) (Key, error)

func (*Memory) GetTeam added in v0.2.0 

func (m *Memory) GetTeam(_ context.Context, name string) (Team, error)

func (*Memory) GetUser added in v0.2.0 

func (m *Memory) GetUser(_ context.Context, name string) (User, error)

func (*Memory) Grant added in v0.2.0 

func (m *Memory) Grant(_ context.Context, b Binding) error

func (*Memory) IncreaseBudget added in v0.2.0 

func (m *Memory) IncreaseBudget(_ context.Context, k BudgetKey, amount float64, expiry time.Time) error

func (*Memory) InsertRecords 

func (m *Memory) InsertRecords(_ context.Context, recs []Record) error

func (*Memory) KeyByHash 

func (m *Memory) KeyByHash(_ context.Context, h [32]byte) (Key, error)

func (*Memory) ListKeys 

func (m *Memory) ListKeys(_ context.Context, team string) ([]Key, error)

func (*Memory) ListOrgs added in v0.2.0 

func (m *Memory) ListOrgs(context.Context) ([]Org, error)

func (*Memory) ListProjects added in v0.2.0 

func (m *Memory) ListProjects(_ context.Context, team string) ([]Project, error)

func (*Memory) ListTeams 

func (m *Memory) ListTeams(_ context.Context, org string) ([]Team, error)

func (*Memory) ListUserTokens added in v0.2.0 

func (m *Memory) ListUserTokens(_ context.Context, user string) ([]UserToken, error)

func (*Memory) ListUsers added in v0.2.0 

func (m *Memory) ListUsers(context.Context) ([]User, error)

func (*Memory) MarkSoftAlert added in v0.2.0 

func (m *Memory) MarkSoftAlert(_ context.Context, k BudgetKey, period time.Time) (bool, error)

func (*Memory) Ping 

func (m *Memory) Ping(context.Context) error

func (*Memory) Records added in v0.2.0 

func (m *Memory) Records() []Record

Records returns a copy of the stored usage records (tests, trials).

func (*Memory) Report added in v0.2.0 

func (m *Memory) Report(ctx context.Context, from, to time.Time, dims []string) ([]ReportRow, error)

func (*Memory) RevokeKey 

func (m *Memory) RevokeKey(_ context.Context, id string) error

func (*Memory) RevokeUserToken added in v0.2.0 

func (m *Memory) RevokeUserToken(_ context.Context, user, id string) error

func (*Memory) RotateKey added in v0.4.0 

func (m *Memory) RotateKey(_ context.Context, id string, newHash, expect [32]byte, prevExpires, at time.Time) error

func (*Memory) SCIMCreate added in v0.4.0 

func (m *Memory) SCIMCreate(_ context.Context, d SCIMDoc) error

func (*Memory) SCIMDelete added in v0.4.0 

func (m *Memory) SCIMDelete(_ context.Context, kind, id string) error

func (*Memory) SCIMGet added in v0.4.0 

func (m *Memory) SCIMGet(_ context.Context, kind, id string) (SCIMDoc, error)

func (*Memory) SCIMList added in v0.4.0 

func (m *Memory) SCIMList(_ context.Context, kind, name string) ([]SCIMDoc, error)

func (*Memory) SCIMReplace added in v0.4.0 

func (m *Memory) SCIMReplace(_ context.Context, d SCIMDoc, ifVersion string) error

func (*Memory) SetAllowlist added in v0.4.0 

func (m *Memory) SetAllowlist(_ context.Context, a Allowlist) error

func (*Memory) SetBudget 

func (m *Memory) SetBudget(_ context.Context, b Budget) error

func (*Memory) SetFail 

func (m *Memory) SetFail(err error)

SetFail makes every later call return err (nil restores the store).

func (*Memory) SetKeyRotation added in v0.4.0 

func (m *Memory) SetKeyRotation(_ context.Context, id string, r Rotation) error

func (*Memory) SetTeamModels added in v0.2.0 

func (m *Memory) SetTeamModels(_ context.Context, team string, models []string) error

func (*Memory) SetTeamTools added in v0.8.0 

func (m *Memory) SetTeamTools(_ context.Context, team string, tools []string) error

func (*Memory) SpendRows added in v0.2.0 

func (m *Memory) SpendRows(_ context.Context, since time.Time) ([]SpendRow, error)

func (*Memory) TamperAudit added in v0.5.0 

func (m *Memory) TamperAudit(seq int64, f func(*AuditEntry))

TamperAudit rewrites one stored entry in place, bypassing the chain: for tests that must show the verifier catches it.

func (*Memory) TeamModels added in v0.2.0 

func (m *Memory) TeamModels(context.Context) (map[string][]string, error)

func (*Memory) TeamTools added in v0.8.0 

func (m *Memory) TeamTools(context.Context) (map[string][]string, error)

func (*Memory) Ungrant added in v0.2.0 

func (m *Memory) Ungrant(_ context.Context, b Binding) error

func (*Memory) UserByTokenHash added in v0.2.0 

func (m *Memory) UserByTokenHash(_ context.Context, h [32]byte) (User, UserToken, error)

type Org added in v0.2.0 

type Org struct {
	Name      string    `json:"name"`
	CreatedAt time.Time `json:"created_at"`
	Disabled  bool      `json:"disabled"`
}

Org is an organisation (an institution), the top of the hierarchy.

type OutboxAlert added in v0.13.1 

type OutboxAlert struct {
	Key, Channel string
	Payload      []byte
	Attempts     int
}

OutboxAlert is a claimed alert.

type Period 

type Period string

Period is a budget period, aligned on UTC calendar boundaries. 

const (
	Day   Period = "day"
	Week  Period = "week"
	Month Period = "month"
)

Periods.

func (Period) Start 

func (p Period) Start(now time.Time) time.Time

Start returns the start of the period containing now (UTC; weeks start on Monday).

func (Period) Valid 

func (p Period) Valid() bool

Valid reports whether p is a known period.

type PoolUsage added in v0.13.1 

type PoolUsage struct {
	Org      string
	Usage    float64
	InFlight int
	// Age is how long ago it was written, by PostgreSQL's clock.
	Age time.Duration
}

PoolUsage is one organisation's fair-share state in a pool, as the owner persisted it.

type Postgres 

type Postgres struct {
	// contains filtered or unexported fields
}

Postgres is a Store on PostgreSQL through pgx.

func OpenPostgres 

func OpenPostgres(ctx context.Context, dsn string) (*Postgres, error)

OpenPostgres connects and creates or upgrades the schema. The error never contains the DSN.

func (*Postgres) Allowlists added in v0.4.0 

func (p *Postgres) Allowlists(ctx context.Context) ([]Allowlist, error)

func (*Postgres) AuditAddCheckpoint added in v0.5.0 

func (p *Postgres) AuditAddCheckpoint(ctx context.Context, c AuditCheckpoint) error

func (*Postgres) AuditAppend added in v0.5.0 

func (p *Postgres) AuditAppend(ctx context.Context, e AuditEntry, hash AuditHasher) (AuditEntry, error)

func (*Postgres) AuditCheckpoints added in v0.5.0 

func (p *Postgres) AuditCheckpoints(ctx context.Context) ([]AuditCheckpoint, error)

func (*Postgres) AuditPrune added in v0.5.0 

func (p *Postgres) AuditPrune(ctx context.Context, through int64) (int64, error)

func (*Postgres) AuditRange added in v0.5.0 

func (p *Postgres) AuditRange(ctx context.Context, from int64, limit int) ([]AuditEntry, error)

func (*Postgres) Bindings added in v0.2.0 

func (p *Postgres) Bindings(ctx context.Context, user string) ([]Binding, error)

func (*Postgres) Budgets 

func (p *Postgres) Budgets(ctx context.Context) ([]Budget, error)

func (*Postgres) ClaimAlerts added in v0.13.1 

func (p *Postgres) ClaimAlerts(ctx context.Context, n int, lease time.Duration) ([]OutboxAlert, error)

func (*Postgres) ClaimLease added in v0.13.1 

func (p *Postgres) ClaimLease(ctx context.Context, c Claim) (Lease, error)

func (*Postgres) Close 

func (p *Postgres) Close()

func (*Postgres) CreateKey 

func (p *Postgres) CreateKey(ctx context.Context, k Key) error

func (*Postgres) CreateOrg added in v0.2.0 

func (p *Postgres) CreateOrg(ctx context.Context, name string) error

func (*Postgres) CreateProject added in v0.2.0 

func (p *Postgres) CreateProject(ctx context.Context, team, name string) error

func (*Postgres) CreateTeam 

func (p *Postgres) CreateTeam(ctx context.Context, name, org string) error

func (*Postgres) CreateUser added in v0.2.0 

func (p *Postgres) CreateUser(ctx context.Context, name string) error

func (*Postgres) CreateUserToken added in v0.2.0 

func (p *Postgres) CreateUserToken(ctx context.Context, t UserToken) error

func (*Postgres) DeleteBudget 

func (p *Postgres) DeleteBudget(ctx context.Context, k BudgetKey) error

func (*Postgres) DisableOrg added in v0.2.0 

func (p *Postgres) DisableOrg(ctx context.Context, name string) error

func (*Postgres) DisableProject added in v0.2.0 

func (p *Postgres) DisableProject(ctx context.Context, team, name string) error

func (*Postgres) DisableTeam 

func (p *Postgres) DisableTeam(ctx context.Context, name string) error

func (*Postgres) DisableUser added in v0.2.0 

func (p *Postgres) DisableUser(ctx context.Context, name string) error

func (*Postgres) DropSealed added in v0.13.1 

func (p *Postgres) DropSealed(ctx context.Context, kind string, id [32]byte) error

func (*Postgres) DueKeys added in v0.4.0 

func (p *Postgres) DueKeys(ctx context.Context, now time.Time) ([]Key, error)

func (*Postgres) EnqueueAlert added in v0.13.1 

func (p *Postgres) EnqueueAlert(ctx context.Context, key string, channels []string, payload []byte) error

func (*Postgres) FinishAlert added in v0.13.1 

func (p *Postgres) FinishAlert(ctx context.Context, key, channel, state string, retryIn time.Duration, lastErr string) error

func (*Postgres) GetKey added in v0.2.0 

func (p *Postgres) GetKey(ctx context.Context, id string) (Key, error)

func (*Postgres) GetSealed added in v0.13.1 

func (p *Postgres) GetSealed(ctx context.Context, kind string, id [32]byte) ([]byte, bool, error)

func (*Postgres) GetTeam added in v0.2.0 

func (p *Postgres) GetTeam(ctx context.Context, name string) (Team, error)

func (*Postgres) GetUser added in v0.2.0 

func (p *Postgres) GetUser(ctx context.Context, name string) (User, error)

func (*Postgres) Grant added in v0.2.0 

func (p *Postgres) Grant(ctx context.Context, b Binding) error

func (*Postgres) IncreaseBudget added in v0.2.0 

func (p *Postgres) IncreaseBudget(ctx context.Context, k BudgetKey, amount float64, expiry time.Time) error

func (*Postgres) InsertRecords 

func (p *Postgres) InsertRecords(ctx context.Context, recs []Record) error

InsertRecords writes a batch: COPY into a temporary table, then one INSERT that skips the records whose ID is already stored, so a batch replayed from the accounting spool is never billed twice.

func (*Postgres) KeyByHash 

func (p *Postgres) KeyByHash(ctx context.Context, h [32]byte) (Key, error)

func (*Postgres) ListKeys 

func (p *Postgres) ListKeys(ctx context.Context, team string) ([]Key, error)

func (*Postgres) ListOrgs added in v0.2.0 

func (p *Postgres) ListOrgs(ctx context.Context) ([]Org, error)

func (*Postgres) ListProjects added in v0.2.0 

func (p *Postgres) ListProjects(ctx context.Context, team string) ([]Project, error)

func (*Postgres) ListTeams 

func (p *Postgres) ListTeams(ctx context.Context, org string) ([]Team, error)

func (*Postgres) ListUserTokens added in v0.2.0 

func (p *Postgres) ListUserTokens(ctx context.Context, user string) ([]UserToken, error)

func (*Postgres) ListUsers added in v0.2.0 

func (p *Postgres) ListUsers(ctx context.Context) ([]User, error)

func (*Postgres) ListenKeys added in v0.13.1 

func (p *Postgres) ListenKeys(ctx context.Context, connected func(), on func(payload string))

func (*Postgres) LoadUsage added in v0.13.1 

func (p *Postgres) LoadUsage(ctx context.Context, pool string) ([]PoolUsage, error)

func (*Postgres) MarkSoftAlert added in v0.2.0 

func (p *Postgres) MarkSoftAlert(ctx context.Context, k BudgetKey, period time.Time) (bool, error)

func (*Postgres) NotifyKeys added in v0.13.1 

func (p *Postgres) NotifyKeys(ctx context.Context, payload string) error

func (*Postgres) Ping 

func (p *Postgres) Ping(ctx context.Context) error

func (*Postgres) PurgeSealed added in v0.13.1 

func (p *Postgres) PurgeSealed(ctx context.Context) error

func (*Postgres) PutSealed added in v0.13.1 

func (p *Postgres) PutSealed(ctx context.Context, kind string, id [32]byte, net string, sealed []byte, ttl time.Duration, lim SealedLimits) error

func (*Postgres) ReleaseLease added in v0.13.1 

func (p *Postgres) ReleaseLease(ctx context.Context, pool, owner string, epoch int64) error

func (*Postgres) Report added in v0.2.0 

func (p *Postgres) Report(ctx context.Context, from, to time.Time, dims []string) ([]ReportRow, error)

func (*Postgres) RevokeKey 

func (p *Postgres) RevokeKey(ctx context.Context, id string) error

func (*Postgres) RevokeUserToken added in v0.2.0 

func (p *Postgres) RevokeUserToken(ctx context.Context, user, id string) error

func (*Postgres) RotateKey added in v0.4.0 

func (p *Postgres) RotateKey(ctx context.Context, id string, newHash, expect [32]byte, prevExpires, at time.Time) error

func (*Postgres) SCIMCreate added in v0.4.0 

func (p *Postgres) SCIMCreate(ctx context.Context, d SCIMDoc) error

func (*Postgres) SCIMDelete added in v0.4.0 

func (p *Postgres) SCIMDelete(ctx context.Context, kind, id string) error

func (*Postgres) SCIMGet added in v0.4.0 

func (p *Postgres) SCIMGet(ctx context.Context, kind, id string) (SCIMDoc, error)

func (*Postgres) SCIMList added in v0.4.0 

func (p *Postgres) SCIMList(ctx context.Context, kind, name string) ([]SCIMDoc, error)

func (*Postgres) SCIMReplace added in v0.4.0 

func (p *Postgres) SCIMReplace(ctx context.Context, d SCIMDoc, ifVersion string) error

func (*Postgres) SaveUsage added in v0.13.1 

func (p *Postgres) SaveUsage(ctx context.Context, pool, owner string, epoch int64, rows []PoolUsage) error

func (*Postgres) SetAllowlist added in v0.4.0 

func (p *Postgres) SetAllowlist(ctx context.Context, a Allowlist) error

func (*Postgres) SetBudget 

func (p *Postgres) SetBudget(ctx context.Context, b Budget) error

func (*Postgres) SetKeyRotation added in v0.4.0 

func (p *Postgres) SetKeyRotation(ctx context.Context, id string, r Rotation) error

func (*Postgres) SetTeamModels added in v0.2.0 

func (p *Postgres) SetTeamModels(ctx context.Context, team string, models []string) error

func (*Postgres) SetTeamTools added in v0.8.0 

func (p *Postgres) SetTeamTools(ctx context.Context, team string, tools []string) error

func (*Postgres) SpendRows added in v0.2.0 

func (p *Postgres) SpendRows(ctx context.Context, since time.Time) ([]SpendRow, error)

func (*Postgres) TakeSealed added in v0.13.1 

func (p *Postgres) TakeSealed(ctx context.Context, kind string, id [32]byte) ([]byte, bool, error)

func (*Postgres) TeamModels added in v0.2.0 

func (p *Postgres) TeamModels(ctx context.Context) (map[string][]string, error)

func (*Postgres) TeamTools added in v0.8.0 

func (p *Postgres) TeamTools(ctx context.Context) (map[string][]string, error)

func (*Postgres) Ungrant added in v0.2.0 

func (p *Postgres) Ungrant(ctx context.Context, b Binding) error

func (*Postgres) UsageRecordIDs added in v0.10.0 

func (p *Postgres) UsageRecordIDs(ctx context.Context) ([]string, error)

UsageRecordIDs lists the record id of every usage row (empty for rows written before ids), for verifying that nothing was billed twice.

func (*Postgres) UserByTokenHash added in v0.2.0 

func (p *Postgres) UserByTokenHash(ctx context.Context, h [32]byte) (User, UserToken, error)

type Project added in v0.2.0 

type Project struct {
	Team      string    `json:"team"`
	Name      string    `json:"name"`
	CreatedAt time.Time `json:"created_at"`
	Disabled  bool      `json:"disabled"`
}

Project is a unit of work inside a team, with its own keys and budget.

type Record 

type Record struct {
	// ID identifies the record: inserting a record whose ID is already
	// stored does nothing, so the accounting spool can replay a batch that
	// may have been written before a crash without billing it twice. Empty
	// for records written before v0.10.0.
	ID      string
	Time    time.Time
	Org     string
	Team    string
	Project string
	KeyID   string
	// User is the gateway user the request is attributed to: the key's
	// owner, or a JWT's username; "" for a service key. Never a token.
	User string
	// Lab is the team's lab at the time of the request ("" for none).
	Lab              string
	Model            string
	Pool             string
	Upstream         string
	PromptTokens     int64
	CompletionTokens int64
	CachedTokens     int64
	// Cost is in Unit: "" for GPU-hours (vLLM models), a currency code
	// such as "USD" for an external provider. Costs of different units are
	// never added together (budgets and reports split by unit).
	Cost      float64
	Unit      string
	LatencyMS float64
	Status    int
	Priority  int
	Borrowing bool
	Stream    bool
	// Tags are the key's tags and the allow-listed request tags.
	Tags []string
	// Estimated is true when vLLM's usage report never arrived (client
	// gone before the final usage chunk or before a non-streaming answer)
	// and the tokens were estimated by the gateway.
	Estimated bool
}

Record is one finished request.

type ReportRow added in v0.2.0 

type ReportRow struct {
	Org               string  `json:"org,omitempty"`
	Lab               string  `json:"lab,omitempty"`
	Team              string  `json:"team,omitempty"`
	Project           string  `json:"project,omitempty"`
	User              string  `json:"user,omitempty"`
	Key               string  `json:"key,omitempty"`
	Model             string  `json:"model,omitempty"`
	Pool              string  `json:"pool,omitempty"`
	Tag               string  `json:"tag,omitempty"`
	Day               string  `json:"day,omitempty"`
	Unit              string  `json:"-"`
	Requests          int64   `json:"requests"`
	PromptTokens      int64   `json:"prompt_tokens"`
	CompletionTokens  int64   `json:"completion_tokens"`
	CachedTokens      int64   `json:"cached_tokens"`
	Cost              float64 `json:"cost"`
	GPUHours          float64 `json:"gpu_hours"`
	EstimatedRequests int64   `json:"estimated_requests"`
}

ReportRow is one group of a spend report. Only the dimensions asked for are filled. Every report is split by cost unit, whatever the dimensions: a row holds the costs of one unit only (Unit "" is GPU-hours), so GPU-hours and a currency are never summed. GPUHours repeats Cost for a GPU-hours row and is 0 otherwise (the v0.11 field, kept for clients that read it).

func Regroup added in v0.2.0 

func Regroup(rows []ReportRow, dims []string) []ReportRow

Regroup aggregates rows onto fewer dimensions, sorted, keeping the cost units apart.

func (ReportRow) Dim added in v0.2.0 

func (r ReportRow) Dim(d string) string

Dim returns the value of one dimension.

func (ReportRow) MarshalJSON added in v0.12.0 

func (r ReportRow) MarshalJSON() ([]byte, error)

MarshalJSON writes the unit by its name ("gpu_hours" for "").

type Role added in v0.2.0 

type Role string

Role is an administrative role. Roles are bound to a user on a scope: the whole proxy, an organisation or a team. The set follows LiteLLM's access-control roles (proxy admin, proxy admin viewer, org admin, team admin, member) and Slurm's account coordinators, who manage the sub-accounts and users of the accounts they coordinate. 

const (
	// RoleProxyAdmin administers everything (like the master key).
	RoleProxyAdmin Role = "proxy_admin"
	// RoleProxyViewer reads everything and changes nothing.
	RoleProxyViewer Role = "proxy_viewer"
	// RoleOrgAdmin administers one organisation: its teams, projects,
	// keys, budgets and the roles of its members.
	RoleOrgAdmin Role = "org_admin"
	// RoleTeamAdmin administers one team: its projects, keys and members.
	RoleTeamAdmin Role = "team_admin"
	// RoleMember uses one team: creates, lists and revokes its own keys.
	RoleMember Role = "member"
	// RoleViewer reads one organisation or team: listings and spend.
	RoleViewer Role = "viewer"
)

Roles, from the most to the least privileged on a given scope.

type Rotation added in v0.4.0 

type Rotation struct {
	Every  Duration `json:"every,omitzero"`
	Grace  Duration `json:"grace_period,omitzero"`
	Target string   `json:"target,omitempty"`
}

Rotation is a scheduled rotation: every Every, the key gets a new secret, written to Target (a vault: secret reference); the replaced secret stays valid for Grace.

type SCIMDoc added in v0.4.0 

type SCIMDoc struct {
	Kind    string // "User" or "Group"
	ID      string
	Name    string
	Data    []byte // the resource as JSON
	Version string
}

SCIMDoc is one provisioned resource.

type ScopeKind added in v0.2.0 

type ScopeKind string

ScopeKind is what a budget caps. 

const (
	ScopeOrg     ScopeKind = "org"
	ScopeTeam    ScopeKind = "team"
	ScopeProject ScopeKind = "project" // Name is "team/project"
	ScopeKey     ScopeKind = "key"     // Name is the key ID
	ScopeTag     ScopeKind = "tag"
)

Budget scopes.

func (ScopeKind) Valid added in v0.2.0 

func (k ScopeKind) Valid() bool

Valid reports whether k is a known scope kind.

type SealedLimits added in v0.13.1 

type SealedLimits struct{ Total, PerNet int }

SealedLimits bound what PutSealed stores: in all for the kind, and per client network (0: no bound).

type SpendRow added in v0.2.0 

type SpendRow struct {
	Org, Team, Project, KeyID, Model string
	Tags                             []string
	Cost                             float64
	Unit                             string
}

SpendRow is the cost of one combination of attributes since a time, in one unit.

type Store 

type Store interface {
	Ping(ctx context.Context) error

	CreateOrg(ctx context.Context, name string) error
	ListOrgs(ctx context.Context) ([]Org, error)
	// DisableOrg disables an organisation, its teams and projects, and
	// revokes all their keys.
	DisableOrg(ctx context.Context, name string) error

	// CreateTeam creates a team in an enabled organisation, re-enables a
	// disabled team of the same organisation, adopts a team created
	// without organisation (v0.1 databases), and returns ErrExists for an
	// enabled team of the same organisation, ErrConflict for a team of
	// another organisation, ErrNotFound for an unknown or disabled org.
	CreateTeam(ctx context.Context, name, org string) error
	GetTeam(ctx context.Context, name string) (Team, error)
	// ListTeams lists the teams of an organisation, or all with org "".
	ListTeams(ctx context.Context, org string) ([]Team, error)
	// DisableTeam disables a team and its projects and revokes its keys.
	DisableTeam(ctx context.Context, name string) error

	CreateProject(ctx context.Context, team, name string) error
	ListProjects(ctx context.Context, team string) ([]Project, error)
	// DisableProject disables a project and revokes its keys.
	DisableProject(ctx context.Context, team, name string) error

	// CreateKey stores a key of an enabled team (and project, when set).
	CreateKey(ctx context.Context, k Key) error
	// ListKeys lists keys, filtered by team when team is not "".
	ListKeys(ctx context.Context, team string) ([]Key, error)
	GetKey(ctx context.Context, id string) (Key, error)
	RevokeKey(ctx context.Context, id string) error
	// RotateKey gives a key a new secret hash. The current one becomes
	// the previous secret, valid until prevExpires (a zero or past time
	// drops it at once), and RotatedAt becomes at (zero: none). With a
	// non-zero expect, the key's current hash must be expect (ErrConflict
	// otherwise), so that two replicas cannot rotate one key twice.
	// ErrNotFound for an unknown or revoked key, ErrExists when the new
	// hash is already another key's, or this key's current one.
	RotateKey(ctx context.Context, id string, newHash, expect [32]byte, prevExpires, at time.Time) error
	// SCIM provisioning documents (see SCIMDoc). SCIMCreate: ErrExists
	// when the id, or a User's name without case, is taken. SCIMReplace:
	// ErrConflict when the stored version is not ifVersion, ErrExists for a
	// taken name. SCIMList filters by name without case when name != "".
	SCIMCreate(ctx context.Context, d SCIMDoc) error
	SCIMGet(ctx context.Context, kind, id string) (SCIMDoc, error)
	SCIMReplace(ctx context.Context, d SCIMDoc, ifVersion string) error
	SCIMDelete(ctx context.Context, kind, id string) error
	SCIMList(ctx context.Context, kind, name string) ([]SCIMDoc, error)

	// The audit log (see AuditEntry). AuditAppend assigns Seq and PrevHash
	// under a lock held across replicas, calls hash, and stores the entry.
	// AuditPrune deletes the entries through seq (the caller checks that a
	// checkpoint anchors what remains).
	AuditAppend(ctx context.Context, e AuditEntry, hash AuditHasher) (AuditEntry, error)
	AuditRange(ctx context.Context, from int64, limit int) ([]AuditEntry, error)
	AuditAddCheckpoint(ctx context.Context, c AuditCheckpoint) error
	AuditCheckpoints(ctx context.Context) ([]AuditCheckpoint, error)
	AuditPrune(ctx context.Context, through int64) (int64, error)

	// SetKeyRotation sets a key's rotation schedule (every 0: none).
	SetKeyRotation(ctx context.Context, id string, r Rotation) error
	// DueKeys lists the keys, not revoked, whose scheduled rotation is due
	// at now: RotatedAt (or CreatedAt) + Every <= now.
	DueKeys(ctx context.Context, now time.Time) ([]Key, error)
	// KeyByHash returns the key whose current or previous hash is hash,
	// revoked or not, with its Org filled in, or ErrNotFound. The caller
	// must still check it with Key.Matches.
	KeyByHash(ctx context.Context, hash [32]byte) (Key, error)

	CreateUser(ctx context.Context, name string) error
	GetUser(ctx context.Context, name string) (User, error)
	ListUsers(ctx context.Context) ([]User, error)
	// DisableUser disables a user and revokes their admin tokens.
	DisableUser(ctx context.Context, name string) error
	CreateUserToken(ctx context.Context, t UserToken) error
	ListUserTokens(ctx context.Context, user string) ([]UserToken, error)
	RevokeUserToken(ctx context.Context, user, id string) error
	// UserByTokenHash returns the token with that hash and its user, or
	// ErrNotFound. The caller compares hashes in constant time and checks
	// Revoked and Disabled.
	UserByTokenHash(ctx context.Context, hash [32]byte) (User, UserToken, error)
	// Grant stores a binding; ErrNotFound when the user, org or team does
	// not exist, ErrExists when it is already granted.
	Grant(ctx context.Context, b Binding) error
	Ungrant(ctx context.Context, b Binding) error
	// Bindings lists the bindings of a user, or all with user "".
	Bindings(ctx context.Context, user string) ([]Binding, error)

	// SetTeamModels replaces the list of model groups a team may use; an
	// empty list allows every configured model.
	SetTeamModels(ctx context.Context, team string, models []string) error
	// TeamModels returns every team's non-empty model list.
	TeamModels(ctx context.Context) (map[string][]string, error)

	// SetTeamTools replaces the MCP tool patterns a team may use
	// ("server.tool", "server.*" or "*"); an empty list allows none.
	SetTeamTools(ctx context.Context, team string, tools []string) error
	// TeamTools returns every team's non-empty tool list.
	TeamTools(ctx context.Context) (map[string][]string, error)

	// SetBudget creates or replaces a budget's period, maximum and soft
	// limit; a temporary increase already granted is kept.
	SetBudget(ctx context.Context, b Budget) error
	// IncreaseBudget grants a temporary increase until expiry, replacing
	// any earlier one; ErrNotFound when the budget does not exist.
	IncreaseBudget(ctx context.Context, k BudgetKey, amount float64, expiry time.Time) error
	DeleteBudget(ctx context.Context, k BudgetKey) error
	Budgets(ctx context.Context) ([]Budget, error)
	// MarkSoftAlert records that the soft-limit alert of a budget was sent
	// for the period starting at period; it returns false when it already
	// was, so that several gateways or restarts send it once.
	MarkSoftAlert(ctx context.Context, k BudgetKey, period time.Time) (bool, error)
	// SpendRows aggregates the cost of records at or after since by org,
	// team, project, key, model and tag set.
	SpendRows(ctx context.Context, since time.Time) ([]SpendRow, error)
	InsertRecords(ctx context.Context, recs []Record) error
	// SetAllowlist replaces the client networks allowed for an
	// organisation, a team or a key; an empty list removes the allowlist.
	SetAllowlist(ctx context.Context, a Allowlist) error
	Allowlists(ctx context.Context) ([]Allowlist, error)
	// Report aggregates the records in [from, to) by the given dimensions
	// (see ReportDims), sorted.
	Report(ctx context.Context, from, to time.Time, dims []string) ([]ReportRow, error)
	Close()
}

Store is the persistence contract.

type Team 

type Team struct {
	Name      string    `json:"name"`
	Org       string    `json:"org"`
	CreatedAt time.Time `json:"created_at"`
	Disabled  bool      `json:"disabled"`
}

Team is a group (a lab) inside an organisation.

type User added in v0.2.0 

type User struct {
	Name      string    `json:"name"`
	CreatedAt time.Time `json:"created_at"`
	Disabled  bool      `json:"disabled"`
}

User is a person who administers or uses the gateway.

type UserToken added in v0.2.0 

type UserToken struct {
	ID        string    `json:"id"`
	Hash      [32]byte  `json:"-"`
	User      string    `json:"user"`
	CreatedAt time.Time `json:"created_at"`
	Revoked   bool      `json:"revoked"`
}

UserToken is an admin API credential of a user; only its hash is kept.

Source Files

View all Source files

Jump to

Keyboard shortcuts

? : This menu
/ : Search site
f or F : Jump to
y or Y : Canonical URL
go.dev uses cookies from Google to deliver and enhance the quality of its services and to analyze traffic. Learn more.