mirror of
https://github.com/coder/coder.git
synced 2026-09-24 15:04:27 +08:00
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:
@@ -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)
|
||||
}
|
||||
|
||||
@@ -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{})
|
||||
|
||||
Reference in New Issue
Block a user