From ce4ee923c2bb2ca319a6af8a6e92899a79f198a0 Mon Sep 17 00:00:00 2001 From: Yevhenii Shcherbina Date: Mon, 27 Jul 2026 12:09:21 -0400 Subject: [PATCH] feat: notify users when AI spend crosses the budget threshold (#27346) Implements: https://linear.app/codercom/issue/AIGOV-289/notify-users-and-admins-on-budget-warning-and-limit-reached Notify users when their AI spend crosses a budget threshold for their effective group. Two thresholds are covered: a warning at 85%, and a limit-reached notification at 100%. Detection runs on the post-response path, right after the interception's cost is added to the user's daily spend. It reads the user's AI spend on the same transaction where token usage is recorded and AI daily spend is incremented, and derives the pre-interception total by subtracting this interception's cost. In case of `oldSpend < threshold && newSpend >= threshold` - notification is sent. A single interception that crosses both thresholds enqueues both notifications. Detection and delivery are best-effort: a failure is logged and never fails usage recording. The payload uses only stable values (the threshold percentage and the spend limit, not the exact spend), so duplicate enqueues are deduplicated by the notification system. The two templates are added via migration and appear in each user's notification settings under the "AI Budget" group. Admin notifications (owners and user admins) are a follow-up: #27415. ## Screenshots: image image --------- Co-authored-by: Cian Johnston --- coderd/aibridged.go | 1 + coderd/aibridgedserver/aibridgedserver.go | 47 +- .../aibridgedserver/aibridgedserver_test.go | 405 ++++++++++++++++++ coderd/aibridgedserver/cost.go | 6 + coderd/aibridgedserver/notifications.go | 141 ++++++ coderd/database/dbauthz/dbauthz.go | 1 + .../000552_ai_budget_notifications.down.sql | 5 + .../000552_ai_budget_notifications.up.sql | 53 +++ coderd/notifications/events.go | 6 + coderd/notifications/notifications_test.go | 36 ++ ...mplateAIBudgetLimitReachedUser.html.golden | 78 ++++ .../TemplateAIBudgetWarningUser.html.golden | 77 ++++ ...mplateAIBudgetLimitReachedUser.json.golden | 28 ++ .../TemplateAIBudgetWarningUser.json.golden | 28 ++ codersdk/deployment.go | 16 + codersdk/deployment_test.go | 13 + enterprise/coderd/aibridgeserve.go | 1 + .../NotificationsPage/NotificationsPage.tsx | 1 + 18 files changed, 940 insertions(+), 3 deletions(-) create mode 100644 coderd/aibridgedserver/notifications.go create mode 100644 coderd/database/migrations/000552_ai_budget_notifications.down.sql create mode 100644 coderd/database/migrations/000552_ai_budget_notifications.up.sql create mode 100644 coderd/notifications/testdata/rendered-templates/smtp/TemplateAIBudgetLimitReachedUser.html.golden create mode 100644 coderd/notifications/testdata/rendered-templates/smtp/TemplateAIBudgetWarningUser.html.golden create mode 100644 coderd/notifications/testdata/rendered-templates/webhook/TemplateAIBudgetLimitReachedUser.json.golden create mode 100644 coderd/notifications/testdata/rendered-templates/webhook/TemplateAIBudgetWarningUser.json.golden diff --git a/coderd/aibridged.go b/coderd/aibridged.go index d163fbfe09..2a88fc53fc 100644 --- a/coderd/aibridged.go +++ b/coderd/aibridged.go @@ -69,6 +69,7 @@ func (api *API) CreateInMemoryAIBridgeServer(dialCtx context.Context) (client ai Store: api.Database, Pubsub: api.Pubsub, AISeatTracker: api.AISeatTracker, + Enqueuer: api.NotificationsEnqueuer, AccessURL: api.AccessURL.String(), GatewayCfg: api.DeploymentValues.AI.BridgeConfig, ExternalAuthConfigs: api.ExternalAuthConfigs, diff --git a/coderd/aibridgedserver/aibridgedserver.go b/coderd/aibridgedserver/aibridgedserver.go index cf750944e9..62f8e1132e 100644 --- a/coderd/aibridgedserver/aibridgedserver.go +++ b/coderd/aibridgedserver/aibridgedserver.go @@ -30,6 +30,7 @@ import ( "github.com/coder/coder/v2/coderd/externalauth" "github.com/coder/coder/v2/coderd/httpmw" codermcp "github.com/coder/coder/v2/coderd/mcp" + "github.com/coder/coder/v2/coderd/notifications" coderdpubsub "github.com/coder/coder/v2/coderd/pubsub" "github.com/coder/coder/v2/coderd/util/ptr" "github.com/coder/coder/v2/codersdk" @@ -82,6 +83,7 @@ type store interface { GetHighestGroupAIBudgetByUser(ctx context.Context, userID uuid.UUID) (database.GetHighestGroupAIBudgetByUserRow, error) 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) // MCPConfigurator-related queries. GetExternalAuthLinksByUserID(ctx context.Context, userID uuid.UUID) ([]database.ExternalAuthLink, error) @@ -117,8 +119,9 @@ type Server struct { budgetPolicy codersdk.AIBudgetPolicy // budgetPeriod is the deployment-configured budgeting period used to // derive the window over which user AI spend is aggregated. - budgetPeriod codersdk.AIBudgetPeriod - clock quartz.Clock + budgetPeriod codersdk.AIBudgetPeriod + clock quartz.Clock + notifEnqueuer notifications.Enqueuer } // Options carries the dependencies required to construct an aibridged Server. @@ -126,6 +129,9 @@ type Options struct { Store store Pubsub pubsub.Pubsub AISeatTracker aiseats.SeatTracker + // Enqueuer enqueues notifications. When nil, NewServer substitutes a no-op + // enqueuer. + Enqueuer notifications.Enqueuer AccessURL string GatewayCfg codersdk.AIBridgeConfig @@ -137,6 +143,11 @@ type Options struct { } func NewServer(lifecycleCtx context.Context, opts Options) (*Server, error) { + enqueuer := opts.Enqueuer + if enqueuer == nil { + enqueuer = notifications.NewNoopEnqueuer() + } + eac := make(map[string]*externalauth.Config, len(opts.ExternalAuthConfigs)) for _, cfg := range opts.ExternalAuthConfigs { @@ -158,6 +169,7 @@ func NewServer(lifecycleCtx context.Context, opts Options) (*Server, error) { budgetPolicy: codersdk.NewAIBudgetPolicyFromString(opts.GatewayCfg.BudgetPolicy), budgetPeriod: codersdk.NewAIBudgetPeriodFromString(opts.GatewayCfg.BudgetPeriod), clock: opts.Clock, + notifEnqueuer: enqueuer, } if opts.GatewayCfg.InjectCoderMCPTools { @@ -357,7 +369,11 @@ func (s *Server) RecordTokenUsage(ctx context.Context, in *proto.RecordTokenUsag // positive, accumulates that cost into the user's daily spend. func (s *Server) recordTokenUsageAndSpend(ctx context.Context, intc database.AIBridgeInterception, cost tokenUsageCost, in *proto.RecordTokenUsageRequest, metadataJSON []byte) error { createdAt := in.GetCreatedAt().AsTime() - return s.store.InTx(func(tx database.Store) error { + + // Populated inside the transaction with any budget thresholds this + // interception crossed. + var crossings []budgetThresholdCrossing + err := s.store.InTx(func(tx database.Store) error { if _, err := tx.InsertAIBridgeTokenUsage(ctx, database.InsertAIBridgeTokenUsageParams{ ID: uuid.New(), InterceptionID: intc.ID, @@ -400,8 +416,33 @@ func (s *Server) recordTokenUsageAndSpend(ctx context.Context, intc database.AIB }); err != nil { return xerrors.Errorf("increment user daily spend: %w", err) } + + // Threshold detection is best-effort: a failed read must not roll back + // the committed spend, so the error is logged rather than propagated. + var detectErr error + crossings, detectErr = s.detectBudgetThresholdCrossings(ctx, tx, intc, cost, createdAt) + if detectErr != nil { + s.logger.Error(ctx, "failed to detect AI budget threshold crossing", + slog.F("interception_id", intc.ID), + slog.F("initiator_id", intc.InitiatorID), + slog.Error(detectErr)) + } return nil }, nil) + if err != nil { + 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)) + } + } + return nil } func (s *Server) RecordPromptUsage(ctx context.Context, in *proto.RecordPromptUsageRequest) (*proto.RecordPromptUsageResponse, error) { diff --git a/coderd/aibridgedserver/aibridgedserver_test.go b/coderd/aibridgedserver/aibridgedserver_test.go index ab6a479b05..1e54445b7a 100644 --- a/coderd/aibridgedserver/aibridgedserver_test.go +++ b/coderd/aibridgedserver/aibridgedserver_test.go @@ -44,6 +44,8 @@ import ( "github.com/coder/coder/v2/coderd/database/pubsub" "github.com/coder/coder/v2/coderd/externalauth" codermcp "github.com/coder/coder/v2/coderd/mcp" + "github.com/coder/coder/v2/coderd/notifications" + "github.com/coder/coder/v2/coderd/notifications/notificationstest" coderdpubsub "github.com/coder/coder/v2/coderd/pubsub" "github.com/coder/coder/v2/coderd/rbac" "github.com/coder/coder/v2/coderd/util/ptr" @@ -1707,6 +1709,9 @@ func TestRecordTokenUsage(t *testing.T) { Day: now.UTC().Truncate(24 * time.Hour), CostMicros: wantCost, }).Return(database.AIUserDailySpend{}, nil) + + db.EXPECT().GetUserAISpendSince(gomock.Any(), gomock.Any()). + Return(database.GetUserAISpendSinceRow{SpendMicros: wantCost}, nil) }, }, { @@ -1759,6 +1764,9 @@ func TestRecordTokenUsage(t *testing.T) { Day: now.UTC().Truncate(24 * time.Hour), CostMicros: wantCost, }).Return(database.AIUserDailySpend{}, nil) + + db.EXPECT().GetUserAISpendSince(gomock.Any(), gomock.Any()). + Return(database.GetUserAISpendSinceRow{SpendMicros: wantCost}, nil) }, }, { @@ -2290,6 +2298,403 @@ func TestRecordTokenUsageAuthorized(t *testing.T) { require.Equal(t, wantCost, spend.SpendMicros, "spend micros") } +// TestRecordTokenUsageBudgetNotifications verifies that recording token usage +// enqueues the right budget notifications: the warning template when spend +// crosses the warning threshold, the limit-reached template at 100%, both when +// a single interception crosses both, and nothing when no threshold is +// freshly crossed. +func TestRecordTokenUsageBudgetNotifications(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) + + // Each case sets its cost via inputTokens (tokensPerDollar tokens = $1). + // newSpend is the post-increment period total; the code derives the + // pre-increment total as newSpend - cost and fires a threshold only when the + // pre-increment total is below the threshold and the post-increment total is + // at or above it. + price := &database.AIModelPrice{ + InputPrice: sql.NullInt64{Int64: inputPrice, Valid: true}, + } + + // Threshold percentage label expected for each template. + wantThreshold := map[uuid.UUID]string{ + notifications.TemplateAIBudgetWarningUser: "85", + notifications.TemplateAIBudgetLimitReachedUser: "100", + } + + testCases := []struct { + name string + inputTokens int64 + newSpend int64 // post-increment period total + wantTemplates []uuid.UUID + }{ + { + name: "crosses warning threshold", + // pre = $84.50 (< $85), post = $85.50 (>= $85) -> warning. + inputTokens: tokensPerDollar, + newSpend: warnAt + dollar/2, + wantTemplates: []uuid.UUID{notifications.TemplateAIBudgetWarningUser}, + }, + { + name: "crosses warning threshold exactly", + // pre = $84 (< $85), post = $85 (>= $85) -> warning. + inputTokens: tokensPerDollar, + newSpend: warnAt, + wantTemplates: []uuid.UUID{notifications.TemplateAIBudgetWarningUser}, + }, + { + name: "stays below warning threshold", + // post = $84.50 (< $85) -> no crossing. + inputTokens: tokensPerDollar, + newSpend: warnAt - dollar/2, + wantTemplates: nil, + }, + { + name: "already at warning threshold", + // pre = $85 (not < $85), post = $86 -> no fresh crossing. + inputTokens: tokensPerDollar, + newSpend: warnAt + dollar, + wantTemplates: nil, + }, + { + name: "already above warning threshold", + // pre = $89 (>= $85, < $100), post = $90 -> no crossing. + inputTokens: tokensPerDollar, + newSpend: warnAt + 5*dollar, + wantTemplates: nil, + }, + { + name: "crosses limit", + // pre = $99.50 (>= $85, so no warning; < $100), post = $100.50 (>= $100) -> limit. + inputTokens: tokensPerDollar, + newSpend: spendLimit + dollar/2, + wantTemplates: []uuid.UUID{notifications.TemplateAIBudgetLimitReachedUser}, + }, + { + name: "crosses warning and limit in one interception", + // pre = $80 (< $85), post = $100 (>= $85 and >= $100) -> warning + limit. + inputTokens: 20 * tokensPerDollar, + newSpend: spendLimit, + wantTemplates: []uuid.UUID{ + notifications.TemplateAIBudgetWarningUser, + notifications.TemplateAIBudgetLimitReachedUser, + }, + }, + { + name: "already above limit", + // pre = $109 (>= $100), post = $110 -> no fresh crossing. + inputTokens: tokensPerDollar, + newSpend: spendLimit + 10*dollar, + wantTemplates: nil, + }, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + ctrl := gomock.NewController(t) + db := dbmock.NewMockStore(ctrl) + enq := ¬ificationstest.FakeEnqueuer{} + + intc := newTestInterception(uuid.New()) + groupID := uuid.New() + 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) + // 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)) + + 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) + + require.Len(t, enq.Sent(), len(tc.wantTemplates), "unexpected number of notifications") + for _, tmpl := range tc.wantTemplates { + sent := enq.Sent(notificationstest.WithTemplateID(tmpl)) + require.Len(t, sent, 1, "expected one notification for template %s", tmpl) + require.Equal(t, intc.InitiatorID, sent[0].UserID) + require.Equal(t, wantThreshold[tmpl], sent[0].Labels["threshold"]) + require.Equal(t, "$100.00", sent[0].Labels["limit"]) + require.Equal(t, "Engineering", sent[0].Labels["effective_group_name"]) + // The interception is recorded at 2026-06-25, so its budget + // period runs June 1 - July 1, 2026. + require.Equal(t, "June 1, 2026", sent[0].Labels["period_start"]) + require.Equal(t, "July 1, 2026", sent[0].Labels["period_end"]) + } + }) + } +} + +// TestRecordTokenUsageBudgetNotificationAcrossPeriodBoundary verifies that an +// interception created in one budget period but processed after the period has +// rolled over is still evaluated against the period it belongs to, so a genuine +// threshold crossing is detected rather than lost across the boundary. +func TestRecordTokenUsageBudgetNotificationAcrossPeriodBoundary(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 + ) + + // The interception was created in the final second of January but is + // processed just after the rollover into February. The spend is bucketed + // into January, so detection must sum against January's period. + createdAt := time.Date(2026, 1, 31, 23, 59, 59, 0, time.UTC) + processedAt := time.Date(2026, 2, 1, 0, 0, 1, 0, time.UTC) + januaryStart := time.Date(2026, 1, 1, 0, 0, 0, 0, time.UTC) + + ctrl := gomock.NewController(t) + db := dbmock.NewMockStore(ctrl) + enq := ¬ificationstest.FakeEnqueuer{} + + intc := newTestInterception(uuid.New()) + groupID := uuid.New() + group := &database.GetHighestGroupAIBudgetByUserRow{GroupID: groupID, SpendLimitMicros: spendLimit} + price := &database.AIModelPrice{InputPrice: sql.NullInt64{Int64: inputPrice, Valid: true}} + + 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) + + // The spend query must run against the period the interception belongs to + // (January), not the period it was processed in (February). + var gotPeriodStart time.Time + db.EXPECT().GetUserAISpendSince(gomock.Any(), gomock.Any()). + DoAndReturn(func(_ context.Context, p database.GetUserAISpendSinceParams) (database.GetUserAISpendSinceRow, error) { + gotPeriodStart = p.PeriodStart + return database.GetUserAISpendSinceRow{SpendMicros: warnAt}, nil + }) + db.EXPECT().GetGroupByID(gomock.Any(), groupID). + Return(database.Group{ID: groupID, Name: "Engineering"}, nil) + + clock := quartz.NewMock(t) + clock.Set(processedAt) + + 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: clock, + }) + require.NoError(t, err) + + _, err = srv.RecordTokenUsage(ctx, &proto.RecordTokenUsageRequest{ + InterceptionId: intc.ID.String(), + MsgId: "msg_boundary", + InputTokens: tokensPerDollar, + CreatedAt: timestamppb.New(createdAt), + }) + require.NoError(t, err) + + require.Equal(t, januaryStart, gotPeriodStart, + "spend must be summed against the period the interception belongs to, not the processing period") + + sent := enq.Sent(notificationstest.WithTemplateID(notifications.TemplateAIBudgetWarningUser)) + require.Len(t, sent, 1, "expected the crossing to be detected against the interception's period") +} + +// TestRecordTokenUsageBudgetNotificationBestEffort verifies that a failure while +// detecting or sending a budget notification is swallowed: the token usage and +// spend are still recorded (RecordTokenUsage returns no error) and no +// notification is enqueued. This guards the best-effort contract, e.g. that a +// detection error is not propagated out of the transaction (which would roll +// back the committed spend). +func TestRecordTokenUsageBudgetNotificationBestEffort(t *testing.T) { + t.Parallel() + + const ( + dollar int64 = 1_000_000 + spendLimit int64 = 100 * dollar + warnAt int64 = 85 * dollar + inputPrice int64 = dollar + tokensPerDollar int64 = 1_000_000 + ) + now := time.Date(2026, 6, 25, 14, 30, 0, 0, time.UTC) + price := &database.AIModelPrice{InputPrice: sql.NullInt64{Int64: inputPrice, Valid: true}} + + testCases := []struct { + name string + // spendSinceErr fails detection (the read inside the transaction); + // groupLookupErr fails the notification after a crossing is detected. + spendSinceErr error + groupLookupErr error + }{ + {name: "detection read fails", spendSinceErr: sql.ErrConnDone}, + {name: "group lookup fails", groupLookupErr: sql.ErrConnDone}, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + ctrl := gomock.NewController(t) + db := dbmock.NewMockStore(ctrl) + enq := ¬ificationstest.FakeEnqueuer{} + + intc := newTestInterception(uuid.New()) + groupID := uuid.New() + 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) }, + ) + // The token usage and spend are recorded before detection runs. + 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) + + switch { + case tc.spendSinceErr != nil: + // Detection fails; the group is never looked up. + db.EXPECT().GetUserAISpendSince(gomock.Any(), gomock.Any()). + Return(database.GetUserAISpendSinceRow{}, tc.spendSinceErr) + case tc.groupLookupErr != nil: + // A crossing is detected ($84 -> $85), but resolving the group + // for the notification fails. + db.EXPECT().GetUserAISpendSince(gomock.Any(), gomock.Any()). + Return(database.GetUserAISpendSinceRow{SpendMicros: warnAt}, nil) + db.EXPECT().GetGroupByID(gomock.Any(), groupID). + Return(database.Group{}, tc.groupLookupErr) + default: + t.Fatal("test case must set spendSinceErr or groupLookupErr") + } + + 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, + // The detect/notify failure is logged; ignore it here since + // triggering it is the point of the test. + Logger: slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}), + Clock: quartz.NewReal(), + }) + require.NoError(t, err) + + // The failure must not surface as an error from RecordTokenUsage. + _, err = srv.RecordTokenUsage(ctx, &proto.RecordTokenUsageRequest{ + InterceptionId: intc.ID.String(), + MsgId: "msg_123", + InputTokens: tokensPerDollar, + CreatedAt: timestamppb.New(now), + }) + require.NoError(t, err) + require.Empty(t, enq.Sent(), "no notification should be enqueued when detection or lookup fails") + }) + } +} + +// TestRecordTokenUsageBudgetNotificationZeroLimit verifies that a zero spend +// limit (used to block a group entirely) produces no budget notification: there +// is no meaningful threshold to cross, and such users are already blocked by +// pre-request enforcement. The token usage is still recorded. +func TestRecordTokenUsageBudgetNotificationZeroLimit(t *testing.T) { + t.Parallel() + + ctrl := gomock.NewController(t) + db := dbmock.NewMockStore(ctrl) + enq := ¬ificationstest.FakeEnqueuer{} + + intc := newTestInterception(uuid.New()) + groupID := uuid.New() + // A zero limit blocks the group; there is no threshold to cross. + group := &database.GetHighestGroupAIBudgetByUserRow{GroupID: groupID, SpendLimitMicros: 0} + price := &database.AIModelPrice{InputPrice: sql.NullInt64{Int64: 1_000_000, Valid: true}} + + 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) }, + ) + // The spend is still recorded; detection then short-circuits on the zero + // limit without reading spend or looking up the group. + 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) + + 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: 1_000_000, // $1 at $1 per million tokens + CreatedAt: timestamppb.New(time.Date(2026, 6, 25, 14, 30, 0, 0, time.UTC)), + }) + require.NoError(t, err) + require.Empty(t, enq.Sent(), "a zero spend limit must not produce a notification") +} + // newTestInterception returns an interception with a fixed initiator, provider, // and model for cost-attribution test setup. func newTestInterception(id uuid.UUID) database.AIBridgeInterception { diff --git a/coderd/aibridgedserver/cost.go b/coderd/aibridgedserver/cost.go index 87ba4ab646..79b59cd1b6 100644 --- a/coderd/aibridgedserver/cost.go +++ b/coderd/aibridgedserver/cost.go @@ -23,6 +23,7 @@ const tokensPerMillion = 1_000_000 // price or cost of 0 is recorded as 0, which is distinct from NULL. type tokenUsageCost struct { effectiveGroupID uuid.NullUUID + spendLimitMicros sql.NullInt64 inputPriceMicros sql.NullInt64 outputPriceMicros sql.NullInt64 cacheReadPriceMicros sql.NullInt64 @@ -52,6 +53,11 @@ func (s *Server) resolveTokenUsageCost(ctx context.Context, intc database.AIBrid slog.F("user_id", intc.InitiatorID)) } 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. + if effectiveGroup.Limit != nil { + result.spendLimitMicros = sql.NullInt64{Int64: effectiveGroup.Limit.SpendLimitMicros, Valid: true} + } } // Snapshot the price for this (provider, model) and compute cost. diff --git a/coderd/aibridgedserver/notifications.go b/coderd/aibridgedserver/notifications.go new file mode 100644 index 0000000000..4481401408 --- /dev/null +++ b/coderd/aibridgedserver/notifications.go @@ -0,0 +1,141 @@ +package aibridgedserver + +import ( + "context" + "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" +) + +// 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 notification +// template sent when a user's spend crosses it. +type budgetThreshold struct { + percent int + notificationTemplate 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}, +} + +// 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 + notificationTemplate uuid.UUID + // 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, + notificationTemplate: t.notificationTemplate, + 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) + } + + 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"), + } + + //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) + } + return nil +} + +// formatSpendLimit renders a spend limit as a USD string. +func formatSpendLimit(micros int64) string { + return fmt.Sprintf("$%.2f", float64(micros)/1_000_000) +} diff --git a/coderd/database/dbauthz/dbauthz.go b/coderd/database/dbauthz/dbauthz.go index 578fb3ab18..6d24264d47 100644 --- a/coderd/database/dbauthz/dbauthz.go +++ b/coderd/database/dbauthz/dbauthz.go @@ -696,6 +696,7 @@ var ( rbac.ResourceAiModelPrice.Type: {policy.ActionRead, policy.ActionUpdate}, // Read: per-interception cost lookup. Update: startup price seeder. rbac.ResourceAiSeat.Type: {policy.ActionCreate}, // Required for UpsertAISeatState. rbac.ResourceAIProvider.Type: {policy.ActionRead}, // Required to load the provider snapshot (and per-provider keys) at startup. + rbac.ResourceGroup.Type: {policy.ActionRead}, // Required to read the effective group. }), User: []rbac.Permission{}, ByOrgID: map[string]rbac.OrgPermissions{}, diff --git a/coderd/database/migrations/000552_ai_budget_notifications.down.sql b/coderd/database/migrations/000552_ai_budget_notifications.down.sql new file mode 100644 index 0000000000..33d74fdb99 --- /dev/null +++ b/coderd/database/migrations/000552_ai_budget_notifications.down.sql @@ -0,0 +1,5 @@ +DELETE FROM notification_templates +WHERE id IN ( + 'b5db9597-de2a-4dea-87e9-25cee6906b86', + 'cdcf2ecd-f003-4169-9800-abb2661ea522' +); diff --git a/coderd/database/migrations/000552_ai_budget_notifications.up.sql b/coderd/database/migrations/000552_ai_budget_notifications.up.sql new file mode 100644 index 0000000000..61535f4f72 --- /dev/null +++ b/coderd/database/migrations/000552_ai_budget_notifications.up.sql @@ -0,0 +1,53 @@ +INSERT INTO notification_templates ( + id, + name, + title_template, + body_template, + actions, + "group", + method, + kind, + enabled_by_default +) +VALUES ( + 'b5db9597-de2a-4dea-87e9-25cee6906b86', + 'AI Budget Warning', + E'You''re approaching your {{.Labels.period}} AI budget limit', + $$You have used more than {{.Labels.threshold}}% of your {{.Labels.period}} AI budget ({{.Labels.limit}}). + +Effective group: **{{.Labels.effective_group_name}}** + +AI budget period: {{.Labels.period_start}} - {{.Labels.period_end}}$$, + '[]'::jsonb, + 'AI Cost Control 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 ( + 'cdcf2ecd-f003-4169-9800-abb2661ea522', + 'AI Budget Limit Reached', + E'You''ve reached your {{.Labels.period}} AI budget limit', + $$You have reached your {{.Labels.period}} AI budget limit ({{.Labels.limit}}). Subsequent requests will be blocked. + +Effective group: **{{.Labels.effective_group_name}}** + +AI budget period: {{.Labels.period_start}} - {{.Labels.period_end}}$$, + '[]'::jsonb, + 'AI Cost Control Events', + NULL, + 'system'::notification_template_kind, + true +); diff --git a/coderd/notifications/events.go b/coderd/notifications/events.go index a2da202702..c2885cc986 100644 --- a/coderd/notifications/events.go +++ b/coderd/notifications/events.go @@ -69,3 +69,9 @@ var ( TemplateChatAutoArchiveDigest = uuid.MustParse("764031be-4863-4220-867b-6ce1a1b7a5f5") TemplateChatShared = uuid.MustParse("b789bd75-d7c6-4cab-9757-1147ab184903") ) + +// AI cost control related events. +var ( + TemplateAIBudgetWarningUser = uuid.MustParse("b5db9597-de2a-4dea-87e9-25cee6906b86") + TemplateAIBudgetLimitReachedUser = uuid.MustParse("cdcf2ecd-f003-4169-9800-abb2661ea522") +) diff --git a/coderd/notifications/notifications_test.go b/coderd/notifications/notifications_test.go index 7b7b5daa32..ac26ecca9f 100644 --- a/coderd/notifications/notifications_test.go +++ b/coderd/notifications/notifications_test.go @@ -1461,6 +1461,42 @@ func TestNotificationTemplates_Golden(t *testing.T) { }, }, }, + { + name: "TemplateAIBudgetWarningUser", + id: notifications.TemplateAIBudgetWarningUser, + payload: types.MessagePayload{ + UserName: "Bobby", + UserEmail: "bobby@coder.com", + UserUsername: "bobby", + Labels: map[string]string{ + "threshold": "85", + "limit": "$1000.00", + "period": "monthly", + "effective_group_name": "Engineering", + "period_start": "July 1, 2026", + "period_end": "August 1, 2026", + }, + Data: map[string]any{}, + }, + }, + { + name: "TemplateAIBudgetLimitReachedUser", + id: notifications.TemplateAIBudgetLimitReachedUser, + payload: types.MessagePayload{ + UserName: "Bobby", + UserEmail: "bobby@coder.com", + UserUsername: "bobby", + Labels: map[string]string{ + "threshold": "100", + "limit": "$1000.00", + "period": "monthly", + "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: diff --git a/coderd/notifications/testdata/rendered-templates/smtp/TemplateAIBudgetLimitReachedUser.html.golden b/coderd/notifications/testdata/rendered-templates/smtp/TemplateAIBudgetLimitReachedUser.html.golden new file mode 100644 index 0000000000..37db6f733c --- /dev/null +++ b/coderd/notifications/testdata/rendered-templates/smtp/TemplateAIBudgetLimitReachedUser.html.golden @@ -0,0 +1,78 @@ +From: system@coder.com +To: bobby@coder.com +Subject: You've reached your 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, + +You have reached your monthly AI budget limit ($1000.00). Subsequent reques= +ts will be blocked. + +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 + + + + + + + You've reached your monthly AI budget limit + + +
+
+ 3D"Cod= +
+

+ You've reached your monthly AI budget limit +

+
+

Hi Bobby,

+

You have reached your monthly AI budget limit ($1000.00). Subseq= +uent requests will be blocked.

+ +

Effective group: Engineering

+ +

AI budget period: July 1, 2026 - August 1, 2026

+
+
+ =20 +
+
+

© 2024 Coder. All rights reserved - h= +ttp://test.com

+

Click here to manage your notification = +settings

+

Stop receiving emails like this

+
+
+ + + +--bbe61b741255b6098bb6b3c1f41b885773df633cb18d2a3002b68e4bc9c4-- diff --git a/coderd/notifications/testdata/rendered-templates/smtp/TemplateAIBudgetWarningUser.html.golden b/coderd/notifications/testdata/rendered-templates/smtp/TemplateAIBudgetWarningUser.html.golden new file mode 100644 index 0000000000..3927ab28e3 --- /dev/null +++ b/coderd/notifications/testdata/rendered-templates/smtp/TemplateAIBudgetWarningUser.html.golden @@ -0,0 +1,77 @@ +From: system@coder.com +To: bobby@coder.com +Subject: You're approaching your 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, + +You have used more than 85% of your 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 + + + + + + + You're approaching your monthly AI budget limit + + +
+
+ 3D"Cod= +
+

+ You're approaching your monthly AI budget limit +

+
+

Hi Bobby,

+

You have used more than 85% of your monthly AI budget ($1000.00)= +.

+ +

Effective group: Engineering

+ +

AI budget period: July 1, 2026 - August 1, 2026

+
+
+ =20 +
+
+

© 2024 Coder. All rights reserved - h= +ttp://test.com

+

Click here to manage your notification = +settings

+

Stop receiving emails like this

+
+
+ + + +--bbe61b741255b6098bb6b3c1f41b885773df633cb18d2a3002b68e4bc9c4-- diff --git a/coderd/notifications/testdata/rendered-templates/webhook/TemplateAIBudgetLimitReachedUser.json.golden b/coderd/notifications/testdata/rendered-templates/webhook/TemplateAIBudgetLimitReachedUser.json.golden new file mode 100644 index 0000000000..4c3470574d --- /dev/null +++ b/coderd/notifications/testdata/rendered-templates/webhook/TemplateAIBudgetLimitReachedUser.json.golden @@ -0,0 +1,28 @@ +{ + "_version": "1.1", + "msg_id": "00000000-0000-0000-0000-000000000000", + "payload": { + "_version": "1.2", + "notification_name": "AI Budget Limit Reached", + "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", + "period": "monthly", + "period_end": "August 1, 2026", + "period_start": "July 1, 2026", + "threshold": "100" + }, + "data": {}, + "targets": null + }, + "title": "You've reached your monthly AI budget limit", + "title_markdown": "You've reached your monthly AI budget limit", + "body": "You have reached your monthly AI budget limit ($1000.00). Subsequent requests will be blocked.\n\nEffective group: Engineering\n\nAI budget period: July 1, 2026 - August 1, 2026", + "body_markdown": "You have reached your monthly AI budget limit ($1000.00). Subsequent requests will be blocked.\n\nEffective group: **Engineering**\n\nAI budget period: July 1, 2026 - August 1, 2026" +} \ No newline at end of file diff --git a/coderd/notifications/testdata/rendered-templates/webhook/TemplateAIBudgetWarningUser.json.golden b/coderd/notifications/testdata/rendered-templates/webhook/TemplateAIBudgetWarningUser.json.golden new file mode 100644 index 0000000000..b9b082b9cb --- /dev/null +++ b/coderd/notifications/testdata/rendered-templates/webhook/TemplateAIBudgetWarningUser.json.golden @@ -0,0 +1,28 @@ +{ + "_version": "1.1", + "msg_id": "00000000-0000-0000-0000-000000000000", + "payload": { + "_version": "1.2", + "notification_name": "AI Budget Warning", + "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", + "period": "monthly", + "period_end": "August 1, 2026", + "period_start": "July 1, 2026", + "threshold": "85" + }, + "data": {}, + "targets": null + }, + "title": "You're approaching your monthly AI budget limit", + "title_markdown": "You're approaching your monthly AI budget limit", + "body": "You have used more than 85% of your monthly AI budget ($1000.00).\n\nEffective group: Engineering\n\nAI budget period: July 1, 2026 - August 1, 2026", + "body_markdown": "You have used more than 85% of your monthly AI budget ($1000.00).\n\nEffective group: **Engineering**\n\nAI budget period: July 1, 2026 - August 1, 2026" +} \ No newline at end of file diff --git a/codersdk/deployment.go b/codersdk/deployment.go index 2c4a0bffa9..4a749c21fc 100644 --- a/codersdk/deployment.go +++ b/codersdk/deployment.go @@ -612,6 +612,22 @@ var AIBudgetPeriods = []string{ string(AIBudgetPeriodMonth), } +// Adjective renders the period as the adjective used in user-facing text (e.g. "monthly"). +func (p AIBudgetPeriod) Adjective() string { + switch p { + case "day": + return "daily" + case "week": + return "weekly" + case AIBudgetPeriodMonth: + return "monthly" + case "year": + return "yearly" + default: + return string(p) + } +} + // NewAIBudgetPeriodFromString converts s to an AIBudgetPeriod, falling back to // AIBudgetPeriodMonth when s is empty or not a recognized period. func NewAIBudgetPeriodFromString(s string) AIBudgetPeriod { diff --git a/codersdk/deployment_test.go b/codersdk/deployment_test.go index a70fef938a..f8eb405017 100644 --- a/codersdk/deployment_test.go +++ b/codersdk/deployment_test.go @@ -149,6 +149,19 @@ func TestDeploymentValues_HighlyConfigurable(t *testing.T) { } } +func TestAIBudgetPeriodAdjective(t *testing.T) { + t.Parallel() + + // Every selectable period must have a real adjective. + for _, p := range codersdk.AIBudgetPeriods { + period := codersdk.AIBudgetPeriod(p) + require.NotEqual(t, p, period.Adjective(), + "add an adjective for AI budget period %q in AIBudgetPeriod.Adjective", p) + } + + require.Equal(t, "monthly", codersdk.AIBudgetPeriodMonth.Adjective()) +} + func TestParseSSHConfigOption(t *testing.T) { t.Parallel() diff --git a/enterprise/coderd/aibridgeserve.go b/enterprise/coderd/aibridgeserve.go index 14663c4610..ff6435f041 100644 --- a/enterprise/coderd/aibridgeserve.go +++ b/enterprise/coderd/aibridgeserve.go @@ -138,6 +138,7 @@ func (api *API) aiGatewayServe(rw http.ResponseWriter, r *http.Request) { Store: api.Database, Pubsub: api.AGPL.Pubsub, AISeatTracker: api.AGPL.AISeatTracker, + Enqueuer: api.AGPL.NotificationsEnqueuer, AccessURL: api.AccessURL.String(), GatewayCfg: api.DeploymentValues.AI.BridgeConfig, ExternalAuthConfigs: api.ExternalAuthConfigs, diff --git a/site/src/pages/UserSettingsPage/NotificationsPage/NotificationsPage.tsx b/site/src/pages/UserSettingsPage/NotificationsPage/NotificationsPage.tsx index 967df0c058..8910e875a0 100644 --- a/site/src/pages/UserSettingsPage/NotificationsPage/NotificationsPage.tsx +++ b/site/src/pages/UserSettingsPage/NotificationsPage/NotificationsPage.tsx @@ -285,6 +285,7 @@ function canSeeNotificationGroup( case "Task Events": case "Chat Events": case "Custom Events": + case "AI Cost Control Events": return true; default: return false;