diff --git a/coderd/aibridged.go b/coderd/aibridged.go index d163fbfe09185..2a88fc53fc017 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 cf750944e93c8..62f8e1132e19d 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 ab6a479b05289..1e54445b7ab83 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 87ba4ab64643a..79b59cd1b65fc 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 0000000000000..4481401408f16 --- /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 578fb3ab183d1..6d24264d470e2 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 0000000000000..33d74fdb99054 --- /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 0000000000000..61535f4f72d34 --- /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 a2da202702b77..c2885cc98660e 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 7b7b5daa32bae..ac26ecca9fa61 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 0000000000000..37db6f733cc5a --- /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 + + + + + + + Codestin Search App + + +
+
+ 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 0000000000000..3927ab28e31dd --- /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 + + + + + + + Codestin Search App + + +
+
+ 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 0000000000000..4c3470574d093 --- /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 0000000000000..b9b082b9cbef2 --- /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 2c4a0bffa9d24..4a749c21fc5b1 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 a70fef938a64c..f8eb405017822 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 14663c4610d6c..ff6435f04115a 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 967df0c05890d..8910e875a0bbb 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;