feat(coderd): resolve effective user AI budget (#26142)

Closes
https://linear.app/codercom/issue/AIGOV-287/add-effective-group-resolution

Implements the effective AI budget resolution from the AI Governance
cost-controls RFC: for a given user, a `user_ai_budget_overrides` row
wins if present, otherwise the deployment budget policy (`highest`)
picks the largest group budget across the user's groups, ties broken
alphabetically.

For now, I keep the logic under `coderd/aibridge/budget`, but that may
change during the implementation of budget enforcement.
This commit is contained in:
Yevhenii Shcherbina
2026-06-10 14:01:20 +00:00
committed by GitHub
parent f9dfa18c46
commit 1cada0649c
9 changed files with 384 additions and 0 deletions
+7
View File
@@ -3625,6 +3625,13 @@ func (q *querier) GetHealthSettings(ctx context.Context) (string, error) {
return q.db.GetHealthSettings(ctx)
}
func (q *querier) GetHighestGroupAIBudgetByUser(ctx context.Context, userID uuid.UUID) (database.GetHighestGroupAIBudgetByUserRow, error) {
if _, err := q.GetUserByID(ctx, userID); err != nil { // AuthZ check
return database.GetHighestGroupAIBudgetByUserRow{}, err
}
return q.db.GetHighestGroupAIBudgetByUser(ctx, userID)
}
func (q *querier) GetInboxNotificationByID(ctx context.Context, id uuid.UUID) (database.InboxNotification, error) {
return fetchWithAction(q.log, q.auth, policy.ActionRead, q.db.GetInboxNotificationByID)(ctx, id)
}
+8
View File
@@ -6502,6 +6502,14 @@ func (s *MethodTestSuite) TestAIBridge() {
check.Args(user.ID).Asserts(user, policy.ActionRead).Returns(override)
}))
s.Run("GetHighestGroupAIBudgetByUser", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
user := testutil.Fake(s.T(), faker, database.User{})
row := testutil.Fake(s.T(), faker, database.GetHighestGroupAIBudgetByUserRow{})
dbm.EXPECT().GetUserByID(gomock.Any(), user.ID).Return(user, nil).AnyTimes()
dbm.EXPECT().GetHighestGroupAIBudgetByUser(gomock.Any(), user.ID).Return(row, nil).AnyTimes()
check.Args(user.ID).Asserts(user, policy.ActionRead).Returns(row)
}))
s.Run("UpsertUserAIBudgetOverride", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
user := testutil.Fake(s.T(), faker, database.User{})
group := testutil.Fake(s.T(), faker, database.Group{})