mirror of
https://github.com/coder/coder.git
synced 2026-09-21 20:51:01 +08:00
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: <img width="1102" height="252" alt="image" src="https://github.com/user-attachments/assets/62291510-09ca-4cdf-a1f5-4bdc11a1db4b" /> <img width="466" height="384" alt="image" src="https://github.com/user-attachments/assets/030460ff-6fe2-4d59-b247-3550c543ef30" /> --------- Co-authored-by: Cian Johnston <cian@coder.com>
This commit is contained in:
Yevhenii Shcherbina
co-authored by
Cian Johnston
Cian Johnston
parent
c9e68987c1
commit
ce4ee923c2
@@ -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,
|
||||
|
||||
@@ -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) {
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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{},
|
||||
|
||||
@@ -0,0 +1,5 @@
|
||||
DELETE FROM notification_templates
|
||||
WHERE id IN (
|
||||
'b5db9597-de2a-4dea-87e9-25cee6906b86',
|
||||
'cdcf2ecd-f003-4169-9800-abb2661ea522'
|
||||
);
|
||||
@@ -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
|
||||
);
|
||||
@@ -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")
|
||||
)
|
||||
|
||||
@@ -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:
|
||||
|
||||
Vendored
+78
@@ -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
|
||||
|
||||
<!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>You've reached your 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;">
|
||||
You've reached your monthly AI budget limit
|
||||
</h1>
|
||||
<div style=3D"line-height: 1.5;">
|
||||
<p>Hi Bobby,</p>
|
||||
<p>You have reached your monthly AI budget limit ($1000.00). Subseq=
|
||||
uent requests will be blocked.</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=3Dcdc=
|
||||
f2ecd-f003-4169-9800-abb2661ea522" 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: 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
|
||||
|
||||
<!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>You're approaching your 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;">
|
||||
You're approaching your monthly AI budget limit
|
||||
</h1>
|
||||
<div style=3D"line-height: 1.5;">
|
||||
<p>Hi Bobby,</p>
|
||||
<p>You have used more than 85% of your monthly 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=3Db5d=
|
||||
b9597-de2a-4dea-87e9-25cee6906b86" style=3D"color: #2563eb; text-decoration=
|
||||
: none;">Stop receiving emails like this</a></p>
|
||||
</div>
|
||||
</div>
|
||||
</body>
|
||||
</html>
|
||||
|
||||
--bbe61b741255b6098bb6b3c1f41b885773df633cb18d2a3002b68e4bc9c4--
|
||||
+28
@@ -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"
|
||||
}
|
||||
Vendored
+28
@@ -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"
|
||||
}
|
||||
@@ -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 {
|
||||
|
||||
@@ -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()
|
||||
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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;
|
||||
|
||||
Reference in New Issue
Block a user