budget

package
v0.2.2 Latest Latest
Warning

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

 Go to latest Published: Nov 6, 2026 License: MIT Imports: 20 Imported by: 0

Documentation

Index

Constants

View Source 
const (
	PeriodHourlySeconds int64 = 3600
	PeriodDailySeconds  int64 = 86400
	PeriodWeeklySeconds int64 = 604800
	// PeriodMonthlySeconds is a sentinel key for calendar-month windows, not a literal 30-day duration.
	PeriodMonthlySeconds int64 = 2592000
)
View Source 
const (
	// SourceConfig marks budgets seeded from static configuration.
	SourceConfig = "config"
	// SourceManual marks budgets created or changed through admin APIs.
	SourceManual = "manual"
)

Variables

View Source 
var ErrBilledBudgetsUnavailable = errors.New("budgets on billed cost require a GoModel Pro license with the billing entitlement and pricing.markups_enabled")

ErrBilledBudgetsUnavailable indicates that a budget on billed cost was found while the deployment lacks the billed-budgets capability or has markups disabled.

View Source 
var ErrNotFound = errors.New("budget not found")
View Source 
var ErrQuotaTemplatesUnavailable = errors.New("per-child quota templates are not enabled")

ErrQuotaTemplatesUnavailable indicates that a per-child template was found while the deployment does not have the quota-templates capability.

View Source 
var ErrUnavailable = errors.New("budget service is unavailable")

ErrUnavailable indicates a budget service was used without an initialized store.

Functions

func NormalizeSubject added in v0.1.60 

func NormalizeSubject(scope Scope, subject string) (string, error)

NormalizeSubject canonicalizes a budget subject for its scope. Label subjects keep their case: they are compared verbatim against the labels recorded on usage entries, the same way the usage label filter and by-label breakdown compare them. Provider subjects are lowercased, as rate limit provider rules are, so "OpenAI" and "openai" cannot persist as two budgets over the same traffic.

func NormalizeUserPath 

func NormalizeUserPath(raw string) (string, error)

func PeriodBounds 

func PeriodBounds(now time.Time, seconds int64, settings Settings) (time.Time, time.Time)

func PeriodLabel 

func PeriodLabel(seconds int64) string

func PeriodSeconds 

func PeriodSeconds(period string) (int64, bool)

func ValidateSettings 

func ValidateSettings(settings Settings) error

Types

type Budget 

type Budget struct {
	Scope     Scope     `json:"scope" bson:"scope"`
	Subject   string    `json:"subject" bson:"subject"`
	PerChild  bool      `json:"per_child" bson:"per_child"`
	CostBasis CostBasis `json:"cost_basis" bson:"cost_basis,omitempty"`
	// EnforcedCostBasis is the basis spend is measured on when it differs
	// from CostBasis: provider while a stored billed budget cannot be enforced
	// on billed cost. It is set when budgets are loaded and never persisted.
	EnforcedCostBasis CostBasis `json:"-" bson:"-"`
	// EffectiveSubject is populated only while evaluating a per-child budget.
	// It is not part of the persisted definition or public JSON representation.
	EffectiveSubject string     `json:"-" bson:"-"`
	PeriodSeconds    int64      `json:"period_seconds" bson:"period_seconds"`
	Amount           float64    `json:"amount" bson:"amount"`
	Source           string     `json:"source,omitempty" bson:"source,omitempty"`
	LastResetAt      *time.Time `json:"last_reset_at,omitempty" bson:"last_reset_at,omitempty"`
	CreatedAt        time.Time  `json:"created_at" bson:"created_at"`
	UpdatedAt        time.Time  `json:"updated_at" bson:"updated_at"`
}

Budget stores one spend limit for one scope, subject, and reset period. CostBasis is an attribute, not part of the identity: one scope, subject, and period holds one budget on either basis.

func NormalizeBudget 

func NormalizeBudget(b Budget) (Budget, error)

func (Budget) EnforcedBasis added in v0.1.109 

func (b Budget) EnforcedBasis() CostBasis

EnforcedBasis returns the basis the budget's spend is measured on.

func (Budget) SubjectLabel added in v0.1.60 

func (b Budget) SubjectLabel() string

SubjectLabel names the budget subject for error messages and logs.

type CheckResult 

type CheckResult struct {
	Budget      Budget    `json:"budget"`
	PeriodStart time.Time `json:"period_start"`
	PeriodEnd   time.Time `json:"period_end"`
	Spent       float64   `json:"spent"`
	HasUsage    bool      `json:"has_usage"`
	Remaining   float64   `json:"remaining"`
}

CheckResult describes one evaluated budget limit.

func (CheckResult) PeriodRatio 

func (r CheckResult) PeriodRatio(now time.Time) float64

PeriodRatio returns the elapsed fraction of the budget period at now, clamped to [0, 1].

func (CheckResult) UsageRatio 

func (r CheckResult) UsageRatio() float64

UsageRatio returns spent/amount for the period, or 0 when the budget amount is not positive. It is deliberately not clamped: values above 1 indicate an exceeded budget.

type CostBasis added in v0.1.109 

type CostBasis string

CostBasis names which recorded cost a budget sums. 

const (
	// CostBasisProvider sums the upstream provider cost (total_cost) of
	// uncached requests. It is the default.
	CostBasisProvider CostBasis = "provider"
	// CostBasisBilled sums what the caller is billed (billed_cost): provider
	// cost plus markup, including billed response-cache hits. Usage recorded
	// without a billed cost counts at its provider cost.
	CostBasisBilled CostBasis = "billed"
)

func NormalizeCostBasis added in v0.1.109 

func NormalizeCostBasis(raw string) (CostBasis, error)

NormalizeCostBasis canonicalizes a cost basis name. An empty basis means provider, keeping budgets stored before cost bases existed unchanged.

type ExceededError 

type ExceededError struct {
	Result CheckResult
}

ExceededError indicates a budget has already been exhausted.

func (*ExceededError) Error 

func (e *ExceededError) Error() string

type MongoDBStore 

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

func NewMongoDBStore 

func NewMongoDBStore(ctx context.Context, database *mongo.Database) (*MongoDBStore, error)

func (*MongoDBStore) Close 

func (s *MongoDBStore) Close() error

func (*MongoDBStore) DeleteBudget 

func (s *MongoDBStore) DeleteBudget(ctx context.Context, scope Scope, subject string, periodSeconds int64) error

func (*MongoDBStore) GetSettings 

func (s *MongoDBStore) GetSettings(ctx context.Context) (Settings, error)

func (*MongoDBStore) ListBudgets 

func (s *MongoDBStore) ListBudgets(ctx context.Context) ([]Budget, error)

func (*MongoDBStore) ReplaceConfigBudgets 

func (s *MongoDBStore) ReplaceConfigBudgets(ctx context.Context, budgets []Budget) error

func (*MongoDBStore) ResetAllBudgets 

func (s *MongoDBStore) ResetAllBudgets(ctx context.Context, at time.Time) error

func (*MongoDBStore) ResetBudget 

func (s *MongoDBStore) ResetBudget(ctx context.Context, scope Scope, subject string, periodSeconds int64, at time.Time) error

func (*MongoDBStore) SaveSettings 

func (s *MongoDBStore) SaveSettings(ctx context.Context, settings Settings) (Settings, error)

func (*MongoDBStore) SumSpend added in v0.1.60 

func (s *MongoDBStore) SumSpend(ctx context.Context, windows []SpendWindow) ([]Spend, error)

SumSpend totals spend for every window in one aggregation.

A budget check matches several budgets at once — a user-path subtree plus a budget per request label — so the pipeline matches the union of their windows once and accumulates a conditional sum per window instead of running an aggregation each.

func (*MongoDBStore) UpsertBudgets 

func (s *MongoDBStore) UpsertBudgets(ctx context.Context, budgets []Budget) error

type Result 

type Result struct {
	Service *Service
	Store   Store
	// contains filtered or unexported fields
}

func New 

func New(ctx context.Context, cfg *config.Config, shared storage.Storage, quotaTemplatesEnabled, billedBudgetsEnabled bool) (*Result, error)

New builds the budget service. quotaTemplatesEnabled and billedBudgetsEnabled report the ext capabilities; budgets on billed cost also need pricing.markups_enabled, which New reads from cfg.

func (*Result) Close 

func (r *Result) Close() error

type SQLStore added in v0.1.60 

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

SQLStore stores budgets and budget settings in a SQL database.

func NewSQLStore added in v0.1.60 

func NewSQLStore(ctx context.Context, db sqlx.DB) (*SQLStore, error)

NewSQLStore creates the budget tables and indexes if needed.

func (*SQLStore) Close added in v0.1.60 

func (s *SQLStore) Close() error

func (*SQLStore) DeleteBudget added in v0.1.60 

func (s *SQLStore) DeleteBudget(ctx context.Context, scope Scope, subject string, periodSeconds int64) error

func (*SQLStore) GetSettings added in v0.1.60 

func (s *SQLStore) GetSettings(ctx context.Context) (Settings, error)

func (*SQLStore) ListBudgets added in v0.1.60 

func (s *SQLStore) ListBudgets(ctx context.Context) ([]Budget, error)

func (*SQLStore) ReplaceConfigBudgets added in v0.1.60 

func (s *SQLStore) ReplaceConfigBudgets(ctx context.Context, budgets []Budget) error

func (*SQLStore) ResetAllBudgets added in v0.1.60 

func (s *SQLStore) ResetAllBudgets(ctx context.Context, at time.Time) error

func (*SQLStore) ResetBudget added in v0.1.60 

func (s *SQLStore) ResetBudget(ctx context.Context, scope Scope, subject string, periodSeconds int64, at time.Time) error

func (*SQLStore) SaveSettings added in v0.1.60 

func (s *SQLStore) SaveSettings(ctx context.Context, settings Settings) (Settings, error)

func (*SQLStore) SumSpend added in v0.1.60 

func (s *SQLStore) SumSpend(ctx context.Context, windows []SpendWindow) ([]Spend, error)

SumSpend totals spend for every window with one scan per chunk.

A budget check matches several budgets at once — a user-path subtree plus a budget per request label — and running one statement each would scan the usage table once per budget. Conditional aggregation collapses them into a single pass: the outer WHERE spans the union of the windows, and each window gets its own SUM(CASE ...) column.

func (*SQLStore) UpsertBudgets added in v0.1.60 

func (s *SQLStore) UpsertBudgets(ctx context.Context, budgets []Budget) error

type Scope added in v0.1.60 

type Scope string

Scope names what a budget limits: a consumer user-path subtree, a request label, or a provider instance. 

const (
	// ScopeUserPath limits a consumer subtree; the subject is a user path and
	// covers all its descendants.
	ScopeUserPath Scope = "user_path"
	// ScopeLabel limits everything carrying one request label; the subject is
	// the label verbatim. A request carrying several labels is charged against
	// every matching label budget.
	ScopeLabel Scope = "label"
	// ScopeProvider limits everything routed to one configured provider
	// instance; the subject is the provider name (e.g. "openai-prod"), matched
	// case-insensitively.
	ScopeProvider Scope = "provider"
)

func NormalizeScope added in v0.1.60 

func NormalizeScope(raw string) (Scope, error)

NormalizeScope canonicalizes a budget scope name. An empty scope means user_path, keeping pre-scope budget definitions and requests valid.

type Service 

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

func NewService 

func NewService(ctx context.Context, store Store, options ...ServiceOption) (*Service, error)

func (*Service) Budgets 

func (s *Service) Budgets() []Budget

func (*Service) Check 

func (s *Service) Check(ctx context.Context, subjects Subjects, now time.Time) error

func (*Service) CheckCacheHit added in v0.1.109 

func (s *Service) CheckCacheHit(ctx context.Context, subjects Subjects, now time.Time) error

CheckCacheHit evaluates the budgets a response-cache hit spends against before the cached response is replayed. Only billed user-path and label budgets qualify, and only while cache hits are billed: provider cost never counts cache hits, and a hit is not routed to any provider. subjects.Provider names the provider the request resolved to, so a hit of a billable: false provider is admitted like the request itself would be. It returns without a store round trip when no such budget exists.

func (*Service) CheckWithResults 

func (s *Service) CheckWithResults(ctx context.Context, subjects Subjects, now time.Time) ([]CheckResult, error)

CheckWithResults evaluates every budget covering the request subjects and stops at the first exhausted one, returning it as an ExceededError alongside the results evaluated so far.

func (*Service) DeleteBudget 

func (s *Service) DeleteBudget(ctx context.Context, scope Scope, subject string, periodSeconds int64) error

func (*Service) FailoverWithinBudget added in v0.2.2 

func (s *Service) FailoverWithinBudget(ctx context.Context, primary, target string) bool

FailoverWithinBudget reports whether a request admitted on the primary provider may fail over to target: no provider budget of target is exhausted and, when the primary was not billable, no user-path or label budget covering the request is either. Admission skips consumer budgets for a non-billable primary, so the first billable target has to pass them; after a billable primary they already passed, and switching targets cannot relieve them, so they are not asked again.

func (*Service) ProviderWithinBudget added in v0.1.104 

func (s *Service) ProviderWithinBudget(ctx context.Context, providerName string) bool

ProviderWithinBudget reports whether no provider budget for providerName is exhausted. Failover uses it to skip targets whose provider budget is spent; consumer budgets are deliberately ignored because switching targets cannot relieve them. A store failure counts as exhausted, matching admission, which refuses a request when the budget check fails.

func (*Service) Refresh 

func (s *Service) Refresh(ctx context.Context) error

func (*Service) ReplaceConfigBudgets 

func (s *Service) ReplaceConfigBudgets(ctx context.Context, budgets []Budget) error

func (*Service) ResetAll 

func (s *Service) ResetAll(ctx context.Context, at time.Time) error

func (*Service) ResetBudget 

func (s *Service) ResetBudget(ctx context.Context, scope Scope, subject string, periodSeconds int64, at time.Time) error

func (*Service) SaveSettings 

func (s *Service) SaveSettings(ctx context.Context, settings Settings) (Settings, error)

func (*Service) SetChangeNotifier added in v0.1.109 

func (s *Service) SetChangeNotifier(fn func(context.Context))

SetChangeNotifier installs fn, called after every successful write of the budgets or their settings so other gateway instances reload them. Call it before serving; writes made before it was set (startup seeding) are not announced.

func (*Service) SetProviderBillable added in v0.2.2 

func (s *Service) SetProviderBillable(fn func(providerName string) bool)

SetProviderBillable installs the lookup that tells a non-billable provider (billable: false) from a billable one. A request routed to a non-billable provider is checked against provider budgets only: its cost is not owed by the caller, so user-path and label budgets neither gate nor count it. Call it before serving.

func (*Service) Settings 

func (s *Service) Settings() Settings

func (*Service) Statuses 

func (s *Service) Statuses(ctx context.Context, now time.Time) ([]CheckResult, error)

Statuses evaluates every configured budget without enforcing limits.

func (*Service) StatusesFor added in v0.1.60 

func (s *Service) StatusesFor(ctx context.Context, subjects Subjects, now time.Time) ([]CheckResult, error)

StatusesFor evaluates every budget covering the request subjects without enforcing limits. Unlike CheckWithResults it never stops at an exhausted budget, so callers get the full status picture even when several are exceeded.

func (*Service) StoredAsBilled added in v0.1.109 

func (s *Service) StoredAsBilled(scope Scope, subject string, periodSeconds int64) bool

StoredAsBilled reports whether the budget with this key is stored on billed cost, whether or not billed budgets are currently available.

func (*Service) UpsertBudgets 

func (s *Service) UpsertBudgets(ctx context.Context, budgets []Budget) error

type ServiceOption added in v0.1.77 

type ServiceOption func(*Service)

ServiceOption configures a budget service.

func WithBilledBudgets added in v0.1.109 

func WithBilledBudgets(enabled bool) ServiceOption

WithBilledBudgets controls whether budgets on billed cost are accepted.

func WithBilledCacheHits added in v0.1.109 

func WithBilledCacheHits(enabled bool) ServiceOption

WithBilledCacheHits reports whether response-cache hits add billed cost (markups enabled and bill_cache_hits on), which makes CheckCacheHit enforce billed budgets before a cached response is replayed.

func WithProviderBilledCacheHits added in v0.1.109 

func WithProviderBilledCacheHits(enabled bool) ServiceOption

WithProviderBilledCacheHits controls whether provider budgets on billed cost count billed response-cache hits. The default is true.

func WithQuotaTemplates added in v0.1.77 

func WithQuotaTemplates(enabled bool) ServiceOption

WithQuotaTemplates controls whether per-child user-path budgets are accepted.

type Settings 

type Settings struct {
	DailyResetHour     int       `json:"daily_reset_hour" bson:"daily_reset_hour"`
	DailyResetMinute   int       `json:"daily_reset_minute" bson:"daily_reset_minute"`
	WeeklyResetWeekday int       `json:"weekly_reset_weekday" bson:"weekly_reset_weekday"`
	WeeklyResetHour    int       `json:"weekly_reset_hour" bson:"weekly_reset_hour"`
	WeeklyResetMinute  int       `json:"weekly_reset_minute" bson:"weekly_reset_minute"`
	MonthlyResetDay    int       `json:"monthly_reset_day" bson:"monthly_reset_day"`
	MonthlyResetHour   int       `json:"monthly_reset_hour" bson:"monthly_reset_hour"`
	MonthlyResetMinute int       `json:"monthly_reset_minute" bson:"monthly_reset_minute"`
	UpdatedAt          time.Time `json:"updated_at" bson:"updated_at"`
}

Settings controls the calendar anchors used to find the active budget period. Values are interpreted in UTC.

func DefaultSettings 

func DefaultSettings() Settings

DefaultSettings returns the reset anchors used when no DB setting exists.

type Spend added in v0.1.60 

type Spend struct {
	Total    float64
	HasUsage bool
}

Spend is the total tracked cost for one SpendWindow. HasUsage is false when the window matched no priced usage at all, which keeps a budget with no traffic from ever blocking.

type SpendWindow added in v0.1.60 

type SpendWindow struct {
	Scope     Scope
	Subject   string
	CostBasis CostBasis
	Start     time.Time
	End       time.Time
	// ExcludeCacheHits keeps billed response-cache hits out of a billed
	// window. Provider windows never count cache hits.
	ExcludeCacheHits bool
}

SpendWindow asks for the tracked spend of one budget subject over one time window.

type Store 

type Store interface {
	ListBudgets(ctx context.Context) ([]Budget, error)
	UpsertBudgets(ctx context.Context, budgets []Budget) error
	DeleteBudget(ctx context.Context, scope Scope, subject string, periodSeconds int64) error
	ReplaceConfigBudgets(ctx context.Context, budgets []Budget) error
	GetSettings(ctx context.Context) (Settings, error)
	SaveSettings(ctx context.Context, settings Settings) (Settings, error)
	ResetBudget(ctx context.Context, scope Scope, subject string, periodSeconds int64, at time.Time) error
	ResetAllBudgets(ctx context.Context, at time.Time) error

	// SumSpend totals spend for every window in one round trip, in the order
	// given: uncached provider cost, or billed cost for a billed window. Budget enforcement runs on every request, so the whole
	// matching set is asked for at once rather than one query per budget.
	SumSpend(ctx context.Context, windows []SpendWindow) ([]Spend, error)

	Close() error
}

Store persists budget definitions, reset settings, and spend lookups.

type Subjects added in v0.1.60 

type Subjects struct {
	UserPath string
	Labels   []string
	Provider string
}

Subjects identifies the dimensions one request can be budgeted by. UserPath is always known at ingress; Labels are whatever tagging extracted for the request; Provider is the routed provider name, empty when the route is not known yet (batch submissions, status lookups).

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.