mirror of
https://github.com/coder/coder.git
synced 2026-09-24 15:04:27 +08:00
feat: add user AI budget override endpoints (#25439)
Implements https://linear.app/codercom/issue/AIGOV-285 Follow the structure established in https://github.com/coder/coder/pull/25203 ## Summary Adds the `user_ai_budget_overrides` table and CRUD API at `/api/v2/users/{user}/ai/budget`. An override sets a custom per-user spend cap that supersedes group-budget resolution, attributing spend to a specific group. ## Schema ```sql CREATE TABLE user_ai_budget_overrides ( user_id UUID PRIMARY KEY REFERENCES users(id) ON DELETE CASCADE, group_id UUID NOT NULL REFERENCES groups(id) ON DELETE CASCADE, spend_limit_micros BIGINT NOT NULL CHECK (spend_limit_micros >= 0), created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW() ); ``` ## Membership lifecycle The membership invariant — a user must be a member of the attributed group, including when that group is "Everyone" — would naturally be expressed as a composite FK on `(user_id, group_id) → group_members_expanded(user_id, group_id)`. PostgreSQL doesn't allow foreign keys to reference views, so enforcement is split across two mechanisms: - **Write-time check.** A CHECK constraint on the table (`user_ai_budget_overrides_must_be_group_member`) calls a `STABLE` function `is_group_member(user_id, group_id)` that queries `group_members_expanded`. The view surfaces both regular group memberships and the implicit "Everyone" group memberships from `organization_members`. Any INSERT or UPDATE that violates the predicate is rejected with a Postgres `check_violation`, which the handler maps to a 400. `is_group_member` is defined as a general predicate, reusable by any future table that needs the same check. - **Cascade on removal.** Two `BEFORE DELETE` triggers handle membership loss: - `trigger_delete_user_ai_budget_overrides_on_group_member_delete` on `group_members` — covers regular group removals (admin action, OIDC sync). - `trigger_delete_user_ai_budget_overrides_on_org_member_delete` on `organization_members` — covers the "Everyone" group, whose membership lives in `organization_members`. The single-column FKs on `users(id)` and `groups(id)` remain to cascade on user or group deletion (those paths don't pass through `group_members`). ## Authorization The dbauthz layer gates each operation against the `User` and (for writes) `Group` resources: | Operation | User resource | Group resource | |-----------|----------------|----------------| | `GET` | `ActionRead` | — | | `PUT` | `ActionUpdate` | `ActionUpdate` | | `DELETE` | `ActionUpdate` | `ActionUpdate` | For `DELETE`, the dbauthz layer fetches the existing override first to learn the attributed `group_id`, then runs both checks. ### Role matrix | Role | GET | PUT | DELETE | |--------------|-----|-----|--------| | Owner | ✅ | ✅ | ✅ | | UserAdmin | ✅ | ✅ | ✅ | | OrgAdmin | ✅ | ❌ | ❌ | | OrgUserAdmin | ✅ | ❌ | ❌ | Internal discussion: https://codercom.slack.com/archives/C096PFVBZKN/p1779392747885359 ## Audit logs Audit logs will be addressed in a follow-up PR.
This commit is contained in:
@@ -2323,6 +2323,32 @@ func (q *querier) DeleteTask(ctx context.Context, arg database.DeleteTaskParams)
|
||||
return q.db.DeleteTask(ctx, arg)
|
||||
}
|
||||
|
||||
func (q *querier) DeleteUserAIBudgetOverride(ctx context.Context, userID uuid.UUID) (database.UserAiBudgetOverride, error) {
|
||||
// Removing a user's AI budget override affects both the user (clearing
|
||||
// their per-user spend cap) and the group it was attributed to.
|
||||
u, err := q.db.GetUserByID(ctx, userID)
|
||||
if err != nil {
|
||||
return database.UserAiBudgetOverride{}, err
|
||||
}
|
||||
if err := q.authorizeContext(ctx, policy.ActionUpdate, u); err != nil {
|
||||
return database.UserAiBudgetOverride{}, err
|
||||
}
|
||||
// Fetch the existing override to learn which group it attributes spend to,
|
||||
// so we can authorize the caller against that group as well.
|
||||
userOverride, err := q.db.GetUserAIBudgetOverride(ctx, userID)
|
||||
if err != nil {
|
||||
return database.UserAiBudgetOverride{}, err
|
||||
}
|
||||
g, err := q.db.GetGroupByID(ctx, userOverride.GroupID)
|
||||
if err != nil {
|
||||
return database.UserAiBudgetOverride{}, err
|
||||
}
|
||||
if err := q.authorizeContext(ctx, policy.ActionUpdate, g); err != nil {
|
||||
return database.UserAiBudgetOverride{}, err
|
||||
}
|
||||
return q.db.DeleteUserAIBudgetOverride(ctx, userID)
|
||||
}
|
||||
|
||||
func (q *querier) DeleteUserAIProviderKey(ctx context.Context, arg database.DeleteUserAIProviderKeyParams) error {
|
||||
u, err := q.db.GetUserByID(ctx, arg.UserID)
|
||||
if err != nil {
|
||||
@@ -4536,6 +4562,13 @@ func (q *querier) GetUnexpiredLicenses(ctx context.Context) ([]database.License,
|
||||
return q.db.GetUnexpiredLicenses(ctx)
|
||||
}
|
||||
|
||||
func (q *querier) GetUserAIBudgetOverride(ctx context.Context, userID uuid.UUID) (database.UserAiBudgetOverride, error) {
|
||||
if _, err := q.GetUserByID(ctx, userID); err != nil { // AuthZ check
|
||||
return database.UserAiBudgetOverride{}, err
|
||||
}
|
||||
return q.db.GetUserAIBudgetOverride(ctx, userID)
|
||||
}
|
||||
|
||||
func (q *querier) GetUserAIProviderKeyByProviderID(ctx context.Context, arg database.GetUserAIProviderKeyByProviderIDParams) (database.UserAiProviderKey, error) {
|
||||
u, err := q.db.GetUserByID(ctx, arg.UserID)
|
||||
if err != nil {
|
||||
@@ -8342,6 +8375,26 @@ func (q *querier) UpsertTemplateUsageStats(ctx context.Context) error {
|
||||
return q.db.UpsertTemplateUsageStats(ctx)
|
||||
}
|
||||
|
||||
func (q *querier) UpsertUserAIBudgetOverride(ctx context.Context, arg database.UpsertUserAIBudgetOverrideParams) (database.UserAiBudgetOverride, error) {
|
||||
// Setting a user's AI budget override affects both the user (their
|
||||
// per-user spend cap) and the group (spend attribution).
|
||||
u, err := q.db.GetUserByID(ctx, arg.UserID)
|
||||
if err != nil {
|
||||
return database.UserAiBudgetOverride{}, err
|
||||
}
|
||||
if err := q.authorizeContext(ctx, policy.ActionUpdate, u); err != nil {
|
||||
return database.UserAiBudgetOverride{}, err
|
||||
}
|
||||
g, err := q.db.GetGroupByID(ctx, arg.GroupID)
|
||||
if err != nil {
|
||||
return database.UserAiBudgetOverride{}, err
|
||||
}
|
||||
if err := q.authorizeContext(ctx, policy.ActionUpdate, g); err != nil {
|
||||
return database.UserAiBudgetOverride{}, err
|
||||
}
|
||||
return q.db.UpsertUserAIBudgetOverride(ctx, arg)
|
||||
}
|
||||
|
||||
func (q *querier) UpsertUserAIProviderKey(ctx context.Context, arg database.UpsertUserAIProviderKeyParams) (database.UserAiProviderKey, error) {
|
||||
u, err := q.db.GetUserByID(ctx, arg.UserID)
|
||||
if err != nil {
|
||||
|
||||
@@ -6475,6 +6475,36 @@ func (s *MethodTestSuite) TestAIBridge() {
|
||||
check.Args(g.ID).Asserts(g, policy.ActionUpdate).Returns(b)
|
||||
}))
|
||||
|
||||
s.Run("GetUserAIBudgetOverride", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
|
||||
user := testutil.Fake(s.T(), faker, database.User{})
|
||||
override := testutil.Fake(s.T(), faker, database.UserAiBudgetOverride{UserID: user.ID})
|
||||
dbm.EXPECT().GetUserByID(gomock.Any(), user.ID).Return(user, nil).AnyTimes()
|
||||
dbm.EXPECT().GetUserAIBudgetOverride(gomock.Any(), user.ID).Return(override, nil).AnyTimes()
|
||||
check.Args(user.ID).Asserts(user, policy.ActionRead).Returns(override)
|
||||
}))
|
||||
|
||||
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{})
|
||||
override := testutil.Fake(s.T(), faker, database.UserAiBudgetOverride{UserID: user.ID, GroupID: group.ID})
|
||||
arg := database.UpsertUserAIBudgetOverrideParams{UserID: user.ID, GroupID: group.ID, SpendLimitMicros: override.SpendLimitMicros}
|
||||
dbm.EXPECT().GetUserByID(gomock.Any(), user.ID).Return(user, nil).AnyTimes()
|
||||
dbm.EXPECT().GetGroupByID(gomock.Any(), group.ID).Return(group, nil).AnyTimes()
|
||||
dbm.EXPECT().UpsertUserAIBudgetOverride(gomock.Any(), arg).Return(override, nil).AnyTimes()
|
||||
check.Args(arg).Asserts(user, policy.ActionUpdate, group, policy.ActionUpdate).Returns(override)
|
||||
}))
|
||||
|
||||
s.Run("DeleteUserAIBudgetOverride", 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{})
|
||||
override := testutil.Fake(s.T(), faker, database.UserAiBudgetOverride{UserID: user.ID, GroupID: group.ID})
|
||||
dbm.EXPECT().GetUserByID(gomock.Any(), user.ID).Return(user, nil).AnyTimes()
|
||||
dbm.EXPECT().GetUserAIBudgetOverride(gomock.Any(), user.ID).Return(override, nil).AnyTimes()
|
||||
dbm.EXPECT().GetGroupByID(gomock.Any(), group.ID).Return(group, nil).AnyTimes()
|
||||
dbm.EXPECT().DeleteUserAIBudgetOverride(gomock.Any(), user.ID).Return(override, nil).AnyTimes()
|
||||
check.Args(user.ID).Asserts(user, policy.ActionUpdate, group, policy.ActionUpdate).Returns(override)
|
||||
}))
|
||||
|
||||
s.Run("GetAIProviderByID", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
|
||||
provider := testutil.Fake(s.T(), faker, database.AIProvider{})
|
||||
dbm.EXPECT().GetAIProviderByID(gomock.Any(), provider.ID).Return(provider, nil).AnyTimes()
|
||||
|
||||
Reference in New Issue
Block a user