internal/budget/types.go

Rama main · réplica de lectura del repositorio original.

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()
}