From 1a91d3179366055b6f74b5b5940fcae1d9729abd Mon Sep 17 00:00:00 2001 From: Yevhenii Shcherbina Date: Fri, 29 May 2026 10:08:25 -0400 Subject: [PATCH] feat: add user AI budget override endpoints (#25439) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- coderd/apidoc/docs.go | 145 ++++++ coderd/apidoc/swagger.json | 131 ++++++ coderd/database/check_constraint.go | 1 + coderd/database/db2sdk/db2sdk.go | 10 + coderd/database/dbauthz/dbauthz.go | 53 +++ coderd/database/dbauthz/dbauthz_test.go | 30 ++ coderd/database/dbmetrics/querymetrics.go | 24 + coderd/database/dbmock/dbmock.go | 45 ++ coderd/database/dump.sql | 62 +++ coderd/database/foreign_key_constraint.go | 2 + .../000513_user_ai_budget_overrides.down.sql | 7 + .../000513_user_ai_budget_overrides.up.sql | 76 +++ .../000513_user_ai_budget_overrides.up.sql | 15 + coderd/database/models.go | 9 + coderd/database/querier.go | 3 + coderd/database/queries.sql.go | 65 +++ coderd/database/queries/aicostcontrol.sql | 17 + coderd/database/unique_constraint.go | 1 + codersdk/aibridge.go | 68 +++ docs/reference/api/enterprise.md | 119 +++++ docs/reference/api/schemas.md | 38 ++ enterprise/coderd/aibridge.go | 118 +++++ enterprise/coderd/aibridge_test.go | 441 ++++++++++++++++++ enterprise/coderd/coderd.go | 11 + site/src/api/typesGenerated.ts | 19 + 25 files changed, 1510 insertions(+) create mode 100644 coderd/database/migrations/000513_user_ai_budget_overrides.down.sql create mode 100644 coderd/database/migrations/000513_user_ai_budget_overrides.up.sql create mode 100644 coderd/database/migrations/testdata/fixtures/000513_user_ai_budget_overrides.up.sql diff --git a/coderd/apidoc/docs.go b/coderd/apidoc/docs.go index 7bca7d7988..1bbb216b1a 100644 --- a/coderd/apidoc/docs.go +++ b/coderd/apidoc/docs.go @@ -9171,6 +9171,110 @@ const docTemplate = `{ ] } }, + "/api/v2/users/{user}/ai/budget": { + "get": { + "produces": [ + "application/json" + ], + "tags": [ + "Enterprise" + ], + "summary": "Get user AI budget override", + "operationId": "get-user-ai-budget-override", + "parameters": [ + { + "type": "string", + "description": "User ID, username, or me", + "name": "user", + "in": "path", + "required": true + } + ], + "responses": { + "200": { + "description": "OK", + "schema": { + "$ref": "#/definitions/codersdk.UserAIBudgetOverride" + } + } + }, + "security": [ + { + "CoderSessionToken": [] + } + ] + }, + "put": { + "consumes": [ + "application/json" + ], + "produces": [ + "application/json" + ], + "tags": [ + "Enterprise" + ], + "summary": "Upsert user AI budget override", + "operationId": "upsert-user-ai-budget-override", + "parameters": [ + { + "type": "string", + "description": "User ID, username, or me", + "name": "user", + "in": "path", + "required": true + }, + { + "description": "Upsert user AI budget override request", + "name": "request", + "in": "body", + "required": true, + "schema": { + "$ref": "#/definitions/codersdk.UpsertUserAIBudgetOverrideRequest" + } + } + ], + "responses": { + "200": { + "description": "OK", + "schema": { + "$ref": "#/definitions/codersdk.UserAIBudgetOverride" + } + } + }, + "security": [ + { + "CoderSessionToken": [] + } + ] + }, + "delete": { + "tags": [ + "Enterprise" + ], + "summary": "Delete user AI budget override", + "operationId": "delete-user-ai-budget-override", + "parameters": [ + { + "type": "string", + "description": "User ID, username, or me", + "name": "user", + "in": "path", + "required": true + } + ], + "responses": { + "204": { + "description": "No Content" + } + }, + "security": [ + { + "CoderSessionToken": [] + } + ] + } + }, "/api/v2/users/{user}/appearance": { "get": { "produces": [ @@ -24663,6 +24767,23 @@ const docTemplate = `{ } } }, + "codersdk.UpsertUserAIBudgetOverrideRequest": { + "type": "object", + "required": [ + "group_id" + ], + "properties": { + "group_id": { + "description": "GroupID is the group the user's spend is attributed to. The user must\nbe a member of this group.", + "type": "string", + "format": "uuid" + }, + "spend_limit_micros": { + "type": "integer", + "minimum": 0 + } + } + }, "codersdk.UpsertWorkspaceAgentPortShareRequest": { "type": "object", "properties": { @@ -24817,6 +24938,30 @@ const docTemplate = `{ } } }, + "codersdk.UserAIBudgetOverride": { + "type": "object", + "properties": { + "created_at": { + "type": "string", + "format": "date-time" + }, + "group_id": { + "type": "string", + "format": "uuid" + }, + "spend_limit_micros": { + "type": "integer" + }, + "updated_at": { + "type": "string", + "format": "date-time" + }, + "user_id": { + "type": "string", + "format": "uuid" + } + } + }, "codersdk.UserActivity": { "type": "object", "properties": { diff --git a/coderd/apidoc/swagger.json b/coderd/apidoc/swagger.json index d641744117..6f7224e972 100644 --- a/coderd/apidoc/swagger.json +++ b/coderd/apidoc/swagger.json @@ -8132,6 +8132,98 @@ ] } }, + "/api/v2/users/{user}/ai/budget": { + "get": { + "produces": ["application/json"], + "tags": ["Enterprise"], + "summary": "Get user AI budget override", + "operationId": "get-user-ai-budget-override", + "parameters": [ + { + "type": "string", + "description": "User ID, username, or me", + "name": "user", + "in": "path", + "required": true + } + ], + "responses": { + "200": { + "description": "OK", + "schema": { + "$ref": "#/definitions/codersdk.UserAIBudgetOverride" + } + } + }, + "security": [ + { + "CoderSessionToken": [] + } + ] + }, + "put": { + "consumes": ["application/json"], + "produces": ["application/json"], + "tags": ["Enterprise"], + "summary": "Upsert user AI budget override", + "operationId": "upsert-user-ai-budget-override", + "parameters": [ + { + "type": "string", + "description": "User ID, username, or me", + "name": "user", + "in": "path", + "required": true + }, + { + "description": "Upsert user AI budget override request", + "name": "request", + "in": "body", + "required": true, + "schema": { + "$ref": "#/definitions/codersdk.UpsertUserAIBudgetOverrideRequest" + } + } + ], + "responses": { + "200": { + "description": "OK", + "schema": { + "$ref": "#/definitions/codersdk.UserAIBudgetOverride" + } + } + }, + "security": [ + { + "CoderSessionToken": [] + } + ] + }, + "delete": { + "tags": ["Enterprise"], + "summary": "Delete user AI budget override", + "operationId": "delete-user-ai-budget-override", + "parameters": [ + { + "type": "string", + "description": "User ID, username, or me", + "name": "user", + "in": "path", + "required": true + } + ], + "responses": { + "204": { + "description": "No Content" + } + }, + "security": [ + { + "CoderSessionToken": [] + } + ] + } + }, "/api/v2/users/{user}/appearance": { "get": { "produces": ["application/json"], @@ -22696,6 +22788,21 @@ } } }, + "codersdk.UpsertUserAIBudgetOverrideRequest": { + "type": "object", + "required": ["group_id"], + "properties": { + "group_id": { + "description": "GroupID is the group the user's spend is attributed to. The user must\nbe a member of this group.", + "type": "string", + "format": "uuid" + }, + "spend_limit_micros": { + "type": "integer", + "minimum": 0 + } + } + }, "codersdk.UpsertWorkspaceAgentPortShareRequest": { "type": "object", "properties": { @@ -22829,6 +22936,30 @@ } } }, + "codersdk.UserAIBudgetOverride": { + "type": "object", + "properties": { + "created_at": { + "type": "string", + "format": "date-time" + }, + "group_id": { + "type": "string", + "format": "uuid" + }, + "spend_limit_micros": { + "type": "integer" + }, + "updated_at": { + "type": "string", + "format": "date-time" + }, + "user_id": { + "type": "string", + "format": "uuid" + } + } + }, "codersdk.UserActivity": { "type": "object", "properties": { diff --git a/coderd/database/check_constraint.go b/coderd/database/check_constraint.go index 5682341ef9..1c20622e58 100644 --- a/coderd/database/check_constraint.go +++ b/coderd/database/check_constraint.go @@ -44,6 +44,7 @@ const ( CheckTelemetryLockEventTypeConstraint CheckConstraint = "telemetry_lock_event_type_constraint" // telemetry_locks CheckValidationMonotonicOrder CheckConstraint = "validation_monotonic_order" // template_version_parameters CheckUsageEventTypeCheck CheckConstraint = "usage_event_type_check" // usage_events + CheckUserAiBudgetOverridesSpendLimitMicrosCheck CheckConstraint = "user_ai_budget_overrides_spend_limit_micros_check" // user_ai_budget_overrides CheckUserAiProviderKeysAPIKeyCheck CheckConstraint = "user_ai_provider_keys_api_key_check" // user_ai_provider_keys CheckUserSkillsContentSize CheckConstraint = "user_skills_content_size" // user_skills CheckUserSkillsDescriptionSize CheckConstraint = "user_skills_description_size" // user_skills diff --git a/coderd/database/db2sdk/db2sdk.go b/coderd/database/db2sdk/db2sdk.go index 36b081b86b..bc93df7cd3 100644 --- a/coderd/database/db2sdk/db2sdk.go +++ b/coderd/database/db2sdk/db2sdk.go @@ -1509,6 +1509,16 @@ func GroupAIBudget(b database.GroupAiBudget) codersdk.GroupAIBudget { } } +func UserAIBudgetOverride(o database.UserAiBudgetOverride) codersdk.UserAIBudgetOverride { + return codersdk.UserAIBudgetOverride{ + UserID: o.UserID, + GroupID: o.GroupID, + SpendLimitMicros: o.SpendLimitMicros, + CreatedAt: o.CreatedAt, + UpdatedAt: o.UpdatedAt, + } +} + func InvalidatedPresets(invalidatedPresets []database.UpdatePresetsLastInvalidatedAtRow) []codersdk.InvalidatedPreset { var presets []codersdk.InvalidatedPreset for _, p := range invalidatedPresets { diff --git a/coderd/database/dbauthz/dbauthz.go b/coderd/database/dbauthz/dbauthz.go index 5e98dfcf30..a1a7497153 100644 --- a/coderd/database/dbauthz/dbauthz.go +++ b/coderd/database/dbauthz/dbauthz.go @@ -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 { diff --git a/coderd/database/dbauthz/dbauthz_test.go b/coderd/database/dbauthz/dbauthz_test.go index cfc87d3ed4..f788fa71e2 100644 --- a/coderd/database/dbauthz/dbauthz_test.go +++ b/coderd/database/dbauthz/dbauthz_test.go @@ -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() diff --git a/coderd/database/dbmetrics/querymetrics.go b/coderd/database/dbmetrics/querymetrics.go index c963d41d3f..e7120ec588 100644 --- a/coderd/database/dbmetrics/querymetrics.go +++ b/coderd/database/dbmetrics/querymetrics.go @@ -793,6 +793,14 @@ func (m queryMetricsStore) DeleteTask(ctx context.Context, arg database.DeleteTa return r0, r1 } +func (m queryMetricsStore) DeleteUserAIBudgetOverride(ctx context.Context, userID uuid.UUID) (database.UserAiBudgetOverride, error) { + start := time.Now() + r0, r1 := m.s.DeleteUserAIBudgetOverride(ctx, userID) + m.queryLatencies.WithLabelValues("DeleteUserAIBudgetOverride").Observe(time.Since(start).Seconds()) + m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "DeleteUserAIBudgetOverride").Inc() + return r0, r1 +} + func (m queryMetricsStore) DeleteUserAIProviderKey(ctx context.Context, arg database.DeleteUserAIProviderKeyParams) error { start := time.Now() r0 := m.s.DeleteUserAIProviderKey(ctx, arg) @@ -2905,6 +2913,14 @@ func (m queryMetricsStore) GetUnexpiredLicenses(ctx context.Context) ([]database return r0, r1 } +func (m queryMetricsStore) GetUserAIBudgetOverride(ctx context.Context, userID uuid.UUID) (database.UserAiBudgetOverride, error) { + start := time.Now() + r0, r1 := m.s.GetUserAIBudgetOverride(ctx, userID) + m.queryLatencies.WithLabelValues("GetUserAIBudgetOverride").Observe(time.Since(start).Seconds()) + m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "GetUserAIBudgetOverride").Inc() + return r0, r1 +} + func (m queryMetricsStore) GetUserAIProviderKeyByProviderID(ctx context.Context, arg database.GetUserAIProviderKeyByProviderIDParams) (database.UserAiProviderKey, error) { start := time.Now() r0, r1 := m.s.GetUserAIProviderKeyByProviderID(ctx, arg) @@ -6049,6 +6065,14 @@ func (m queryMetricsStore) UpsertTemplateUsageStats(ctx context.Context) error { return r0 } +func (m queryMetricsStore) UpsertUserAIBudgetOverride(ctx context.Context, arg database.UpsertUserAIBudgetOverrideParams) (database.UserAiBudgetOverride, error) { + start := time.Now() + r0, r1 := m.s.UpsertUserAIBudgetOverride(ctx, arg) + m.queryLatencies.WithLabelValues("UpsertUserAIBudgetOverride").Observe(time.Since(start).Seconds()) + m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "UpsertUserAIBudgetOverride").Inc() + return r0, r1 +} + func (m queryMetricsStore) UpsertUserAIProviderKey(ctx context.Context, arg database.UpsertUserAIProviderKeyParams) (database.UserAiProviderKey, error) { start := time.Now() r0, r1 := m.s.UpsertUserAIProviderKey(ctx, arg) diff --git a/coderd/database/dbmock/dbmock.go b/coderd/database/dbmock/dbmock.go index 75ac8f44ad..0f6799e638 100644 --- a/coderd/database/dbmock/dbmock.go +++ b/coderd/database/dbmock/dbmock.go @@ -1349,6 +1349,21 @@ func (mr *MockStoreMockRecorder) DeleteTask(ctx, arg any) *gomock.Call { return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DeleteTask", reflect.TypeOf((*MockStore)(nil).DeleteTask), ctx, arg) } +// DeleteUserAIBudgetOverride mocks base method. +func (m *MockStore) DeleteUserAIBudgetOverride(ctx context.Context, userID uuid.UUID) (database.UserAiBudgetOverride, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "DeleteUserAIBudgetOverride", ctx, userID) + ret0, _ := ret[0].(database.UserAiBudgetOverride) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// DeleteUserAIBudgetOverride indicates an expected call of DeleteUserAIBudgetOverride. +func (mr *MockStoreMockRecorder) DeleteUserAIBudgetOverride(ctx, userID any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DeleteUserAIBudgetOverride", reflect.TypeOf((*MockStore)(nil).DeleteUserAIBudgetOverride), ctx, userID) +} + // DeleteUserAIProviderKey mocks base method. func (m *MockStore) DeleteUserAIProviderKey(ctx context.Context, arg database.DeleteUserAIProviderKeyParams) error { m.ctrl.T.Helper() @@ -5445,6 +5460,21 @@ func (mr *MockStoreMockRecorder) GetUnexpiredLicenses(ctx any) *gomock.Call { return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetUnexpiredLicenses", reflect.TypeOf((*MockStore)(nil).GetUnexpiredLicenses), ctx) } +// GetUserAIBudgetOverride mocks base method. +func (m *MockStore) GetUserAIBudgetOverride(ctx context.Context, userID uuid.UUID) (database.UserAiBudgetOverride, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "GetUserAIBudgetOverride", ctx, userID) + ret0, _ := ret[0].(database.UserAiBudgetOverride) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// GetUserAIBudgetOverride indicates an expected call of GetUserAIBudgetOverride. +func (mr *MockStoreMockRecorder) GetUserAIBudgetOverride(ctx, userID any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetUserAIBudgetOverride", reflect.TypeOf((*MockStore)(nil).GetUserAIBudgetOverride), ctx, userID) +} + // GetUserAIProviderKeyByProviderID mocks base method. func (m *MockStore) GetUserAIProviderKeyByProviderID(ctx context.Context, arg database.GetUserAIProviderKeyByProviderIDParams) (database.UserAiProviderKey, error) { m.ctrl.T.Helper() @@ -11344,6 +11374,21 @@ func (mr *MockStoreMockRecorder) UpsertTemplateUsageStats(ctx any) *gomock.Call return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpsertTemplateUsageStats", reflect.TypeOf((*MockStore)(nil).UpsertTemplateUsageStats), ctx) } +// UpsertUserAIBudgetOverride mocks base method. +func (m *MockStore) UpsertUserAIBudgetOverride(ctx context.Context, arg database.UpsertUserAIBudgetOverrideParams) (database.UserAiBudgetOverride, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "UpsertUserAIBudgetOverride", ctx, arg) + ret0, _ := ret[0].(database.UserAiBudgetOverride) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// UpsertUserAIBudgetOverride indicates an expected call of UpsertUserAIBudgetOverride. +func (mr *MockStoreMockRecorder) UpsertUserAIBudgetOverride(ctx, arg any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpsertUserAIBudgetOverride", reflect.TypeOf((*MockStore)(nil).UpsertUserAIBudgetOverride), ctx, arg) +} + // UpsertUserAIProviderKey mocks base method. func (m *MockStore) UpsertUserAIProviderKey(ctx context.Context, arg database.UpsertUserAIProviderKeyParams) (database.UserAiProviderKey, error) { m.ctrl.T.Helper() diff --git a/coderd/database/dump.sql b/coderd/database/dump.sql index fad35755b9..82aa376d34 100644 --- a/coderd/database/dump.sql +++ b/coderd/database/dump.sql @@ -841,6 +841,42 @@ BEGIN END; $$; +CREATE FUNCTION delete_user_ai_budget_overrides_on_group_member_delete() RETURNS trigger + LANGUAGE plpgsql + AS $$ +BEGIN + DELETE FROM user_ai_budget_overrides + WHERE user_id = OLD.user_id AND group_id = OLD.group_id; + RETURN OLD; +END; +$$; + +CREATE FUNCTION delete_user_ai_budget_overrides_on_org_member_delete() RETURNS trigger + LANGUAGE plpgsql + AS $$ +BEGIN + DELETE FROM user_ai_budget_overrides + WHERE user_id = OLD.user_id AND group_id = OLD.organization_id; + RETURN OLD; +END; +$$; + +CREATE FUNCTION enforce_user_ai_budget_override_membership() RETURNS trigger + LANGUAGE plpgsql + AS $$ +BEGIN + IF NOT EXISTS ( + SELECT 1 FROM group_members_expanded + WHERE user_id = NEW.user_id AND group_id = NEW.group_id + ) THEN + RAISE EXCEPTION 'user % is not a member of group %', NEW.user_id, NEW.group_id + USING ERRCODE = 'check_violation', + CONSTRAINT = 'user_ai_budget_overrides_must_be_group_member'; + END IF; + RETURN NEW; +END; +$$; + CREATE FUNCTION enforce_user_secrets_per_user_limits() RETURNS trigger LANGUAGE plpgsql AS $$ @@ -3137,6 +3173,17 @@ COMMENT ON TABLE usage_events_daily IS 'usage_events_daily is a daily rollup of COMMENT ON COLUMN usage_events_daily.day IS 'The date of the summed usage events, always in UTC.'; +CREATE TABLE user_ai_budget_overrides ( + user_id uuid NOT NULL, + group_id uuid NOT NULL, + spend_limit_micros bigint NOT NULL, + created_at timestamp with time zone DEFAULT now() NOT NULL, + updated_at timestamp with time zone DEFAULT now() NOT NULL, + CONSTRAINT user_ai_budget_overrides_spend_limit_micros_check CHECK ((spend_limit_micros >= 0)) +); + +COMMENT ON TABLE user_ai_budget_overrides IS 'Per-user AI spend override that supersedes group budget resolution.'; + CREATE TABLE user_ai_provider_keys ( id uuid DEFAULT gen_random_uuid() NOT NULL, user_id uuid NOT NULL, @@ -3986,6 +4033,9 @@ ALTER TABLE ONLY usage_events_daily ALTER TABLE ONLY usage_events ADD CONSTRAINT usage_events_pkey PRIMARY KEY (id); +ALTER TABLE ONLY user_ai_budget_overrides + ADD CONSTRAINT user_ai_budget_overrides_pkey PRIMARY KEY (user_id); + ALTER TABLE ONLY user_ai_provider_keys ADD CONSTRAINT user_ai_provider_keys_pkey PRIMARY KEY (id); @@ -4467,6 +4517,12 @@ CREATE TRIGGER trigger_delete_group_members_on_org_member_delete BEFORE DELETE O CREATE TRIGGER trigger_delete_oauth2_provider_app_token AFTER DELETE ON oauth2_provider_app_tokens FOR EACH ROW EXECUTE FUNCTION delete_deleted_oauth2_provider_app_token_api_key(); +CREATE TRIGGER trigger_delete_user_ai_budget_overrides_on_group_member_delete BEFORE DELETE ON group_members FOR EACH ROW EXECUTE FUNCTION delete_user_ai_budget_overrides_on_group_member_delete(); + +CREATE TRIGGER trigger_delete_user_ai_budget_overrides_on_org_member_delete BEFORE DELETE ON organization_members FOR EACH ROW EXECUTE FUNCTION delete_user_ai_budget_overrides_on_org_member_delete(); + +CREATE TRIGGER trigger_enforce_user_ai_budget_override_membership BEFORE INSERT OR UPDATE ON user_ai_budget_overrides FOR EACH ROW EXECUTE FUNCTION enforce_user_ai_budget_override_membership(); + CREATE TRIGGER trigger_insert_apikeys BEFORE INSERT ON api_keys FOR EACH ROW EXECUTE FUNCTION insert_apikey_fail_if_user_deleted(); CREATE TRIGGER trigger_insert_organization_system_roles AFTER INSERT ON organizations FOR EACH ROW EXECUTE FUNCTION insert_organization_system_roles(); @@ -4792,6 +4848,12 @@ ALTER TABLE ONLY templates ALTER TABLE ONLY templates ADD CONSTRAINT templates_organization_id_fkey FOREIGN KEY (organization_id) REFERENCES organizations(id) ON DELETE CASCADE; +ALTER TABLE ONLY user_ai_budget_overrides + ADD CONSTRAINT user_ai_budget_overrides_group_id_fkey FOREIGN KEY (group_id) REFERENCES groups(id) ON DELETE CASCADE; + +ALTER TABLE ONLY user_ai_budget_overrides + ADD CONSTRAINT user_ai_budget_overrides_user_id_fkey FOREIGN KEY (user_id) REFERENCES users(id) ON DELETE CASCADE; + ALTER TABLE ONLY user_ai_provider_keys ADD CONSTRAINT user_ai_provider_keys_ai_provider_id_fkey FOREIGN KEY (ai_provider_id) REFERENCES ai_providers(id) ON DELETE CASCADE; diff --git a/coderd/database/foreign_key_constraint.go b/coderd/database/foreign_key_constraint.go index a31346b423..5eeb24587a 100644 --- a/coderd/database/foreign_key_constraint.go +++ b/coderd/database/foreign_key_constraint.go @@ -105,6 +105,8 @@ const ( ForeignKeyTemplateVersionsTemplateID ForeignKeyConstraint = "template_versions_template_id_fkey" // ALTER TABLE ONLY template_versions ADD CONSTRAINT template_versions_template_id_fkey FOREIGN KEY (template_id) REFERENCES templates(id) ON DELETE CASCADE; ForeignKeyTemplatesCreatedBy ForeignKeyConstraint = "templates_created_by_fkey" // ALTER TABLE ONLY templates ADD CONSTRAINT templates_created_by_fkey FOREIGN KEY (created_by) REFERENCES users(id) ON DELETE RESTRICT; ForeignKeyTemplatesOrganizationID ForeignKeyConstraint = "templates_organization_id_fkey" // ALTER TABLE ONLY templates ADD CONSTRAINT templates_organization_id_fkey FOREIGN KEY (organization_id) REFERENCES organizations(id) ON DELETE CASCADE; + ForeignKeyUserAiBudgetOverridesGroupID ForeignKeyConstraint = "user_ai_budget_overrides_group_id_fkey" // ALTER TABLE ONLY user_ai_budget_overrides ADD CONSTRAINT user_ai_budget_overrides_group_id_fkey FOREIGN KEY (group_id) REFERENCES groups(id) ON DELETE CASCADE; + ForeignKeyUserAiBudgetOverridesUserID ForeignKeyConstraint = "user_ai_budget_overrides_user_id_fkey" // ALTER TABLE ONLY user_ai_budget_overrides ADD CONSTRAINT user_ai_budget_overrides_user_id_fkey FOREIGN KEY (user_id) REFERENCES users(id) ON DELETE CASCADE; ForeignKeyUserAiProviderKeysAiProviderID ForeignKeyConstraint = "user_ai_provider_keys_ai_provider_id_fkey" // ALTER TABLE ONLY user_ai_provider_keys ADD CONSTRAINT user_ai_provider_keys_ai_provider_id_fkey FOREIGN KEY (ai_provider_id) REFERENCES ai_providers(id) ON DELETE CASCADE; ForeignKeyUserAiProviderKeysAPIKeyKeyID ForeignKeyConstraint = "user_ai_provider_keys_api_key_key_id_fkey" // ALTER TABLE ONLY user_ai_provider_keys ADD CONSTRAINT user_ai_provider_keys_api_key_key_id_fkey FOREIGN KEY (api_key_key_id) REFERENCES dbcrypt_keys(active_key_digest); ForeignKeyUserAiProviderKeysUserID ForeignKeyConstraint = "user_ai_provider_keys_user_id_fkey" // ALTER TABLE ONLY user_ai_provider_keys ADD CONSTRAINT user_ai_provider_keys_user_id_fkey FOREIGN KEY (user_id) REFERENCES users(id) ON DELETE CASCADE; diff --git a/coderd/database/migrations/000513_user_ai_budget_overrides.down.sql b/coderd/database/migrations/000513_user_ai_budget_overrides.down.sql new file mode 100644 index 0000000000..1a1a8e2160 --- /dev/null +++ b/coderd/database/migrations/000513_user_ai_budget_overrides.down.sql @@ -0,0 +1,7 @@ +DROP TRIGGER IF EXISTS trigger_delete_user_ai_budget_overrides_on_org_member_delete ON organization_members; +DROP FUNCTION IF EXISTS delete_user_ai_budget_overrides_on_org_member_delete; +DROP TRIGGER IF EXISTS trigger_delete_user_ai_budget_overrides_on_group_member_delete ON group_members; +DROP FUNCTION IF EXISTS delete_user_ai_budget_overrides_on_group_member_delete; +DROP TRIGGER IF EXISTS trigger_enforce_user_ai_budget_override_membership ON user_ai_budget_overrides; +DROP FUNCTION IF EXISTS enforce_user_ai_budget_override_membership; +DROP TABLE IF EXISTS user_ai_budget_overrides CASCADE; diff --git a/coderd/database/migrations/000513_user_ai_budget_overrides.up.sql b/coderd/database/migrations/000513_user_ai_budget_overrides.up.sql new file mode 100644 index 0000000000..b1ab1cd9d2 --- /dev/null +++ b/coderd/database/migrations/000513_user_ai_budget_overrides.up.sql @@ -0,0 +1,76 @@ +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 applied to the user, in micro-units (1 unit = 1,000,000). + 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() + -- The membership invariant (user must be a member of the attributed + -- group, including when that group is "Everyone") would naturally be + -- a composite FK to group_members_expanded, but PostgreSQL does not + -- allow FKs to views. It's enforced instead by a write-time trigger + -- on this table and removal-time triggers on the underlying + -- membership tables. +); + +COMMENT ON TABLE user_ai_budget_overrides IS 'Per-user AI spend override that supersedes group budget resolution.'; + +-- Write-time membership check. Reads from group_members_expanded so +-- the "Everyone" group (whose membership lives in organization_members) +-- is correctly handled. Raises check_violation with a constraint name +-- so callers can match it via database.IsCheckViolation in Go. +CREATE FUNCTION enforce_user_ai_budget_override_membership() RETURNS TRIGGER + LANGUAGE plpgsql +AS $$ +BEGIN + IF NOT EXISTS ( + SELECT 1 FROM group_members_expanded + WHERE user_id = NEW.user_id AND group_id = NEW.group_id + ) THEN + RAISE EXCEPTION 'user % is not a member of group %', NEW.user_id, NEW.group_id + USING ERRCODE = 'check_violation', + CONSTRAINT = 'user_ai_budget_overrides_must_be_group_member'; + END IF; + RETURN NEW; +END; +$$; + +CREATE TRIGGER trigger_enforce_user_ai_budget_override_membership + BEFORE INSERT OR UPDATE ON user_ai_budget_overrides + FOR EACH ROW +EXECUTE PROCEDURE enforce_user_ai_budget_override_membership(); + +-- When a user is removed from a regular group (any group except +-- "Everyone"), delete any override attributed to that group. +CREATE FUNCTION delete_user_ai_budget_overrides_on_group_member_delete() RETURNS TRIGGER + LANGUAGE plpgsql +AS $$ +BEGIN + DELETE FROM user_ai_budget_overrides + WHERE user_id = OLD.user_id AND group_id = OLD.group_id; + RETURN OLD; +END; +$$; + +CREATE TRIGGER trigger_delete_user_ai_budget_overrides_on_group_member_delete + BEFORE DELETE ON group_members + FOR EACH ROW +EXECUTE PROCEDURE delete_user_ai_budget_overrides_on_group_member_delete(); + +-- When a user is removed from an organization, delete any override +-- attributed to that organization's "Everyone" group (which has +-- id == organization_id). +CREATE FUNCTION delete_user_ai_budget_overrides_on_org_member_delete() RETURNS TRIGGER + LANGUAGE plpgsql +AS $$ +BEGIN + DELETE FROM user_ai_budget_overrides + WHERE user_id = OLD.user_id AND group_id = OLD.organization_id; + RETURN OLD; +END; +$$; + +CREATE TRIGGER trigger_delete_user_ai_budget_overrides_on_org_member_delete + BEFORE DELETE ON organization_members + FOR EACH ROW +EXECUTE PROCEDURE delete_user_ai_budget_overrides_on_org_member_delete(); diff --git a/coderd/database/migrations/testdata/fixtures/000513_user_ai_budget_overrides.up.sql b/coderd/database/migrations/testdata/fixtures/000513_user_ai_budget_overrides.up.sql new file mode 100644 index 0000000000..787b808b7d --- /dev/null +++ b/coderd/database/migrations/testdata/fixtures/000513_user_ai_budget_overrides.up.sql @@ -0,0 +1,15 @@ +-- Seed a group_members row so the override below references a real +-- membership. +INSERT INTO group_members ( + user_id, + group_id +) VALUES + ('30095c71-380b-457a-8995-97b8ee6e5307', 'bb640d07-ca8a-4869-b6bc-ae61ebb2fda1') +ON CONFLICT DO NOTHING; + +INSERT INTO user_ai_budget_overrides ( + user_id, + group_id, + spend_limit_micros +) VALUES + ('30095c71-380b-457a-8995-97b8ee6e5307', 'bb640d07-ca8a-4869-b6bc-ae61ebb2fda1', 500000000); diff --git a/coderd/database/models.go b/coderd/database/models.go index 6149e12a10..ebfaa7a051 100644 --- a/coderd/database/models.go +++ b/coderd/database/models.go @@ -5730,6 +5730,15 @@ type User struct { ChatSpendLimitMicros sql.NullInt64 `db:"chat_spend_limit_micros" json:"chat_spend_limit_micros"` } +// Per-user AI spend override that supersedes group budget resolution. +type UserAiBudgetOverride struct { + UserID uuid.UUID `db:"user_id" json:"user_id"` + GroupID uuid.UUID `db:"group_id" json:"group_id"` + SpendLimitMicros int64 `db:"spend_limit_micros" json:"spend_limit_micros"` + CreatedAt time.Time `db:"created_at" json:"created_at"` + UpdatedAt time.Time `db:"updated_at" json:"updated_at"` +} + // User-owned API keys associated with AI providers. These keys are used only when BYOK is enabled. type UserAiProviderKey struct { ID uuid.UUID `db:"id" json:"id"` diff --git a/coderd/database/querier.go b/coderd/database/querier.go index 0b3d565b69..a6c8f3e7db 100644 --- a/coderd/database/querier.go +++ b/coderd/database/querier.go @@ -198,6 +198,7 @@ type sqlcQuerier interface { DeleteTailnetPeer(ctx context.Context, arg DeleteTailnetPeerParams) (DeleteTailnetPeerRow, error) DeleteTailnetTunnel(ctx context.Context, arg DeleteTailnetTunnelParams) (DeleteTailnetTunnelRow, error) DeleteTask(ctx context.Context, arg DeleteTaskParams) (uuid.UUID, error) + DeleteUserAIBudgetOverride(ctx context.Context, userID uuid.UUID) (UserAiBudgetOverride, error) DeleteUserAIProviderKey(ctx context.Context, arg DeleteUserAIProviderKeyParams) error DeleteUserAIProviderKeysByProviderID(ctx context.Context, aiProviderID uuid.UUID) error DeleteUserChatCompactionThreshold(ctx context.Context, arg DeleteUserChatCompactionThresholdParams) error @@ -738,6 +739,7 @@ type sqlcQuerier interface { // inclusive. GetTotalUsageDCManagedAgentsV1(ctx context.Context, arg GetTotalUsageDCManagedAgentsV1Params) (int64, error) GetUnexpiredLicenses(ctx context.Context) ([]License, error) + GetUserAIBudgetOverride(ctx context.Context, userID uuid.UUID) (UserAiBudgetOverride, error) GetUserAIProviderKeyByProviderID(ctx context.Context, arg GetUserAIProviderKeyByProviderIDParams) (UserAiProviderKey, error) // GetUserAIProviderKeys is used by dbcrypt key rotation. Request paths should use // user-scoped lookups instead of this bulk accessor. @@ -1407,6 +1409,7 @@ type sqlcQuerier interface { // used to store the data, and the minutes are summed for each user and template // combination. The result is stored in the template_usage_stats table. UpsertTemplateUsageStats(ctx context.Context) error + UpsertUserAIBudgetOverride(ctx context.Context, arg UpsertUserAIBudgetOverrideParams) (UserAiBudgetOverride, error) // UpsertUserAIProviderKey preserves the original id and created_at when the // user/provider pair already exists. On conflict, callers provide id and // created_at for the insert path only. diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go index b49e757549..dc646121dc 100644 --- a/coderd/database/queries.sql.go +++ b/coderd/database/queries.sql.go @@ -2449,6 +2449,23 @@ func (q *sqlQuerier) DeleteGroupAIBudget(ctx context.Context, groupID uuid.UUID) return i, err } +const deleteUserAIBudgetOverride = `-- name: DeleteUserAIBudgetOverride :one +DELETE FROM user_ai_budget_overrides WHERE user_id = $1 RETURNING user_id, group_id, spend_limit_micros, created_at, updated_at +` + +func (q *sqlQuerier) DeleteUserAIBudgetOverride(ctx context.Context, userID uuid.UUID) (UserAiBudgetOverride, error) { + row := q.db.QueryRowContext(ctx, deleteUserAIBudgetOverride, userID) + var i UserAiBudgetOverride + err := row.Scan( + &i.UserID, + &i.GroupID, + &i.SpendLimitMicros, + &i.CreatedAt, + &i.UpdatedAt, + ) + return i, err +} + const getAIModelPriceByProviderModel = `-- name: GetAIModelPriceByProviderModel :one SELECT provider, model, input_price, output_price, cache_read_price, cache_write_price, created_at, updated_at FROM ai_model_prices @@ -2494,6 +2511,25 @@ func (q *sqlQuerier) GetGroupAIBudget(ctx context.Context, groupID uuid.UUID) (G return i, err } +const getUserAIBudgetOverride = `-- name: GetUserAIBudgetOverride :one +SELECT user_id, group_id, spend_limit_micros, created_at, updated_at +FROM user_ai_budget_overrides +WHERE user_id = $1 +` + +func (q *sqlQuerier) GetUserAIBudgetOverride(ctx context.Context, userID uuid.UUID) (UserAiBudgetOverride, error) { + row := q.db.QueryRowContext(ctx, getUserAIBudgetOverride, userID) + var i UserAiBudgetOverride + err := row.Scan( + &i.UserID, + &i.GroupID, + &i.SpendLimitMicros, + &i.CreatedAt, + &i.UpdatedAt, + ) + return i, err +} + const upsertAIModelPrices = `-- name: UpsertAIModelPrices :exec INSERT INTO ai_model_prices ( provider, model, input_price, output_price, cache_read_price, cache_write_price @@ -2548,6 +2584,35 @@ func (q *sqlQuerier) UpsertGroupAIBudget(ctx context.Context, arg UpsertGroupAIB return i, err } +const upsertUserAIBudgetOverride = `-- name: UpsertUserAIBudgetOverride :one +INSERT INTO user_ai_budget_overrides (user_id, group_id, spend_limit_micros) +VALUES ($1, $2, $3) +ON CONFLICT (user_id) DO UPDATE SET + group_id = EXCLUDED.group_id, + spend_limit_micros = EXCLUDED.spend_limit_micros, + updated_at = NOW() +RETURNING user_id, group_id, spend_limit_micros, created_at, updated_at +` + +type UpsertUserAIBudgetOverrideParams struct { + UserID uuid.UUID `db:"user_id" json:"user_id"` + GroupID uuid.UUID `db:"group_id" json:"group_id"` + SpendLimitMicros int64 `db:"spend_limit_micros" json:"spend_limit_micros"` +} + +func (q *sqlQuerier) UpsertUserAIBudgetOverride(ctx context.Context, arg UpsertUserAIBudgetOverrideParams) (UserAiBudgetOverride, error) { + row := q.db.QueryRowContext(ctx, upsertUserAIBudgetOverride, arg.UserID, arg.GroupID, arg.SpendLimitMicros) + var i UserAiBudgetOverride + err := row.Scan( + &i.UserID, + &i.GroupID, + &i.SpendLimitMicros, + &i.CreatedAt, + &i.UpdatedAt, + ) + return i, err +} + const getActiveAISeatCount = `-- name: GetActiveAISeatCount :one SELECT COUNT(*) diff --git a/coderd/database/queries/aicostcontrol.sql b/coderd/database/queries/aicostcontrol.sql index 6740b2568c..188ec7357e 100644 --- a/coderd/database/queries/aicostcontrol.sql +++ b/coderd/database/queries/aicostcontrol.sql @@ -40,3 +40,20 @@ RETURNING *; -- name: DeleteGroupAIBudget :one DELETE FROM group_ai_budgets WHERE group_id = @group_id RETURNING *; + +-- name: GetUserAIBudgetOverride :one +SELECT * +FROM user_ai_budget_overrides +WHERE user_id = @user_id; + +-- name: UpsertUserAIBudgetOverride :one +INSERT INTO user_ai_budget_overrides (user_id, group_id, spend_limit_micros) +VALUES (@user_id, @group_id, @spend_limit_micros) +ON CONFLICT (user_id) DO UPDATE SET + group_id = EXCLUDED.group_id, + spend_limit_micros = EXCLUDED.spend_limit_micros, + updated_at = NOW() +RETURNING *; + +-- name: DeleteUserAIBudgetOverride :one +DELETE FROM user_ai_budget_overrides WHERE user_id = @user_id RETURNING *; diff --git a/coderd/database/unique_constraint.go b/coderd/database/unique_constraint.go index 8ef517a9cb..3d5e5dabcf 100644 --- a/coderd/database/unique_constraint.go +++ b/coderd/database/unique_constraint.go @@ -97,6 +97,7 @@ const ( UniqueTemplatesPkey UniqueConstraint = "templates_pkey" // ALTER TABLE ONLY templates ADD CONSTRAINT templates_pkey PRIMARY KEY (id); UniqueUsageEventsDailyPkey UniqueConstraint = "usage_events_daily_pkey" // ALTER TABLE ONLY usage_events_daily ADD CONSTRAINT usage_events_daily_pkey PRIMARY KEY (day, event_type); UniqueUsageEventsPkey UniqueConstraint = "usage_events_pkey" // ALTER TABLE ONLY usage_events ADD CONSTRAINT usage_events_pkey PRIMARY KEY (id); + UniqueUserAiBudgetOverridesPkey UniqueConstraint = "user_ai_budget_overrides_pkey" // ALTER TABLE ONLY user_ai_budget_overrides ADD CONSTRAINT user_ai_budget_overrides_pkey PRIMARY KEY (user_id); UniqueUserAiProviderKeysPkey UniqueConstraint = "user_ai_provider_keys_pkey" // ALTER TABLE ONLY user_ai_provider_keys ADD CONSTRAINT user_ai_provider_keys_pkey PRIMARY KEY (id); UniqueUserAiProviderKeysUserIDAiProviderIDKey UniqueConstraint = "user_ai_provider_keys_user_id_ai_provider_id_key" // ALTER TABLE ONLY user_ai_provider_keys ADD CONSTRAINT user_ai_provider_keys_user_id_ai_provider_id_key UNIQUE (user_id, ai_provider_id); UniqueUserConfigsPkey UniqueConstraint = "user_configs_pkey" // ALTER TABLE ONLY user_configs ADD CONSTRAINT user_configs_pkey PRIMARY KEY (user_id, key); diff --git a/codersdk/aibridge.go b/codersdk/aibridge.go index 9bb7df0aac..d04359acb3 100644 --- a/codersdk/aibridge.go +++ b/codersdk/aibridge.go @@ -435,3 +435,71 @@ func (c *Client) DeleteGroupAIBudget(ctx context.Context, group uuid.UUID) error } return nil } + +type UserAIBudgetOverride struct { + UserID uuid.UUID `json:"user_id" format:"uuid"` + GroupID uuid.UUID `json:"group_id" format:"uuid"` + SpendLimitMicros int64 `json:"spend_limit_micros"` + CreatedAt time.Time `json:"created_at" format:"date-time"` + UpdatedAt time.Time `json:"updated_at" format:"date-time"` +} + +type UpsertUserAIBudgetOverrideRequest struct { + // GroupID is the group the user's spend is attributed to. The user must + // be a member of this group. + GroupID uuid.UUID `json:"group_id" format:"uuid" validate:"required"` + SpendLimitMicros int64 `json:"spend_limit_micros" validate:"gte=0"` +} + +// UserAIBudgetOverride returns the AI spend budget override configured for the given user. +func (c *Client) UserAIBudgetOverride(ctx context.Context, user uuid.UUID) (UserAIBudgetOverride, error) { + res, err := c.Request(ctx, http.MethodGet, + fmt.Sprintf("/api/v2/users/%s/ai/budget", user.String()), + nil, + ) + if err != nil { + return UserAIBudgetOverride{}, xerrors.Errorf("make request: %w", err) + } + defer res.Body.Close() + + if res.StatusCode != http.StatusOK { + return UserAIBudgetOverride{}, ReadBodyAsError(res) + } + var resp UserAIBudgetOverride + return resp, json.NewDecoder(res.Body).Decode(&resp) +} + +// UpsertUserAIBudgetOverride creates or updates the AI spend budget override for the given user. +func (c *Client) UpsertUserAIBudgetOverride(ctx context.Context, user uuid.UUID, req UpsertUserAIBudgetOverrideRequest) (UserAIBudgetOverride, error) { + res, err := c.Request(ctx, http.MethodPut, + fmt.Sprintf("/api/v2/users/%s/ai/budget", user.String()), + req, + ) + if err != nil { + return UserAIBudgetOverride{}, xerrors.Errorf("make request: %w", err) + } + defer res.Body.Close() + + if res.StatusCode != http.StatusOK { + return UserAIBudgetOverride{}, ReadBodyAsError(res) + } + var resp UserAIBudgetOverride + return resp, json.NewDecoder(res.Body).Decode(&resp) +} + +// DeleteUserAIBudgetOverride removes the AI spend budget override for the given user. +func (c *Client) DeleteUserAIBudgetOverride(ctx context.Context, user uuid.UUID) error { + res, err := c.Request(ctx, http.MethodDelete, + fmt.Sprintf("/api/v2/users/%s/ai/budget", user.String()), + nil, + ) + if err != nil { + return xerrors.Errorf("make request: %w", err) + } + defer res.Body.Close() + + if res.StatusCode != http.StatusNoContent { + return ReadBodyAsError(res) + } + return nil +} diff --git a/docs/reference/api/enterprise.md b/docs/reference/api/enterprise.md index eca72025c0..ed1ce268e7 100644 --- a/docs/reference/api/enterprise.md +++ b/docs/reference/api/enterprise.md @@ -3418,6 +3418,125 @@ curl -X POST http://coder-server:8080/api/v2/templates/{template}/prebuilds/inva To perform this operation, you must be authenticated. [Learn more](authentication.md). +## Get user AI budget override + +### Code samples + +```shell +# Example request using curl +curl -X GET http://coder-server:8080/api/v2/users/{user}/ai/budget \ + -H 'Accept: application/json' \ + -H 'Coder-Session-Token: API_KEY' +``` + +`GET /api/v2/users/{user}/ai/budget` + +### Parameters + +| Name | In | Type | Required | Description | +|--------|------|--------|----------|--------------------------| +| `user` | path | string | true | User ID, username, or me | + +### Example responses + +> 200 Response + +```json +{ + "created_at": "2019-08-24T14:15:22Z", + "group_id": "306db4e0-7449-4501-b76f-075576fe2d8f", + "spend_limit_micros": 0, + "updated_at": "2019-08-24T14:15:22Z", + "user_id": "a169451c-8525-4352-b8ca-070dd449a1a5" +} +``` + +### Responses + +| Status | Meaning | Description | Schema | +|--------|---------------------------------------------------------|-------------|--------------------------------------------------------------------------| +| 200 | [OK](https://tools.ietf.org/html/rfc7231#section-6.3.1) | OK | [codersdk.UserAIBudgetOverride](schemas.md#codersdkuseraibudgetoverride) | + +To perform this operation, you must be authenticated. [Learn more](authentication.md). + +## Upsert user AI budget override + +### Code samples + +```shell +# Example request using curl +curl -X PUT http://coder-server:8080/api/v2/users/{user}/ai/budget \ + -H 'Content-Type: application/json' \ + -H 'Accept: application/json' \ + -H 'Coder-Session-Token: API_KEY' +``` + +`PUT /api/v2/users/{user}/ai/budget` + +> Body parameter + +```json +{ + "group_id": "306db4e0-7449-4501-b76f-075576fe2d8f", + "spend_limit_micros": 0 +} +``` + +### Parameters + +| Name | In | Type | Required | Description | +|--------|------|----------------------------------------------------------------------------------------------------|----------|----------------------------------------| +| `user` | path | string | true | User ID, username, or me | +| `body` | body | [codersdk.UpsertUserAIBudgetOverrideRequest](schemas.md#codersdkupsertuseraibudgetoverriderequest) | true | Upsert user AI budget override request | + +### Example responses + +> 200 Response + +```json +{ + "created_at": "2019-08-24T14:15:22Z", + "group_id": "306db4e0-7449-4501-b76f-075576fe2d8f", + "spend_limit_micros": 0, + "updated_at": "2019-08-24T14:15:22Z", + "user_id": "a169451c-8525-4352-b8ca-070dd449a1a5" +} +``` + +### Responses + +| Status | Meaning | Description | Schema | +|--------|---------------------------------------------------------|-------------|--------------------------------------------------------------------------| +| 200 | [OK](https://tools.ietf.org/html/rfc7231#section-6.3.1) | OK | [codersdk.UserAIBudgetOverride](schemas.md#codersdkuseraibudgetoverride) | + +To perform this operation, you must be authenticated. [Learn more](authentication.md). + +## Delete user AI budget override + +### Code samples + +```shell +# Example request using curl +curl -X DELETE http://coder-server:8080/api/v2/users/{user}/ai/budget \ + -H 'Coder-Session-Token: API_KEY' +``` + +`DELETE /api/v2/users/{user}/ai/budget` + +### Parameters + +| Name | In | Type | Required | Description | +|--------|------|--------|----------|--------------------------| +| `user` | path | string | true | User ID, username, or me | + +### Responses + +| Status | Meaning | Description | Schema | +|--------|-----------------------------------------------------------------|-------------|--------| +| 204 | [No Content](https://tools.ietf.org/html/rfc7231#section-6.3.5) | No Content | | + +To perform this operation, you must be authenticated. [Learn more](authentication.md). + ## Get user quiet hours schedule ### Code samples diff --git a/docs/reference/api/schemas.md b/docs/reference/api/schemas.md index 771657d3c0..ea8f19c4bf 100644 --- a/docs/reference/api/schemas.md +++ b/docs/reference/api/schemas.md @@ -13601,6 +13601,22 @@ If the schedule is empty, the user will be updated to use the default schedule.| |----------------------|---------|----------|--------------|-------------| | `spend_limit_micros` | integer | false | | | +## codersdk.UpsertUserAIBudgetOverrideRequest + +```json +{ + "group_id": "306db4e0-7449-4501-b76f-075576fe2d8f", + "spend_limit_micros": 0 +} +``` + +### Properties + +| Name | Type | Required | Restrictions | Description | +|----------------------|---------|----------|--------------|---------------------------------------------------------------------------------------------------| +| `group_id` | string | true | | Group ID is the group the user's spend is attributed to. The user must be a member of this group. | +| `spend_limit_micros` | integer | false | | | + ## codersdk.UpsertWorkspaceAgentPortShareRequest ```json @@ -13730,6 +13746,28 @@ If the schedule is empty, the user will be updated to use the default schedule.| |----------|-----------------------| | `status` | `active`, `suspended` | +## codersdk.UserAIBudgetOverride + +```json +{ + "created_at": "2019-08-24T14:15:22Z", + "group_id": "306db4e0-7449-4501-b76f-075576fe2d8f", + "spend_limit_micros": 0, + "updated_at": "2019-08-24T14:15:22Z", + "user_id": "a169451c-8525-4352-b8ca-070dd449a1a5" +} +``` + +### Properties + +| Name | Type | Required | Restrictions | Description | +|----------------------|---------|----------|--------------|-------------| +| `created_at` | string | false | | | +| `group_id` | string | false | | | +| `spend_limit_micros` | integer | false | | | +| `updated_at` | string | false | | | +| `user_id` | string | false | | | + ## codersdk.UserActivity ```json diff --git a/enterprise/coderd/aibridge.go b/enterprise/coderd/aibridge.go index 9773dba352..8a220760de 100644 --- a/enterprise/coderd/aibridge.go +++ b/enterprise/coderd/aibridge.go @@ -43,6 +43,11 @@ const ( // reference a valid resource in the expected scope. var errInvalidCursor = xerrors.New("invalid pagination cursor") +// This name is raised by a trigger function with USING CONSTRAINT. +// It is not a table CHECK constraint, so dbgen does not emit it in +// check_constraint.go. +const userAIBudgetOverridesMustBeGroupMemberConstraint database.CheckConstraint = "user_ai_budget_overrides_must_be_group_member" + // aibridgeHandler handles all aibridged-related endpoints. func aibridgeHandler(api *API, middlewares ...func(http.Handler) http.Handler) func(r chi.Router) { // Build the overload protection middleware chain for the aibridged handler. @@ -821,3 +826,116 @@ func (api *API) deleteGroupAIBudget(rw http.ResponseWriter, r *http.Request) { rw.WriteHeader(http.StatusNoContent) } + +// @Summary Get user AI budget override +// @ID get-user-ai-budget-override +// @Security CoderSessionToken +// @Produce json +// @Tags Enterprise +// @Param user path string true "User ID, username, or me" +// @Success 200 {object} codersdk.UserAIBudgetOverride +// @Router /api/v2/users/{user}/ai/budget [get] +func (api *API) userAIBudgetOverride(rw http.ResponseWriter, r *http.Request) { + ctx := r.Context() + user := httpmw.UserParam(r) + + override, err := api.Database.GetUserAIBudgetOverride(ctx, user.ID) + if httpapi.Is404Error(err) { + httpapi.ResourceNotFound(rw) + return + } + if err != nil { + api.Logger.Error(ctx, "get user AI budget override", slog.Error(err)) + httpapi.InternalServerError(rw, err) + return + } + + httpapi.Write(ctx, rw, http.StatusOK, db2sdk.UserAIBudgetOverride(override)) +} + +// @Summary Upsert user AI budget override +// @ID upsert-user-ai-budget-override +// @Security CoderSessionToken +// @Accept json +// @Produce json +// @Tags Enterprise +// @Param user path string true "User ID, username, or me" +// @Param request body codersdk.UpsertUserAIBudgetOverrideRequest true "Upsert user AI budget override request" +// @Success 200 {object} codersdk.UserAIBudgetOverride +// @Router /api/v2/users/{user}/ai/budget [put] +func (api *API) upsertUserAIBudgetOverride(rw http.ResponseWriter, r *http.Request) { + ctx := r.Context() + user := httpmw.UserParam(r) + + var req codersdk.UpsertUserAIBudgetOverrideRequest + if !httpapi.Read(ctx, rw, r, &req) { + return + } + + // Look up the group first so a missing or forbidden group_id returns + // 404, distinct from the 400 "not a member" case handled below. + if _, err := api.Database.GetGroupByID(ctx, req.GroupID); err != nil { + if httpapi.Is404Error(err) { + httpapi.ResourceNotFound(rw) + return + } + api.Logger.Error(ctx, "get group for user AI budget override", slog.Error(err)) + httpapi.InternalServerError(rw, err) + return + } + + override, err := api.Database.UpsertUserAIBudgetOverride(ctx, database.UpsertUserAIBudgetOverrideParams{ + UserID: user.ID, + GroupID: req.GroupID, + SpendLimitMicros: req.SpendLimitMicros, + }) + // A trigger enforces that the user must be a member of the attributed + // group; it raises check_violation with this constraint name. Map + // the violation to a structured 400. + if database.IsCheckViolation(err, userAIBudgetOverridesMustBeGroupMemberConstraint) { + httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{ + Message: "User is not a member of the referenced group.", + Validations: []codersdk.ValidationError{{ + Field: "group_id", + Detail: "user must be a member of this group", + }}, + }) + return + } + if httpapi.Is404Error(err) { + httpapi.ResourceNotFound(rw) + return + } + if err != nil { + api.Logger.Error(ctx, "upsert user AI budget override", slog.Error(err)) + httpapi.InternalServerError(rw, err) + return + } + + httpapi.Write(ctx, rw, http.StatusOK, db2sdk.UserAIBudgetOverride(override)) +} + +// @Summary Delete user AI budget override +// @ID delete-user-ai-budget-override +// @Security CoderSessionToken +// @Tags Enterprise +// @Param user path string true "User ID, username, or me" +// @Success 204 +// @Router /api/v2/users/{user}/ai/budget [delete] +func (api *API) deleteUserAIBudgetOverride(rw http.ResponseWriter, r *http.Request) { + ctx := r.Context() + user := httpmw.UserParam(r) + + _, err := api.Database.DeleteUserAIBudgetOverride(ctx, user.ID) + if httpapi.Is404Error(err) { + httpapi.ResourceNotFound(rw) + return + } + if err != nil { + api.Logger.Error(ctx, "delete user AI budget override", slog.Error(err)) + httpapi.InternalServerError(rw, err) + return + } + + rw.WriteHeader(http.StatusNoContent) +} diff --git a/enterprise/coderd/aibridge_test.go b/enterprise/coderd/aibridge_test.go index 158f682842..1faadd1f53 100644 --- a/enterprise/coderd/aibridge_test.go +++ b/enterprise/coderd/aibridge_test.go @@ -2871,6 +2871,447 @@ func TestGroupAIBudget(t *testing.T) { }) } +func TestUserAIBudgetOverride(t *testing.T) { + t.Parallel() + + t.Run("Upsert/CreatesAndUpdates", func(t *testing.T) { + t.Parallel() + + adminClient, targetUser, group := setupUserAIBudgetOverrideTest(t) + ctx := testutil.Context(t, testutil.WaitLong) + + // First upsert creates the override. + newOverride, err := adminClient.UpsertUserAIBudgetOverride(ctx, targetUser.ID, codersdk.UpsertUserAIBudgetOverrideRequest{ + GroupID: group.ID, + SpendLimitMicros: 500_000_000, + }) + require.NoError(t, err) + require.Equal(t, targetUser.ID, newOverride.UserID) + require.Equal(t, group.ID, newOverride.GroupID) + require.EqualValues(t, 500_000_000, newOverride.SpendLimitMicros) + + // Second upsert updates the existing override. + updatedOverride, err := adminClient.UpsertUserAIBudgetOverride(ctx, targetUser.ID, codersdk.UpsertUserAIBudgetOverrideRequest{ + GroupID: group.ID, + SpendLimitMicros: 1_000_000_000, + }) + require.NoError(t, err) + require.EqualValues(t, 1_000_000_000, updatedOverride.SpendLimitMicros) + + // GET returns the latest value. + currentOverride, err := adminClient.UserAIBudgetOverride(ctx, targetUser.ID) + require.NoError(t, err) + require.EqualValues(t, 1_000_000_000, currentOverride.SpendLimitMicros) + }) + + t.Run("Upsert/ReassignsGroup", func(t *testing.T) { + t.Parallel() + + adminClient, targetUser, groupA := setupUserAIBudgetOverrideTest(t) + ctx := testutil.Context(t, testutil.WaitLong) + + // First upsert: attribute spend to groupA. + _, err := adminClient.UpsertUserAIBudgetOverride(ctx, targetUser.ID, codersdk.UpsertUserAIBudgetOverrideRequest{ + GroupID: groupA.ID, + SpendLimitMicros: 500_000_000, + }) + require.NoError(t, err) + + // Create groupB in the same org and add the target user. + groupB, err := adminClient.CreateGroup(ctx, targetUser.OrganizationIDs[0], codersdk.CreateGroupRequest{ + Name: "reassign-test-group-b", + }) + require.NoError(t, err) + _, err = adminClient.PatchGroup(ctx, groupB.ID, codersdk.PatchGroupRequest{ + AddUsers: []string{targetUser.ID.String()}, + }) + require.NoError(t, err) + + // Reassign the override's attribution to groupB. + updated, err := adminClient.UpsertUserAIBudgetOverride(ctx, targetUser.ID, codersdk.UpsertUserAIBudgetOverrideRequest{ + GroupID: groupB.ID, + SpendLimitMicros: 500_000_000, + }) + require.NoError(t, err) + require.Equal(t, groupB.ID, updated.GroupID, "upsert should change attributed group") + + // GET reflects the new group. + got, err := adminClient.UserAIBudgetOverride(ctx, targetUser.ID) + require.NoError(t, err) + require.Equal(t, groupB.ID, got.GroupID, "GET should reflect new group") + }) + + t.Run("Upsert/EveryoneGroup", func(t *testing.T) { + t.Parallel() + + adminClient, targetUser, _ := setupUserAIBudgetOverrideTest(t) + ctx := testutil.Context(t, testutil.WaitLong) + + // The Everyone group has id == organization_id, and the target user + // is implicitly a member via organization_members rather than + // group_members. The membership trigger queries + // group_members_expanded (a UNION of both tables), so this case + // exercises the organization_members branch. + everyoneGroupID := targetUser.OrganizationIDs[0] + + override, err := adminClient.UpsertUserAIBudgetOverride(ctx, targetUser.ID, codersdk.UpsertUserAIBudgetOverrideRequest{ + GroupID: everyoneGroupID, + SpendLimitMicros: 500_000_000, + }) + require.NoError(t, err, "should be able to attribute override to Everyone group") + require.Equal(t, targetUser.ID, override.UserID) + require.Equal(t, everyoneGroupID, override.GroupID) + require.EqualValues(t, 500_000_000, override.SpendLimitMicros) + }) + + t.Run("Upsert/AcceptsZeroSpendLimit", func(t *testing.T) { + t.Parallel() + + adminClient, targetUser, group := setupUserAIBudgetOverrideTest(t) + ctx := testutil.Context(t, testutil.WaitLong) + + // 0 is a valid value: it blocks all spend for the user. + override, err := adminClient.UpsertUserAIBudgetOverride(ctx, targetUser.ID, codersdk.UpsertUserAIBudgetOverrideRequest{ + GroupID: group.ID, + SpendLimitMicros: 0, + }) + require.NoError(t, err) + require.EqualValues(t, 0, override.SpendLimitMicros) + }) + + t.Run("Upsert/RejectsNegativeSpend", func(t *testing.T) { + t.Parallel() + + adminClient, targetUser, group := setupUserAIBudgetOverrideTest(t) + ctx := testutil.Context(t, testutil.WaitLong) + + _, err := adminClient.UpsertUserAIBudgetOverride(ctx, targetUser.ID, codersdk.UpsertUserAIBudgetOverrideRequest{ + GroupID: group.ID, + SpendLimitMicros: -1, + }) + var sdkErr *codersdk.Error + require.ErrorAs(t, err, &sdkErr) + require.Equal(t, http.StatusBadRequest, sdkErr.StatusCode()) + }) + + t.Run("Upsert/RejectsUnknownGroup", func(t *testing.T) { + t.Parallel() + + adminClient, targetUser, _ := setupUserAIBudgetOverrideTest(t) + ctx := testutil.Context(t, testutil.WaitLong) + + // A group_id that doesn't exist (or that the caller can't see) + // is rejected by the visibility check before the membership check. + _, err := adminClient.UpsertUserAIBudgetOverride(ctx, targetUser.ID, codersdk.UpsertUserAIBudgetOverrideRequest{ + GroupID: uuid.New(), + SpendLimitMicros: 500_000_000, + }) + var sdkErr *codersdk.Error + require.ErrorAs(t, err, &sdkErr) + require.Equal(t, http.StatusNotFound, sdkErr.StatusCode()) + }) + + t.Run("Upsert/RejectsNonMemberGroup", func(t *testing.T) { + t.Parallel() + + adminClient, targetUser, _ := setupUserAIBudgetOverrideTest(t) + ctx := testutil.Context(t, testutil.WaitLong) + + // Create a second group the target is NOT a member of. + outsiderGroup, err := adminClient.CreateGroup(ctx, targetUser.OrganizationIDs[0], codersdk.CreateGroupRequest{ + Name: "outsider-group", + }) + require.NoError(t, err) + + _, err = adminClient.UpsertUserAIBudgetOverride(ctx, targetUser.ID, codersdk.UpsertUserAIBudgetOverrideRequest{ + GroupID: outsiderGroup.ID, + SpendLimitMicros: 500_000_000, + }) + var sdkErr *codersdk.Error + require.ErrorAs(t, err, &sdkErr) + require.Equal(t, http.StatusBadRequest, sdkErr.StatusCode()) + }) + + t.Run("Get/AbsentReturns404", func(t *testing.T) { + t.Parallel() + + adminClient, targetUser, _ := setupUserAIBudgetOverrideTest(t) + ctx := testutil.Context(t, testutil.WaitLong) + + _, err := adminClient.UserAIBudgetOverride(ctx, targetUser.ID) + var sdkErr *codersdk.Error + require.ErrorAs(t, err, &sdkErr) + require.Equal(t, http.StatusNotFound, sdkErr.StatusCode()) + }) + + t.Run("Get/UnknownUserReturns404", func(t *testing.T) { + t.Parallel() + + adminClient, _, _ := setupUserAIBudgetOverrideTest(t) + ctx := testutil.Context(t, testutil.WaitLong) + + _, err := adminClient.UserAIBudgetOverride(ctx, uuid.New()) + var sdkErr *codersdk.Error + require.ErrorAs(t, err, &sdkErr) + require.Equal(t, http.StatusNotFound, sdkErr.StatusCode()) + }) + + t.Run("Delete/RoundTrip", func(t *testing.T) { + t.Parallel() + + adminClient, targetUser, group := setupUserAIBudgetOverrideTest(t) + ctx := testutil.Context(t, testutil.WaitLong) + + _, err := adminClient.UpsertUserAIBudgetOverride(ctx, targetUser.ID, codersdk.UpsertUserAIBudgetOverrideRequest{ + GroupID: group.ID, + SpendLimitMicros: 500_000_000, + }) + require.NoError(t, err) + + require.NoError(t, adminClient.DeleteUserAIBudgetOverride(ctx, targetUser.ID)) + + _, err = adminClient.UserAIBudgetOverride(ctx, targetUser.ID) + var sdkErr *codersdk.Error + require.ErrorAs(t, err, &sdkErr) + require.Equal(t, http.StatusNotFound, sdkErr.StatusCode()) + }) + + t.Run("Delete/AbsentReturns404", func(t *testing.T) { + t.Parallel() + + adminClient, targetUser, _ := setupUserAIBudgetOverrideTest(t) + ctx := testutil.Context(t, testutil.WaitLong) + + err := adminClient.DeleteUserAIBudgetOverride(ctx, targetUser.ID) + var sdkErr *codersdk.Error + require.ErrorAs(t, err, &sdkErr) + require.Equal(t, http.StatusNotFound, sdkErr.StatusCode()) + }) +} + +// TestUserAIBudgetOverrideRoleAccess verifies the authz matrix for the roles +// expected to interact with user budget overrides: +// +// - Owner / UserAdmin: full CRUD. +// - OrgAdmin / OrgUserAdmin: read-only. Writes require ActionUpdate on the +// User resource (site-scoped), which neither role has. +// +//nolint:tparallel // Subtests run sequentially: they share the same deployment and group, and parallel PatchGroup calls on the same group race. +func TestUserAIBudgetOverrideRoleAccess(t *testing.T) { + t.Parallel() + + dv := coderdtest.DeploymentValues(t) + dv.AI.BridgeConfig.Enabled = serpent.Bool(true) + ownerClient, owner := coderdenttest.New(t, &coderdenttest.Options{ + Options: &coderdtest.Options{DeploymentValues: dv}, + LicenseOptions: &coderdenttest.LicenseOptions{ + Features: license.Features{ + codersdk.FeatureTemplateRBAC: 1, + codersdk.FeatureAIBridge: 1, + }, + }, + }) + userAdminClient, _ := coderdtest.CreateAnotherUser(t, ownerClient, owner.OrganizationID, rbac.RoleUserAdmin()) + orgAdminClient, _ := coderdtest.CreateAnotherUser(t, ownerClient, owner.OrganizationID, rbac.ScopedRoleOrgAdmin(owner.OrganizationID)) + orgUserAdminClient, _ := coderdtest.CreateAnotherUser(t, ownerClient, owner.OrganizationID, rbac.ScopedRoleOrgUserAdmin(owner.OrganizationID)) + + setupCtx := testutil.Context(t, testutil.WaitLong) + group, err := userAdminClient.CreateGroup(setupCtx, owner.OrganizationID, codersdk.CreateGroupRequest{ + Name: "role-access-group", + }) + require.NoError(t, err) + + cases := []struct { + Name string + Client *codersdk.Client + CanWrite bool + }{ + {Name: "Owner", Client: ownerClient, CanWrite: true}, + {Name: "UserAdmin", Client: userAdminClient, CanWrite: true}, + {Name: "OrgAdmin", Client: orgAdminClient, CanWrite: false}, + {Name: "OrgUserAdmin", Client: orgUserAdminClient, CanWrite: false}, + } + + //nolint:paralleltest // Subtests run sequentially: they share the same deployment and group, and parallel PatchGroup calls on the same group race. + for _, tc := range cases { + t.Run(tc.Name, func(t *testing.T) { + ctx := testutil.Context(t, testutil.WaitLong) + + // Each case gets a fresh target user. + _, targetUser := coderdtest.CreateAnotherUser(t, ownerClient, owner.OrganizationID) + _, err := userAdminClient.PatchGroup(ctx, group.ID, codersdk.PatchGroupRequest{ + AddUsers: []string{targetUser.ID.String()}, + }) + require.NoError(t, err) + + upsertReq := codersdk.UpsertUserAIBudgetOverrideRequest{ + GroupID: group.ID, + SpendLimitMicros: 500_000_000, + } + + if tc.CanWrite { + // Full CRUD lifecycle. + override, err := tc.Client.UpsertUserAIBudgetOverride(ctx, targetUser.ID, upsertReq) + require.NoError(t, err, "PUT") + require.Equal(t, group.ID, override.GroupID) + + got, err := tc.Client.UserAIBudgetOverride(ctx, targetUser.ID) + require.NoError(t, err, "GET") + require.EqualValues(t, 500_000_000, got.SpendLimitMicros) + + err = tc.Client.DeleteUserAIBudgetOverride(ctx, targetUser.ID) + require.NoError(t, err, "DELETE") + } else { + // PUT rejected. + _, err := tc.Client.UpsertUserAIBudgetOverride(ctx, targetUser.ID, upsertReq) + var sdkErr *codersdk.Error + require.ErrorAs(t, err, &sdkErr) + require.Equal(t, http.StatusNotFound, sdkErr.StatusCode(), "PUT") + + // Seed a row via UserAdmin so we can verify read access still works. + _, err = userAdminClient.UpsertUserAIBudgetOverride(ctx, targetUser.ID, upsertReq) + require.NoError(t, err) + + // GET still works (all roles have ActionRead on User). + got, err := tc.Client.UserAIBudgetOverride(ctx, targetUser.ID) + require.NoError(t, err, "GET") + require.EqualValues(t, 500_000_000, got.SpendLimitMicros) + + // DELETE rejected. + err = tc.Client.DeleteUserAIBudgetOverride(ctx, targetUser.ID) + require.ErrorAs(t, err, &sdkErr) + require.Equal(t, http.StatusNotFound, sdkErr.StatusCode(), "DELETE") + } + }) + } +} + +// TestUserAIBudgetOverrideDeletedOnMembershipRemoval verifies that a per-user +// override is deleted automatically when the user loses membership in the +// attributed group. Two paths are exercised: +// +// - RegularGroup: membership stored in group_members; removed via +// PatchGroup with RemoveUsers. +// - EveryoneGroup: membership stored in organization_members; removed +// via DeleteOrganizationMember. +func TestUserAIBudgetOverrideDeletedOnMembershipRemoval(t *testing.T) { + t.Parallel() + + dv := coderdtest.DeploymentValues(t) + dv.AI.BridgeConfig.Enabled = serpent.Bool(true) + ownerClient, owner := coderdenttest.New(t, &coderdenttest.Options{ + Options: &coderdtest.Options{DeploymentValues: dv}, + LicenseOptions: &coderdenttest.LicenseOptions{ + Features: license.Features{ + codersdk.FeatureTemplateRBAC: 1, + codersdk.FeatureAIBridge: 1, + }, + }, + }) + adminClient, _ := coderdtest.CreateAnotherUser(t, ownerClient, owner.OrganizationID, rbac.RoleUserAdmin()) + + // "Regular group" means any group except "Everyone". + t.Run("RegularGroup", func(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + + _, targetUser := coderdtest.CreateAnotherUser(t, ownerClient, owner.OrganizationID) + + group, err := adminClient.CreateGroup(ctx, owner.OrganizationID, codersdk.CreateGroupRequest{ + Name: "cascade-regular-group", + }) + require.NoError(t, err) + + _, err = adminClient.PatchGroup(ctx, group.ID, codersdk.PatchGroupRequest{ + AddUsers: []string{targetUser.ID.String()}, + }) + require.NoError(t, err) + + _, err = adminClient.UpsertUserAIBudgetOverride(ctx, targetUser.ID, codersdk.UpsertUserAIBudgetOverrideRequest{ + GroupID: group.ID, + SpendLimitMicros: 500_000_000, + }) + require.NoError(t, err, "set override") + + // Sanity-check the override exists. + _, err = adminClient.UserAIBudgetOverride(ctx, targetUser.ID) + require.NoError(t, err, "override should exist before removal") + + _, err = adminClient.PatchGroup(ctx, group.ID, codersdk.PatchGroupRequest{ + RemoveUsers: []string{targetUser.ID.String()}, + }) + require.NoError(t, err, "remove user from group") + + _, err = adminClient.UserAIBudgetOverride(ctx, targetUser.ID) + var sdkErr *codersdk.Error + require.ErrorAs(t, err, &sdkErr) + require.Equal(t, http.StatusNotFound, sdkErr.StatusCode(), + "override should be deleted after user is removed from the attributed group") + }) + + t.Run("EveryoneGroup", func(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + + _, targetUser := coderdtest.CreateAnotherUser(t, ownerClient, owner.OrganizationID) + + // The Everyone group has id == organization_id. + everyoneGroupID := owner.OrganizationID + + _, err := adminClient.UpsertUserAIBudgetOverride(ctx, targetUser.ID, codersdk.UpsertUserAIBudgetOverrideRequest{ + GroupID: everyoneGroupID, + SpendLimitMicros: 500_000_000, + }) + require.NoError(t, err, "set override") + + // Sanity-check the override exists. + _, err = adminClient.UserAIBudgetOverride(ctx, targetUser.ID) + require.NoError(t, err, "override should exist before removal") + + err = adminClient.DeleteOrganizationMember(ctx, owner.OrganizationID, targetUser.ID.String()) + require.NoError(t, err, "remove user from organization") + + _, err = adminClient.UserAIBudgetOverride(ctx, targetUser.ID) + var sdkErr *codersdk.Error + require.ErrorAs(t, err, &sdkErr) + require.Equal(t, http.StatusNotFound, sdkErr.StatusCode(), + "override should be deleted after user is removed from the organization") + }) +} + +// setupUserAIBudgetOverrideTest returns an Admin client, a target user, and a +// group the target user is a member of. +func setupUserAIBudgetOverrideTest(t *testing.T) (adminClient *codersdk.Client, targetUser codersdk.User, group codersdk.Group) { + t.Helper() + + dv := coderdtest.DeploymentValues(t) + dv.AI.BridgeConfig.Enabled = serpent.Bool(true) + ownerClient, owner := coderdenttest.New(t, &coderdenttest.Options{ + Options: &coderdtest.Options{DeploymentValues: dv}, + LicenseOptions: &coderdenttest.LicenseOptions{ + Features: license.Features{ + codersdk.FeatureTemplateRBAC: 1, + codersdk.FeatureAIBridge: 1, + }, + }, + }) + adminClient, _ = coderdtest.CreateAnotherUser(t, ownerClient, owner.OrganizationID, rbac.RoleUserAdmin()) + _, targetUser = coderdtest.CreateAnotherUser(t, ownerClient, owner.OrganizationID) + + ctx := testutil.Context(t, testutil.WaitLong) + g, err := adminClient.CreateGroup(ctx, owner.OrganizationID, codersdk.CreateGroupRequest{ + Name: "override-test-group", + }) + require.NoError(t, err) + g, err = adminClient.PatchGroup(ctx, g.ID, codersdk.PatchGroupRequest{ + AddUsers: []string{targetUser.ID.String()}, + }) + require.NoError(t, err) + return adminClient, targetUser, g +} + // setupGroupAIBudgetTest returns an Admin client along with a newly created group inside it. func setupGroupAIBudgetTest(t *testing.T) (adminClient *codersdk.Client, group codersdk.Group) { t.Helper() diff --git a/enterprise/coderd/coderd.go b/enterprise/coderd/coderd.go index 2a55c6218f..33732eea3d 100644 --- a/enterprise/coderd/coderd.go +++ b/enterprise/coderd/coderd.go @@ -596,6 +596,17 @@ func New(ctx context.Context, options *Options) (_ *API, err error) { r.Get("/", api.userQuietHoursSchedule) r.Put("/", api.putUserQuietHoursSchedule) }) + r.Route("/users/{user}/ai/budget", func(r chi.Router) { + // AI cost controls are a paid feature (AI Governance add-on). + r.Use( + api.RequireFeatureMW(codersdk.FeatureAIBridge), + apiKeyMiddleware, + httpmw.ExtractUserParam(options.Database), + ) + r.Get("/", api.userAIBudgetOverride) + r.Put("/", api.upsertUserAIBudgetOverride) + r.Delete("/", api.deleteUserAIBudgetOverride) + }) r.Route("/prebuilds", func(r chi.Router) { r.Use( apiKeyMiddleware, diff --git a/site/src/api/typesGenerated.ts b/site/src/api/typesGenerated.ts index 8a465a3073..c4fae7a3d0 100644 --- a/site/src/api/typesGenerated.ts +++ b/site/src/api/typesGenerated.ts @@ -9160,6 +9160,16 @@ export interface UpsertGroupAIBudgetRequest { readonly spend_limit_micros: number; } +// From codersdk/aibridge.go +export interface UpsertUserAIBudgetOverrideRequest { + /** + * GroupID is the group the user's spend is attributed to. The user must + * be a member of this group. + */ + readonly group_id: string; + readonly spend_limit_micros: number; +} + // From codersdk/workspaceagentportshare.go export interface UpsertWorkspaceAgentPortShareRequest { readonly agent_name: string; @@ -9204,6 +9214,15 @@ export interface User extends ReducedUser { readonly has_ai_seat: boolean; } +// From codersdk/aibridge.go +export interface UserAIBudgetOverride { + readonly user_id: string; + readonly group_id: string; + readonly spend_limit_micros: number; + readonly created_at: string; + readonly updated_at: string; +} + // From codersdk/chats.go /** * UserAIProviderKeyConfig is a provider summary from the current user's