diff --git a/coderd/aibridgedserver/aibridgedserver.go b/coderd/aibridgedserver/aibridgedserver.go index bf1beb83048..4a856bc8aa0 100644 --- a/coderd/aibridgedserver/aibridgedserver.go +++ b/coderd/aibridgedserver/aibridgedserver.go @@ -85,6 +85,7 @@ type store interface { GetUserEveryoneFallbackGroup(ctx context.Context, userID uuid.UUID) (uuid.UUID, error) GetUserAISpendSince(ctx context.Context, arg database.GetUserAISpendSinceParams) (database.GetUserAISpendSinceRow, error) GetGroupByID(ctx context.Context, id uuid.UUID) (database.Group, error) + GetOrganizationByID(ctx context.Context, id uuid.UUID) (database.Organization, error) GetUsers(ctx context.Context, arg database.GetUsersParams) ([]database.GetUsersRow, error) // MCPConfigurator-related queries. diff --git a/coderd/aibridgedserver/aibridgedserver_test.go b/coderd/aibridgedserver/aibridgedserver_test.go index a68f2b18252..d1eae1a1bcb 100644 --- a/coderd/aibridgedserver/aibridgedserver_test.go +++ b/coderd/aibridgedserver/aibridgedserver_test.go @@ -2574,6 +2574,90 @@ func TestRecordTokenUsageAuthorized(t *testing.T) { require.Equal(t, wantCost, spend.SpendMicros, "spend micros") } +// TestBudgetNotificationAuthorized exercises the budget threshold notification +// path end-to-end against a real database through the dbauthz layer as +// subjectAibridged. This catches missing RBAC grants on the aibridged subject +// and verifies the group and organization labels round-trip from storage. +func TestBudgetNotificationAuthorized(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + logger := testutil.Logger(t) + + rawDB, _ := dbtestutil.NewDB(t) + authzDB := dbauthz.New(rawDB, rbac.NewStrictAuthorizer(prometheus.NewRegistry()), logger, coderdtest.AccessControlStorePointer()) + + // Seed prerequisites via the raw (unauthorized) store. The user belongs to a + // group with a budget, so the effective group resolves to that group. + org := dbgen.Organization(t, rawDB, database.Organization{}) + user := dbgen.User(t, rawDB, database.User{}) + dbgen.OrganizationMember(t, rawDB, database.OrganizationMember{OrganizationID: org.ID, UserID: user.ID}) + group := dbgen.Group(t, rawDB, database.Group{OrganizationID: org.ID}) + dbgen.GroupMember(t, rawDB, database.GroupMemberTable{UserID: user.ID, GroupID: group.ID}) + + // The interception below costs 1555 micros, which crosses 85% of this limit + // (1530 micros) without reaching the limit itself. + _, err := rawDB.UpsertGroupAIBudget(ctx, database.UpsertGroupAIBudgetParams{ + GroupID: group.ID, + SpendLimitMicros: 1_800, + }) + require.NoError(t, err, "upsert group AI budget") + + const provider, model = "anthropic", "claude-sonnet-4-6" + priceSeed, err := json.Marshal([]map[string]any{{ + "provider": provider, + "model": model, + "input_price": 3_000_000, + "output_price": 6_000_000, + "cache_read_price": 300_000, + "cache_write_price": 4_000_000, + }}) + require.NoError(t, err) + require.NoError(t, rawDB.UpsertAIModelPrices(ctx, database.UpsertAIModelPricesParams{Seed: priceSeed, Source: database.AIModelPriceSourceDefault}), "seed model prices") + + aiProvider := dbgen.AIProvider(t, rawDB, database.AIProvider{ + Name: "anthropic-eu", + Type: database.AIProviderTypeAnthropic, + }) + intc := dbgen.AIBridgeInterception(t, rawDB, database.InsertAIBridgeInterceptionParams{ + InitiatorID: user.ID, + Provider: provider, + ProviderName: aiProvider.Name, + Model: model, + }, nil) + + enq := ¬ificationstest.FakeEnqueuer{} + srv, err := aibridgedserver.NewServer(ctx, aibridgedserver.Options{ + Store: authzDB, + AISeatTracker: agplaiseats.Noop{}, + AccessURL: "/", + GatewayCfg: codersdk.AIBridgeConfig{}, + Experiments: requiredExperiments, + Enqueuer: enq, + Logger: logger, + Clock: quartz.NewReal(), + }) + require.NoError(t, err) + + _, err = srv.RecordTokenUsage(ctx, &proto.RecordTokenUsageRequest{ + InterceptionId: intc.ID.String(), + MsgId: "msg_budget_authz", + InputTokens: 100, + OutputTokens: 200, + CacheReadInputTokens: 50, + CacheWriteInputTokens: 10, + CreatedAt: timestamppb.New(time.Date(2026, 6, 25, 14, 30, 0, 0, time.UTC)), + }) + require.NoError(t, err, "record token usage") + + // Notification failures are logged rather than returned, so the enqueued + // message is the only signal that both lookups were authorized. + sent := enq.Sent(notificationstest.WithTemplateID(notifications.TemplateAIBudgetWarningUser)) + require.Len(t, sent, 1, "expected one budget warning notification") + require.Equal(t, group.Name, sent[0].Labels["effective_group_name"]) + require.Equal(t, org.Name, sent[0].Labels["organization_name"]) +} + // TestRecordTokenUsageModelPriceResolution covers which price an interception // snapshots when a model carries a price from the embedded book, a price set // through the API, or both. @@ -2985,11 +3069,15 @@ func TestRecordTokenUsageBudgetNotifications(t *testing.T) { Return(database.AIUserDailySpend{}, nil) db.EXPECT().GetUserAISpendSince(gomock.Any(), gomock.Any()). Return(database.GetUserAISpendSinceRow{SpendMicros: tc.newSpend}, nil) - // The group and user are resolved once per interception that - // notifies, regardless of how many thresholds it crosses. + // The group, organization, and user are resolved once per + // interception that notifies, regardless of how many thresholds it + // crosses. if len(tc.wantTemplates) > 0 { + orgID := uuid.New() db.EXPECT().GetGroupByID(gomock.Any(), groupID). - Return(database.Group{ID: groupID, Name: "Engineering"}, nil) + Return(database.Group{ID: groupID, Name: "Engineering", OrganizationID: orgID}, nil) + db.EXPECT().GetOrganizationByID(gomock.Any(), orgID). + Return(database.Organization{ID: orgID, Name: "coder"}, nil) db.EXPECT().GetUserByID(gomock.Any(), intc.InitiatorID). Return(database.User{ID: intc.InitiatorID, Username: "bob"}, nil) // No admins configured, so only the user is notified. @@ -3083,8 +3171,11 @@ func TestRecordTokenUsageBudgetNotificationAcrossPeriodBoundary(t *testing.T) { gotPeriodStart = p.PeriodStart return database.GetUserAISpendSinceRow{SpendMicros: warnAt}, nil }) + orgID := uuid.New() db.EXPECT().GetGroupByID(gomock.Any(), groupID). - Return(database.Group{ID: groupID, Name: "Engineering"}, nil) + Return(database.Group{ID: groupID, Name: "Engineering", OrganizationID: orgID}, nil) + db.EXPECT().GetOrganizationByID(gomock.Any(), orgID). + Return(database.Organization{ID: orgID, Name: "coder"}, nil) db.EXPECT().GetUserByID(gomock.Any(), intc.InitiatorID). Return(database.User{ID: intc.InitiatorID, Username: "bob"}, nil) // No admins configured, so only the user is notified. @@ -3366,8 +3457,11 @@ func TestRecordTokenUsageBudgetAdminNotification(t *testing.T) { Return(database.AIUserDailySpend{}, nil) db.EXPECT().GetUserAISpendSince(gomock.Any(), gomock.Any()). Return(database.GetUserAISpendSinceRow{SpendMicros: tc.newSpend}, nil) + orgID := uuid.New() db.EXPECT().GetGroupByID(gomock.Any(), groupID). - Return(database.Group{ID: groupID, Name: "Engineering"}, nil) + Return(database.Group{ID: groupID, Name: "Engineering", OrganizationID: orgID}, nil) + db.EXPECT().GetOrganizationByID(gomock.Any(), orgID). + Return(database.Organization{ID: orgID, Name: "coder"}, nil) db.EXPECT().GetUserByID(gomock.Any(), intc.InitiatorID). Return(database.User{ID: intc.InitiatorID, Username: "bob"}, nil) db.EXPECT().GetUsers(gomock.Any(), database.GetUsersParams{ @@ -3409,6 +3503,7 @@ func TestRecordTokenUsageBudgetAdminNotification(t *testing.T) { require.Equal(t, tc.wantThreshold, adminSent[0].Labels["threshold"]) require.Equal(t, "$100.00", adminSent[0].Labels["limit"]) require.Equal(t, "Engineering", adminSent[0].Labels["effective_group_name"]) + require.Equal(t, "coder", adminSent[0].Labels["organization_name"]) require.Equal(t, tc.wantLimitSource, adminSent[0].Labels["limit_source"]) }) } diff --git a/coderd/aibridgedserver/notifications.go b/coderd/aibridgedserver/notifications.go index aa1fbd0d0ea..6484b288162 100644 --- a/coderd/aibridgedserver/notifications.go +++ b/coderd/aibridgedserver/notifications.go @@ -137,6 +137,10 @@ func (s *Server) notifyBudgetThresholdCrossings(ctx context.Context, crossings [ if err != nil { return xerrors.Errorf("look up group %q: %w", effectiveGroupID, err) } + org, err := s.store.GetOrganizationByID(ctx, group.OrganizationID) + if err != nil { + return xerrors.Errorf("look up organization %q: %w", group.OrganizationID, err) + } user, err := s.store.GetUserByID(ctx, userID) if err != nil { return xerrors.Errorf("look up user %q: %w", userID, err) @@ -158,6 +162,7 @@ func (s *Server) notifyBudgetThresholdCrossings(ctx context.Context, crossings [ "limit_source": string(c.limitSource), "username": user.Username, "effective_group_name": group.Name, + "organization_name": org.Name, // Both bounds carry the year so a period straddling a year boundary // (e.g. December 1, 2026 - January 1, 2027) is unambiguous. "period_start": c.periodStart.UTC().Format("January 2, 2006"), diff --git a/coderd/database/dbauthz/dbauthz.go b/coderd/database/dbauthz/dbauthz.go index 104f188d951..4a45c0fb9b1 100644 --- a/coderd/database/dbauthz/dbauthz.go +++ b/coderd/database/dbauthz/dbauthz.go @@ -725,6 +725,7 @@ var ( 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. + rbac.ResourceOrganization.Type: {policy.ActionRead}, // Required to name the organization in budget notification links. }), User: []rbac.Permission{}, ByOrgID: map[string]rbac.OrgPermissions{}, diff --git a/coderd/database/migrations/000579_ai_budget_admin_notification_actions.down.sql b/coderd/database/migrations/000579_ai_budget_admin_notification_actions.down.sql new file mode 100644 index 00000000000..017b0a66534 --- /dev/null +++ b/coderd/database/migrations/000579_ai_budget_admin_notification_actions.down.sql @@ -0,0 +1,29 @@ +UPDATE notification_templates +SET + body_template = $$User **{{.Labels.username}}** has used more than {{.Labels.threshold}}% of their {{.Labels.period}} AI budget ({{.Labels.limit}}). + +Effective group: **{{.Labels.effective_group_name}}** +{{- if eq .Labels.limit_source "user_override"}} + +This limit is a per-user override. +{{- end}} + +AI budget period: {{.Labels.period_start}} - {{.Labels.period_end}}$$, + actions = '[]'::jsonb +WHERE + id = '2a7b0ac1-00e1-4625-9cd5-1e5933972c77'; + +UPDATE notification_templates +SET + body_template = $$User **{{.Labels.username}}** has reached their {{.Labels.period}} AI budget limit ({{.Labels.limit}}). Subsequent requests will be blocked. + +Effective group: **{{.Labels.effective_group_name}}** +{{- if eq .Labels.limit_source "user_override"}} + +This limit is a per-user override. +{{- end}} + +AI budget period: {{.Labels.period_start}} - {{.Labels.period_end}}$$, + actions = '[]'::jsonb +WHERE + id = '0bafe0ea-a78b-4217-ad05-1ef12e92e025'; diff --git a/coderd/database/migrations/000579_ai_budget_admin_notification_actions.up.sql b/coderd/database/migrations/000579_ai_budget_admin_notification_actions.up.sql new file mode 100644 index 00000000000..de4a33cc67e --- /dev/null +++ b/coderd/database/migrations/000579_ai_budget_admin_notification_actions.up.sql @@ -0,0 +1,43 @@ +UPDATE notification_templates +SET + body_template = $$User **{{.Labels.username}}** has used more than {{.Labels.threshold}}% of their {{.Labels.period}} AI budget ({{.Labels.limit}}). + +Effective group: **{{.Labels.effective_group_name}}** +{{- if eq .Labels.limit_source "user_override"}} + +This limit is a per-user override. +{{- end}} + +AI budget period: {{.Labels.period_start}} - {{.Labels.period_end}} + +[How to configure AI budgets](https://coder.com/docs/ai-coder/ai-gateway/cost-controls#budget)$$, + actions = '[ + { + "label": "Manage user AI budget", + "url": "{{base_url}}/organizations/{{.Labels.organization_name}}/groups/{{.Labels.effective_group_name}}?filter={{.Labels.username}}" + } + ]'::jsonb +WHERE + id = '2a7b0ac1-00e1-4625-9cd5-1e5933972c77'; + +UPDATE notification_templates +SET + body_template = $$User **{{.Labels.username}}** has reached their {{.Labels.period}} AI budget limit ({{.Labels.limit}}). Subsequent requests will be blocked until {{.Labels.period_end}}. + +Effective group: **{{.Labels.effective_group_name}}** +{{- if eq .Labels.limit_source "user_override"}} + +This limit is a per-user override. +{{- end}} + +AI budget period: {{.Labels.period_start}} - {{.Labels.period_end}} + +Raise this user's limit to unblock them sooner. [How to configure AI budgets](https://coder.com/docs/ai-coder/ai-gateway/cost-controls#budget).$$, + actions = '[ + { + "label": "Manage user AI budget", + "url": "{{base_url}}/organizations/{{.Labels.organization_name}}/groups/{{.Labels.effective_group_name}}?filter={{.Labels.username}}" + } + ]'::jsonb +WHERE + id = '0bafe0ea-a78b-4217-ad05-1ef12e92e025'; diff --git a/coderd/notifications/notifications_test.go b/coderd/notifications/notifications_test.go index 5cb42b52005..78b988ab0e2 100644 --- a/coderd/notifications/notifications_test.go +++ b/coderd/notifications/notifications_test.go @@ -1591,6 +1591,7 @@ func TestNotificationTemplates_Golden(t *testing.T) { "period": "monthly", "limit_source": "group", "effective_group_name": "Engineering", + "organization_name": "coder", "period_start": "July 1, 2026", "period_end": "August 1, 2026", }, @@ -1610,6 +1611,7 @@ func TestNotificationTemplates_Golden(t *testing.T) { "period": "monthly", "limit_source": "user_override", "effective_group_name": "Engineering", + "organization_name": "coder", "period_start": "July 1, 2026", "period_end": "August 1, 2026", }, diff --git a/coderd/notifications/testdata/rendered-templates/smtp/TemplateAIBudgetLimitReachedAdmin.html.golden b/coderd/notifications/testdata/rendered-templates/smtp/TemplateAIBudgetLimitReachedAdmin.html.golden index 550fd9b56fd..7e968ca02eb 100644 --- a/coderd/notifications/testdata/rendered-templates/smtp/TemplateAIBudgetLimitReachedAdmin.html.golden +++ b/coderd/notifications/testdata/rendered-templates/smtp/TemplateAIBudgetLimitReachedAdmin.html.golden @@ -13,7 +13,7 @@ Content-Type: text/plain; charset=UTF-8 Hi Bobby, User alice has reached their monthly AI budget limit ($1000.00). Subsequent= - requests will be blocked. + requests will be blocked until August 1, 2026. Effective group: Engineering @@ -21,6 +21,12 @@ This limit is a per-user override. AI budget period: July 1, 2026 - August 1, 2026 +Raise this user's limit to unblock them sooner. How to configure AI budgets= + (https://coder.com/docs/ai-coder/ai-gateway/cost-controls#budget). + + +Manage user AI budget: http://test.com/organizations/coder/groups/Engineeri= +ng?filter=3Dalice --bbe61b741255b6098bb6b3c1f41b885773df633cb18d2a3002b68e4bc9c4 Content-Transfer-Encoding: quoted-printable @@ -52,16 +58,28 @@ argin: 8px 0 32px; line-height: 1.5;">

Hi Bobby,

User alice has reached their monthly AI budget = -limit ($1000.00). Subsequent requests will be blocked.

+limit ($1000.00). Subsequent requests will be blocked until August 1, 2026.= +

Effective group: Engineering

This limit is a per-user override.

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

+ +

Raise this user’s limit to unblock them sooner. How to configure = +AI budgets.

=20 + + Manage user AI budget + + =20
diff --git a/coderd/notifications/testdata/rendered-templates/smtp/TemplateAIBudgetWarningAdmin.html.golden b/coderd/notifications/testdata/rendered-templates/smtp/TemplateAIBudgetWarningAdmin.html.golden index 87937d7fcd4..f5fe1bd55f6 100644 --- a/coderd/notifications/testdata/rendered-templates/smtp/TemplateAIBudgetWarningAdmin.html.golden +++ b/coderd/notifications/testdata/rendered-templates/smtp/TemplateAIBudgetWarningAdmin.html.golden @@ -18,6 +18,12 @@ Effective group: Engineering AI budget period: July 1, 2026 - August 1, 2026 +How to configure AI budgets (https://coder.com/docs/ai-coder/ai-gateway/cos= +t-controls#budget) + + +Manage user AI budget: http://test.com/organizations/coder/groups/Engineeri= +ng?filter=3Dalice --bbe61b741255b6098bb6b3c1f41b885773df633cb18d2a3002b68e4bc9c4 Content-Transfer-Encoding: quoted-printable @@ -54,9 +60,19 @@ hly AI budget ($1000.00).

Effective group: Engineering

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

+ +

How to configure AI budgets

=20 + + Manage user AI budget + + =20
diff --git a/coderd/notifications/testdata/rendered-templates/webhook/TemplateAIBudgetLimitReachedAdmin.json.golden b/coderd/notifications/testdata/rendered-templates/webhook/TemplateAIBudgetLimitReachedAdmin.json.golden index 4315def7665..ff95682ee73 100644 --- a/coderd/notifications/testdata/rendered-templates/webhook/TemplateAIBudgetLimitReachedAdmin.json.golden +++ b/coderd/notifications/testdata/rendered-templates/webhook/TemplateAIBudgetLimitReachedAdmin.json.golden @@ -9,11 +9,17 @@ "user_email": "bobby@coder.com", "user_name": "Bobby", "user_username": "bobby", - "actions": [], + "actions": [ + { + "label": "Manage user AI budget", + "url": "http://test.com/organizations/coder/groups/Engineering?filter=alice" + } + ], "labels": { "effective_group_name": "Engineering", "limit": "$1000.00", "limit_source": "user_override", + "organization_name": "coder", "period": "monthly", "period_end": "August 1, 2026", "period_start": "July 1, 2026", @@ -24,6 +30,6 @@ }, "title": "alice has reached their monthly AI budget limit", "title_markdown": "alice has reached their monthly AI budget limit", - "body": "User alice has reached their monthly AI budget limit ($1000.00). Subsequent requests will be blocked.\n\nEffective group: Engineering\n\nThis limit is a per-user override.\n\nAI budget period: July 1, 2026 - August 1, 2026", - "body_markdown": "User **alice** has reached their monthly AI budget limit ($1000.00). Subsequent requests will be blocked.\n\nEffective group: **Engineering**\n\nThis limit is a per-user override.\n\nAI budget period: July 1, 2026 - August 1, 2026" + "body": "User alice has reached their monthly AI budget limit ($1000.00). Subsequent requests will be blocked until August 1, 2026.\n\nEffective group: Engineering\n\nThis limit is a per-user override.\n\nAI budget period: July 1, 2026 - August 1, 2026\n\nRaise this user's limit to unblock them sooner. How to configure AI budgets (https://coder.com/docs/ai-coder/ai-gateway/cost-controls#budget).", + "body_markdown": "User **alice** has reached their monthly AI budget limit ($1000.00). Subsequent requests will be blocked until August 1, 2026.\n\nEffective group: **Engineering**\n\nThis limit is a per-user override.\n\nAI budget period: July 1, 2026 - August 1, 2026\n\nRaise this user's limit to unblock them sooner. [How to configure AI budgets](https://coder.com/docs/ai-coder/ai-gateway/cost-controls#budget)." } \ No newline at end of file diff --git a/coderd/notifications/testdata/rendered-templates/webhook/TemplateAIBudgetWarningAdmin.json.golden b/coderd/notifications/testdata/rendered-templates/webhook/TemplateAIBudgetWarningAdmin.json.golden index 6e83fd5d71c..004ee506cde 100644 --- a/coderd/notifications/testdata/rendered-templates/webhook/TemplateAIBudgetWarningAdmin.json.golden +++ b/coderd/notifications/testdata/rendered-templates/webhook/TemplateAIBudgetWarningAdmin.json.golden @@ -9,11 +9,17 @@ "user_email": "bobby@coder.com", "user_name": "Bobby", "user_username": "bobby", - "actions": [], + "actions": [ + { + "label": "Manage user AI budget", + "url": "http://test.com/organizations/coder/groups/Engineering?filter=alice" + } + ], "labels": { "effective_group_name": "Engineering", "limit": "$1000.00", "limit_source": "group", + "organization_name": "coder", "period": "monthly", "period_end": "August 1, 2026", "period_start": "July 1, 2026", @@ -25,6 +31,6 @@ }, "title": "alice is approaching their monthly AI budget limit", "title_markdown": "alice is approaching their monthly AI budget limit", - "body": "User alice has used more than 85% of their monthly AI budget ($1000.00).\n\nEffective group: Engineering\n\nAI budget period: July 1, 2026 - August 1, 2026", - "body_markdown": "User **alice** has used more than 85% of their monthly AI budget ($1000.00).\n\nEffective group: **Engineering**\n\nAI budget period: July 1, 2026 - August 1, 2026" + "body": "User alice has used more than 85% of their monthly AI budget ($1000.00).\n\nEffective group: Engineering\n\nAI budget period: July 1, 2026 - August 1, 2026\n\nHow to configure AI budgets (https://coder.com/docs/ai-coder/ai-gateway/cost-controls#budget)", + "body_markdown": "User **alice** has used more than 85% of their monthly AI budget ($1000.00).\n\nEffective group: **Engineering**\n\nAI budget period: July 1, 2026 - August 1, 2026\n\n[How to configure AI budgets](https://coder.com/docs/ai-coder/ai-gateway/cost-controls#budget)" } \ No newline at end of file