package budget import ( "fmt" "math" "slices" "strings" "time" "github.com/enterpilot/gomodel/internal/core" ) 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 ) const ( // SourceConfig marks budgets seeded from static configuration. SourceConfig = "config" // SourceManual marks budgets created or changed through admin APIs. SourceManual = "manual" ) // Scope names what a budget limits: a consumer user-path subtree, a request // label, or a provider instance. type Scope string 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" ) // CostBasis names which recorded cost a budget sums. type CostBasis string 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" ) // NormalizeCostBasis canonicalizes a cost basis name. An empty basis means // provider, keeping budgets stored before cost bases existed unchanged. func NormalizeCostBasis(raw string) (CostBasis, error) { switch CostBasis(strings.ToLower(strings.TrimSpace(raw))) { case "", CostBasisProvider: return CostBasisProvider, nil case CostBasisBilled: return CostBasisBilled, nil default: return "", fmt.Errorf("cost_basis must be one of provider, billed") } } // 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. 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"` } // EnforcedBasis returns the basis the budget's spend is measured on. func (b Budget) EnforcedBasis() CostBasis { if b.EnforcedCostBasis != "" { return b.EnforcedCostBasis } return b.CostBasis } // 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). type Subjects struct { UserPath string Labels []string Provider string } // appliesTo reports whether the budget covers the request subjects. func (b Budget) appliesTo(s Subjects) bool { switch b.Scope { case ScopeLabel: return slices.Contains(s.Labels, b.Subject) case ScopeProvider: return s.Provider != "" && strings.EqualFold(b.Subject, strings.TrimSpace(s.Provider)) default: return budgetAppliesToPath(b.Subject, s.UserPath) } } // resolve returns the budget instance that applies to the request. Per-child // definitions retain Subject as their configured template and carry the // derived child separately for spend queries and diagnostics. func (b Budget) resolve(s Subjects) (Budget, bool) { if !b.PerChild { return b, b.appliesTo(s) } if b.Scope != ScopeUserPath { return Budget{}, false } child, ok := core.UserPathChild(b.Subject, s.UserPath) if !ok { return Budget{}, false } b.EffectiveSubject = child return b, true } // SubjectLabel names the budget subject for error messages and logs. func (b Budget) SubjectLabel() string { switch b.Scope { case ScopeLabel: return "label " + b.Subject case ScopeProvider: return "provider " + b.Subject } if b.EffectiveSubject != "" { return b.EffectiveSubject } return b.Subject } func (b Budget) evaluationSubject() string { if b.EffectiveSubject != "" { return b.EffectiveSubject } return b.Subject } // budgetAppliesToPath reports whether a budget path covers the request path, // using the same subtree semantics as rate limits. func budgetAppliesToPath(budgetPath, requestPath string) bool { budgetPath = strings.TrimSpace(budgetPath) requestPath = strings.TrimSpace(requestPath) if budgetPath == "/" { return true } return requestPath == budgetPath || strings.HasPrefix(requestPath, budgetPath+"/") } // Settings controls the calendar anchors used to find the active budget period. // Values are interpreted in UTC. 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"` } // CheckResult describes one evaluated budget limit. 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"` } // 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. func (r CheckResult) UsageRatio() float64 { if r.Budget.Amount <= 0 { return 0 } return r.Spent / r.Budget.Amount } // PeriodRatio returns the elapsed fraction of the budget period at now, // clamped to [0, 1]. func (r CheckResult) PeriodRatio(now time.Time) float64 { duration := r.PeriodEnd.Sub(r.PeriodStart).Seconds() if duration <= 0 { return 0 } ratio := now.Sub(r.PeriodStart).Seconds() / duration if ratio < 0 { return 0 } if ratio > 1 { return 1 } return ratio } // ExceededError indicates a budget has already been exhausted. type ExceededError struct { Result CheckResult } func (e *ExceededError) Error() string { if e == nil { return "" } return fmt.Sprintf( "budget exceeded for %s %s limit: spent %.6f of %.6f", e.Result.Budget.SubjectLabel(), PeriodLabel(e.Result.Budget.PeriodSeconds), e.Result.Spent, e.Result.Budget.Amount, ) } // DefaultSettings returns the reset anchors used when no DB setting exists. func DefaultSettings() Settings { return Settings{ DailyResetHour: 0, DailyResetMinute: 0, WeeklyResetWeekday: int(time.Monday), WeeklyResetHour: 0, WeeklyResetMinute: 0, MonthlyResetDay: 1, MonthlyResetHour: 0, MonthlyResetMinute: 0, } } func NormalizeUserPath(raw string) (string, error) { path, err := core.NormalizeUserPath(raw) if err != nil { return "", err } if path == "" { return "/", nil } return path, nil } // NormalizeScope canonicalizes a budget scope name. An empty scope means // user_path, keeping pre-scope budget definitions and requests valid. func NormalizeScope(raw string) (Scope, error) { switch Scope(strings.ToLower(strings.TrimSpace(raw))) { case "", ScopeUserPath: return ScopeUserPath, nil case ScopeLabel: return ScopeLabel, nil case ScopeProvider: return ScopeProvider, nil default: return "", fmt.Errorf("scope must be one of user_path, label, provider") } } // 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 NormalizeSubject(scope Scope, subject string) (string, error) { switch scope { case ScopeLabel: subject = strings.TrimSpace(subject) if subject == "" { return "", fmt.Errorf("label budget subject is required") } return subject, nil case ScopeProvider: subject = strings.ToLower(strings.TrimSpace(subject)) if subject == "" { return "", fmt.Errorf("provider budget subject is required") } if strings.ContainsAny(subject, "/ \t") { return "", fmt.Errorf("provider budget subject must be a provider name without slashes or spaces") } return subject, nil default: return NormalizeUserPath(subject) } } func NormalizeBudget(b Budget) (Budget, error) { scope, err := NormalizeScope(string(b.Scope)) if err != nil { return Budget{}, err } b.Scope = scope subject, err := NormalizeSubject(scope, b.Subject) if err != nil { return Budget{}, err } b.Subject = subject b.EffectiveSubject = "" costBasis, err := NormalizeCostBasis(string(b.CostBasis)) if err != nil { return Budget{}, err } b.CostBasis = costBasis if b.PerChild && scope != ScopeUserPath { return Budget{}, fmt.Errorf("per_child is only valid for user_path budgets") } if b.PeriodSeconds <= 0 { return Budget{}, fmt.Errorf("period_seconds must be greater than 0") } if math.IsNaN(b.Amount) || math.IsInf(b.Amount, 0) || b.Amount <= 0 { return Budget{}, fmt.Errorf("amount must be a finite number greater than 0") } b.Source = strings.TrimSpace(b.Source) if b.LastResetAt != nil { t := b.LastResetAt.UTC() b.LastResetAt = &t } now := time.Now().UTC() if b.CreatedAt.IsZero() { b.CreatedAt = now } b.UpdatedAt = now return b, nil } func ValidateSettings(settings Settings) error { if settings.DailyResetHour < 0 || settings.DailyResetHour > 23 { return fmt.Errorf("daily_reset_hour must be between 0 and 23") } if settings.DailyResetMinute < 0 || settings.DailyResetMinute > 59 { return fmt.Errorf("daily_reset_minute must be between 0 and 59") } if settings.WeeklyResetWeekday < int(time.Sunday) || settings.WeeklyResetWeekday > int(time.Saturday) { return fmt.Errorf("weekly_reset_weekday must be between 0 and 6") } if settings.WeeklyResetHour < 0 || settings.WeeklyResetHour > 23 { return fmt.Errorf("weekly_reset_hour must be between 0 and 23") } if settings.WeeklyResetMinute < 0 || settings.WeeklyResetMinute > 59 { return fmt.Errorf("weekly_reset_minute must be between 0 and 59") } if settings.MonthlyResetDay < 1 || settings.MonthlyResetDay > 31 { return fmt.Errorf("monthly_reset_day must be between 1 and 31") } if settings.MonthlyResetHour < 0 || settings.MonthlyResetHour > 23 { return fmt.Errorf("monthly_reset_hour must be between 0 and 23") } if settings.MonthlyResetMinute < 0 || settings.MonthlyResetMinute > 59 { return fmt.Errorf("monthly_reset_minute must be between 0 and 59") } return nil } func PeriodSeconds(period string) (int64, bool) { switch strings.ToLower(strings.TrimSpace(period)) { case "hour", "hourly", "hours": return PeriodHourlySeconds, true case "day", "daily", "days": return PeriodDailySeconds, true case "week", "weekly", "weeks": return PeriodWeeklySeconds, true case "month", "monthly", "months": return PeriodMonthlySeconds, true default: return 0, false } } func PeriodLabel(seconds int64) string { switch seconds { case PeriodHourlySeconds: return "hourly" case PeriodDailySeconds: return "daily" case PeriodWeeklySeconds: return "weekly" case PeriodMonthlySeconds: return "monthly" default: return fmt.Sprintf("%ds", seconds) } } func PeriodBounds(now time.Time, seconds int64, settings Settings) (time.Time, time.Time) { now = now.UTC() switch seconds { case PeriodHourlySeconds: start := now.Truncate(time.Hour) return start, start.Add(time.Hour) case PeriodDailySeconds: start := anchoredDayStart(now, settings.DailyResetHour, settings.DailyResetMinute) return start, start.AddDate(0, 0, 1) case PeriodWeeklySeconds: start := anchoredWeekStart(now, time.Weekday(settings.WeeklyResetWeekday), settings.WeeklyResetHour, settings.WeeklyResetMinute) return start, start.AddDate(0, 0, 7) case PeriodMonthlySeconds: start := anchoredMonthStart(now, settings.MonthlyResetDay, settings.MonthlyResetHour, settings.MonthlyResetMinute) nextMonth := time.Date(start.Year(), start.Month()+1, 1, 0, 0, 0, 0, time.UTC) return start, monthAnchor(nextMonth.Year(), nextMonth.Month(), settings.MonthlyResetDay, settings.MonthlyResetHour, settings.MonthlyResetMinute) default: if seconds <= 0 { return now, now } startUnix := now.Unix() - (now.Unix() % seconds) start := time.Unix(startUnix, 0).UTC() return start, start.Add(time.Duration(seconds) * time.Second) } } func anchoredDayStart(now time.Time, hour, minute int) time.Time { start := time.Date(now.Year(), now.Month(), now.Day(), hour, minute, 0, 0, time.UTC) if now.Before(start) { start = start.AddDate(0, 0, -1) } return start } func anchoredWeekStart(now time.Time, weekday time.Weekday, hour, minute int) time.Time { todayAnchor := time.Date(now.Year(), now.Month(), now.Day(), hour, minute, 0, 0, time.UTC) daysBack := (int(todayAnchor.Weekday()) - int(weekday) + 7) % 7 start := todayAnchor.AddDate(0, 0, -daysBack) if now.Before(start) { start = start.AddDate(0, 0, -7) } return start } func anchoredMonthStart(now time.Time, day, hour, minute int) time.Time { start := monthAnchor(now.Year(), now.Month(), day, hour, minute) if now.Before(start) { prevFirst := time.Date(now.Year(), now.Month(), 1, now.Hour(), now.Minute(), now.Second(), now.Nanosecond(), now.Location()) prev := prevFirst.AddDate(0, -1, 0) start = monthAnchor(prev.Year(), prev.Month(), day, hour, minute) } return start } func monthAnchor(year int, month time.Month, day, hour, minute int) time.Time { anchorDay := min(day, daysInMonth(year, month)) return time.Date(year, month, anchorDay, hour, minute, 0, 0, time.UTC) } func daysInMonth(year int, month time.Month) int { return time.Date(year, month+1, 0, 0, 0, 0, 0, time.UTC).Day() }