package aibridgedserver import ( "context" "errors" "fmt" "strconv" "time" "github.com/google/uuid" "golang.org/x/xerrors" "github.com/coder/coder/v2/coderd/aibridge/budget" "github.com/coder/coder/v2/coderd/database" "github.com/coder/coder/v2/coderd/database/dbauthz" "github.com/coder/coder/v2/coderd/notifications" "github.com/coder/coder/v2/codersdk" ) // warningThresholdPercent triggers a warning notification; limitThresholdPercent // triggers the limit-reached notification (after which requests are blocked). const ( warningThresholdPercent = 85 limitThresholdPercent = 100 ) // budgetNotificationsCreatedBy records what enqueued AI budget notifications. const budgetNotificationsCreatedBy = "aigateway" // budgetThreshold pairs a percentage of the spend limit with the notifications // sent when a user's spend crosses it. type budgetThreshold struct { percent int userNotificationTemplate uuid.UUID adminNotificationTemplate uuid.UUID } // budgetThresholds are the thresholds evaluated on every priced interception, // ordered ascending. A single interception can cross more than one. var budgetThresholds = []budgetThreshold{ { percent: warningThresholdPercent, userNotificationTemplate: notifications.TemplateAIBudgetWarningUser, adminNotificationTemplate: notifications.TemplateAIBudgetWarningAdmin, }, { percent: limitThresholdPercent, userNotificationTemplate: notifications.TemplateAIBudgetLimitReachedUser, adminNotificationTemplate: notifications.TemplateAIBudgetLimitReachedAdmin, }, } // budgetThresholdCrossing describes a user crossing a budget threshold on a // single interception. It carries only stable values so that if the same // crossing is ever enqueued more than once, the payloads match and the // duplicate is dropped. type budgetThresholdCrossing struct { userID uuid.UUID effectiveGroupID uuid.UUID spendLimitMicros int64 thresholdPercent int userNotificationTemplate uuid.UUID adminNotificationTemplate uuid.UUID limitSource codersdk.AIBudgetLimitSource // periodStart and periodEnd bound the budget period [start, end) the // crossing occurred in. periodStart time.Time periodEnd time.Time } // detectBudgetThresholdCrossings checks whether this interception's cost pushed // the user's period spend across any budget thresholds, returning each one // crossed. A single interception can cross several at once (e.g. straight past // both the warning and limit thresholds). // // The period is derived from the interception's recorded time (the same // timestamp the spend row is bucketed by) rather than the current wall clock, // so an interception recorded near a period boundary but processed after it is // evaluated against the period it belongs to. func (s *Server) detectBudgetThresholdCrossings(ctx context.Context, tx database.Store, intc database.AIBridgeInterception, cost tokenUsageCost, createdAt time.Time) ([]budgetThresholdCrossing, error) { if !cost.effectiveGroupID.Valid || !cost.spendLimitMicros.Valid || cost.spendLimitMicros.Int64 <= 0 { return nil, nil } period, err := budget.CurrentPeriod(createdAt, s.budgetPeriod) if err != nil { return nil, xerrors.Errorf("compute AI budget period: %w", err) } spend, err := tx.GetUserAISpendSince(ctx, database.GetUserAISpendSinceParams{ UserID: intc.InitiatorID, EffectiveGroupID: cost.effectiveGroupID.UUID, PeriodStart: period.Start, }) if err != nil { return nil, xerrors.Errorf("get user AI spend for user %q in group %q: %w", intc.InitiatorID, cost.effectiveGroupID.UUID, err) } limit := cost.spendLimitMicros.Int64 newSpend := spend.SpendMicros // Pre-interception total is the current total minus this interception's cost. oldSpend := newSpend - cost.costMicros.Int64 var crossings []budgetThresholdCrossing for _, t := range budgetThresholds { at := limit * int64(t.percent) / 100 if oldSpend < at && newSpend >= at { crossings = append(crossings, budgetThresholdCrossing{ userID: intc.InitiatorID, effectiveGroupID: cost.effectiveGroupID.UUID, spendLimitMicros: limit, thresholdPercent: t.percent, userNotificationTemplate: t.userNotificationTemplate, adminNotificationTemplate: t.adminNotificationTemplate, limitSource: cost.limitSource, periodStart: period.Start, periodEnd: period.End, }) } } return crossings, nil } // notifyBudgetThresholdCrossings enqueues user and admin notifications for the // thresholds crossed by a single interception. All crossings share the same // user and effective group, so the group, username, and admin recipients are // resolved once. func (s *Server) notifyBudgetThresholdCrossings(ctx context.Context, crossings []budgetThresholdCrossing) error { if len(crossings) == 0 { return nil } userID := crossings[0].userID effectiveGroupID := crossings[0].effectiveGroupID group, err := s.store.GetGroupByID(ctx, effectiveGroupID) if err != nil { return xerrors.Errorf("look up group %q: %w", effectiveGroupID, err) } org, err := s.store.GetOrganizationByID(ctx, group.OrganizationID) if err != nil { return xerrors.Errorf("look up organization %q: %w", group.OrganizationID, err) } user, err := s.store.GetUserByID(ctx, userID) if err != nil { return xerrors.Errorf("look up user %q: %w", userID, err) } admins, err := s.budgetNotificationAdmins(ctx, userID) if err != nil { return xerrors.Errorf("look up budget notification admins: %w", err) } //nolint:gocritic // Enqueuing notifications requires the notifier actor. notifCtx := dbauthz.AsNotifier(ctx) var errs []error for _, c := range crossings { labels := map[string]string{ "threshold": strconv.Itoa(c.thresholdPercent), "limit": formatSpendLimit(c.spendLimitMicros), "period": s.budgetPeriod.Adjective(), "limit_source": string(c.limitSource), "username": user.Username, "effective_group_name": group.Name, "organization_name": org.Name, // Both bounds carry the year so a period straddling a year boundary // (e.g. December 1, 2026 - January 1, 2027) is unambiguous. "period_start": c.periodStart.UTC().Format("January 2, 2006"), "period_end": c.periodEnd.UTC().Format("January 2, 2006"), } // Notify the user who crossed the threshold. if _, err := s.notifEnqueuer.EnqueueWithData(notifCtx, userID, c.userNotificationTemplate, labels, nil, budgetNotificationsCreatedBy, effectiveGroupID); err != nil { errs = append(errs, xerrors.Errorf("enqueue user notification (threshold %d%%): %w", c.thresholdPercent, err)) } // Notify admins, naming the affected user. for _, admin := range admins { if _, err := s.notifEnqueuer.EnqueueWithData(notifCtx, admin.ID, c.adminNotificationTemplate, labels, nil, budgetNotificationsCreatedBy, userID, effectiveGroupID); err != nil { errs = append(errs, xerrors.Errorf("enqueue admin notification (threshold %d%%, admin %q): %w", c.thresholdPercent, admin.ID, err)) } } } return errors.Join(errs...) } // budgetNotificationAdmins returns the users who should receive admin budget // notifications: deployment owners and user admins, excluding the affected // user, who receives the user-facing notification instead. func (s *Server) budgetNotificationAdmins(ctx context.Context, excludeUserID uuid.UUID) ([]database.GetUsersRow, error) { admins, err := s.store.GetUsers(ctx, database.GetUsersParams{ RbacRole: []string{codersdk.RoleOwner, codersdk.RoleUserAdmin}, }) if err != nil { return nil, err } recipients := make([]database.GetUsersRow, 0, len(admins)) for _, admin := range admins { if admin.ID == excludeUserID { continue } recipients = append(recipients, admin) } return recipients, nil } // formatSpendLimit renders a spend limit as a USD string. func formatSpendLimit(micros int64) string { return fmt.Sprintf("$%.2f", float64(micros)/1_000_000) }