mirror of
https://github.com/coder/coder.git
synced 2026-09-22 05:05:20 +08:00
feat: notify admins when a user crosses an AI budget threshold (#27415)
Implements: https://linear.app/codercom/issue/AIGOV-289/notify-users-and-admins-on-budget-warning-and-limit-reached Notify admins when a user crosses an AI budget threshold, complementing the user-facing notifications from https://github.com/coder/coder/pull/27346 When a priced interception pushes a user's period spend across the warning (85%) or limit (100%) threshold, the Owners and User Admins now receive an admin notification naming the affected user, alongside the user's own notification. The affected user is excluded from the admin recipients since they already get the user-facing copy. Delivery is best-effort: a failure to enqueue is logged and never blocks recording the interception. The admin templates always show the effective group the spend is attributed to, and note when the limit comes from a per-user override rather than the group budget. Depends on https://github.com/coder/coder/pull/27346 ## Screenshots: <img width="1101" height="440" alt="image" src="https://github.com/user-attachments/assets/eb731088-05c8-47bd-9d06-fc9d07f63a08" /> <img width="468" height="391" alt="image" src="https://github.com/user-attachments/assets/b89b76a6-3fa8-4735-99a2-43e119a7a7e3" />
This commit is contained in:
@@ -84,6 +84,7 @@ type store interface {
|
||||
GetUserEveryoneFallbackGroup(ctx context.Context, userID uuid.UUID) (uuid.UUID, error)
|
||||
GetUserAISpendSince(ctx context.Context, arg database.GetUserAISpendSinceParams) (database.GetUserAISpendSinceRow, error)
|
||||
GetGroupByID(ctx context.Context, id uuid.UUID) (database.Group, error)
|
||||
GetUsers(ctx context.Context, arg database.GetUsersParams) ([]database.GetUsersRow, error)
|
||||
|
||||
// MCPConfigurator-related queries.
|
||||
GetExternalAuthLinksByUserID(ctx context.Context, userID uuid.UUID) ([]database.ExternalAuthLink, error)
|
||||
@@ -433,14 +434,10 @@ func (s *Server) recordTokenUsageAndSpend(ctx context.Context, intc database.AIB
|
||||
return err
|
||||
}
|
||||
|
||||
for _, crossing := range crossings {
|
||||
if err := s.notifyBudgetThresholdCrossing(ctx, crossing); err != nil {
|
||||
s.logger.Error(ctx, "failed to send AI budget notification",
|
||||
slog.F("user_id", crossing.userID),
|
||||
slog.F("group_id", crossing.effectiveGroupID),
|
||||
slog.F("threshold_percent", crossing.thresholdPercent),
|
||||
slog.Error(err))
|
||||
}
|
||||
if err := s.notifyBudgetThresholdCrossings(ctx, crossings); err != nil {
|
||||
s.logger.Error(ctx, "failed to send AI budget notifications",
|
||||
slog.F("initiator_id", intc.InitiatorID),
|
||||
slog.Error(err))
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
@@ -2420,10 +2420,16 @@ func TestRecordTokenUsageBudgetNotifications(t *testing.T) {
|
||||
Return(database.AIUserDailySpend{}, nil)
|
||||
db.EXPECT().GetUserAISpendSince(gomock.Any(), gomock.Any()).
|
||||
Return(database.GetUserAISpendSinceRow{SpendMicros: tc.newSpend}, nil)
|
||||
// The group is looked up once per crossing that notifies.
|
||||
db.EXPECT().GetGroupByID(gomock.Any(), groupID).
|
||||
Return(database.Group{ID: groupID, Name: "Engineering"}, nil).
|
||||
Times(len(tc.wantTemplates))
|
||||
// The group and user are resolved once per interception that
|
||||
// notifies, regardless of how many thresholds it crosses.
|
||||
if len(tc.wantTemplates) > 0 {
|
||||
db.EXPECT().GetGroupByID(gomock.Any(), groupID).
|
||||
Return(database.Group{ID: groupID, Name: "Engineering"}, nil)
|
||||
db.EXPECT().GetUserByID(gomock.Any(), intc.InitiatorID).
|
||||
Return(database.User{ID: intc.InitiatorID, Username: "bob"}, nil)
|
||||
// No admins configured, so only the user is notified.
|
||||
db.EXPECT().GetUsers(gomock.Any(), gomock.Any()).Return(nil, nil)
|
||||
}
|
||||
|
||||
ctx := testutil.Context(t, testutil.WaitLong)
|
||||
srv, err := aibridgedserver.NewServer(ctx, aibridgedserver.Options{
|
||||
@@ -2514,6 +2520,10 @@ func TestRecordTokenUsageBudgetNotificationAcrossPeriodBoundary(t *testing.T) {
|
||||
})
|
||||
db.EXPECT().GetGroupByID(gomock.Any(), groupID).
|
||||
Return(database.Group{ID: groupID, Name: "Engineering"}, nil)
|
||||
db.EXPECT().GetUserByID(gomock.Any(), intc.InitiatorID).
|
||||
Return(database.User{ID: intc.InitiatorID, Username: "bob"}, nil)
|
||||
// No admins configured, so only the user is notified.
|
||||
db.EXPECT().GetUsers(gomock.Any(), gomock.Any()).Return(nil, nil)
|
||||
|
||||
clock := quartz.NewMock(t)
|
||||
clock.Set(processedAt)
|
||||
@@ -2695,6 +2705,150 @@ func TestRecordTokenUsageBudgetNotificationZeroLimit(t *testing.T) {
|
||||
require.Empty(t, enq.Sent(), "a zero spend limit must not produce a notification")
|
||||
}
|
||||
|
||||
// TestRecordTokenUsageBudgetAdminNotification verifies that crossing a budget
|
||||
// threshold also notifies deployment owners and user admins, and that the
|
||||
// affected user is not double-notified as an admin.
|
||||
func TestRecordTokenUsageBudgetAdminNotification(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
const (
|
||||
dollar int64 = 1_000_000 // micros per USD dollar
|
||||
spendLimit int64 = 100 * dollar // $100 limit (100% threshold)
|
||||
warnAt int64 = 85 * dollar // $85 (85% of the limit)
|
||||
inputPrice int64 = dollar // $1 per million tokens
|
||||
tokensPerDollar int64 = 1_000_000 // 1,000,000 tokens = $1 at the price above
|
||||
)
|
||||
now := time.Date(2026, 6, 25, 14, 30, 0, 0, time.UTC)
|
||||
|
||||
cases := []struct {
|
||||
name string
|
||||
useOverride bool
|
||||
inputTokens int64
|
||||
newSpend int64
|
||||
wantThreshold string
|
||||
wantLimitSource string
|
||||
wantUserTemplate uuid.UUID
|
||||
wantAdminTemplate uuid.UUID
|
||||
}{
|
||||
{
|
||||
// A $1 interception takes spend to warnAt ($85): pre ($84) < warnAt <= post.
|
||||
name: "warning",
|
||||
inputTokens: tokensPerDollar,
|
||||
newSpend: warnAt,
|
||||
wantThreshold: "85",
|
||||
wantLimitSource: string(codersdk.AIBudgetLimitSourceGroup),
|
||||
wantUserTemplate: notifications.TemplateAIBudgetWarningUser,
|
||||
wantAdminTemplate: notifications.TemplateAIBudgetWarningAdmin,
|
||||
},
|
||||
{
|
||||
// A $5 interception takes spend from $95 to the $100 limit, so the limit threshold is crossed.
|
||||
name: "limit reached",
|
||||
inputTokens: 5 * tokensPerDollar,
|
||||
newSpend: spendLimit,
|
||||
wantThreshold: "100",
|
||||
wantLimitSource: string(codersdk.AIBudgetLimitSourceGroup),
|
||||
wantUserTemplate: notifications.TemplateAIBudgetLimitReachedUser,
|
||||
wantAdminTemplate: notifications.TemplateAIBudgetLimitReachedAdmin,
|
||||
},
|
||||
{
|
||||
// A per-user override supplies the limit, so limit_source is user_override.
|
||||
name: "warning, per-user override",
|
||||
useOverride: true,
|
||||
inputTokens: tokensPerDollar,
|
||||
newSpend: warnAt,
|
||||
wantThreshold: "85",
|
||||
wantLimitSource: string(codersdk.AIBudgetLimitSourceUserOverride),
|
||||
wantUserTemplate: notifications.TemplateAIBudgetWarningUser,
|
||||
wantAdminTemplate: notifications.TemplateAIBudgetWarningAdmin,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tc := range cases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
price := &database.AIModelPrice{
|
||||
InputPrice: sql.NullInt64{Int64: inputPrice, Valid: true},
|
||||
}
|
||||
|
||||
ctrl := gomock.NewController(t)
|
||||
db := dbmock.NewMockStore(ctrl)
|
||||
enq := ¬ificationstest.FakeEnqueuer{}
|
||||
|
||||
intc := newTestInterception(uuid.New())
|
||||
groupID := uuid.New()
|
||||
|
||||
admin := database.GetUsersRow{ID: uuid.New(), Username: "admin1"}
|
||||
// The affected user is also an admin; they must not receive the admin copy.
|
||||
selfAdmin := database.GetUsersRow{ID: intc.InitiatorID, Username: "bob"}
|
||||
|
||||
// Whether the limit comes from a group budget or a per-user override, the
|
||||
// spend is attributed to the same effective group, so the group lookup and
|
||||
// notifications are identical apart from the limit_source label.
|
||||
if tc.useOverride {
|
||||
override := &database.UserAIBudgetOverride{GroupID: groupID, SpendLimitMicros: spendLimit}
|
||||
expectTokenUsageCostLookups(db, intc, override, nil, nil, price)
|
||||
} else {
|
||||
group := &database.GetHighestGroupAIBudgetByUserRow{GroupID: groupID, SpendLimitMicros: spendLimit}
|
||||
expectTokenUsageCostLookups(db, intc, nil, group, nil, price)
|
||||
}
|
||||
db.EXPECT().InTx(gomock.Any(), nil).DoAndReturn(
|
||||
func(fn func(database.Store) error, _ *database.TxOptions) error { return fn(db) },
|
||||
)
|
||||
db.EXPECT().InsertAIBridgeTokenUsage(gomock.Any(), gomock.Any()).
|
||||
Return(database.AIBridgeTokenUsage{ID: uuid.New(), InterceptionID: intc.ID}, nil)
|
||||
db.EXPECT().IncrementUserAIDailySpend(gomock.Any(), gomock.Any()).
|
||||
Return(database.AIUserDailySpend{}, nil)
|
||||
db.EXPECT().GetUserAISpendSince(gomock.Any(), gomock.Any()).
|
||||
Return(database.GetUserAISpendSinceRow{SpendMicros: tc.newSpend}, nil)
|
||||
db.EXPECT().GetGroupByID(gomock.Any(), groupID).
|
||||
Return(database.Group{ID: groupID, Name: "Engineering"}, nil)
|
||||
db.EXPECT().GetUserByID(gomock.Any(), intc.InitiatorID).
|
||||
Return(database.User{ID: intc.InitiatorID, Username: "bob"}, nil)
|
||||
db.EXPECT().GetUsers(gomock.Any(), database.GetUsersParams{
|
||||
RbacRole: []string{codersdk.RoleOwner, codersdk.RoleUserAdmin},
|
||||
}).Return([]database.GetUsersRow{admin, selfAdmin}, nil)
|
||||
|
||||
ctx := testutil.Context(t, testutil.WaitLong)
|
||||
srv, err := aibridgedserver.NewServer(ctx, aibridgedserver.Options{
|
||||
Store: db,
|
||||
AISeatTracker: agplaiseats.Noop{},
|
||||
AccessURL: "/",
|
||||
GatewayCfg: codersdk.AIBridgeConfig{},
|
||||
Experiments: requiredExperiments,
|
||||
Enqueuer: enq,
|
||||
Logger: testutil.Logger(t),
|
||||
Clock: quartz.NewReal(),
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
_, err = srv.RecordTokenUsage(ctx, &proto.RecordTokenUsageRequest{
|
||||
InterceptionId: intc.ID.String(),
|
||||
MsgId: "msg_123",
|
||||
InputTokens: tc.inputTokens,
|
||||
CreatedAt: timestamppb.New(now),
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
// The user who crossed the threshold gets the user-facing notification.
|
||||
userSent := enq.Sent(notificationstest.WithTemplateID(tc.wantUserTemplate))
|
||||
require.Len(t, userSent, 1)
|
||||
require.Equal(t, intc.InitiatorID, userSent[0].UserID)
|
||||
|
||||
// The admin (but not the affected user, who is also an admin) gets the
|
||||
// admin notification naming the affected user.
|
||||
adminSent := enq.Sent(notificationstest.WithTemplateID(tc.wantAdminTemplate))
|
||||
require.Len(t, adminSent, 1)
|
||||
require.Equal(t, admin.ID, adminSent[0].UserID)
|
||||
require.Equal(t, "bob", adminSent[0].Labels["username"])
|
||||
require.Equal(t, tc.wantThreshold, adminSent[0].Labels["threshold"])
|
||||
require.Equal(t, "$100.00", adminSent[0].Labels["limit"])
|
||||
require.Equal(t, "Engineering", adminSent[0].Labels["effective_group_name"])
|
||||
require.Equal(t, tc.wantLimitSource, adminSent[0].Labels["limit_source"])
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// newTestInterception returns an interception with a fixed initiator, provider,
|
||||
// and model for cost-attribution test setup.
|
||||
func newTestInterception(id uuid.UUID) database.AIBridgeInterception {
|
||||
|
||||
@@ -12,6 +12,7 @@ import (
|
||||
"github.com/coder/coder/v2/coderd/aibridge/budget"
|
||||
"github.com/coder/coder/v2/coderd/aibridged/proto"
|
||||
"github.com/coder/coder/v2/coderd/database"
|
||||
"github.com/coder/coder/v2/codersdk"
|
||||
)
|
||||
|
||||
// tokensPerMillion is the divisor for prices, which are quoted per million
|
||||
@@ -24,6 +25,7 @@ const tokensPerMillion = 1_000_000
|
||||
type tokenUsageCost struct {
|
||||
effectiveGroupID uuid.NullUUID
|
||||
spendLimitMicros sql.NullInt64
|
||||
limitSource codersdk.AIBudgetLimitSource
|
||||
inputPriceMicros sql.NullInt64
|
||||
outputPriceMicros sql.NullInt64
|
||||
cacheReadPriceMicros sql.NullInt64
|
||||
@@ -54,9 +56,10 @@ func (s *Server) resolveTokenUsageCost(ctx context.Context, intc database.AIBrid
|
||||
} else {
|
||||
result.effectiveGroupID = uuid.NullUUID{UUID: effectiveGroup.GroupID, Valid: true}
|
||||
// Limit is nil for the unlimited Everyone fallback; only a budgeted
|
||||
// group carries the spend limit.
|
||||
// group carries the spend limit and its source.
|
||||
if effectiveGroup.Limit != nil {
|
||||
result.spendLimitMicros = sql.NullInt64{Int64: effectiveGroup.Limit.SpendLimitMicros, Valid: true}
|
||||
result.limitSource = effectiveGroup.Limit.Source
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -2,6 +2,7 @@ package aibridgedserver
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"strconv"
|
||||
"time"
|
||||
@@ -13,6 +14,7 @@ import (
|
||||
"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
|
||||
@@ -25,18 +27,27 @@ const (
|
||||
// budgetNotificationsCreatedBy records what enqueued AI budget notifications.
|
||||
const budgetNotificationsCreatedBy = "aigateway"
|
||||
|
||||
// budgetThreshold pairs a percentage of the spend limit with the notification
|
||||
// template sent when a user's spend crosses it.
|
||||
// budgetThreshold pairs a percentage of the spend limit with the notifications
|
||||
// sent when a user's spend crosses it.
|
||||
type budgetThreshold struct {
|
||||
percent int
|
||||
notificationTemplate uuid.UUID
|
||||
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, notificationTemplate: notifications.TemplateAIBudgetWarningUser},
|
||||
{percent: limitThresholdPercent, notificationTemplate: notifications.TemplateAIBudgetLimitReachedUser},
|
||||
{
|
||||
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
|
||||
@@ -44,11 +55,13 @@ var budgetThresholds = []budgetThreshold{
|
||||
// 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
|
||||
notificationTemplate uuid.UUID
|
||||
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
|
||||
@@ -93,46 +106,100 @@ func (s *Server) detectBudgetThresholdCrossings(ctx context.Context, tx database
|
||||
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,
|
||||
notificationTemplate: t.notificationTemplate,
|
||||
periodStart: period.Start,
|
||||
periodEnd: period.End,
|
||||
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
|
||||
}
|
||||
|
||||
// notifyBudgetThresholdCrossing enqueues the notification for the user who
|
||||
// crossed the threshold.
|
||||
func (s *Server) notifyBudgetThresholdCrossing(ctx context.Context, crossing budgetThresholdCrossing) error {
|
||||
group, err := s.store.GetGroupByID(ctx, crossing.effectiveGroupID)
|
||||
if err != nil {
|
||||
return xerrors.Errorf("look up group %q: %w", crossing.effectiveGroupID, err)
|
||||
// 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
|
||||
}
|
||||
|
||||
labels := map[string]string{
|
||||
"threshold": strconv.Itoa(crossing.thresholdPercent),
|
||||
"limit": formatSpendLimit(crossing.spendLimitMicros),
|
||||
"period": s.budgetPeriod.Adjective(),
|
||||
"effective_group_name": group.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": crossing.periodStart.UTC().Format("January 2, 2006"),
|
||||
"period_end": crossing.periodEnd.UTC().Format("January 2, 2006"),
|
||||
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)
|
||||
}
|
||||
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.
|
||||
if _, err := s.notifEnqueuer.EnqueueWithData(dbauthz.AsNotifier(ctx), crossing.userID, crossing.notificationTemplate,
|
||||
labels, nil, budgetNotificationsCreatedBy,
|
||||
crossing.effectiveGroupID,
|
||||
); err != nil {
|
||||
return xerrors.Errorf("enqueue notification: %w", err)
|
||||
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,
|
||||
// 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 nil
|
||||
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.
|
||||
|
||||
@@ -0,0 +1,5 @@
|
||||
DELETE FROM notification_templates
|
||||
WHERE id IN (
|
||||
'2a7b0ac1-00e1-4625-9cd5-1e5933972c77',
|
||||
'0bafe0ea-a78b-4217-ad05-1ef12e92e025'
|
||||
);
|
||||
@@ -0,0 +1,61 @@
|
||||
INSERT INTO notification_templates (
|
||||
id,
|
||||
name,
|
||||
title_template,
|
||||
body_template,
|
||||
actions,
|
||||
"group",
|
||||
method,
|
||||
kind,
|
||||
enabled_by_default
|
||||
)
|
||||
VALUES (
|
||||
'2a7b0ac1-00e1-4625-9cd5-1e5933972c77',
|
||||
'User Approaching AI Budget Limit',
|
||||
E'{{.Labels.username}} is approaching their {{.Labels.period}} AI budget limit',
|
||||
$$User **{{.Labels.username}}** has used more than {{.Labels.threshold}}% of their {{.Labels.period}} AI budget ({{.Labels.limit}}).
|
||||
|
||||
Effective group: **{{.Labels.effective_group_name}}**
|
||||
{{- if eq .Labels.limit_source "user_override"}}
|
||||
|
||||
This limit is a per-user override.
|
||||
{{- end}}
|
||||
|
||||
AI budget period: {{.Labels.period_start}} - {{.Labels.period_end}}$$,
|
||||
'[]'::jsonb,
|
||||
'AI Cost Control Admin Events',
|
||||
NULL,
|
||||
'system'::notification_template_kind,
|
||||
true
|
||||
);
|
||||
|
||||
INSERT INTO notification_templates (
|
||||
id,
|
||||
name,
|
||||
title_template,
|
||||
body_template,
|
||||
actions,
|
||||
"group",
|
||||
method,
|
||||
kind,
|
||||
enabled_by_default
|
||||
)
|
||||
VALUES (
|
||||
'0bafe0ea-a78b-4217-ad05-1ef12e92e025',
|
||||
'User Reached AI Budget Limit',
|
||||
E'{{.Labels.username}} has reached their {{.Labels.period}} AI budget limit',
|
||||
$$User **{{.Labels.username}}** has reached their {{.Labels.period}} AI budget limit ({{.Labels.limit}}). Subsequent requests will be blocked.
|
||||
|
||||
Effective group: **{{.Labels.effective_group_name}}**
|
||||
{{- if eq .Labels.limit_source "user_override"}}
|
||||
|
||||
This limit is a per-user override.
|
||||
{{- end}}
|
||||
|
||||
AI budget period: {{.Labels.period_start}} - {{.Labels.period_end}}$$,
|
||||
'[]'::jsonb,
|
||||
'AI Cost Control Admin Events',
|
||||
NULL,
|
||||
'system'::notification_template_kind,
|
||||
true
|
||||
);
|
||||
@@ -72,6 +72,8 @@ var (
|
||||
|
||||
// AI cost control related events.
|
||||
var (
|
||||
TemplateAIBudgetWarningUser = uuid.MustParse("b5db9597-de2a-4dea-87e9-25cee6906b86")
|
||||
TemplateAIBudgetLimitReachedUser = uuid.MustParse("cdcf2ecd-f003-4169-9800-abb2661ea522")
|
||||
TemplateAIBudgetWarningUser = uuid.MustParse("b5db9597-de2a-4dea-87e9-25cee6906b86")
|
||||
TemplateAIBudgetLimitReachedUser = uuid.MustParse("cdcf2ecd-f003-4169-9800-abb2661ea522")
|
||||
TemplateAIBudgetWarningAdmin = uuid.MustParse("2a7b0ac1-00e1-4625-9cd5-1e5933972c77")
|
||||
TemplateAIBudgetLimitReachedAdmin = uuid.MustParse("0bafe0ea-a78b-4217-ad05-1ef12e92e025")
|
||||
)
|
||||
|
||||
@@ -1497,6 +1497,45 @@ func TestNotificationTemplates_Golden(t *testing.T) {
|
||||
Data: map[string]any{},
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "TemplateAIBudgetWarningAdmin",
|
||||
id: notifications.TemplateAIBudgetWarningAdmin,
|
||||
payload: types.MessagePayload{
|
||||
UserName: "Bobby",
|
||||
UserEmail: "bobby@coder.com",
|
||||
UserUsername: "bobby",
|
||||
Labels: map[string]string{
|
||||
"username": "alice",
|
||||
"threshold": "85",
|
||||
"limit": "$1000.00",
|
||||
"period": "monthly",
|
||||
"limit_source": "group",
|
||||
"effective_group_name": "Engineering",
|
||||
"period_start": "July 1, 2026",
|
||||
"period_end": "August 1, 2026",
|
||||
},
|
||||
Data: map[string]any{},
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "TemplateAIBudgetLimitReachedAdmin",
|
||||
id: notifications.TemplateAIBudgetLimitReachedAdmin,
|
||||
payload: types.MessagePayload{
|
||||
UserName: "Bobby",
|
||||
UserEmail: "bobby@coder.com",
|
||||
UserUsername: "bobby",
|
||||
Labels: map[string]string{
|
||||
"username": "alice",
|
||||
"limit": "$1000.00",
|
||||
"period": "monthly",
|
||||
"limit_source": "user_override",
|
||||
"effective_group_name": "Engineering",
|
||||
"period_start": "July 1, 2026",
|
||||
"period_end": "August 1, 2026",
|
||||
},
|
||||
Data: map[string]any{},
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
// We must have a test case for every notification_template. This is enforced below:
|
||||
|
||||
Vendored
+82
@@ -0,0 +1,82 @@
|
||||
From: system@coder.com
|
||||
To: bobby@coder.com
|
||||
Subject: alice has reached their monthly AI budget limit
|
||||
Message-Id: 02ee4935-73be-4fa1-a290-ff9999026b13@blush-whale-48
|
||||
Date: Fri, 11 Oct 2024 09:03:06 +0000
|
||||
Content-Type: multipart/alternative; boundary=bbe61b741255b6098bb6b3c1f41b885773df633cb18d2a3002b68e4bc9c4
|
||||
MIME-Version: 1.0
|
||||
|
||||
--bbe61b741255b6098bb6b3c1f41b885773df633cb18d2a3002b68e4bc9c4
|
||||
Content-Transfer-Encoding: quoted-printable
|
||||
Content-Type: text/plain; charset=UTF-8
|
||||
|
||||
Hi Bobby,
|
||||
|
||||
User alice has reached their monthly AI budget limit ($1000.00). Subsequent=
|
||||
requests will be blocked.
|
||||
|
||||
Effective group: Engineering
|
||||
|
||||
This limit is a per-user override.
|
||||
|
||||
AI budget period: July 1, 2026 - August 1, 2026
|
||||
|
||||
|
||||
--bbe61b741255b6098bb6b3c1f41b885773df633cb18d2a3002b68e4bc9c4
|
||||
Content-Transfer-Encoding: quoted-printable
|
||||
Content-Type: text/html; charset=UTF-8
|
||||
|
||||
<!doctype html>
|
||||
<html lang=3D"en">
|
||||
<head>
|
||||
<meta charset=3D"UTF-8" />
|
||||
<meta name=3D"viewport" content=3D"width=3Ddevice-width, initial-scale=
|
||||
=3D1.0" />
|
||||
<title>alice has reached their monthly AI budget limit</title>
|
||||
</head>
|
||||
<body style=3D"margin: 0; padding: 0; font-family: -apple-system, system-=
|
||||
ui, BlinkMacSystemFont, 'Segoe UI', 'Roboto', 'Oxygen', 'Ubuntu', 'Cantarel=
|
||||
l', 'Fira Sans', 'Droid Sans', 'Helvetica Neue', sans-serif; color: #020617=
|
||||
; background: #f8fafc;">
|
||||
<div style=3D"max-width: 600px; margin: 20px auto; padding: 60px; borde=
|
||||
r: 1px solid #e2e8f0; border-radius: 8px; background-color: #fff; text-alig=
|
||||
n: left; font-size: 14px; line-height: 1.5;">
|
||||
<div style=3D"text-align: center;">
|
||||
<img src=3D"https://coder.com/coder-logo-horizontal.png" alt=3D"Cod=
|
||||
er Logo" style=3D"height: 40px;" />
|
||||
</div>
|
||||
<h1 style=3D"text-align: center; font-size: 24px; font-weight: 400; m=
|
||||
argin: 8px 0 32px; line-height: 1.5;">
|
||||
alice has reached their monthly AI budget limit
|
||||
</h1>
|
||||
<div style=3D"line-height: 1.5;">
|
||||
<p>Hi Bobby,</p>
|
||||
<p>User <strong>alice</strong> has reached their monthly AI budget =
|
||||
limit ($1000.00). Subsequent requests will be blocked.</p>
|
||||
|
||||
<p>Effective group: <strong>Engineering</strong></p>
|
||||
|
||||
<p>This limit is a per-user override.</p>
|
||||
|
||||
<p>AI budget period: July 1, 2026 - August 1, 2026</p>
|
||||
</div>
|
||||
<div style=3D"text-align: center; margin-top: 32px;">
|
||||
=20
|
||||
</div>
|
||||
<div style=3D"border-top: 1px solid #e2e8f0; color: #475569; font-siz=
|
||||
e: 12px; margin-top: 64px; padding-top: 24px; line-height: 1.6;">
|
||||
<p>© 2024 Coder. All rights reserved - <a =
|
||||
href=3D"http://test.com" style=3D"color: #2563eb; text-decoration: none;">h=
|
||||
ttp://test.com</a></p>
|
||||
<p><a href=3D"http://test.com/settings/notifications" style=3D"colo=
|
||||
r: #2563eb; text-decoration: none;">Click here to manage your notification =
|
||||
settings</a></p>
|
||||
<p><a href=3D"http://test.com/settings/notifications?disabled=3D0ba=
|
||||
fe0ea-a78b-4217-ad05-1ef12e92e025" style=3D"color: #2563eb; text-decoration=
|
||||
: none;">Stop receiving emails like this</a></p>
|
||||
</div>
|
||||
</div>
|
||||
</body>
|
||||
</html>
|
||||
|
||||
--bbe61b741255b6098bb6b3c1f41b885773df633cb18d2a3002b68e4bc9c4--
|
||||
Vendored
+77
@@ -0,0 +1,77 @@
|
||||
From: system@coder.com
|
||||
To: bobby@coder.com
|
||||
Subject: alice is approaching their monthly AI budget limit
|
||||
Message-Id: 02ee4935-73be-4fa1-a290-ff9999026b13@blush-whale-48
|
||||
Date: Fri, 11 Oct 2024 09:03:06 +0000
|
||||
Content-Type: multipart/alternative; boundary=bbe61b741255b6098bb6b3c1f41b885773df633cb18d2a3002b68e4bc9c4
|
||||
MIME-Version: 1.0
|
||||
|
||||
--bbe61b741255b6098bb6b3c1f41b885773df633cb18d2a3002b68e4bc9c4
|
||||
Content-Transfer-Encoding: quoted-printable
|
||||
Content-Type: text/plain; charset=UTF-8
|
||||
|
||||
Hi Bobby,
|
||||
|
||||
User alice has used more than 85% of their monthly AI budget ($1000.00).
|
||||
|
||||
Effective group: Engineering
|
||||
|
||||
AI budget period: July 1, 2026 - August 1, 2026
|
||||
|
||||
|
||||
--bbe61b741255b6098bb6b3c1f41b885773df633cb18d2a3002b68e4bc9c4
|
||||
Content-Transfer-Encoding: quoted-printable
|
||||
Content-Type: text/html; charset=UTF-8
|
||||
|
||||
<!doctype html>
|
||||
<html lang=3D"en">
|
||||
<head>
|
||||
<meta charset=3D"UTF-8" />
|
||||
<meta name=3D"viewport" content=3D"width=3Ddevice-width, initial-scale=
|
||||
=3D1.0" />
|
||||
<title>alice is approaching their monthly AI budget limit</title>
|
||||
</head>
|
||||
<body style=3D"margin: 0; padding: 0; font-family: -apple-system, system-=
|
||||
ui, BlinkMacSystemFont, 'Segoe UI', 'Roboto', 'Oxygen', 'Ubuntu', 'Cantarel=
|
||||
l', 'Fira Sans', 'Droid Sans', 'Helvetica Neue', sans-serif; color: #020617=
|
||||
; background: #f8fafc;">
|
||||
<div style=3D"max-width: 600px; margin: 20px auto; padding: 60px; borde=
|
||||
r: 1px solid #e2e8f0; border-radius: 8px; background-color: #fff; text-alig=
|
||||
n: left; font-size: 14px; line-height: 1.5;">
|
||||
<div style=3D"text-align: center;">
|
||||
<img src=3D"https://coder.com/coder-logo-horizontal.png" alt=3D"Cod=
|
||||
er Logo" style=3D"height: 40px;" />
|
||||
</div>
|
||||
<h1 style=3D"text-align: center; font-size: 24px; font-weight: 400; m=
|
||||
argin: 8px 0 32px; line-height: 1.5;">
|
||||
alice is approaching their monthly AI budget limit
|
||||
</h1>
|
||||
<div style=3D"line-height: 1.5;">
|
||||
<p>Hi Bobby,</p>
|
||||
<p>User <strong>alice</strong> has used more than 85% of their mont=
|
||||
hly AI budget ($1000.00).</p>
|
||||
|
||||
<p>Effective group: <strong>Engineering</strong></p>
|
||||
|
||||
<p>AI budget period: July 1, 2026 - August 1, 2026</p>
|
||||
</div>
|
||||
<div style=3D"text-align: center; margin-top: 32px;">
|
||||
=20
|
||||
</div>
|
||||
<div style=3D"border-top: 1px solid #e2e8f0; color: #475569; font-siz=
|
||||
e: 12px; margin-top: 64px; padding-top: 24px; line-height: 1.6;">
|
||||
<p>© 2024 Coder. All rights reserved - <a =
|
||||
href=3D"http://test.com" style=3D"color: #2563eb; text-decoration: none;">h=
|
||||
ttp://test.com</a></p>
|
||||
<p><a href=3D"http://test.com/settings/notifications" style=3D"colo=
|
||||
r: #2563eb; text-decoration: none;">Click here to manage your notification =
|
||||
settings</a></p>
|
||||
<p><a href=3D"http://test.com/settings/notifications?disabled=3D2a7=
|
||||
b0ac1-00e1-4625-9cd5-1e5933972c77" style=3D"color: #2563eb; text-decoration=
|
||||
: none;">Stop receiving emails like this</a></p>
|
||||
</div>
|
||||
</div>
|
||||
</body>
|
||||
</html>
|
||||
|
||||
--bbe61b741255b6098bb6b3c1f41b885773df633cb18d2a3002b68e4bc9c4--
|
||||
+29
@@ -0,0 +1,29 @@
|
||||
{
|
||||
"_version": "1.1",
|
||||
"msg_id": "00000000-0000-0000-0000-000000000000",
|
||||
"payload": {
|
||||
"_version": "1.2",
|
||||
"notification_name": "User Reached AI Budget Limit",
|
||||
"notification_template_id": "00000000-0000-0000-0000-000000000000",
|
||||
"user_id": "00000000-0000-0000-0000-000000000000",
|
||||
"user_email": "bobby@coder.com",
|
||||
"user_name": "Bobby",
|
||||
"user_username": "bobby",
|
||||
"actions": [],
|
||||
"labels": {
|
||||
"effective_group_name": "Engineering",
|
||||
"limit": "$1000.00",
|
||||
"limit_source": "user_override",
|
||||
"period": "monthly",
|
||||
"period_end": "August 1, 2026",
|
||||
"period_start": "July 1, 2026",
|
||||
"username": "alice"
|
||||
},
|
||||
"data": {},
|
||||
"targets": null
|
||||
},
|
||||
"title": "alice has reached their monthly AI budget limit",
|
||||
"title_markdown": "alice has reached their monthly AI budget limit",
|
||||
"body": "User alice has reached their monthly AI budget limit ($1000.00). Subsequent requests will be blocked.\n\nEffective group: Engineering\n\nThis limit is a per-user override.\n\nAI budget period: July 1, 2026 - August 1, 2026",
|
||||
"body_markdown": "User **alice** has reached their monthly AI budget limit ($1000.00). Subsequent requests will be blocked.\n\nEffective group: **Engineering**\n\nThis limit is a per-user override.\n\nAI budget period: July 1, 2026 - August 1, 2026"
|
||||
}
|
||||
Vendored
+30
@@ -0,0 +1,30 @@
|
||||
{
|
||||
"_version": "1.1",
|
||||
"msg_id": "00000000-0000-0000-0000-000000000000",
|
||||
"payload": {
|
||||
"_version": "1.2",
|
||||
"notification_name": "User Approaching AI Budget Limit",
|
||||
"notification_template_id": "00000000-0000-0000-0000-000000000000",
|
||||
"user_id": "00000000-0000-0000-0000-000000000000",
|
||||
"user_email": "bobby@coder.com",
|
||||
"user_name": "Bobby",
|
||||
"user_username": "bobby",
|
||||
"actions": [],
|
||||
"labels": {
|
||||
"effective_group_name": "Engineering",
|
||||
"limit": "$1000.00",
|
||||
"limit_source": "group",
|
||||
"period": "monthly",
|
||||
"period_end": "August 1, 2026",
|
||||
"period_start": "July 1, 2026",
|
||||
"threshold": "85",
|
||||
"username": "alice"
|
||||
},
|
||||
"data": {},
|
||||
"targets": null
|
||||
},
|
||||
"title": "alice is approaching their monthly AI budget limit",
|
||||
"title_markdown": "alice is approaching their monthly AI budget limit",
|
||||
"body": "User alice has used more than 85% of their monthly AI budget ($1000.00).\n\nEffective group: Engineering\n\nAI budget period: July 1, 2026 - August 1, 2026",
|
||||
"body_markdown": "User **alice** has used more than 85% of their monthly AI budget ($1000.00).\n\nEffective group: **Engineering**\n\nAI budget period: July 1, 2026 - August 1, 2026"
|
||||
}
|
||||
Reference in New Issue
Block a user