From 238968cfa09c525c9bcb25b22b6c776960d0b7e3 Mon Sep 17 00:00:00 2001 From: Yevhenii Shcherbina Date: Thu, 14 May 2026 15:54:37 -0400 Subject: [PATCH] feat: add per-group AI budget table and endpoints (#25203) Closes https://linear.app/codercom/issue/AIGOV-284/add-group-budgets-table-and-crud-api ## Summary Adds the `group_ai_budgets` table and the following endpoints: - `GET /api/v2/groups/{group}/ai/budget` - `PUT /api/v2/groups/{group}/ai/budget` - `DELETE /api/v2/groups/{group}/ai/budget` Each group may have at most one budget row. If no row exists, no budget is enforced. ### Feature gate Added `RequireFeatureMW(FeatureAIBridge)` on the `/ai/budget` sub-route. ## RBAC Authorization reuses `rbac.ResourceGroup` with the existing `.InOrganization(...).WithID(...)` scoping model. The `dbauthz` wrappers load the parent `groups` row and authorize against it. No new resource type is introduced. As a result, anyone with `group:update` permissions (Owner, OrgAdmin, or UserAdmin within the organization) can manage AI budgets for that group. ## Read access for group members `database.Group.RBACObject()` grants `policy.ActionRead` to all members of the group through the group ACL: ```go func (g Group) RBACObject() rbac.Object { return rbac.ResourceGroup.WithID(g.ID). InOrg(g.OrganizationID). // Group members can read the group. WithGroupACL(map[string][]policy.Action{ g.ID.String(): { policy.ActionRead, }, }) } ``` Because the `GET` endpoint authorizes against the same loaded `Group` object, any group member can call: ```text GET /api/v2/groups/{group}/ai/budget ``` `PUT` and `DELETE` remain admin-only. The group ACL grants only `ActionRead`, so write operations continue to require role-based `group:update` permissions. ## Alternative considered A dedicated `rbac.ResourceGroupAiBudget` resource would allow budget management to be separated from general group administration. We decided not to add that complexity for now. --- coderd/apidoc/docs.go | 136 ++++++++++++ coderd/apidoc/swagger.json | 124 +++++++++++ coderd/database/check_constraint.go | 1 + coderd/database/db2sdk/db2sdk.go | 9 + coderd/database/dbauthz/dbauthz.go | 36 ++++ coderd/database/dbauthz/dbauthz_test.go | 25 +++ coderd/database/dbmetrics/querymetrics.go | 24 +++ coderd/database/dbmock/dbmock.go | 45 ++++ coderd/database/dump.sql | 16 ++ coderd/database/foreign_key_constraint.go | 1 + .../000497_group_ai_budgets.down.sql | 1 + .../migrations/000497_group_ai_budgets.up.sql | 9 + .../fixtures/000497_group_ai_budgets.up.sql | 5 + coderd/database/models.go | 8 + coderd/database/querier.go | 3 + coderd/database/queries.sql.go | 60 ++++++ coderd/database/queries/aicostcontrol.sql | 16 ++ coderd/database/unique_constraint.go | 1 + codersdk/aibridge.go | 65 ++++++ docs/reference/api/enterprise.md | 116 ++++++++++ docs/reference/api/schemas.md | 34 +++ enterprise/coderd/aibridge.go | 87 ++++++++ enterprise/coderd/aibridge_test.go | 198 ++++++++++++++++++ enterprise/coderd/coderd.go | 7 + site/src/api/typesGenerated.ts | 13 ++ 25 files changed, 1040 insertions(+) create mode 100644 coderd/database/migrations/000497_group_ai_budgets.down.sql create mode 100644 coderd/database/migrations/000497_group_ai_budgets.up.sql create mode 100644 coderd/database/migrations/testdata/fixtures/000497_group_ai_budgets.up.sql diff --git a/coderd/apidoc/docs.go b/coderd/apidoc/docs.go index 01647cbd34..b77aec28f6 100644 --- a/coderd/apidoc/docs.go +++ b/coderd/apidoc/docs.go @@ -2520,6 +2520,113 @@ const docTemplate = `{ ] } }, + "/api/v2/groups/{group}/ai/budget": { + "get": { + "produces": [ + "application/json" + ], + "tags": [ + "Enterprise" + ], + "summary": "Get group AI budget", + "operationId": "get-group-ai-budget", + "parameters": [ + { + "type": "string", + "format": "uuid", + "description": "Group ID", + "name": "group", + "in": "path", + "required": true + } + ], + "responses": { + "200": { + "description": "OK", + "schema": { + "$ref": "#/definitions/codersdk.GroupAIBudget" + } + } + }, + "security": [ + { + "CoderSessionToken": [] + } + ] + }, + "put": { + "consumes": [ + "application/json" + ], + "produces": [ + "application/json" + ], + "tags": [ + "Enterprise" + ], + "summary": "Upsert group AI budget", + "operationId": "upsert-group-ai-budget", + "parameters": [ + { + "type": "string", + "format": "uuid", + "description": "Group ID", + "name": "group", + "in": "path", + "required": true + }, + { + "description": "Upsert group AI budget request", + "name": "request", + "in": "body", + "required": true, + "schema": { + "$ref": "#/definitions/codersdk.UpsertGroupAIBudgetRequest" + } + } + ], + "responses": { + "200": { + "description": "OK", + "schema": { + "$ref": "#/definitions/codersdk.GroupAIBudget" + } + } + }, + "security": [ + { + "CoderSessionToken": [] + } + ] + }, + "delete": { + "tags": [ + "Enterprise" + ], + "summary": "Delete group AI budget", + "operationId": "delete-group-ai-budget", + "parameters": [ + { + "type": "string", + "format": "uuid", + "description": "Group ID", + "name": "group", + "in": "path", + "required": true + } + ], + "responses": { + "204": { + "description": "No Content" + } + }, + "security": [ + { + "CoderSessionToken": [] + } + ] + } + }, "/api/v2/groups/{group}/members": { "get": { "produces": [ @@ -18569,6 +18676,26 @@ const docTemplate = `{ } } }, + "codersdk.GroupAIBudget": { + "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" + } + } + }, "codersdk.GroupMembersResponse": { "type": "object", "properties": { @@ -23753,6 +23880,15 @@ const docTemplate = `{ } } }, + "codersdk.UpsertGroupAIBudgetRequest": { + "type": "object", + "properties": { + "spend_limit_micros": { + "type": "integer", + "minimum": 0 + } + } + }, "codersdk.UpsertWorkspaceAgentPortShareRequest": { "type": "object", "properties": { diff --git a/coderd/apidoc/swagger.json b/coderd/apidoc/swagger.json index 4cc8ee4374..d6de642bbd 100644 --- a/coderd/apidoc/swagger.json +++ b/coderd/apidoc/swagger.json @@ -2217,6 +2217,101 @@ ] } }, + "/api/v2/groups/{group}/ai/budget": { + "get": { + "produces": ["application/json"], + "tags": ["Enterprise"], + "summary": "Get group AI budget", + "operationId": "get-group-ai-budget", + "parameters": [ + { + "type": "string", + "format": "uuid", + "description": "Group ID", + "name": "group", + "in": "path", + "required": true + } + ], + "responses": { + "200": { + "description": "OK", + "schema": { + "$ref": "#/definitions/codersdk.GroupAIBudget" + } + } + }, + "security": [ + { + "CoderSessionToken": [] + } + ] + }, + "put": { + "consumes": ["application/json"], + "produces": ["application/json"], + "tags": ["Enterprise"], + "summary": "Upsert group AI budget", + "operationId": "upsert-group-ai-budget", + "parameters": [ + { + "type": "string", + "format": "uuid", + "description": "Group ID", + "name": "group", + "in": "path", + "required": true + }, + { + "description": "Upsert group AI budget request", + "name": "request", + "in": "body", + "required": true, + "schema": { + "$ref": "#/definitions/codersdk.UpsertGroupAIBudgetRequest" + } + } + ], + "responses": { + "200": { + "description": "OK", + "schema": { + "$ref": "#/definitions/codersdk.GroupAIBudget" + } + } + }, + "security": [ + { + "CoderSessionToken": [] + } + ] + }, + "delete": { + "tags": ["Enterprise"], + "summary": "Delete group AI budget", + "operationId": "delete-group-ai-budget", + "parameters": [ + { + "type": "string", + "format": "uuid", + "description": "Group ID", + "name": "group", + "in": "path", + "required": true + } + ], + "responses": { + "204": { + "description": "No Content" + } + }, + "security": [ + { + "CoderSessionToken": [] + } + ] + } + }, "/api/v2/groups/{group}/members": { "get": { "produces": ["application/json"], @@ -16914,6 +17009,26 @@ } } }, + "codersdk.GroupAIBudget": { + "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" + } + } + }, "codersdk.GroupMembersResponse": { "type": "object", "properties": { @@ -21867,6 +21982,15 @@ } } }, + "codersdk.UpsertGroupAIBudgetRequest": { + "type": "object", + "properties": { + "spend_limit_micros": { + "type": "integer", + "minimum": 0 + } + } + }, "codersdk.UpsertWorkspaceAgentPortShareRequest": { "type": "object", "properties": { diff --git a/coderd/database/check_constraint.go b/coderd/database/check_constraint.go index 6ba18bd6c6..90750642cd 100644 --- a/coderd/database/check_constraint.go +++ b/coderd/database/check_constraint.go @@ -27,6 +27,7 @@ const ( CheckUsersServiceAccountLoginType CheckConstraint = "users_service_account_login_type" // users CheckUsersUsernameMinLength CheckConstraint = "users_username_min_length" // users CheckOrganizationIDNotZero CheckConstraint = "organization_id_not_zero" // custom_roles + CheckGroupAiBudgetsSpendLimitMicrosCheck CheckConstraint = "group_ai_budgets_spend_limit_micros_check" // group_ai_budgets CheckGroupsChatSpendLimitMicrosCheck CheckConstraint = "groups_chat_spend_limit_micros_check" // groups CheckMcpServerConfigsAuthTypeCheck CheckConstraint = "mcp_server_configs_auth_type_check" // mcp_server_configs CheckMcpServerConfigsAvailabilityCheck CheckConstraint = "mcp_server_configs_availability_check" // mcp_server_configs diff --git a/coderd/database/db2sdk/db2sdk.go b/coderd/database/db2sdk/db2sdk.go index 8e1ed20330..e485ade9ff 100644 --- a/coderd/database/db2sdk/db2sdk.go +++ b/coderd/database/db2sdk/db2sdk.go @@ -1418,6 +1418,15 @@ func flattenAndSum(sums map[string]int64, prefix string, m map[string]json.RawMe } } +func GroupAIBudget(b database.GroupAiBudget) codersdk.GroupAIBudget { + return codersdk.GroupAIBudget{ + GroupID: b.GroupID, + SpendLimitMicros: b.SpendLimitMicros, + CreatedAt: b.CreatedAt, + UpdatedAt: b.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 f251dcf06d..599eb87460 100644 --- a/coderd/database/dbauthz/dbauthz.go +++ b/coderd/database/dbauthz/dbauthz.go @@ -2016,6 +2016,18 @@ func (q *querier) DeleteExternalAuthLink(ctx context.Context, arg database.Delet }, q.db.DeleteExternalAuthLink)(ctx, arg) } +func (q *querier) DeleteGroupAIBudget(ctx context.Context, groupID uuid.UUID) (database.GroupAiBudget, error) { + // Removing a group's AI budget counts as updating the group. + group, err := q.db.GetGroupByID(ctx, groupID) + if err != nil { + return database.GroupAiBudget{}, err + } + if err := q.authorizeContext(ctx, policy.ActionUpdate, group); err != nil { + return database.GroupAiBudget{}, err + } + return q.db.DeleteGroupAIBudget(ctx, groupID) +} + func (q *querier) DeleteGroupByID(ctx context.Context, id uuid.UUID) error { return deleteQ(q.log, q.auth, q.db.GetGroupByID, q.db.DeleteGroupByID)(ctx, id) } @@ -3364,6 +3376,18 @@ func (q *querier) GetGitSSHKey(ctx context.Context, userID uuid.UUID) (database. return fetchWithAction(q.log, q.auth, policy.ActionReadPersonal, q.db.GetGitSSHKey)(ctx, userID) } +func (q *querier) GetGroupAIBudget(ctx context.Context, groupID uuid.UUID) (database.GroupAiBudget, error) { + // Reading a group's AI budget requires read on the parent group. + group, err := q.db.GetGroupByID(ctx, groupID) + if err != nil { + return database.GroupAiBudget{}, err + } + if err := q.authorizeContext(ctx, policy.ActionRead, group); err != nil { + return database.GroupAiBudget{}, err + } + return q.db.GetGroupAIBudget(ctx, groupID) +} + func (q *querier) GetGroupByID(ctx context.Context, id uuid.UUID) (database.Group, error) { return fetch(q.log, q.auth, q.db.GetGroupByID)(ctx, id) } @@ -7927,6 +7951,18 @@ func (q *querier) UpsertDefaultProxy(ctx context.Context, arg database.UpsertDef return q.db.UpsertDefaultProxy(ctx, arg) } +func (q *querier) UpsertGroupAIBudget(ctx context.Context, arg database.UpsertGroupAIBudgetParams) (database.GroupAiBudget, error) { + // Setting a group's AI budget counts as updating the group. + group, err := q.db.GetGroupByID(ctx, arg.GroupID) + if err != nil { + return database.GroupAiBudget{}, err + } + if err := q.authorizeContext(ctx, policy.ActionUpdate, group); err != nil { + return database.GroupAiBudget{}, err + } + return q.db.UpsertGroupAIBudget(ctx, arg) +} + func (q *querier) UpsertHealthSettings(ctx context.Context, value string) error { if err := q.authorizeContext(ctx, policy.ActionUpdate, rbac.ResourceDeploymentConfig); err != nil { return err diff --git a/coderd/database/dbauthz/dbauthz_test.go b/coderd/database/dbauthz/dbauthz_test.go index f6983b0a20..ef65d0f9b6 100644 --- a/coderd/database/dbauthz/dbauthz_test.go +++ b/coderd/database/dbauthz/dbauthz_test.go @@ -6209,6 +6209,31 @@ func (s *MethodTestSuite) TestAIBridge() { check.Args(database.GetAIModelPriceByProviderModelParams{}).Asserts(rbac.ResourceAiModelPrice, policy.ActionRead) })) + s.Run("GetGroupAIBudget", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { + g := testutil.Fake(s.T(), faker, database.Group{}) + b := testutil.Fake(s.T(), faker, database.GroupAiBudget{GroupID: g.ID}) + dbm.EXPECT().GetGroupByID(gomock.Any(), g.ID).Return(g, nil).AnyTimes() + dbm.EXPECT().GetGroupAIBudget(gomock.Any(), g.ID).Return(b, nil).AnyTimes() + check.Args(g.ID).Asserts(g, policy.ActionRead).Returns(b) + })) + + s.Run("UpsertGroupAIBudget", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { + g := testutil.Fake(s.T(), faker, database.Group{}) + b := testutil.Fake(s.T(), faker, database.GroupAiBudget{GroupID: g.ID}) + arg := database.UpsertGroupAIBudgetParams{GroupID: g.ID, SpendLimitMicros: b.SpendLimitMicros} + dbm.EXPECT().GetGroupByID(gomock.Any(), g.ID).Return(g, nil).AnyTimes() + dbm.EXPECT().UpsertGroupAIBudget(gomock.Any(), arg).Return(b, nil).AnyTimes() + check.Args(arg).Asserts(g, policy.ActionUpdate).Returns(b) + })) + + s.Run("DeleteGroupAIBudget", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { + g := testutil.Fake(s.T(), faker, database.Group{}) + b := testutil.Fake(s.T(), faker, database.GroupAiBudget{GroupID: g.ID}) + dbm.EXPECT().GetGroupByID(gomock.Any(), g.ID).Return(g, nil).AnyTimes() + dbm.EXPECT().DeleteGroupAIBudget(gomock.Any(), g.ID).Return(b, nil).AnyTimes() + check.Args(g.ID).Asserts(g, policy.ActionUpdate).Returns(b) + })) + 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 69255ac229..eb5c0950bf 100644 --- a/coderd/database/dbmetrics/querymetrics.go +++ b/coderd/database/dbmetrics/querymetrics.go @@ -537,6 +537,14 @@ func (m queryMetricsStore) DeleteExternalAuthLink(ctx context.Context, arg datab return r0 } +func (m queryMetricsStore) DeleteGroupAIBudget(ctx context.Context, groupID uuid.UUID) (database.GroupAiBudget, error) { + start := time.Now() + r0, r1 := m.s.DeleteGroupAIBudget(ctx, groupID) + m.queryLatencies.WithLabelValues("DeleteGroupAIBudget").Observe(time.Since(start).Seconds()) + m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "DeleteGroupAIBudget").Inc() + return r0, r1 +} + func (m queryMetricsStore) DeleteGroupByID(ctx context.Context, id uuid.UUID) error { start := time.Now() r0 := m.s.DeleteGroupByID(ctx, id) @@ -1857,6 +1865,14 @@ func (m queryMetricsStore) GetGitSSHKey(ctx context.Context, userID uuid.UUID) ( return r0, r1 } +func (m queryMetricsStore) GetGroupAIBudget(ctx context.Context, groupID uuid.UUID) (database.GroupAiBudget, error) { + start := time.Now() + r0, r1 := m.s.GetGroupAIBudget(ctx, groupID) + m.queryLatencies.WithLabelValues("GetGroupAIBudget").Observe(time.Since(start).Seconds()) + m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "GetGroupAIBudget").Inc() + return r0, r1 +} + func (m queryMetricsStore) GetGroupByID(ctx context.Context, id uuid.UUID) (database.Group, error) { start := time.Now() r0, r1 := m.s.GetGroupByID(ctx, id) @@ -5761,6 +5777,14 @@ func (m queryMetricsStore) UpsertDefaultProxy(ctx context.Context, arg database. return r0 } +func (m queryMetricsStore) UpsertGroupAIBudget(ctx context.Context, arg database.UpsertGroupAIBudgetParams) (database.GroupAiBudget, error) { + start := time.Now() + r0, r1 := m.s.UpsertGroupAIBudget(ctx, arg) + m.queryLatencies.WithLabelValues("UpsertGroupAIBudget").Observe(time.Since(start).Seconds()) + m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "UpsertGroupAIBudget").Inc() + return r0, r1 +} + func (m queryMetricsStore) UpsertHealthSettings(ctx context.Context, value string) error { start := time.Now() r0 := m.s.UpsertHealthSettings(ctx, value) diff --git a/coderd/database/dbmock/dbmock.go b/coderd/database/dbmock/dbmock.go index 40396a9b5a..298ef831a0 100644 --- a/coderd/database/dbmock/dbmock.go +++ b/coderd/database/dbmock/dbmock.go @@ -888,6 +888,21 @@ func (mr *MockStoreMockRecorder) DeleteExternalAuthLink(ctx, arg any) *gomock.Ca return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DeleteExternalAuthLink", reflect.TypeOf((*MockStore)(nil).DeleteExternalAuthLink), ctx, arg) } +// DeleteGroupAIBudget mocks base method. +func (m *MockStore) DeleteGroupAIBudget(ctx context.Context, groupID uuid.UUID) (database.GroupAiBudget, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "DeleteGroupAIBudget", ctx, groupID) + ret0, _ := ret[0].(database.GroupAiBudget) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// DeleteGroupAIBudget indicates an expected call of DeleteGroupAIBudget. +func (mr *MockStoreMockRecorder) DeleteGroupAIBudget(ctx, groupID any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DeleteGroupAIBudget", reflect.TypeOf((*MockStore)(nil).DeleteGroupAIBudget), ctx, groupID) +} + // DeleteGroupByID mocks base method. func (m *MockStore) DeleteGroupByID(ctx context.Context, id uuid.UUID) error { m.ctrl.T.Helper() @@ -3436,6 +3451,21 @@ func (mr *MockStoreMockRecorder) GetGitSSHKey(ctx, userID any) *gomock.Call { return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetGitSSHKey", reflect.TypeOf((*MockStore)(nil).GetGitSSHKey), ctx, userID) } +// GetGroupAIBudget mocks base method. +func (m *MockStore) GetGroupAIBudget(ctx context.Context, groupID uuid.UUID) (database.GroupAiBudget, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "GetGroupAIBudget", ctx, groupID) + ret0, _ := ret[0].(database.GroupAiBudget) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// GetGroupAIBudget indicates an expected call of GetGroupAIBudget. +func (mr *MockStoreMockRecorder) GetGroupAIBudget(ctx, groupID any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetGroupAIBudget", reflect.TypeOf((*MockStore)(nil).GetGroupAIBudget), ctx, groupID) +} + // GetGroupByID mocks base method. func (m *MockStore) GetGroupByID(ctx context.Context, id uuid.UUID) (database.Group, error) { m.ctrl.T.Helper() @@ -10804,6 +10834,21 @@ func (mr *MockStoreMockRecorder) UpsertDefaultProxy(ctx, arg any) *gomock.Call { return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpsertDefaultProxy", reflect.TypeOf((*MockStore)(nil).UpsertDefaultProxy), ctx, arg) } +// UpsertGroupAIBudget mocks base method. +func (m *MockStore) UpsertGroupAIBudget(ctx context.Context, arg database.UpsertGroupAIBudgetParams) (database.GroupAiBudget, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "UpsertGroupAIBudget", ctx, arg) + ret0, _ := ret[0].(database.GroupAiBudget) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// UpsertGroupAIBudget indicates an expected call of UpsertGroupAIBudget. +func (mr *MockStoreMockRecorder) UpsertGroupAIBudget(ctx, arg any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpsertGroupAIBudget", reflect.TypeOf((*MockStore)(nil).UpsertGroupAIBudget), ctx, arg) +} + // UpsertHealthSettings mocks base method. func (m *MockStore) UpsertHealthSettings(ctx context.Context, value string) error { m.ctrl.T.Helper() diff --git a/coderd/database/dump.sql b/coderd/database/dump.sql index 792f3712af..7648855838 100644 --- a/coderd/database/dump.sql +++ b/coderd/database/dump.sql @@ -1765,6 +1765,16 @@ CREATE TABLE gitsshkeys ( public_key text NOT NULL ); +CREATE TABLE group_ai_budgets ( + 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 group_ai_budgets_spend_limit_micros_check CHECK ((spend_limit_micros >= 0)) +); + +COMMENT ON TABLE group_ai_budgets IS 'Per-group AI spend limit applied to each member of the group. No row means no budget is enforced.'; + CREATE TABLE group_members ( user_id uuid NOT NULL, group_id uuid NOT NULL @@ -3594,6 +3604,9 @@ ALTER TABLE ONLY external_auth_links ALTER TABLE ONLY gitsshkeys ADD CONSTRAINT gitsshkeys_pkey PRIMARY KEY (user_id); +ALTER TABLE ONLY group_ai_budgets + ADD CONSTRAINT group_ai_budgets_pkey PRIMARY KEY (group_id); + ALTER TABLE ONLY group_members ADD CONSTRAINT group_members_user_id_group_id_key UNIQUE (user_id, group_id); @@ -4358,6 +4371,9 @@ ALTER TABLE ONLY external_auth_links ALTER TABLE ONLY gitsshkeys ADD CONSTRAINT gitsshkeys_user_id_fkey FOREIGN KEY (user_id) REFERENCES users(id); +ALTER TABLE ONLY group_ai_budgets + ADD CONSTRAINT group_ai_budgets_group_id_fkey FOREIGN KEY (group_id) REFERENCES groups(id) ON DELETE CASCADE; + ALTER TABLE ONLY group_members ADD CONSTRAINT group_members_group_id_fkey FOREIGN KEY (group_id) REFERENCES groups(id) ON DELETE CASCADE; diff --git a/coderd/database/foreign_key_constraint.go b/coderd/database/foreign_key_constraint.go index 78acd9b3f1..f989f19046 100644 --- a/coderd/database/foreign_key_constraint.go +++ b/coderd/database/foreign_key_constraint.go @@ -43,6 +43,7 @@ const ( ForeignKeyGitAuthLinksOauthAccessTokenKeyID ForeignKeyConstraint = "git_auth_links_oauth_access_token_key_id_fkey" // ALTER TABLE ONLY external_auth_links ADD CONSTRAINT git_auth_links_oauth_access_token_key_id_fkey FOREIGN KEY (oauth_access_token_key_id) REFERENCES dbcrypt_keys(active_key_digest); ForeignKeyGitAuthLinksOauthRefreshTokenKeyID ForeignKeyConstraint = "git_auth_links_oauth_refresh_token_key_id_fkey" // ALTER TABLE ONLY external_auth_links ADD CONSTRAINT git_auth_links_oauth_refresh_token_key_id_fkey FOREIGN KEY (oauth_refresh_token_key_id) REFERENCES dbcrypt_keys(active_key_digest); ForeignKeyGitSSHKeysUserID ForeignKeyConstraint = "gitsshkeys_user_id_fkey" // ALTER TABLE ONLY gitsshkeys ADD CONSTRAINT gitsshkeys_user_id_fkey FOREIGN KEY (user_id) REFERENCES users(id); + ForeignKeyGroupAiBudgetsGroupID ForeignKeyConstraint = "group_ai_budgets_group_id_fkey" // ALTER TABLE ONLY group_ai_budgets ADD CONSTRAINT group_ai_budgets_group_id_fkey FOREIGN KEY (group_id) REFERENCES groups(id) ON DELETE CASCADE; ForeignKeyGroupMembersGroupID ForeignKeyConstraint = "group_members_group_id_fkey" // ALTER TABLE ONLY group_members ADD CONSTRAINT group_members_group_id_fkey FOREIGN KEY (group_id) REFERENCES groups(id) ON DELETE CASCADE; ForeignKeyGroupMembersUserID ForeignKeyConstraint = "group_members_user_id_fkey" // ALTER TABLE ONLY group_members ADD CONSTRAINT group_members_user_id_fkey FOREIGN KEY (user_id) REFERENCES users(id) ON DELETE CASCADE; ForeignKeyGroupsOrganizationID ForeignKeyConstraint = "groups_organization_id_fkey" // ALTER TABLE ONLY groups ADD CONSTRAINT groups_organization_id_fkey FOREIGN KEY (organization_id) REFERENCES organizations(id) ON DELETE CASCADE; diff --git a/coderd/database/migrations/000497_group_ai_budgets.down.sql b/coderd/database/migrations/000497_group_ai_budgets.down.sql new file mode 100644 index 0000000000..afcdf2f7b3 --- /dev/null +++ b/coderd/database/migrations/000497_group_ai_budgets.down.sql @@ -0,0 +1 @@ +DROP TABLE IF EXISTS group_ai_budgets CASCADE; diff --git a/coderd/database/migrations/000497_group_ai_budgets.up.sql b/coderd/database/migrations/000497_group_ai_budgets.up.sql new file mode 100644 index 0000000000..76255f6cd1 --- /dev/null +++ b/coderd/database/migrations/000497_group_ai_budgets.up.sql @@ -0,0 +1,9 @@ +CREATE TABLE group_ai_budgets ( + group_id UUID PRIMARY KEY REFERENCES groups(id) ON DELETE CASCADE, + -- Spend limit applied to each member, 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() +); + +COMMENT ON TABLE group_ai_budgets IS 'Per-group AI spend limit applied to each member of the group. No row means no budget is enforced.'; diff --git a/coderd/database/migrations/testdata/fixtures/000497_group_ai_budgets.up.sql b/coderd/database/migrations/testdata/fixtures/000497_group_ai_budgets.up.sql new file mode 100644 index 0000000000..140e9f7305 --- /dev/null +++ b/coderd/database/migrations/testdata/fixtures/000497_group_ai_budgets.up.sql @@ -0,0 +1,5 @@ +INSERT INTO group_ai_budgets ( + group_id, + spend_limit_micros +) VALUES + ('bb640d07-ca8a-4869-b6bc-ae61ebb2fda1', 500000000); diff --git a/coderd/database/models.go b/coderd/database/models.go index 296c8232bc..81e7cb9713 100644 --- a/coderd/database/models.go +++ b/coderd/database/models.go @@ -4819,6 +4819,14 @@ type Group struct { ChatSpendLimitMicros sql.NullInt64 `db:"chat_spend_limit_micros" json:"chat_spend_limit_micros"` } +// Per-group AI spend limit applied to each member of the group. No row means no budget is enforced. +type GroupAiBudget struct { + 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"` +} + type GroupMember struct { UserID uuid.UUID `db:"user_id" json:"user_id"` UserEmail string `db:"user_email" json:"user_email"` diff --git a/coderd/database/querier.go b/coderd/database/querier.go index e5d1ade48b..728463c8d5 100644 --- a/coderd/database/querier.go +++ b/coderd/database/querier.go @@ -131,6 +131,7 @@ type sqlcQuerier interface { DeleteCustomRole(ctx context.Context, arg DeleteCustomRoleParams) error DeleteExpiredAPIKeys(ctx context.Context, arg DeleteExpiredAPIKeysParams) (int64, error) DeleteExternalAuthLink(ctx context.Context, arg DeleteExternalAuthLinkParams) error + DeleteGroupAIBudget(ctx context.Context, groupID uuid.UUID) (GroupAiBudget, error) DeleteGroupByID(ctx context.Context, id uuid.UUID) error DeleteGroupMemberFromGroup(ctx context.Context, arg DeleteGroupMemberFromGroupParams) error DeleteLicense(ctx context.Context, id int32) (int32, error) @@ -457,6 +458,7 @@ type sqlcQuerier interface { GetFilteredInboxNotificationsByUserID(ctx context.Context, arg GetFilteredInboxNotificationsByUserIDParams) ([]InboxNotification, error) GetForcedMCPServerConfigs(ctx context.Context) ([]MCPServerConfig, error) GetGitSSHKey(ctx context.Context, userID uuid.UUID) (GitSSHKey, error) + GetGroupAIBudget(ctx context.Context, groupID uuid.UUID) (GroupAiBudget, error) GetGroupByID(ctx context.Context, id uuid.UUID) (Group, error) GetGroupByOrgAndName(ctx context.Context, arg GetGroupByOrgAndNameParams) (Group, error) GetGroupMembers(ctx context.Context, includeSystem bool) ([]GroupMember, error) @@ -1334,6 +1336,7 @@ type sqlcQuerier interface { // So we need to store it's configuration here for display purposes. // The functional values are immutable and controlled implicitly. UpsertDefaultProxy(ctx context.Context, arg UpsertDefaultProxyParams) error + UpsertGroupAIBudget(ctx context.Context, arg UpsertGroupAIBudgetParams) (GroupAiBudget, error) UpsertHealthSettings(ctx context.Context, value string) error UpsertLastUpdateCheck(ctx context.Context, value string) error UpsertLogoURL(ctx context.Context, value string) error diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go index 513105f5cd..92394600a2 100644 --- a/coderd/database/queries.sql.go +++ b/coderd/database/queries.sql.go @@ -2280,6 +2280,22 @@ func (q *sqlQuerier) UpdateAIBridgeInterceptionEnded(ctx context.Context, arg Up return i, err } +const deleteGroupAIBudget = `-- name: DeleteGroupAIBudget :one +DELETE FROM group_ai_budgets WHERE group_id = $1 RETURNING group_id, spend_limit_micros, created_at, updated_at +` + +func (q *sqlQuerier) DeleteGroupAIBudget(ctx context.Context, groupID uuid.UUID) (GroupAiBudget, error) { + row := q.db.QueryRowContext(ctx, deleteGroupAIBudget, groupID) + var i GroupAiBudget + err := row.Scan( + &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 @@ -2307,6 +2323,24 @@ func (q *sqlQuerier) GetAIModelPriceByProviderModel(ctx context.Context, arg Get return i, err } +const getGroupAIBudget = `-- name: GetGroupAIBudget :one +SELECT group_id, spend_limit_micros, created_at, updated_at +FROM group_ai_budgets +WHERE group_id = $1 +` + +func (q *sqlQuerier) GetGroupAIBudget(ctx context.Context, groupID uuid.UUID) (GroupAiBudget, error) { + row := q.db.QueryRowContext(ctx, getGroupAIBudget, groupID) + var i GroupAiBudget + err := row.Scan( + &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 @@ -2335,6 +2369,32 @@ func (q *sqlQuerier) UpsertAIModelPrices(ctx context.Context, seed json.RawMessa return err } +const upsertGroupAIBudget = `-- name: UpsertGroupAIBudget :one +INSERT INTO group_ai_budgets (group_id, spend_limit_micros) +VALUES ($1, $2) +ON CONFLICT (group_id) DO UPDATE SET + spend_limit_micros = EXCLUDED.spend_limit_micros, + updated_at = NOW() +RETURNING group_id, spend_limit_micros, created_at, updated_at +` + +type UpsertGroupAIBudgetParams struct { + GroupID uuid.UUID `db:"group_id" json:"group_id"` + SpendLimitMicros int64 `db:"spend_limit_micros" json:"spend_limit_micros"` +} + +func (q *sqlQuerier) UpsertGroupAIBudget(ctx context.Context, arg UpsertGroupAIBudgetParams) (GroupAiBudget, error) { + row := q.db.QueryRowContext(ctx, upsertGroupAIBudget, arg.GroupID, arg.SpendLimitMicros) + var i GroupAiBudget + err := row.Scan( + &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 d2b66c4d3b..6740b2568c 100644 --- a/coderd/database/queries/aicostcontrol.sql +++ b/coderd/database/queries/aicostcontrol.sql @@ -24,3 +24,19 @@ ON CONFLICT (provider, model) DO UPDATE SET SELECT * FROM ai_model_prices WHERE provider = @provider AND model = @model; + +-- name: GetGroupAIBudget :one +SELECT * +FROM group_ai_budgets +WHERE group_id = @group_id; + +-- name: UpsertGroupAIBudget :one +INSERT INTO group_ai_budgets (group_id, spend_limit_micros) +VALUES (@group_id, @spend_limit_micros) +ON CONFLICT (group_id) DO UPDATE SET + spend_limit_micros = EXCLUDED.spend_limit_micros, + updated_at = NOW() +RETURNING *; + +-- name: DeleteGroupAIBudget :one +DELETE FROM group_ai_budgets WHERE group_id = @group_id RETURNING *; diff --git a/coderd/database/unique_constraint.go b/coderd/database/unique_constraint.go index 9a6a828691..0ca0283b1d 100644 --- a/coderd/database/unique_constraint.go +++ b/coderd/database/unique_constraint.go @@ -41,6 +41,7 @@ const ( UniqueFilesPkey UniqueConstraint = "files_pkey" // ALTER TABLE ONLY files ADD CONSTRAINT files_pkey PRIMARY KEY (id); UniqueGitAuthLinksProviderIDUserIDKey UniqueConstraint = "git_auth_links_provider_id_user_id_key" // ALTER TABLE ONLY external_auth_links ADD CONSTRAINT git_auth_links_provider_id_user_id_key UNIQUE (provider_id, user_id); UniqueGitSSHKeysPkey UniqueConstraint = "gitsshkeys_pkey" // ALTER TABLE ONLY gitsshkeys ADD CONSTRAINT gitsshkeys_pkey PRIMARY KEY (user_id); + UniqueGroupAiBudgetsPkey UniqueConstraint = "group_ai_budgets_pkey" // ALTER TABLE ONLY group_ai_budgets ADD CONSTRAINT group_ai_budgets_pkey PRIMARY KEY (group_id); UniqueGroupMembersUserIDGroupIDKey UniqueConstraint = "group_members_user_id_group_id_key" // ALTER TABLE ONLY group_members ADD CONSTRAINT group_members_user_id_group_id_key UNIQUE (user_id, group_id); UniqueGroupsNameOrganizationIDKey UniqueConstraint = "groups_name_organization_id_key" // ALTER TABLE ONLY groups ADD CONSTRAINT groups_name_organization_id_key UNIQUE (name, organization_id); UniqueGroupsPkey UniqueConstraint = "groups_pkey" // ALTER TABLE ONLY groups ADD CONSTRAINT groups_pkey PRIMARY KEY (id); diff --git a/codersdk/aibridge.go b/codersdk/aibridge.go index 2f00f8e80c..4f7f314330 100644 --- a/codersdk/aibridge.go +++ b/codersdk/aibridge.go @@ -9,6 +9,7 @@ import ( "time" "github.com/google/uuid" + "golang.org/x/xerrors" ) type AIBridgeInterception struct { @@ -350,3 +351,67 @@ func (c *Client) AIBridgeListClients(ctx context.Context) ([]string, error) { var clients []string return clients, json.NewDecoder(res.Body).Decode(&clients) } + +type GroupAIBudget struct { + 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 UpsertGroupAIBudgetRequest struct { + SpendLimitMicros int64 `json:"spend_limit_micros" validate:"gte=0"` +} + +// GroupAIBudget returns the AI spend budget configured for the given group. +func (c *Client) GroupAIBudget(ctx context.Context, group uuid.UUID) (GroupAIBudget, error) { + res, err := c.Request(ctx, http.MethodGet, + fmt.Sprintf("/api/v2/groups/%s/ai/budget", group.String()), + nil, + ) + if err != nil { + return GroupAIBudget{}, xerrors.Errorf("make request: %w", err) + } + defer res.Body.Close() + + if res.StatusCode != http.StatusOK { + return GroupAIBudget{}, ReadBodyAsError(res) + } + var resp GroupAIBudget + return resp, json.NewDecoder(res.Body).Decode(&resp) +} + +// UpsertGroupAIBudget creates or updates the AI spend budget for the given group. +func (c *Client) UpsertGroupAIBudget(ctx context.Context, group uuid.UUID, req UpsertGroupAIBudgetRequest) (GroupAIBudget, error) { + res, err := c.Request(ctx, http.MethodPut, + fmt.Sprintf("/api/v2/groups/%s/ai/budget", group.String()), + req, + ) + if err != nil { + return GroupAIBudget{}, xerrors.Errorf("make request: %w", err) + } + defer res.Body.Close() + + if res.StatusCode != http.StatusOK { + return GroupAIBudget{}, ReadBodyAsError(res) + } + var resp GroupAIBudget + return resp, json.NewDecoder(res.Body).Decode(&resp) +} + +// DeleteGroupAIBudget removes the AI spend budget for the given group. +func (c *Client) DeleteGroupAIBudget(ctx context.Context, group uuid.UUID) error { + res, err := c.Request(ctx, http.MethodDelete, + fmt.Sprintf("/api/v2/groups/%s/ai/budget", group.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 965241cd85..def5c219a2 100644 --- a/docs/reference/api/enterprise.md +++ b/docs/reference/api/enterprise.md @@ -679,6 +679,122 @@ curl -X PATCH http://coder-server:8080/api/v2/groups/{group} \ To perform this operation, you must be authenticated. [Learn more](authentication.md). +## Get group AI budget + +### Code samples + +```shell +# Example request using curl +curl -X GET http://coder-server:8080/api/v2/groups/{group}/ai/budget \ + -H 'Accept: application/json' \ + -H 'Coder-Session-Token: API_KEY' +``` + +`GET /api/v2/groups/{group}/ai/budget` + +### Parameters + +| Name | In | Type | Required | Description | +|---------|------|--------------|----------|-------------| +| `group` | path | string(uuid) | true | Group ID | + +### 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" +} +``` + +### Responses + +| Status | Meaning | Description | Schema | +|--------|---------------------------------------------------------|-------------|------------------------------------------------------------| +| 200 | [OK](https://tools.ietf.org/html/rfc7231#section-6.3.1) | OK | [codersdk.GroupAIBudget](schemas.md#codersdkgroupaibudget) | + +To perform this operation, you must be authenticated. [Learn more](authentication.md). + +## Upsert group AI budget + +### Code samples + +```shell +# Example request using curl +curl -X PUT http://coder-server:8080/api/v2/groups/{group}/ai/budget \ + -H 'Content-Type: application/json' \ + -H 'Accept: application/json' \ + -H 'Coder-Session-Token: API_KEY' +``` + +`PUT /api/v2/groups/{group}/ai/budget` + +> Body parameter + +```json +{ + "spend_limit_micros": 0 +} +``` + +### Parameters + +| Name | In | Type | Required | Description | +|---------|------|--------------------------------------------------------------------------------------|----------|--------------------------------| +| `group` | path | string(uuid) | true | Group ID | +| `body` | body | [codersdk.UpsertGroupAIBudgetRequest](schemas.md#codersdkupsertgroupaibudgetrequest) | true | Upsert group AI budget 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" +} +``` + +### Responses + +| Status | Meaning | Description | Schema | +|--------|---------------------------------------------------------|-------------|------------------------------------------------------------| +| 200 | [OK](https://tools.ietf.org/html/rfc7231#section-6.3.1) | OK | [codersdk.GroupAIBudget](schemas.md#codersdkgroupaibudget) | + +To perform this operation, you must be authenticated. [Learn more](authentication.md). + +## Delete group AI budget + +### Code samples + +```shell +# Example request using curl +curl -X DELETE http://coder-server:8080/api/v2/groups/{group}/ai/budget \ + -H 'Coder-Session-Token: API_KEY' +``` + +`DELETE /api/v2/groups/{group}/ai/budget` + +### Parameters + +| Name | In | Type | Required | Description | +|---------|------|--------------|----------|-------------| +| `group` | path | string(uuid) | true | Group ID | + +### 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 group members by group ID ### Code samples diff --git a/docs/reference/api/schemas.md b/docs/reference/api/schemas.md index db26714321..4b53900dc9 100644 --- a/docs/reference/api/schemas.md +++ b/docs/reference/api/schemas.md @@ -7313,6 +7313,26 @@ Only certain features set these fields: - FeatureManagedAgentLimit| | `source` | [codersdk.GroupSource](#codersdkgroupsource) | false | | | | `total_member_count` | integer | false | | How many members are in this group. Shows the total count, even if the user is not authorized to read group member details. May be greater than `len(Group.Members)`. | +## codersdk.GroupAIBudget + +```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" +} +``` + +### Properties + +| Name | Type | Required | Restrictions | Description | +|----------------------|---------|----------|--------------|-------------| +| `created_at` | string | false | | | +| `group_id` | string | false | | | +| `spend_limit_micros` | integer | false | | | +| `updated_at` | string | false | | | + ## codersdk.GroupMembersResponse ```json @@ -13188,6 +13208,20 @@ If the schedule is empty, the user will be updated to use the default schedule.| |--------|--------|----------|--------------|-------------| | `hash` | string | false | | | +## codersdk.UpsertGroupAIBudgetRequest + +```json +{ + "spend_limit_micros": 0 +} +``` + +### Properties + +| Name | Type | Required | Restrictions | Description | +|----------------------|---------|----------|--------------|-------------| +| `spend_limit_micros` | integer | false | | | + ## codersdk.UpsertWorkspaceAgentPortShareRequest ```json diff --git a/enterprise/coderd/aibridge.go b/enterprise/coderd/aibridge.go index b1a8d8838a..3f27ff5169 100644 --- a/enterprise/coderd/aibridge.go +++ b/enterprise/coderd/aibridge.go @@ -697,3 +697,90 @@ func populatedAndConvertAIBridgeInterceptions(ctx context.Context, db database.S return items, nil } + +// @Summary Get group AI budget +// @ID get-group-ai-budget +// @Security CoderSessionToken +// @Produce json +// @Tags Enterprise +// @Param group path string true "Group ID" format(uuid) +// @Success 200 {object} codersdk.GroupAIBudget +// @Router /api/v2/groups/{group}/ai/budget [get] +func (api *API) groupAIBudget(rw http.ResponseWriter, r *http.Request) { + ctx := r.Context() + group := httpmw.GroupParam(r) + + budget, err := api.Database.GetGroupAIBudget(ctx, group.ID) + if httpapi.Is404Error(err) { + httpapi.ResourceNotFound(rw) + return + } + if err != nil { + api.Logger.Error(ctx, "get group AI budget", slog.Error(err)) + httpapi.InternalServerError(rw, err) + return + } + + httpapi.Write(ctx, rw, http.StatusOK, db2sdk.GroupAIBudget(budget)) +} + +// @Summary Upsert group AI budget +// @ID upsert-group-ai-budget +// @Security CoderSessionToken +// @Accept json +// @Produce json +// @Tags Enterprise +// @Param group path string true "Group ID" format(uuid) +// @Param request body codersdk.UpsertGroupAIBudgetRequest true "Upsert group AI budget request" +// @Success 200 {object} codersdk.GroupAIBudget +// @Router /api/v2/groups/{group}/ai/budget [put] +func (api *API) upsertGroupAIBudget(rw http.ResponseWriter, r *http.Request) { + ctx := r.Context() + group := httpmw.GroupParam(r) + + var req codersdk.UpsertGroupAIBudgetRequest + if !httpapi.Read(ctx, rw, r, &req) { + return + } + + budget, err := api.Database.UpsertGroupAIBudget(ctx, database.UpsertGroupAIBudgetParams{ + GroupID: group.ID, + SpendLimitMicros: req.SpendLimitMicros, + }) + if httpapi.Is404Error(err) { + httpapi.ResourceNotFound(rw) + return + } + if err != nil { + api.Logger.Error(ctx, "upsert group AI budget", slog.Error(err)) + httpapi.InternalServerError(rw, err) + return + } + + httpapi.Write(ctx, rw, http.StatusOK, db2sdk.GroupAIBudget(budget)) +} + +// @Summary Delete group AI budget +// @ID delete-group-ai-budget +// @Security CoderSessionToken +// @Tags Enterprise +// @Param group path string true "Group ID" format(uuid) +// @Success 204 +// @Router /api/v2/groups/{group}/ai/budget [delete] +func (api *API) deleteGroupAIBudget(rw http.ResponseWriter, r *http.Request) { + ctx := r.Context() + group := httpmw.GroupParam(r) + + _, err := api.Database.DeleteGroupAIBudget(ctx, group.ID) + if httpapi.Is404Error(err) { + httpapi.ResourceNotFound(rw) + return + } + if err != nil { + api.Logger.Error(ctx, "delete group AI budget", 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 fa78c61956..e4015b75b7 100644 --- a/enterprise/coderd/aibridge_test.go +++ b/enterprise/coderd/aibridge_test.go @@ -2603,3 +2603,201 @@ func TestAIBridgeAllowBYOK(t *testing.T) { }) } } + +func TestGroupAIBudget(t *testing.T) { + t.Parallel() + + t.Run("Upsert", func(t *testing.T) { + t.Parallel() + + adminClient, group := setupGroupAIBudgetTest(t) + ctx := testutil.Context(t, testutil.WaitLong) + + // First upsert creates the budget. + newBudget, err := adminClient.UpsertGroupAIBudget(ctx, group.ID, codersdk.UpsertGroupAIBudgetRequest{ + SpendLimitMicros: 500_000_000, + }) + require.NoError(t, err) + require.Equal(t, group.ID, newBudget.GroupID) + require.EqualValues(t, 500_000_000, newBudget.SpendLimitMicros) + + // Second upsert updates the existing budget. + updatedBudget, err := adminClient.UpsertGroupAIBudget(ctx, group.ID, codersdk.UpsertGroupAIBudgetRequest{ + SpendLimitMicros: 1_000_000_000, + }) + require.NoError(t, err) + require.EqualValues(t, 1_000_000_000, updatedBudget.SpendLimitMicros) + + // GET returns the latest value. + currentBudget, err := adminClient.GroupAIBudget(ctx, group.ID) + require.NoError(t, err) + require.EqualValues(t, 1_000_000_000, currentBudget.SpendLimitMicros) + }) + + t.Run("GetWhenAbsent_404", func(t *testing.T) { + t.Parallel() + + adminClient, group := setupGroupAIBudgetTest(t) + ctx := testutil.Context(t, testutil.WaitLong) + + _, err := adminClient.GroupAIBudget(ctx, group.ID) + var sdkErr *codersdk.Error + require.ErrorAs(t, err, &sdkErr) + require.Equal(t, http.StatusNotFound, sdkErr.StatusCode()) + }) + + t.Run("DeleteWhenAbsent_404", func(t *testing.T) { + t.Parallel() + + adminClient, group := setupGroupAIBudgetTest(t) + ctx := testutil.Context(t, testutil.WaitLong) + + err := adminClient.DeleteGroupAIBudget(ctx, group.ID) + var sdkErr *codersdk.Error + require.ErrorAs(t, err, &sdkErr) + require.Equal(t, http.StatusNotFound, sdkErr.StatusCode()) + }) + + t.Run("DeleteWhenPresent", func(t *testing.T) { + t.Parallel() + + adminClient, group := setupGroupAIBudgetTest(t) + ctx := testutil.Context(t, testutil.WaitLong) + + _, err := adminClient.UpsertGroupAIBudget(ctx, group.ID, codersdk.UpsertGroupAIBudgetRequest{ + SpendLimitMicros: 500_000_000, + }) + require.NoError(t, err) + + require.NoError(t, adminClient.DeleteGroupAIBudget(ctx, group.ID)) + + _, err = adminClient.GroupAIBudget(ctx, group.ID) + var sdkErr *codersdk.Error + require.ErrorAs(t, err, &sdkErr) + require.Equal(t, http.StatusNotFound, sdkErr.StatusCode()) + }) + + t.Run("RejectsNegativeSpendLimit", func(t *testing.T) { + t.Parallel() + + adminClient, group := setupGroupAIBudgetTest(t) + ctx := testutil.Context(t, testutil.WaitLong) + + _, err := adminClient.UpsertGroupAIBudget(ctx, group.ID, codersdk.UpsertGroupAIBudgetRequest{ + SpendLimitMicros: -1, + }) + var sdkErr *codersdk.Error + require.ErrorAs(t, err, &sdkErr) + require.Equal(t, http.StatusBadRequest, sdkErr.StatusCode()) + }) + + t.Run("AcceptsZeroSpendLimitToBlock", func(t *testing.T) { + t.Parallel() + + adminClient, group := setupGroupAIBudgetTest(t) + ctx := testutil.Context(t, testutil.WaitLong) + + // 0 is a valid value: it blocks all spend for the group's members. + budget, err := adminClient.UpsertGroupAIBudget(ctx, group.ID, codersdk.UpsertGroupAIBudgetRequest{ + SpendLimitMicros: 0, + }) + require.NoError(t, err) + require.EqualValues(t, 0, budget.SpendLimitMicros) + }) + + t.Run("UnknownGroup_404", func(t *testing.T) { + t.Parallel() + + adminClient, _ := setupGroupAIBudgetTest(t) + ctx := testutil.Context(t, testutil.WaitLong) + + _, err := adminClient.GroupAIBudget(ctx, uuid.New()) + var sdkErr *codersdk.Error + require.ErrorAs(t, err, &sdkErr) + require.Equal(t, http.StatusNotFound, sdkErr.StatusCode()) + }) + + t.Run("GroupMemberCanReadButNotWrite", func(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()) + memberClient, member := coderdtest.CreateAnotherUser(t, ownerClient, owner.OrganizationID) + + ctx := testutil.Context(t, testutil.WaitLong) + group, err := adminClient.CreateGroup(ctx, owner.OrganizationID, codersdk.CreateGroupRequest{ + Name: "budget-group", + }) + require.NoError(t, err) + + // Add the member to the group so the Group.RBACObject ACL grants them read. + _, err = adminClient.PatchGroup(ctx, group.ID, codersdk.PatchGroupRequest{ + AddUsers: []string{member.ID.String()}, + }) + require.NoError(t, err) + + // Admin sets the budget so there is a row to read. + _, err = adminClient.UpsertGroupAIBudget(ctx, group.ID, codersdk.UpsertGroupAIBudgetRequest{ + SpendLimitMicros: 500_000_000, + }) + require.NoError(t, err) + + // Group members can read the budget. + got, err := memberClient.GroupAIBudget(ctx, group.ID) + require.NoError(t, err) + require.EqualValues(t, 500_000_000, got.SpendLimitMicros) + + // Group members cannot write the budget. + _, err = memberClient.UpsertGroupAIBudget(ctx, group.ID, codersdk.UpsertGroupAIBudgetRequest{ + SpendLimitMicros: 1_000_000_000, + }) + var sdkErr *codersdk.Error + require.ErrorAs(t, err, &sdkErr) + require.Equal(t, http.StatusNotFound, sdkErr.StatusCode()) + + // Group members cannot delete the budget. + err = memberClient.DeleteGroupAIBudget(ctx, group.ID) + require.ErrorAs(t, err, &sdkErr) + require.Equal(t, http.StatusNotFound, sdkErr.StatusCode()) + + // The failed upsert and delete left the budget untouched. + got, err = memberClient.GroupAIBudget(ctx, group.ID) + require.NoError(t, err) + require.EqualValues(t, 500_000_000, got.SpendLimitMicros) + }) +} + +// 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() + + 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()) + + ctx := testutil.Context(t, testutil.WaitLong) + g, err := adminClient.CreateGroup(ctx, owner.OrganizationID, codersdk.CreateGroupRequest{ + Name: "budget-test-group", + }) + require.NoError(t, err) + return adminClient, g +} diff --git a/enterprise/coderd/coderd.go b/enterprise/coderd/coderd.go index a2c81ae5ba..3d85473801 100644 --- a/enterprise/coderd/coderd.go +++ b/enterprise/coderd/coderd.go @@ -549,6 +549,13 @@ func New(ctx context.Context, options *Options) (_ *API, err error) { r.Patch("/", api.patchGroup) r.Delete("/", api.deleteGroup) r.Get("/members", api.groupMembers) + r.Route("/ai/budget", func(r chi.Router) { + // AI cost controls are a paid feature (AI Governance add-on). + r.Use(api.RequireFeatureMW(codersdk.FeatureAIBridge)) + r.Get("/", api.groupAIBudget) + r.Put("/", api.upsertGroupAIBudget) + r.Delete("/", api.deleteGroupAIBudget) + }) }) }) r.Route("/workspace-quota", func(r chi.Router) { diff --git a/site/src/api/typesGenerated.ts b/site/src/api/typesGenerated.ts index 0ee4f550ea..9b033b058e 100644 --- a/site/src/api/typesGenerated.ts +++ b/site/src/api/typesGenerated.ts @@ -4369,6 +4369,14 @@ export interface Group { readonly organization_display_name: string; } +// From codersdk/aibridge.go +export interface GroupAIBudget { + readonly group_id: string; + readonly spend_limit_micros: number; + readonly created_at: string; + readonly updated_at: string; +} + // From codersdk/groups.go export interface GroupArguments { /** @@ -8693,6 +8701,11 @@ export interface UpsertChatUsageLimitOverrideRequest { readonly spend_limit_micros: number; // Must be greater than 0. } +// From codersdk/aibridge.go +export interface UpsertGroupAIBudgetRequest { + readonly spend_limit_micros: number; +} + // From codersdk/workspaceagentportshare.go export interface UpsertWorkspaceAgentPortShareRequest { readonly agent_name: string;