From deda790c7e8649f1ee2c3a950bf01ce42f040110 Mon Sep 17 00:00:00 2001 From: Susana Cardoso Ferreira Date: Thu, 9 Jul 2026 12:09:23 +0000 Subject: [PATCH 1/5] feat: add GET /organizations/{org}/groups/ai/spend --- coderd/apidoc/docs.go | 80 ++++++ coderd/apidoc/swagger.json | 76 ++++++ coderd/database/dbauthz/dbauthz.go | 4 + coderd/database/dbauthz/dbauthz_test.go | 16 ++ coderd/database/dbmetrics/querymetrics.go | 8 + coderd/database/dbmock/dbmock.go | 15 ++ coderd/database/modelmethods.go | 4 + coderd/database/querier.go | 6 + coderd/database/querier_test.go | 306 ++++++++++++++++++++++ coderd/database/queries.sql.go | 62 +++++ coderd/database/queries/aicostcontrol.sql | 20 ++ codersdk/aibridge.go | 63 ++++- docs/reference/api/enterprise.md | 46 ++++ docs/reference/api/schemas.md | 42 +++ enterprise/coderd/aibridge.go | 88 ++++++- enterprise/coderd/aibridge_test.go | 284 ++++++++++++++++++++ enterprise/coderd/coderd.go | 9 + site/src/api/typesGenerated.ts | 60 ++++- 18 files changed, 1170 insertions(+), 19 deletions(-) diff --git a/coderd/apidoc/docs.go b/coderd/apidoc/docs.go index 54add36baea86..9ac635035b325 100644 --- a/coderd/apidoc/docs.go +++ b/coderd/apidoc/docs.go @@ -4752,6 +4752,48 @@ const docTemplate = `{ ] } }, + "/api/v2/organizations/{organization}/groups/ai/spend": { + "get": { + "produces": [ + "application/json" + ], + "tags": [ + "Enterprise" + ], + "summary": "Get organization groups AI spend", + "operationId": "get-organization-groups-ai-spend", + "parameters": [ + { + "type": "string", + "format": "uuid", + "description": "Organization ID", + "name": "organization", + "in": "path", + "required": true + }, + { + "type": "string", + "description": "Comma-separated list of group IDs", + "name": "group_ids", + "in": "query", + "required": true + } + ], + "responses": { + "200": { + "description": "OK", + "schema": { + "$ref": "#/definitions/codersdk.OrganizationGroupsAISpend" + } + } + }, + "security": [ + { + "CoderSessionToken": [] + } + ] + } + }, "/api/v2/organizations/{organization}/groups/{groupName}": { "get": { "produces": [ @@ -21624,6 +21666,44 @@ const docTemplate = `{ } } }, + "codersdk.OrganizationGroupAISpend": { + "type": "object", + "properties": { + "current_spend_micros": { + "description": "CurrentSpendMicros is the group's spend over the current budget\nperiod.", + "type": "integer" + }, + "group_id": { + "type": "string", + "format": "uuid" + }, + "spend_limit_micros": { + "description": "SpendLimitMicros is the group's configured AI spend limit. Null when\nthe group has no configured budget.", + "type": "integer" + } + } + }, + "codersdk.OrganizationGroupsAISpend": { + "type": "object", + "properties": { + "groups": { + "type": "array", + "items": { + "$ref": "#/definitions/codersdk.OrganizationGroupAISpend" + } + }, + "period_end": { + "description": "PeriodEnd is the exclusive upper bound of the current budget\nperiod.", + "type": "string", + "format": "date-time" + }, + "period_start": { + "description": "PeriodStart is the inclusive lower bound of the current budget\nperiod.", + "type": "string", + "format": "date-time" + } + } + }, "codersdk.OrganizationMember": { "type": "object", "properties": { diff --git a/coderd/apidoc/swagger.json b/coderd/apidoc/swagger.json index 5687248ce4f85..7b5f4592be325 100644 --- a/coderd/apidoc/swagger.json +++ b/coderd/apidoc/swagger.json @@ -4193,6 +4193,44 @@ ] } }, + "/api/v2/organizations/{organization}/groups/ai/spend": { + "get": { + "produces": ["application/json"], + "tags": ["Enterprise"], + "summary": "Get organization groups AI spend", + "operationId": "get-organization-groups-ai-spend", + "parameters": [ + { + "type": "string", + "format": "uuid", + "description": "Organization ID", + "name": "organization", + "in": "path", + "required": true + }, + { + "type": "string", + "description": "Comma-separated list of group IDs", + "name": "group_ids", + "in": "query", + "required": true + } + ], + "responses": { + "200": { + "description": "OK", + "schema": { + "$ref": "#/definitions/codersdk.OrganizationGroupsAISpend" + } + } + }, + "security": [ + { + "CoderSessionToken": [] + } + ] + } + }, "/api/v2/organizations/{organization}/groups/{groupName}": { "get": { "produces": ["application/json"], @@ -19731,6 +19769,44 @@ } } }, + "codersdk.OrganizationGroupAISpend": { + "type": "object", + "properties": { + "current_spend_micros": { + "description": "CurrentSpendMicros is the group's spend over the current budget\nperiod.", + "type": "integer" + }, + "group_id": { + "type": "string", + "format": "uuid" + }, + "spend_limit_micros": { + "description": "SpendLimitMicros is the group's configured AI spend limit. Null when\nthe group has no configured budget.", + "type": "integer" + } + } + }, + "codersdk.OrganizationGroupsAISpend": { + "type": "object", + "properties": { + "groups": { + "type": "array", + "items": { + "$ref": "#/definitions/codersdk.OrganizationGroupAISpend" + } + }, + "period_end": { + "description": "PeriodEnd is the exclusive upper bound of the current budget\nperiod.", + "type": "string", + "format": "date-time" + }, + "period_start": { + "description": "PeriodStart is the inclusive lower bound of the current budget\nperiod.", + "type": "string", + "format": "date-time" + } + } + }, "codersdk.OrganizationMember": { "type": "object", "properties": { diff --git a/coderd/database/dbauthz/dbauthz.go b/coderd/database/dbauthz/dbauthz.go index 6221d042aa14b..48c26244c17e7 100644 --- a/coderd/database/dbauthz/dbauthz.go +++ b/coderd/database/dbauthz/dbauthz.go @@ -4217,6 +4217,10 @@ func (q *querier) GetOrganizationByName(ctx context.Context, name database.GetOr return fetch(q.log, q.auth, q.db.GetOrganizationByName)(ctx, name) } +func (q *querier) GetOrganizationGroupsAISpend(ctx context.Context, arg database.GetOrganizationGroupsAISpendParams) ([]database.GetOrganizationGroupsAISpendRow, error) { + return fetchWithPostFilter(q.auth, policy.ActionRead, q.db.GetOrganizationGroupsAISpend)(ctx, arg) +} + func (q *querier) GetOrganizationIDsByMemberIDs(ctx context.Context, ids []uuid.UUID) ([]database.GetOrganizationIDsByMemberIDsRow, error) { // TODO: This should be rewritten to return a list of database.OrganizationMember for consistent RBAC objects. // Currently this row returns a list of org ids per user, which is challenging to check against the RBAC system. diff --git a/coderd/database/dbauthz/dbauthz_test.go b/coderd/database/dbauthz/dbauthz_test.go index 28e2b91ae3159..dd62a2fb16bb9 100644 --- a/coderd/database/dbauthz/dbauthz_test.go +++ b/coderd/database/dbauthz/dbauthz_test.go @@ -6938,6 +6938,22 @@ func (s *MethodTestSuite) TestAIBridge() { check.Args(database.GetAIModelPriceByProviderModelParams{}).Asserts(rbac.ResourceAiModelPrice, policy.ActionRead) })) + s.Run("GetOrganizationGroupsAISpend", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { + org := testutil.Fake(s.T(), faker, database.Organization{}) + row1 := testutil.Fake(s.T(), faker, database.GetOrganizationGroupsAISpendRow{OrganizationID: org.ID}) + row2 := testutil.Fake(s.T(), faker, database.GetOrganizationGroupsAISpendRow{OrganizationID: org.ID}) + arg := database.GetOrganizationGroupsAISpendParams{ + OrganizationID: org.ID, + GroupIds: []uuid.UUID{row1.GroupID, row2.GroupID}, + PeriodStart: time.Now().UTC().Truncate(24 * time.Hour), + } + dbm.EXPECT().GetOrganizationGroupsAISpend(gomock.Any(), arg). + Return([]database.GetOrganizationGroupsAISpendRow{row1, row2}, nil).AnyTimes() + check.Args(arg). + Asserts(row1, policy.ActionRead, row2, policy.ActionRead). + Returns([]database.GetOrganizationGroupsAISpendRow{row1, row2}) + })) + 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}) diff --git a/coderd/database/dbmetrics/querymetrics.go b/coderd/database/dbmetrics/querymetrics.go index 1e8c654e73bde..6225be9058d87 100644 --- a/coderd/database/dbmetrics/querymetrics.go +++ b/coderd/database/dbmetrics/querymetrics.go @@ -2537,6 +2537,14 @@ func (m queryMetricsStore) GetOrganizationByName(ctx context.Context, arg databa return r0, r1 } +func (m queryMetricsStore) GetOrganizationGroupsAISpend(ctx context.Context, arg database.GetOrganizationGroupsAISpendParams) ([]database.GetOrganizationGroupsAISpendRow, error) { + start := time.Now() + r0, r1 := m.s.GetOrganizationGroupsAISpend(ctx, arg) + m.queryLatencies.WithLabelValues("GetOrganizationGroupsAISpend").Observe(time.Since(start).Seconds()) + m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "GetOrganizationGroupsAISpend").Inc() + return r0, r1 +} + func (m queryMetricsStore) GetOrganizationIDsByMemberIDs(ctx context.Context, ids []uuid.UUID) ([]database.GetOrganizationIDsByMemberIDsRow, error) { start := time.Now() r0, r1 := m.s.GetOrganizationIDsByMemberIDs(ctx, ids) diff --git a/coderd/database/dbmock/dbmock.go b/coderd/database/dbmock/dbmock.go index a4e2eda82ddd4..1c403c54f40f7 100644 --- a/coderd/database/dbmock/dbmock.go +++ b/coderd/database/dbmock/dbmock.go @@ -4708,6 +4708,21 @@ func (mr *MockStoreMockRecorder) GetOrganizationByName(ctx, arg any) *gomock.Cal return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetOrganizationByName", reflect.TypeOf((*MockStore)(nil).GetOrganizationByName), ctx, arg) } +// GetOrganizationGroupsAISpend mocks base method. +func (m *MockStore) GetOrganizationGroupsAISpend(ctx context.Context, arg database.GetOrganizationGroupsAISpendParams) ([]database.GetOrganizationGroupsAISpendRow, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "GetOrganizationGroupsAISpend", ctx, arg) + ret0, _ := ret[0].([]database.GetOrganizationGroupsAISpendRow) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// GetOrganizationGroupsAISpend indicates an expected call of GetOrganizationGroupsAISpend. +func (mr *MockStoreMockRecorder) GetOrganizationGroupsAISpend(ctx, arg any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetOrganizationGroupsAISpend", reflect.TypeOf((*MockStore)(nil).GetOrganizationGroupsAISpend), ctx, arg) +} + // GetOrganizationIDsByMemberIDs mocks base method. func (m *MockStore) GetOrganizationIDsByMemberIDs(ctx context.Context, ids []uuid.UUID) ([]database.GetOrganizationIDsByMemberIDsRow, error) { m.ctrl.T.Helper() diff --git a/coderd/database/modelmethods.go b/coderd/database/modelmethods.go index a86c7ad1e748c..1fba7b4d99ee4 100644 --- a/coderd/database/modelmethods.go +++ b/coderd/database/modelmethods.go @@ -458,6 +458,10 @@ func (g GetGroupsRow) RBACObject() rbac.Object { return g.Group.RBACObject() } +func (g GetOrganizationGroupsAISpendRow) RBACObject() rbac.Object { + return Group{ID: g.GroupID, OrganizationID: g.OrganizationID}.RBACObject() +} + func (gm GroupMember) RBACObject() rbac.Object { return rbac.ResourceGroupMember.WithID(gm.UserID).InOrg(gm.OrganizationID).WithOwner(gm.UserID.String()) } diff --git a/coderd/database/querier.go b/coderd/database/querier.go index 097e7d2ad4de0..6a5f0ce47d73c 100644 --- a/coderd/database/querier.go +++ b/coderd/database/querier.go @@ -651,6 +651,12 @@ type sqlcQuerier interface { GetOAuth2ProviderAppsByUserID(ctx context.Context, userID uuid.UUID) ([]GetOAuth2ProviderAppsByUserIDRow, error) GetOrganizationByID(ctx context.Context, id uuid.UUID) (Organization, error) GetOrganizationByName(ctx context.Context, arg GetOrganizationByNameParams) (Organization, error) + // Returns AI spend limits and aggregate spend for groups in @group_ids that + // belong to @organization_id, on or after period_start until NOW. The spend + // limit is null when the group has no configured budget. + // The period_start parameter is normalized to its UTC calendar day. + // Only return groups from @group_ids that belong to @organization_id. + GetOrganizationGroupsAISpend(ctx context.Context, arg GetOrganizationGroupsAISpendParams) ([]GetOrganizationGroupsAISpendRow, error) GetOrganizationIDsByMemberIDs(ctx context.Context, ids []uuid.UUID) ([]GetOrganizationIDsByMemberIDsRow, error) GetOrganizationResourceCountByID(ctx context.Context, organizationID uuid.UUID) (GetOrganizationResourceCountByIDRow, error) GetOrganizations(ctx context.Context, arg GetOrganizationsParams) ([]Organization, error) diff --git a/coderd/database/querier_test.go b/coderd/database/querier_test.go index c590dc0640723..38c99dd7e0e82 100644 --- a/coderd/database/querier_test.go +++ b/coderd/database/querier_test.go @@ -12800,6 +12800,312 @@ func TestGetUserAISpendSince(t *testing.T) { }) } +func TestGetOrganizationGroupsAISpend(t *testing.T) { + t.Parallel() + + // Use fixed dates to keep the test deterministic. + monthStart := time.Date(2024, 6, 1, 0, 0, 0, 0, time.UTC) + now := monthStart.AddDate(0, 0, 14) // 2024-06-15 + prevMonthLastDay := monthStart.AddDate(0, 0, -1) // 2024-05-31 + + type seedRow struct { + day time.Time + spend int64 + } + + tests := []struct { + name string + setBudget bool + spendLimit int64 + rows []seedRow + wantCurrentSpend int64 + }{ + { + name: "NoBudgetNoSpend", + wantCurrentSpend: 0, + }, + { + name: "ZeroLimitBudget", + setBudget: true, + spendLimit: 0, + wantCurrentSpend: 0, + }, + { + name: "BudgetZeroSpend", + setBudget: true, + spendLimit: 1_000_000, + wantCurrentSpend: 0, + }, + { + name: "BudgetWithSpend", + setBudget: true, + spendLimit: 1_000_000, + rows: []seedRow{{now, 250}}, + wantCurrentSpend: 250, + }, + { + name: "NoBudgetWithSpend", + rows: []seedRow{{now, 100}}, + wantCurrentSpend: 100, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + db, _ := dbtestutil.NewDB(t) + ctx := testutil.Context(t, testutil.WaitShort) + + // Given: an org with a single group, optionally with a budget and seeded spend. + user := dbgen.User(t, db, database.User{}) + org := dbgen.Organization(t, db, database.Organization{}) + group := dbgen.Group(t, db, database.Group{OrganizationID: org.ID}) + if tt.setBudget { + _, err := db.UpsertGroupAIBudget(ctx, database.UpsertGroupAIBudgetParams{ + GroupID: group.ID, + SpendLimitMicros: tt.spendLimit, + }) + require.NoError(t, err) + } + for _, r := range tt.rows { + _, err := db.IncrementUserAIDailySpend(ctx, database.IncrementUserAIDailySpendParams{ + UserID: user.ID, + EffectiveGroupID: group.ID, + Day: r.day, + CostMicros: r.spend, + }) + require.NoError(t, err) + } + + // When: querying spend for the group since monthStart. + got, err := db.GetOrganizationGroupsAISpend(ctx, database.GetOrganizationGroupsAISpendParams{ + OrganizationID: org.ID, + GroupIds: []uuid.UUID{group.ID}, + PeriodStart: monthStart, + }) + require.NoError(t, err) + + // Then: one row is returned with the group's limit and spend. + require.Len(t, got, 1) + require.Equal(t, group.ID, got[0].GroupID) + require.Equal(t, org.ID, got[0].OrganizationID) + if tt.setBudget { + require.True(t, got[0].SpendLimitMicros.Valid, "expected configured budget") + require.Equal(t, tt.spendLimit, got[0].SpendLimitMicros.Int64, "spend_limit_micros") + } else { + require.False(t, got[0].SpendLimitMicros.Valid, "expected no configured budget") + } + require.Equal(t, tt.wantCurrentSpend, got[0].CurrentSpendMicros) + }) + } + + t.Run("MultipleGroupsInSameOrg", func(t *testing.T) { + t.Parallel() + db, _ := dbtestutil.NewDB(t) + ctx := testutil.Context(t, testutil.WaitShort) + + // Given: two groups in the same org with different budget and spend. + user := dbgen.User(t, db, database.User{}) + org := dbgen.Organization(t, db, database.Organization{}) + groupA := dbgen.Group(t, db, database.Group{OrganizationID: org.ID}) + groupB := dbgen.Group(t, db, database.Group{OrganizationID: org.ID}) + _, err := db.UpsertGroupAIBudget(ctx, database.UpsertGroupAIBudgetParams{ + GroupID: groupA.ID, + SpendLimitMicros: 1_000_000, + }) + require.NoError(t, err) + _, err = db.IncrementUserAIDailySpend(ctx, database.IncrementUserAIDailySpendParams{ + UserID: user.ID, EffectiveGroupID: groupA.ID, Day: now, CostMicros: 250, + }) + require.NoError(t, err) + _, err = db.IncrementUserAIDailySpend(ctx, database.IncrementUserAIDailySpendParams{ + UserID: user.ID, EffectiveGroupID: groupB.ID, Day: now, CostMicros: 500, + }) + require.NoError(t, err) + + // When: querying spend for both groups in one call. + got, err := db.GetOrganizationGroupsAISpend(ctx, database.GetOrganizationGroupsAISpendParams{ + OrganizationID: org.ID, + GroupIds: []uuid.UUID{groupA.ID, groupB.ID}, + PeriodStart: monthStart, + }) + require.NoError(t, err) + + // Then: both are returned with their own budget and spend aggregates. + require.Len(t, got, 2) + byID := make(map[uuid.UUID]database.GetOrganizationGroupsAISpendRow, len(got)) + for _, r := range got { + byID[r.GroupID] = r + } + rowA, ok := byID[groupA.ID] + require.True(t, ok, "groupA missing from response") + require.Equal(t, sql.NullInt64{Int64: 1_000_000, Valid: true}, rowA.SpendLimitMicros) + require.Equal(t, int64(250), rowA.CurrentSpendMicros) + rowB, ok := byID[groupB.ID] + require.True(t, ok, "groupB missing from response") + require.Equal(t, sql.NullInt64{}, rowB.SpendLimitMicros) + require.Equal(t, int64(500), rowB.CurrentSpendMicros) + }) + + t.Run("ExcludesGroupsInOtherOrgs", func(t *testing.T) { + t.Parallel() + db, _ := dbtestutil.NewDB(t) + ctx := testutil.Context(t, testutil.WaitShort) + + // Given: a group in a different org with its own budget and spend. + user := dbgen.User(t, db, database.User{}) + org := dbgen.Organization(t, db, database.Organization{}) + otherOrg := dbgen.Organization(t, db, database.Organization{}) + group := dbgen.Group(t, db, database.Group{OrganizationID: org.ID}) + otherOrgGroup := dbgen.Group(t, db, database.Group{OrganizationID: otherOrg.ID}) + _, err := db.UpsertGroupAIBudget(ctx, database.UpsertGroupAIBudgetParams{ + GroupID: otherOrgGroup.ID, + SpendLimitMicros: 9_999_999, + }) + require.NoError(t, err) + _, err = db.IncrementUserAIDailySpend(ctx, database.IncrementUserAIDailySpendParams{ + UserID: user.ID, EffectiveGroupID: otherOrgGroup.ID, Day: now, CostMicros: 999, + }) + require.NoError(t, err) + + // When: querying the primary org with both group IDs. + got, err := db.GetOrganizationGroupsAISpend(ctx, database.GetOrganizationGroupsAISpendParams{ + OrganizationID: org.ID, + GroupIds: []uuid.UUID{group.ID, otherOrgGroup.ID}, + PeriodStart: monthStart, + }) + require.NoError(t, err) + + // Then: only the primary-org group is returned, and the cross-org group's budget and spend are absent. + require.Len(t, got, 1) + require.Equal(t, group.ID, got[0].GroupID) + require.Equal(t, sql.NullInt64{}, got[0].SpendLimitMicros, + "cross-org group's budget must not leak") + require.Equal(t, int64(0), got[0].CurrentSpendMicros, + "cross-org group's spend must not leak") + }) + + t.Run("ExcludesGroupIDsNotInList", func(t *testing.T) { + t.Parallel() + db, _ := dbtestutil.NewDB(t) + ctx := testutil.Context(t, testutil.WaitShort) + + // Given: two groups in the same org. + org := dbgen.Organization(t, db, database.Organization{}) + groupA := dbgen.Group(t, db, database.Group{OrganizationID: org.ID}) + _ = dbgen.Group(t, db, database.Group{OrganizationID: org.ID}) + + // When: querying with only one of the group IDs. + got, err := db.GetOrganizationGroupsAISpend(ctx, database.GetOrganizationGroupsAISpendParams{ + OrganizationID: org.ID, + GroupIds: []uuid.UUID{groupA.ID}, + PeriodStart: monthStart, + }) + require.NoError(t, err) + + // Then: only the requested group is returned. + require.Len(t, got, 1) + require.Equal(t, groupA.ID, got[0].GroupID) + }) + + t.Run("ExcludesSpendBeforePeriodStart", func(t *testing.T) { + t.Parallel() + db, _ := dbtestutil.NewDB(t) + ctx := testutil.Context(t, testutil.WaitShort) + + // Given: spend both in the prior period and in the current period. + user := dbgen.User(t, db, database.User{}) + org := dbgen.Organization(t, db, database.Organization{}) + group := dbgen.Group(t, db, database.Group{OrganizationID: org.ID}) + _, err := db.IncrementUserAIDailySpend(ctx, database.IncrementUserAIDailySpendParams{ + UserID: user.ID, EffectiveGroupID: group.ID, Day: prevMonthLastDay, CostMicros: 999, + }) + require.NoError(t, err) + _, err = db.IncrementUserAIDailySpend(ctx, database.IncrementUserAIDailySpendParams{ + UserID: user.ID, EffectiveGroupID: group.ID, Day: monthStart, CostMicros: 25, + }) + require.NoError(t, err) + + // When: querying since monthStart. + got, err := db.GetOrganizationGroupsAISpend(ctx, database.GetOrganizationGroupsAISpendParams{ + OrganizationID: org.ID, + GroupIds: []uuid.UUID{group.ID}, + PeriodStart: monthStart, + }) + require.NoError(t, err) + + // Then: only the current-period spend is aggregated. + require.Len(t, got, 1) + require.Equal(t, int64(25), got[0].CurrentSpendMicros) + }) + + t.Run("AggregatesSpendAcrossUsers", func(t *testing.T) { + t.Parallel() + db, _ := dbtestutil.NewDB(t) + ctx := testutil.Context(t, testutil.WaitShort) + + // Given: spend from two users attributed to the same group. + userA := dbgen.User(t, db, database.User{}) + userB := dbgen.User(t, db, database.User{}) + org := dbgen.Organization(t, db, database.Organization{}) + group := dbgen.Group(t, db, database.Group{OrganizationID: org.ID}) + _, err := db.IncrementUserAIDailySpend(ctx, database.IncrementUserAIDailySpendParams{ + UserID: userA.ID, EffectiveGroupID: group.ID, Day: now, CostMicros: 100, + }) + require.NoError(t, err) + _, err = db.IncrementUserAIDailySpend(ctx, database.IncrementUserAIDailySpendParams{ + UserID: userB.ID, EffectiveGroupID: group.ID, Day: now, CostMicros: 25, + }) + require.NoError(t, err) + + // When: querying the group's spend. + got, err := db.GetOrganizationGroupsAISpend(ctx, database.GetOrganizationGroupsAISpendParams{ + OrganizationID: org.ID, + GroupIds: []uuid.UUID{group.ID}, + PeriodStart: monthStart, + }) + require.NoError(t, err) + + // Then: the group's aggregate sums both users' spend. + require.Len(t, got, 1) + require.Equal(t, int64(125), got[0].CurrentSpendMicros) + }) + + t.Run("NormalizesNonUTCPeriodStart", func(t *testing.T) { + t.Parallel() + db, _ := dbtestutil.NewDB(t) + ctx := testutil.Context(t, testutil.WaitShort) + + // Given: spend both in the prior UTC day and the first day of the current UTC month. + user := dbgen.User(t, db, database.User{}) + org := dbgen.Organization(t, db, database.Organization{}) + group := dbgen.Group(t, db, database.Group{OrganizationID: org.ID}) + _, err := db.IncrementUserAIDailySpend(ctx, database.IncrementUserAIDailySpendParams{ + UserID: user.ID, EffectiveGroupID: group.ID, Day: prevMonthLastDay, CostMicros: 999, + }) + require.NoError(t, err) + _, err = db.IncrementUserAIDailySpend(ctx, database.IncrementUserAIDailySpendParams{ + UserID: user.ID, EffectiveGroupID: group.ID, Day: monthStart, CostMicros: 25, + }) + require.NoError(t, err) + + // When: querying with a non-UTC period_start that normalizes to June 1 UTC. + // 2024-05-31 23:00 in UTC-5 is 2024-06-01 04:00 UTC. + localLate := time.Date(2024, 5, 31, 23, 0, 0, 0, time.FixedZone("UTC-5", -5*3600)) + got, err := db.GetOrganizationGroupsAISpend(ctx, database.GetOrganizationGroupsAISpendParams{ + OrganizationID: org.ID, + GroupIds: []uuid.UUID{group.ID}, + PeriodStart: localLate, + }) + require.NoError(t, err) + + // Then: the prior UTC day's spend is excluded from the aggregate. + require.Len(t, got, 1) + require.Equal(t, int64(25), got[0].CurrentSpendMicros, + "sum must exclude prevMonthLastDay row after normalization") + }) +} + func TestChatPinOrderQueries(t *testing.T) { t.Parallel() if testing.Short() { diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go index dbb5c8c7a6fac..7ade08c9dd8d5 100644 --- a/coderd/database/queries.sql.go +++ b/coderd/database/queries.sql.go @@ -2568,6 +2568,68 @@ func (q *sqlQuerier) GetHighestGroupAIBudgetByUser(ctx context.Context, userID u return i, err } +const getOrganizationGroupsAISpend = `-- name: GetOrganizationGroupsAISpend :many +SELECT + groups.id AS group_id, + groups.organization_id AS organization_id, + budget.spend_limit_micros AS spend_limit_micros, + COALESCE(SUM(spend.spend_micros), 0)::BIGINT AS current_spend_micros +FROM groups +LEFT JOIN group_ai_budgets budget ON budget.group_id = groups.id +LEFT JOIN ai_user_daily_spend spend + ON spend.effective_group_id = groups.id + AND spend.day >= (($1::timestamptz) AT TIME ZONE 'UTC')::date +WHERE groups.organization_id = $2 + AND groups.id = ANY($3::uuid[]) +GROUP BY groups.id, budget.spend_limit_micros +` + +type GetOrganizationGroupsAISpendParams struct { + PeriodStart time.Time `db:"period_start" json:"period_start"` + OrganizationID uuid.UUID `db:"organization_id" json:"organization_id"` + GroupIds []uuid.UUID `db:"group_ids" json:"group_ids"` +} + +type GetOrganizationGroupsAISpendRow struct { + GroupID uuid.UUID `db:"group_id" json:"group_id"` + OrganizationID uuid.UUID `db:"organization_id" json:"organization_id"` + SpendLimitMicros sql.NullInt64 `db:"spend_limit_micros" json:"spend_limit_micros"` + CurrentSpendMicros int64 `db:"current_spend_micros" json:"current_spend_micros"` +} + +// Returns AI spend limits and aggregate spend for groups in @group_ids that +// belong to @organization_id, on or after period_start until NOW. The spend +// limit is null when the group has no configured budget. +// The period_start parameter is normalized to its UTC calendar day. +// Only return groups from @group_ids that belong to @organization_id. +func (q *sqlQuerier) GetOrganizationGroupsAISpend(ctx context.Context, arg GetOrganizationGroupsAISpendParams) ([]GetOrganizationGroupsAISpendRow, error) { + rows, err := q.db.QueryContext(ctx, getOrganizationGroupsAISpend, arg.PeriodStart, arg.OrganizationID, pq.Array(arg.GroupIds)) + if err != nil { + return nil, err + } + defer rows.Close() + var items []GetOrganizationGroupsAISpendRow + for rows.Next() { + var i GetOrganizationGroupsAISpendRow + if err := rows.Scan( + &i.GroupID, + &i.OrganizationID, + &i.SpendLimitMicros, + &i.CurrentSpendMicros, + ); err != nil { + return nil, err + } + items = append(items, i) + } + if err := rows.Close(); err != nil { + return nil, err + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + const getUserAIBudgetOverride = `-- name: GetUserAIBudgetOverride :one SELECT user_id, group_id, spend_limit_micros, created_at, updated_at FROM user_ai_budget_overrides diff --git a/coderd/database/queries/aicostcontrol.sql b/coderd/database/queries/aicostcontrol.sql index 00c4917371e11..82647081e4f35 100644 --- a/coderd/database/queries/aicostcontrol.sql +++ b/coderd/database/queries/aicostcontrol.sql @@ -101,3 +101,23 @@ FROM ai_user_daily_spend WHERE user_id = @user_id AND effective_group_id = @effective_group_id AND day >= ((@period_start::timestamptz) AT TIME ZONE 'UTC')::date; + +-- name: GetOrganizationGroupsAISpend :many +-- Returns AI spend limits and aggregate spend for groups in @group_ids that +-- belong to @organization_id, on or after period_start until NOW. The spend +-- limit is null when the group has no configured budget. +-- The period_start parameter is normalized to its UTC calendar day. +SELECT + groups.id AS group_id, + groups.organization_id AS organization_id, + budget.spend_limit_micros AS spend_limit_micros, + COALESCE(SUM(spend.spend_micros), 0)::BIGINT AS current_spend_micros +FROM groups +LEFT JOIN group_ai_budgets budget ON budget.group_id = groups.id +LEFT JOIN ai_user_daily_spend spend + ON spend.effective_group_id = groups.id + AND spend.day >= ((@period_start::timestamptz) AT TIME ZONE 'UTC')::date +-- Only return groups from @group_ids that belong to @organization_id. +WHERE groups.organization_id = @organization_id + AND groups.id = ANY(@group_ids::uuid[]) +GROUP BY groups.id, budget.spend_limit_micros; diff --git a/codersdk/aibridge.go b/codersdk/aibridge.go index 7b92638ac5993..6b845addbcb80 100644 --- a/codersdk/aibridge.go +++ b/codersdk/aibridge.go @@ -40,19 +40,44 @@ type UserAIBudgetSummary struct { LimitSource *AIBudgetLimitSource `json:"limit_source"` } +// AISpendPeriodWindow is the [Start, End) window over which AI spend is +// aggregated. +type AISpendPeriodWindow struct { + // PeriodStart is the inclusive lower bound of the current budget + // period. + PeriodStart time.Time `json:"period_start" format:"date-time"` + // PeriodEnd is the exclusive upper bound of the current budget + // period. + PeriodEnd time.Time `json:"period_end" format:"date-time"` +} + // UserAISpendStatus is the current AI spend snapshot for a user within // the active budget period. type UserAISpendStatus struct { UserAIBudgetSummary + AISpendPeriodWindow // CurrentSpendMicros is the user's spend on their effective group over // the current budget period. CurrentSpendMicros int64 `json:"current_spend_micros"` - // PeriodStart is the inclusive lower bound of the current budget - // period. - PeriodStart time.Time `json:"period_start" format:"date-time"` - // PeriodEnd is the exclusive upper bound of the current budget +} + +// OrganizationGroupsAISpend reports AI spend for a set of groups in the +// active budget period. +type OrganizationGroupsAISpend struct { + AISpendPeriodWindow + Groups []OrganizationGroupAISpend `json:"groups"` +} + +// OrganizationGroupAISpend is the current AI spend snapshot for a group +// within the active budget period. +type OrganizationGroupAISpend struct { + GroupID uuid.UUID `json:"group_id" format:"uuid"` + // SpendLimitMicros is the group's configured AI spend limit. Null when + // the group has no configured budget. + SpendLimitMicros *int64 `json:"spend_limit_micros"` + // CurrentSpendMicros is the group's spend over the current budget // period. - PeriodEnd time.Time `json:"period_end" format:"date-time"` + CurrentSpendMicros int64 `json:"current_spend_micros"` } type AIBridgeSession struct { @@ -444,3 +469,31 @@ func (c *Client) UserAISpendStatus(ctx context.Context, user uuid.UUID) (UserAIS var resp UserAISpendStatus return resp, json.NewDecoder(res.Body).Decode(&resp) } + +// OrganizationGroupsAISpend returns AI spend for the given groups within the +// organization for the active budget period. +func (c *Client) OrganizationGroupsAISpend(ctx context.Context, organization uuid.UUID, groupIDs []uuid.UUID) (OrganizationGroupsAISpend, error) { + ids := make([]string, len(groupIDs)) + for i, id := range groupIDs { + ids[i] = id.String() + } + res, err := c.Request(ctx, http.MethodGet, + fmt.Sprintf("/api/v2/organizations/%s/groups/ai/spend", organization.String()), + nil, + func(r *http.Request) { + q := r.URL.Query() + q.Set("group_ids", strings.Join(ids, ",")) + r.URL.RawQuery = q.Encode() + }, + ) + if err != nil { + return OrganizationGroupsAISpend{}, xerrors.Errorf("make request: %w", err) + } + defer res.Body.Close() + + if res.StatusCode != http.StatusOK { + return OrganizationGroupsAISpend{}, ReadBodyAsError(res) + } + var resp OrganizationGroupsAISpend + return resp, json.NewDecoder(res.Body).Decode(&resp) +} diff --git a/docs/reference/api/enterprise.md b/docs/reference/api/enterprise.md index 6368a193130b7..052e836254087 100644 --- a/docs/reference/api/enterprise.md +++ b/docs/reference/api/enterprise.md @@ -1839,6 +1839,52 @@ curl -X POST http://coder-server:8080/api/v2/organizations/{organization}/groups To perform this operation, you must be authenticated. [Learn more](authentication.md). +## Get organization groups AI spend + +### Code samples + +```shell +# Example request using curl +curl -X GET http://coder-server:8080/api/v2/organizations/{organization}/groups/ai/spend?group_ids=string \ + -H 'Accept: application/json' \ + -H 'Coder-Session-Token: API_KEY' +``` + +`GET /api/v2/organizations/{organization}/groups/ai/spend` + +### Parameters + +| Name | In | Type | Required | Description | +|----------------|-------|--------------|----------|-----------------------------------| +| `organization` | path | string(uuid) | true | Organization ID | +| `group_ids` | query | string | true | Comma-separated list of group IDs | + +### Example responses + +> 200 Response + +```json +{ + "groups": [ + { + "current_spend_micros": 0, + "group_id": "306db4e0-7449-4501-b76f-075576fe2d8f", + "spend_limit_micros": 0 + } + ], + "period_end": "2019-08-24T14:15:22Z", + "period_start": "2019-08-24T14:15:22Z" +} +``` + +### Responses + +| Status | Meaning | Description | Schema | +|--------|---------------------------------------------------------|-------------|------------------------------------------------------------------------------------| +| 200 | [OK](https://tools.ietf.org/html/rfc7231#section-6.3.1) | OK | [codersdk.OrganizationGroupsAISpend](schemas.md#codersdkorganizationgroupsaispend) | + +To perform this operation, you must be authenticated. [Learn more](authentication.md). + ## Get group by organization and group name ### Code samples diff --git a/docs/reference/api/schemas.md b/docs/reference/api/schemas.md index fcbe267ec2036..e938a624b7f8e 100644 --- a/docs/reference/api/schemas.md +++ b/docs/reference/api/schemas.md @@ -9142,6 +9142,48 @@ Only certain features set these fields: - FeatureManagedAgentLimit| | `name` | string | false | | | | `updated_at` | string | true | | | +## codersdk.OrganizationGroupAISpend + +```json +{ + "current_spend_micros": 0, + "group_id": "306db4e0-7449-4501-b76f-075576fe2d8f", + "spend_limit_micros": 0 +} +``` + +### Properties + +| Name | Type | Required | Restrictions | Description | +|------------------------|---------|----------|--------------|------------------------------------------------------------------------------------------------------------| +| `current_spend_micros` | integer | false | | Current spend micros is the group's spend over the current budget period. | +| `group_id` | string | false | | | +| `spend_limit_micros` | integer | false | | Spend limit micros is the group's configured AI spend limit. Null when the group has no configured budget. | + +## codersdk.OrganizationGroupsAISpend + +```json +{ + "groups": [ + { + "current_spend_micros": 0, + "group_id": "306db4e0-7449-4501-b76f-075576fe2d8f", + "spend_limit_micros": 0 + } + ], + "period_end": "2019-08-24T14:15:22Z", + "period_start": "2019-08-24T14:15:22Z" +} +``` + +### Properties + +| Name | Type | Required | Restrictions | Description | +|----------------|---------------------------------------------------------------------------------|----------|--------------|-------------------------------------------------------------------------| +| `groups` | array of [codersdk.OrganizationGroupAISpend](#codersdkorganizationgroupaispend) | false | | | +| `period_end` | string | false | | Period end is the exclusive upper bound of the current budget period. | +| `period_start` | string | false | | Period start is the inclusive lower bound of the current budget period. | + ## codersdk.OrganizationMember ```json diff --git a/enterprise/coderd/aibridge.go b/enterprise/coderd/aibridge.go index ac6b3d80fe177..9fff6234ce797 100644 --- a/enterprise/coderd/aibridge.go +++ b/enterprise/coderd/aibridge.go @@ -36,7 +36,8 @@ const ( defaultListClientsLimit = 100 // aiBridgeRateLimitWindow is the fixed duration for rate limiting AI Bridge // requests. This is hardcoded to keep configuration simple. - aiBridgeRateLimitWindow = time.Second + aiBridgeRateLimitWindow = time.Second + maxOrganizationGroupsAISpendGroupIDs = 100 ) // errInvalidCursor is returned when a pagination cursor does not @@ -914,8 +915,10 @@ func (api *API) userAISpendStatus(rw http.ResponseWriter, r *http.Request) { UserAIBudgetSummary: codersdk.UserAIBudgetSummary{ UserID: user.ID, }, - PeriodStart: periodWindow.Start, - PeriodEnd: periodWindow.End, + AISpendPeriodWindow: codersdk.AISpendPeriodWindow{ + PeriodStart: periodWindow.Start, + PeriodEnd: periodWindow.End, + }, } if ok { @@ -939,3 +942,82 @@ func (api *API) userAISpendStatus(rw http.ResponseWriter, r *http.Request) { httpapi.Write(ctx, rw, http.StatusOK, resp) } + +// @Summary Get organization groups AI spend +// @ID get-organization-groups-ai-spend +// @Security CoderSessionToken +// @Produce json +// @Tags Enterprise +// @Param organization path string true "Organization ID" format(uuid) +// @Param group_ids query string true "Comma-separated list of group IDs" +// @Success 200 {object} codersdk.OrganizationGroupsAISpend +// @Router /api/v2/organizations/{organization}/groups/ai/spend [get] +func (api *API) organizationGroupsAISpend(rw http.ResponseWriter, r *http.Request) { + ctx := r.Context() + org := httpmw.OrganizationParam(r) + logger := api.Logger.With(slog.F("organization_id", org.ID)) + + parser := httpapi.NewQueryParamParser() + parser.RequiredNotEmpty("group_ids") + groupIDs := parser.UUIDs(r.URL.Query(), nil, "group_ids") + parser.ErrorExcessParams(r.URL.Query()) + if len(parser.Errors) > 0 { + httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{ + Message: "Query parameters have invalid values.", + Validations: parser.Errors, + }) + return + } + if len(groupIDs) > maxOrganizationGroupsAISpendGroupIDs { + httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{ + Message: fmt.Sprintf( + "group_ids has %d entries, maximum is %d.", + len(groupIDs), maxOrganizationGroupsAISpendGroupIDs, + ), + }) + return + } + + period := codersdk.NewAIBudgetPeriodFromString(api.DeploymentValues.AI.BridgeConfig.BudgetPeriod) + periodWindow, err := budget.CurrentPeriod(api.Clock.Now(), period) + if err != nil { + logger.Error(ctx, "failed to compute AI budget period", slog.Error(err)) + httpapi.InternalServerError(rw, err) + return + } + logger = logger.With( + slog.F("period_start", periodWindow.Start), + slog.F("period_end", periodWindow.End), + ) + + rows, err := api.Database.GetOrganizationGroupsAISpend(ctx, database.GetOrganizationGroupsAISpendParams{ + OrganizationID: org.ID, + GroupIds: groupIDs, + PeriodStart: periodWindow.Start, + }) + if err != nil { + logger.Error(ctx, "failed to get organization groups AI spend", slog.Error(err)) + httpapi.InternalServerError(rw, err) + return + } + + resp := codersdk.OrganizationGroupsAISpend{ + AISpendPeriodWindow: codersdk.AISpendPeriodWindow{ + PeriodStart: periodWindow.Start, + PeriodEnd: periodWindow.End, + }, + Groups: make([]codersdk.OrganizationGroupAISpend, 0, len(rows)), + } + for _, row := range rows { + entry := codersdk.OrganizationGroupAISpend{ + GroupID: row.GroupID, + CurrentSpendMicros: row.CurrentSpendMicros, + } + if row.SpendLimitMicros.Valid { + entry.SpendLimitMicros = &row.SpendLimitMicros.Int64 + } + resp.Groups = append(resp.Groups, entry) + } + + httpapi.Write(ctx, rw, http.StatusOK, resp) +} diff --git a/enterprise/coderd/aibridge_test.go b/enterprise/coderd/aibridge_test.go index ccfe716864b57..b212f96381df6 100644 --- a/enterprise/coderd/aibridge_test.go +++ b/enterprise/coderd/aibridge_test.go @@ -3206,6 +3206,290 @@ func TestUserAISpendStatusRoleAccess(t *testing.T) { } } +func TestOrganizationGroupsAISpend(t *testing.T) { + t.Parallel() + + t.Run("RequiresLicenseFeature", func(t *testing.T) { + t.Parallel() + + dv := coderdtest.DeploymentValues(t) + dv.Experiments = []string{string(codersdk.ExperimentAIGatewayCostControl)} + client, owner := coderdenttest.New(t, &coderdenttest.Options{ + Options: &coderdtest.Options{DeploymentValues: dv}, + LicenseOptions: &coderdenttest.LicenseOptions{ + Features: license.Features{}, + }, + }) + ctx := testutil.Context(t, testutil.WaitLong) + + //nolint:gocritic // Owner role is irrelevant here; the request is blocked before RBAC. + _, err := client.OrganizationGroupsAISpend(ctx, owner.OrganizationID, []uuid.UUID{uuid.New()}) + var sdkErr *codersdk.Error + require.ErrorAs(t, err, &sdkErr) + require.Equal(t, http.StatusForbidden, sdkErr.StatusCode()) + }) + + t.Run("RequiresExperiment", func(t *testing.T) { + t.Parallel() + + dv := coderdtest.DeploymentValues(t) + dv.AI.BridgeConfig.Enabled = serpent.Bool(true) + client, owner := coderdenttest.New(t, &coderdenttest.Options{ + Options: &coderdtest.Options{DeploymentValues: dv}, + LicenseOptions: &coderdenttest.LicenseOptions{ + Features: license.Features{ + codersdk.FeatureAIBridge: 1, + }, + }, + }) + ctx := testutil.Context(t, testutil.WaitLong) + + //nolint:gocritic // Owner role is irrelevant here; the request is blocked before RBAC. + _, err := client.OrganizationGroupsAISpend(ctx, owner.OrganizationID, []uuid.UUID{uuid.New()}) + var sdkErr *codersdk.Error + require.ErrorAs(t, err, &sdkErr) + require.Equal(t, http.StatusForbidden, sdkErr.StatusCode()) + }) + + t.Run("MissingGroupIDs", func(t *testing.T) { + t.Parallel() + + adminClient, _, group := setupAICostControlTest(t, aiCostControlTestOptions{GroupName: "missing-ids-group"}) + ctx := testutil.Context(t, testutil.WaitLong) + + // Given: no group_ids query parameter. + // When: querying spend. + _, err := adminClient.OrganizationGroupsAISpend(ctx, group.OrganizationID, nil) + + // Then: request fails with 400. + var sdkErr *codersdk.Error + require.ErrorAs(t, err, &sdkErr) + require.Equal(t, http.StatusBadRequest, sdkErr.StatusCode()) + }) + + t.Run("TooManyGroupIDs", func(t *testing.T) { + t.Parallel() + + adminClient, _, group := setupAICostControlTest(t, aiCostControlTestOptions{GroupName: "too-many-group-ids-group"}) + ctx := testutil.Context(t, testutil.WaitLong) + + // Given: 101 group_ids, above the cap of 100. + ids := make([]uuid.UUID, 101) + for i := range ids { + ids[i] = uuid.New() + } + + // When: querying spend. + _, err := adminClient.OrganizationGroupsAISpend(ctx, group.OrganizationID, ids) + + // Then: request fails with 400. + var sdkErr *codersdk.Error + require.ErrorAs(t, err, &sdkErr) + require.Equal(t, http.StatusBadRequest, sdkErr.StatusCode()) + }) + + t.Run("MalformedGroupID", func(t *testing.T) { + t.Parallel() + + adminClient, _, group := setupAICostControlTest(t, aiCostControlTestOptions{GroupName: "malformed-group-id-group"}) + ctx := testutil.Context(t, testutil.WaitLong) + + // Given: a malformed UUID passed via raw HTTP. + // When: querying spend. + res, err := adminClient.Request(ctx, http.MethodGet, + "/api/v2/organizations/"+group.OrganizationID.String()+"/groups/ai/spend", + nil, + func(r *http.Request) { + q := r.URL.Query() + q.Set("group_ids", "not-a-uuid") + r.URL.RawQuery = q.Encode() + }, + ) + require.NoError(t, err) + defer res.Body.Close() + + // Then: 400. + require.Equal(t, http.StatusBadRequest, res.StatusCode) + }) + + t.Run("GroupInOtherOrgExcluded", func(t *testing.T) { + t.Parallel() + + // Given: two groups, one in the queried org and one in a different org. + db, ps := dbtestutil.NewDB(t) + adminClient, _, group := setupAICostControlTest(t, aiCostControlTestOptions{ + GroupName: "primary-org-group", + Database: db, + Pubsub: ps, + }) + otherOrg := dbgen.Organization(t, db, database.Organization{}) + otherOrgGroup := dbgen.Group(t, db, database.Group{OrganizationID: otherOrg.ID}) + ctx := testutil.Context(t, testutil.WaitLong) + + // When: querying the primary org with both group IDs. + resp, err := adminClient.OrganizationGroupsAISpend(ctx, group.OrganizationID, []uuid.UUID{group.ID, otherOrgGroup.ID}) + require.NoError(t, err) + + // Then: only the primary-org group is returned. + require.Len(t, resp.Groups, 1) + require.Equal(t, group.ID, resp.Groups[0].GroupID) + }) + + tests := []struct { + name string + setBudget bool + spendLimit int64 + spent int64 + wantSpendLimit *int64 + wantCurrentSpend int64 + }{ + { + name: "NoBudgetNoSpend", + }, + { + name: "ZeroLimitBudget", + setBudget: true, + spendLimit: 0, + wantSpendLimit: ptr.Ref(int64(0)), + wantCurrentSpend: 0, + }, + { + name: "BudgetZeroSpend", + setBudget: true, + spendLimit: 1_000_000_000, + wantSpendLimit: ptr.Ref(int64(1_000_000_000)), + wantCurrentSpend: 0, + }, + { + name: "BudgetWithSpend", + setBudget: true, + spendLimit: 1_000_000_000, + spent: 250_000_000, + wantSpendLimit: ptr.Ref(int64(1_000_000_000)), + wantCurrentSpend: 250_000_000, + }, + { + name: "SpendExceedsLimit", + setBudget: true, + spendLimit: 1_000_000_000, + spent: 1_500_000_000, + wantSpendLimit: ptr.Ref(int64(1_000_000_000)), + wantCurrentSpend: 1_500_000_000, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + // Given: an admin, a group, and optionally a budget and seeded spend. + clock := quartz.NewMock(t) + db, ps := dbtestutil.NewDB(t) + adminClient, targetUser, group := setupAICostControlTest(t, aiCostControlTestOptions{ + GroupName: "spend-test-group", + Clock: clock, + Database: db, + Pubsub: ps, + }) + ctx := testutil.Context(t, testutil.WaitLong) + clock.Set(time.Date(2026, time.March, 15, 12, 0, 0, 0, time.UTC)) + wantPeriodStart := time.Date(2026, time.March, 1, 0, 0, 0, 0, time.UTC) + wantPeriodEnd := time.Date(2026, time.April, 1, 0, 0, 0, 0, time.UTC) + + if tt.setBudget { + _, err := adminClient.UpsertGroupAIBudget(ctx, group.ID, codersdk.UpsertGroupAIBudgetRequest{ + SpendLimitMicros: tt.spendLimit, + }) + require.NoError(t, err) + } + if tt.spent > 0 { + _, err := db.IncrementUserAIDailySpend(ctx, database.IncrementUserAIDailySpendParams{ + UserID: targetUser.ID, + EffectiveGroupID: group.ID, + Day: clock.Now(), + CostMicros: tt.spent, + }) + require.NoError(t, err) + } + + // When: querying the group's spend. + got, err := adminClient.OrganizationGroupsAISpend(ctx, group.OrganizationID, []uuid.UUID{group.ID}) + require.NoError(t, err) + + // Then: the response contains one row with the expected fields. + require.Equal(t, wantPeriodStart, got.PeriodStart) + require.Equal(t, wantPeriodEnd, got.PeriodEnd) + require.Len(t, got.Groups, 1) + require.Equal(t, group.ID, got.Groups[0].GroupID) + require.Equal(t, tt.wantSpendLimit, got.Groups[0].SpendLimitMicros) + require.Equal(t, tt.wantCurrentSpend, got.Groups[0].CurrentSpendMicros) + }) + } +} + +func TestOrganizationGroupsAISpendRoleAccess(t *testing.T) { + t.Parallel() + + dv := coderdtest.DeploymentValues(t) + dv.AI.BridgeConfig.Enabled = serpent.Bool(true) + dv.Experiments = []string{string(codersdk.ExperimentAIGatewayCostControl)} + ownerClient, owner := coderdenttest.New(t, &coderdenttest.Options{ + Options: &coderdtest.Options{DeploymentValues: dv}, + LicenseOptions: &coderdenttest.LicenseOptions{ + Features: license.Features{ + codersdk.FeatureTemplateRBAC: 1, + codersdk.FeatureAIBridge: 1, + codersdk.FeatureMultipleOrganizations: 1, + }, + }, + }) + userAdminClient, _ := coderdtest.CreateAnotherUser(t, ownerClient, owner.OrganizationID, rbac.RoleUserAdmin()) + orgAdminClient, _ := coderdtest.CreateAnotherUser(t, ownerClient, owner.OrganizationID, rbac.ScopedRoleOrgAdmin(owner.OrganizationID)) + orgUserAdminClient, _ := coderdtest.CreateAnotherUser(t, ownerClient, owner.OrganizationID, rbac.ScopedRoleOrgUserAdmin(owner.OrganizationID)) + memberClient, _ := coderdtest.CreateAnotherUser(t, ownerClient, owner.OrganizationID) + + otherOrg := coderdenttest.CreateOrganization(t, ownerClient, coderdenttest.CreateOrganizationOptions{}) + otherOrgMemberClient, _ := coderdtest.CreateAnotherUser(t, ownerClient, otherOrg.ID) + + ctx := testutil.Context(t, testutil.WaitLong) + group, err := userAdminClient.CreateGroup(ctx, owner.OrganizationID, codersdk.CreateGroupRequest{ + Name: "role-access-group", + }) + require.NoError(t, err) + + cases := []struct { + name string + client *codersdk.Client + wantGroup bool + }{ + {name: "Owner", client: ownerClient, wantGroup: true}, + {name: "UserAdmin", client: userAdminClient, wantGroup: true}, + {name: "OrgAdmin", client: orgAdminClient, wantGroup: true}, + {name: "OrgUserAdmin", client: orgUserAdminClient, wantGroup: true}, + {name: "Member", client: memberClient, wantGroup: true}, + {name: "OtherOrgMember", client: otherOrgMemberClient, wantGroup: false}, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + + resp, err := tc.client.OrganizationGroupsAISpend(ctx, owner.OrganizationID, []uuid.UUID{group.ID}) + if !tc.wantGroup { + var sdkErr *codersdk.Error + require.ErrorAs(t, err, &sdkErr) + require.Equal(t, http.StatusNotFound, sdkErr.StatusCode()) + return + } + require.NoError(t, err) + require.Len(t, resp.Groups, 1) + require.Equal(t, group.ID, resp.Groups[0].GroupID) + }) + } +} + // aiCostControlTestOptions configures the setup of an AI cost control test // deployment. GroupName is required. Clock, Database, and Pubsub are // optional overrides (leave nil for defaults). diff --git a/enterprise/coderd/coderd.go b/enterprise/coderd/coderd.go index faf204227d06e..c08680c6fc35b 100644 --- a/enterprise/coderd/coderd.go +++ b/enterprise/coderd/coderd.go @@ -503,6 +503,15 @@ func New(ctx context.Context, options *Options) (_ *API, err error) { ) r.Post("/", api.postGroupByOrganization) r.Get("/", api.groupsByOrganization) + r.Route("/ai/spend", func(r chi.Router) { + // AI cost controls are a paid feature (AI Governance add-on). + r.Use( + // TODO(AIGOV-443): remove once AI Gateway cost control functionality is stable. + httpmw.RequireExperiment(api.AGPL.Experiments, codersdk.ExperimentAIGatewayCostControl), + api.RequireFeatureMW(codersdk.FeatureAIBridge), + ) + r.Get("/", api.organizationGroupsAISpend) + }) r.Route("/{groupName}", func(r chi.Router) { r.Use( httpmw.ExtractGroupByNameParam(api.Database), diff --git a/site/src/api/typesGenerated.ts b/site/src/api/typesGenerated.ts index b4ee9f1ab0b3e..1f7c26a1a4a6a 100644 --- a/site/src/api/typesGenerated.ts +++ b/site/src/api/typesGenerated.ts @@ -504,6 +504,24 @@ export const AIProviderTypes: AIProviderType[] = [ "vercel", ]; +// From codersdk/aibridge.go +/** + * AISpendPeriodWindow is the [Start, End) window over which AI spend is + * aggregated. + */ +export interface AISpendPeriodWindow { + /** + * PeriodStart is the inclusive lower bound of the current budget + * period. + */ + readonly period_start: string; + /** + * PeriodEnd is the exclusive upper bound of the current budget + * period. + */ + readonly period_end: string; +} + // From codersdk/allowlist.go /** * APIAllowListTarget represents a single allow-list entry using the canonical @@ -6536,6 +6554,34 @@ export interface Organization extends MinimalOrganization { readonly default_org_member_roles: readonly string[]; } +// From codersdk/aibridge.go +/** + * OrganizationGroupAISpend is the current AI spend snapshot for a group + * within the active budget period. + */ +export interface OrganizationGroupAISpend { + readonly group_id: string; + /** + * SpendLimitMicros is the group's configured AI spend limit. Null when + * the group has no configured budget. + */ + readonly spend_limit_micros: number | null; + /** + * CurrentSpendMicros is the group's spend over the current budget + * period. + */ + readonly current_spend_micros: number; +} + +// From codersdk/aibridge.go +/** + * OrganizationGroupsAISpend reports AI spend for a set of groups in the + * active budget period. + */ +export interface OrganizationGroupsAISpend extends AISpendPeriodWindow { + readonly groups: readonly OrganizationGroupAISpend[]; +} + // From codersdk/organizations.go export interface OrganizationMember { readonly user_id: string; @@ -9825,22 +9871,14 @@ export interface UserAIProviderKeyConfig { * UserAISpendStatus is the current AI spend snapshot for a user within * the active budget period. */ -export interface UserAISpendStatus extends UserAIBudgetSummary { +export interface UserAISpendStatus + extends UserAIBudgetSummary, + AISpendPeriodWindow { /** * CurrentSpendMicros is the user's spend on their effective group over * the current budget period. */ readonly current_spend_micros: number; - /** - * PeriodStart is the inclusive lower bound of the current budget - * period. - */ - readonly period_start: string; - /** - * PeriodEnd is the exclusive upper bound of the current budget - * period. - */ - readonly period_end: string; } // From codersdk/insights.go From d1b6d1b6efa4cfb3d12db4853c2c0adaba0b707f Mon Sep 17 00:00:00 2001 From: Susana Cardoso Ferreira Date: Wed, 15 Jul 2026 10:51:53 +0000 Subject: [PATCH 2/5] chore: address review comments --- coderd/apidoc/docs.go | 3 ++- coderd/apidoc/swagger.json | 3 ++- coderd/database/querier.go | 1 - coderd/database/queries.sql.go | 2 +- coderd/database/queries/aicostcontrol.sql | 4 ++-- docs/reference/api/enterprise.md | 10 ++++---- enterprise/coderd/aibridge.go | 16 +++++++++---- enterprise/coderd/aibridge_test.go | 28 +++++++++++++++++++++-- 8 files changed, 50 insertions(+), 17 deletions(-) diff --git a/coderd/apidoc/docs.go b/coderd/apidoc/docs.go index 9ac635035b325..bcc1fb240560a 100644 --- a/coderd/apidoc/docs.go +++ b/coderd/apidoc/docs.go @@ -4754,6 +4754,7 @@ const docTemplate = `{ }, "/api/v2/organizations/{organization}/groups/ai/spend": { "get": { + "description": "Returns AI spend limits and aggregate spend for the requested groups. Unknown or unreadable group IDs are silently omitted.", "produces": [ "application/json" ], @@ -4773,7 +4774,7 @@ const docTemplate = `{ }, { "type": "string", - "description": "Comma-separated list of group IDs", + "description": "Comma-separated list of group IDs (maximum 100)", "name": "group_ids", "in": "query", "required": true diff --git a/coderd/apidoc/swagger.json b/coderd/apidoc/swagger.json index 7b5f4592be325..9793f9f18dd71 100644 --- a/coderd/apidoc/swagger.json +++ b/coderd/apidoc/swagger.json @@ -4195,6 +4195,7 @@ }, "/api/v2/organizations/{organization}/groups/ai/spend": { "get": { + "description": "Returns AI spend limits and aggregate spend for the requested groups. Unknown or unreadable group IDs are silently omitted.", "produces": ["application/json"], "tags": ["Enterprise"], "summary": "Get organization groups AI spend", @@ -4210,7 +4211,7 @@ }, { "type": "string", - "description": "Comma-separated list of group IDs", + "description": "Comma-separated list of group IDs (maximum 100)", "name": "group_ids", "in": "query", "required": true diff --git a/coderd/database/querier.go b/coderd/database/querier.go index 6a5f0ce47d73c..b9bb106391310 100644 --- a/coderd/database/querier.go +++ b/coderd/database/querier.go @@ -655,7 +655,6 @@ type sqlcQuerier interface { // belong to @organization_id, on or after period_start until NOW. The spend // limit is null when the group has no configured budget. // The period_start parameter is normalized to its UTC calendar day. - // Only return groups from @group_ids that belong to @organization_id. GetOrganizationGroupsAISpend(ctx context.Context, arg GetOrganizationGroupsAISpendParams) ([]GetOrganizationGroupsAISpendRow, error) GetOrganizationIDsByMemberIDs(ctx context.Context, ids []uuid.UUID) ([]GetOrganizationIDsByMemberIDsRow, error) GetOrganizationResourceCountByID(ctx context.Context, organizationID uuid.UUID) (GetOrganizationResourceCountByIDRow, error) diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go index 7ade08c9dd8d5..ac81aac2075e1 100644 --- a/coderd/database/queries.sql.go +++ b/coderd/database/queries.sql.go @@ -2582,6 +2582,7 @@ LEFT JOIN ai_user_daily_spend spend WHERE groups.organization_id = $2 AND groups.id = ANY($3::uuid[]) GROUP BY groups.id, budget.spend_limit_micros +ORDER BY groups.id ` type GetOrganizationGroupsAISpendParams struct { @@ -2601,7 +2602,6 @@ type GetOrganizationGroupsAISpendRow struct { // belong to @organization_id, on or after period_start until NOW. The spend // limit is null when the group has no configured budget. // The period_start parameter is normalized to its UTC calendar day. -// Only return groups from @group_ids that belong to @organization_id. func (q *sqlQuerier) GetOrganizationGroupsAISpend(ctx context.Context, arg GetOrganizationGroupsAISpendParams) ([]GetOrganizationGroupsAISpendRow, error) { rows, err := q.db.QueryContext(ctx, getOrganizationGroupsAISpend, arg.PeriodStart, arg.OrganizationID, pq.Array(arg.GroupIds)) if err != nil { diff --git a/coderd/database/queries/aicostcontrol.sql b/coderd/database/queries/aicostcontrol.sql index 82647081e4f35..7738e3e09e073 100644 --- a/coderd/database/queries/aicostcontrol.sql +++ b/coderd/database/queries/aicostcontrol.sql @@ -117,7 +117,7 @@ LEFT JOIN group_ai_budgets budget ON budget.group_id = groups.id LEFT JOIN ai_user_daily_spend spend ON spend.effective_group_id = groups.id AND spend.day >= ((@period_start::timestamptz) AT TIME ZONE 'UTC')::date --- Only return groups from @group_ids that belong to @organization_id. WHERE groups.organization_id = @organization_id AND groups.id = ANY(@group_ids::uuid[]) -GROUP BY groups.id, budget.spend_limit_micros; +GROUP BY groups.id, budget.spend_limit_micros +ORDER BY groups.id; diff --git a/docs/reference/api/enterprise.md b/docs/reference/api/enterprise.md index 052e836254087..6657fe0a95ba3 100644 --- a/docs/reference/api/enterprise.md +++ b/docs/reference/api/enterprise.md @@ -1852,12 +1852,14 @@ curl -X GET http://coder-server:8080/api/v2/organizations/{organization}/groups/ `GET /api/v2/organizations/{organization}/groups/ai/spend` +Returns AI spend limits and aggregate spend for the requested groups. Unknown or unreadable group IDs are silently omitted. + ### Parameters -| Name | In | Type | Required | Description | -|----------------|-------|--------------|----------|-----------------------------------| -| `organization` | path | string(uuid) | true | Organization ID | -| `group_ids` | query | string | true | Comma-separated list of group IDs | +| Name | In | Type | Required | Description | +|----------------|-------|--------------|----------|-------------------------------------------------| +| `organization` | path | string(uuid) | true | Organization ID | +| `group_ids` | query | string | true | Comma-separated list of group IDs (maximum 100) | ### Example responses diff --git a/enterprise/coderd/aibridge.go b/enterprise/coderd/aibridge.go index 9fff6234ce797..c29dc29a7bab3 100644 --- a/enterprise/coderd/aibridge.go +++ b/enterprise/coderd/aibridge.go @@ -878,6 +878,13 @@ func (api *API) deleteUserAIBudgetOverride(rw http.ResponseWriter, r *http.Reque rw.WriteHeader(http.StatusNoContent) } +// currentAIBudgetWindow returns the current AI budget period window based on +// the configured budget period. +func (api *API) currentAIBudgetWindow() (budget.PeriodWindow, error) { + period := codersdk.NewAIBudgetPeriodFromString(api.DeploymentValues.AI.BridgeConfig.BudgetPeriod) + return budget.CurrentPeriod(api.Clock.Now(), period) +} + // @Summary Get user AI spend // @ID get-user-ai-spend // @Security CoderSessionToken @@ -891,8 +898,7 @@ func (api *API) userAISpendStatus(rw http.ResponseWriter, r *http.Request) { user := httpmw.UserParam(r) logger := api.Logger.With(slog.F("user_id", user.ID)) - period := codersdk.NewAIBudgetPeriodFromString(api.DeploymentValues.AI.BridgeConfig.BudgetPeriod) - periodWindow, err := budget.CurrentPeriod(api.Clock.Now(), period) + periodWindow, err := api.currentAIBudgetWindow() if err != nil { logger.Error(ctx, "failed to compute AI budget period", slog.Error(err)) httpapi.InternalServerError(rw, err) @@ -944,12 +950,13 @@ func (api *API) userAISpendStatus(rw http.ResponseWriter, r *http.Request) { } // @Summary Get organization groups AI spend +// @Description Returns AI spend limits and aggregate spend for the requested groups. Unknown or unreadable group IDs are silently omitted. // @ID get-organization-groups-ai-spend // @Security CoderSessionToken // @Produce json // @Tags Enterprise // @Param organization path string true "Organization ID" format(uuid) -// @Param group_ids query string true "Comma-separated list of group IDs" +// @Param group_ids query string true "Comma-separated list of group IDs (maximum 100)" // @Success 200 {object} codersdk.OrganizationGroupsAISpend // @Router /api/v2/organizations/{organization}/groups/ai/spend [get] func (api *API) organizationGroupsAISpend(rw http.ResponseWriter, r *http.Request) { @@ -978,8 +985,7 @@ func (api *API) organizationGroupsAISpend(rw http.ResponseWriter, r *http.Reques return } - period := codersdk.NewAIBudgetPeriodFromString(api.DeploymentValues.AI.BridgeConfig.BudgetPeriod) - periodWindow, err := budget.CurrentPeriod(api.Clock.Now(), period) + periodWindow, err := api.currentAIBudgetWindow() if err != nil { logger.Error(ctx, "failed to compute AI budget period", slog.Error(err)) httpapi.InternalServerError(rw, err) diff --git a/enterprise/coderd/aibridge_test.go b/enterprise/coderd/aibridge_test.go index b212f96381df6..e3d2844769d7b 100644 --- a/enterprise/coderd/aibridge_test.go +++ b/enterprise/coderd/aibridge_test.go @@ -3217,7 +3217,9 @@ func TestOrganizationGroupsAISpend(t *testing.T) { client, owner := coderdenttest.New(t, &coderdenttest.Options{ Options: &coderdtest.Options{DeploymentValues: dv}, LicenseOptions: &coderdenttest.LicenseOptions{ - Features: license.Features{}, + Features: license.Features{ + codersdk.FeatureTemplateRBAC: 1, + }, }, }) ctx := testutil.Context(t, testutil.WaitLong) @@ -3227,6 +3229,7 @@ func TestOrganizationGroupsAISpend(t *testing.T) { var sdkErr *codersdk.Error require.ErrorAs(t, err, &sdkErr) require.Equal(t, http.StatusForbidden, sdkErr.StatusCode()) + require.Contains(t, sdkErr.Message, "AI Gateway is a Premium feature") }) t.Run("RequiresExperiment", func(t *testing.T) { @@ -3238,7 +3241,8 @@ func TestOrganizationGroupsAISpend(t *testing.T) { Options: &coderdtest.Options{DeploymentValues: dv}, LicenseOptions: &coderdenttest.LicenseOptions{ Features: license.Features{ - codersdk.FeatureAIBridge: 1, + codersdk.FeatureTemplateRBAC: 1, + codersdk.FeatureAIBridge: 1, }, }, }) @@ -3249,6 +3253,7 @@ func TestOrganizationGroupsAISpend(t *testing.T) { var sdkErr *codersdk.Error require.ErrorAs(t, err, &sdkErr) require.Equal(t, http.StatusForbidden, sdkErr.StatusCode()) + require.Contains(t, sdkErr.Message, "ai-gateway-cost-control") }) t.Run("MissingGroupIDs", func(t *testing.T) { @@ -3267,6 +3272,25 @@ func TestOrganizationGroupsAISpend(t *testing.T) { require.Equal(t, http.StatusBadRequest, sdkErr.StatusCode()) }) + t.Run("InclusiveMaxGroupIDs", func(t *testing.T) { + t.Parallel() + + adminClient, _, group := setupAICostControlTest(t, aiCostControlTestOptions{GroupName: "inclusive-max-group-ids-group"}) + ctx := testutil.Context(t, testutil.WaitLong) + + // Given: 100 group_ids, exactly at the cap. + ids := make([]uuid.UUID, 100) + for i := range ids { + ids[i] = uuid.New() + } + + // When: querying spend. + _, err := adminClient.OrganizationGroupsAISpend(ctx, group.OrganizationID, ids) + + // Then: request succeeds. + require.NoError(t, err) + }) + t.Run("TooManyGroupIDs", func(t *testing.T) { t.Parallel() From 582345bd8dc2b665faaf028f36d1786f283f301c Mon Sep 17 00:00:00 2001 From: Susana Cardoso Ferreira Date: Thu, 16 Jul 2026 09:44:01 +0000 Subject: [PATCH 3/5] chore: minor improvements --- coderd/apidoc/docs.go | 2 +- coderd/apidoc/swagger.json | 2 +- codersdk/aibridge.go | 11 ++++++----- docs/reference/api/enterprise.md | 4 +++- enterprise/coderd/aibridge.go | 4 +++- 5 files changed, 14 insertions(+), 9 deletions(-) diff --git a/coderd/apidoc/docs.go b/coderd/apidoc/docs.go index bcc1fb240560a..d3a9660ce60c2 100644 --- a/coderd/apidoc/docs.go +++ b/coderd/apidoc/docs.go @@ -4754,7 +4754,7 @@ const docTemplate = `{ }, "/api/v2/organizations/{organization}/groups/ai/spend": { "get": { - "description": "Returns AI spend limits and aggregate spend for the requested groups. Unknown or unreadable group IDs are silently omitted.", + "description": "Returns AI spend limits and aggregate spend for the requested groups.\nA maximum of 100 group IDs may be requested per call, and requests with more are rejected, so callers are expected to batch across multiple requests.\nUnknown or unreadable group IDs are silently omitted.", "produces": [ "application/json" ], diff --git a/coderd/apidoc/swagger.json b/coderd/apidoc/swagger.json index 9793f9f18dd71..30881bdc5bfc7 100644 --- a/coderd/apidoc/swagger.json +++ b/coderd/apidoc/swagger.json @@ -4195,7 +4195,7 @@ }, "/api/v2/organizations/{organization}/groups/ai/spend": { "get": { - "description": "Returns AI spend limits and aggregate spend for the requested groups. Unknown or unreadable group IDs are silently omitted.", + "description": "Returns AI spend limits and aggregate spend for the requested groups.\nA maximum of 100 group IDs may be requested per call, and requests with more are rejected, so callers are expected to batch across multiple requests.\nUnknown or unreadable group IDs are silently omitted.", "produces": ["application/json"], "tags": ["Enterprise"], "summary": "Get organization groups AI spend", diff --git a/codersdk/aibridge.go b/codersdk/aibridge.go index 6b845addbcb80..7cec7fbfa3116 100644 --- a/codersdk/aibridge.go +++ b/codersdk/aibridge.go @@ -10,6 +10,8 @@ import ( "github.com/google/uuid" "golang.org/x/xerrors" + + "github.com/coder/coder/v2/coderd/util/slice" ) // AIBudgetLimitSource identifies which tier produced the user's @@ -471,12 +473,11 @@ func (c *Client) UserAISpendStatus(ctx context.Context, user uuid.UUID) (UserAIS } // OrganizationGroupsAISpend returns AI spend for the given groups within the -// organization for the active budget period. +// organization for the active budget period. At most 100 group IDs may be +// requested per call, and callers with more groups are expected to batch +// across multiple requests. func (c *Client) OrganizationGroupsAISpend(ctx context.Context, organization uuid.UUID, groupIDs []uuid.UUID) (OrganizationGroupsAISpend, error) { - ids := make([]string, len(groupIDs)) - for i, id := range groupIDs { - ids[i] = id.String() - } + ids := slice.List(groupIDs, func(id uuid.UUID) string { return id.String() }) res, err := c.Request(ctx, http.MethodGet, fmt.Sprintf("/api/v2/organizations/%s/groups/ai/spend", organization.String()), nil, diff --git a/docs/reference/api/enterprise.md b/docs/reference/api/enterprise.md index 6657fe0a95ba3..646774a14fd8c 100644 --- a/docs/reference/api/enterprise.md +++ b/docs/reference/api/enterprise.md @@ -1852,7 +1852,9 @@ curl -X GET http://coder-server:8080/api/v2/organizations/{organization}/groups/ `GET /api/v2/organizations/{organization}/groups/ai/spend` -Returns AI spend limits and aggregate spend for the requested groups. Unknown or unreadable group IDs are silently omitted. +Returns AI spend limits and aggregate spend for the requested groups. +A maximum of 100 group IDs may be requested per call, and requests with more are rejected, so callers are expected to batch across multiple requests. +Unknown or unreadable group IDs are silently omitted. ### Parameters diff --git a/enterprise/coderd/aibridge.go b/enterprise/coderd/aibridge.go index c29dc29a7bab3..15f0a1ba023fb 100644 --- a/enterprise/coderd/aibridge.go +++ b/enterprise/coderd/aibridge.go @@ -950,7 +950,9 @@ func (api *API) userAISpendStatus(rw http.ResponseWriter, r *http.Request) { } // @Summary Get organization groups AI spend -// @Description Returns AI spend limits and aggregate spend for the requested groups. Unknown or unreadable group IDs are silently omitted. +// @Description Returns AI spend limits and aggregate spend for the requested groups. +// @Description A maximum of 100 group IDs may be requested per call, and requests with more are rejected, so callers are expected to batch across multiple requests. +// @Description Unknown or unreadable group IDs are silently omitted. // @ID get-organization-groups-ai-spend // @Security CoderSessionToken // @Produce json From 0d677e4ea5668ad0c18d5b0518ae42070460a331 Mon Sep 17 00:00:00 2001 From: Susana Cardoso Ferreira Date: Thu, 16 Jul 2026 10:30:37 +0000 Subject: [PATCH 4/5] fix: make gen --- docs/reference/api/enterprise.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/reference/api/enterprise.md b/docs/reference/api/enterprise.md index 646774a14fd8c..5497703829c8f 100644 --- a/docs/reference/api/enterprise.md +++ b/docs/reference/api/enterprise.md @@ -1843,7 +1843,7 @@ To perform this operation, you must be authenticated. [Learn more](authenticatio ### Code samples -```shell +```sh # Example request using curl curl -X GET http://coder-server:8080/api/v2/organizations/{organization}/groups/ai/spend?group_ids=string \ -H 'Accept: application/json' \ From 50a654db3a71dffa8b9a585bc1a734c301d9964b Mon Sep 17 00:00:00 2001 From: Susana Cardoso Ferreira Date: Thu, 16 Jul 2026 13:16:35 +0000 Subject: [PATCH 5/5] chore: minor improvements --- coderd/database/db2sdk/db2sdk.go | 11 +++++++++++ enterprise/coderd/aibridge.go | 9 +-------- 2 files changed, 12 insertions(+), 8 deletions(-) diff --git a/coderd/database/db2sdk/db2sdk.go b/coderd/database/db2sdk/db2sdk.go index 4754bbbe9506b..9a510674ea7cc 100644 --- a/coderd/database/db2sdk/db2sdk.go +++ b/coderd/database/db2sdk/db2sdk.go @@ -1458,6 +1458,17 @@ func UserAIBudgetOverride(o database.UserAIBudgetOverride) codersdk.UserAIBudget } } +func OrganizationGroupAISpend(row database.GetOrganizationGroupsAISpendRow) codersdk.OrganizationGroupAISpend { + group := codersdk.OrganizationGroupAISpend{ + GroupID: row.GroupID, + CurrentSpendMicros: row.CurrentSpendMicros, + } + if row.SpendLimitMicros.Valid { + group.SpendLimitMicros = &row.SpendLimitMicros.Int64 + } + return group +} + func InvalidatedPresets(invalidatedPresets []database.UpdatePresetsLastInvalidatedAtRow) []codersdk.InvalidatedPreset { var presets []codersdk.InvalidatedPreset for _, p := range invalidatedPresets { diff --git a/enterprise/coderd/aibridge.go b/enterprise/coderd/aibridge.go index 15f0a1ba023fb..d128a82aa870a 100644 --- a/enterprise/coderd/aibridge.go +++ b/enterprise/coderd/aibridge.go @@ -1017,14 +1017,7 @@ func (api *API) organizationGroupsAISpend(rw http.ResponseWriter, r *http.Reques Groups: make([]codersdk.OrganizationGroupAISpend, 0, len(rows)), } for _, row := range rows { - entry := codersdk.OrganizationGroupAISpend{ - GroupID: row.GroupID, - CurrentSpendMicros: row.CurrentSpendMicros, - } - if row.SpendLimitMicros.Valid { - entry.SpendLimitMicros = &row.SpendLimitMicros.Int64 - } - resp.Groups = append(resp.Groups, entry) + resp.Groups = append(resp.Groups, db2sdk.OrganizationGroupAISpend(row)) } httpapi.Write(ctx, rw, http.StatusOK, resp)