opensource/coder
3
0
Fork 0
mirror of https://github.com/coder/coder.git synced 2026-09-24 15:04:27 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat: add per-group AI budget table and endpoints (#25203)

Browse Source 
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.
This commit is contained in:
Yevhenii Shcherbina
2026-05-14 15:54:37 -04:00
committed by GitHub
parent d97f5ae2a6
commit 238968cfa0
25 changed files with 1040 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
coderd/apidoc/docs.go
Generated
+136
View File
@@ -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": {
coderd/apidoc/swagger.json
Generated
+124
View File
@@ -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": {
coderd/database/check_constraint.go
+1
View File
@@ -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
coderd/database/db2sdk/db2sdk.go
+9
View File
@@ -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 {
coderd/database/dbauthz/dbauthz.go
+36
View File
@@ -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
coderd/database/dbauthz/dbauthz_test.go
+25
View File
@@ -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()
coderd/database/dbmetrics/querymetrics.go
+24
View File
@@ -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)
coderd/database/dbmock/dbmock.go
+45
View File
@@ -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()
coderd/database/dump.sql
Generated
+16
View File
@@ -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;
coderd/database/foreign_key_constraint.go
+1
View File
@@ -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;
coderd/database/migrations/000497_group_ai_budgets.down.sql
+1
View File
@@ -0,0 +1 @@
DROP TABLE IF EXISTS group_ai_budgets CASCADE;
coderd/database/migrations/000497_group_ai_budgets.up.sql
+9
View File
@@ -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.';
coderd/database/migrations/testdata/fixtures/000497_group_ai_budgets.up.sql
Vendored
+5
View File
@@ -0,0 +1,5 @@
INSERT INTO group_ai_budgets (
group_id,
spend_limit_micros
) VALUES
('bb640d07-ca8a-4869-b6bc-ae61ebb2fda1', 500000000);
coderd/database/models.go
+8
View File
@@ -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"`
coderd/database/querier.go
+3
View File
@@ -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
coderd/database/queries.sql.go
+60
View File
@@ -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(*)
coderd/database/queries/aicostcontrol.sql
+16
View File
@@ -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 *;
coderd/database/unique_constraint.go
+1
View File
@@ -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);
codersdk/aibridge.go
+65
View File
@@ -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
}
docs/reference/api/enterprise.md
Generated
+116
View File
@@ -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
docs/reference/api/schemas.md
Generated
+34
View File
@@ -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
enterprise/coderd/aibridge.go
+87
View File
@@ -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)
}
enterprise/coderd/aibridge_test.go
+198
View File
@@ -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
}
enterprise/coderd/coderd.go
+7
View File
@@ -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) {
site/src/api/typesGenerated.ts
Generated
+13
View File
@@ -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;