feat: add GET /organizations/{org}/groups/ai/spend (#27123)

## Description

Adds `GET /api/v2/organizations/{org}/groups/ai/spend?group_ids=...` to return per-group AI spend and configured limits for a set of groups in an organization.

In the UI, this endpoint is used alongside the existing `/api/v2/organizations/{org}/groups` endpoint. AI spend data is kept separate from that endpoint so that:

- Different concepts stay on different endpoints: identity (groups) vs. cost control (spend). Cost control is an additional feature layered on top of groups/orgs.
- Callers that don't need spend information don't pay for its computation.

UI flow:

1. Request `/api/v2/organizations/{org}/groups` → returns the organization's groups.
2. Request `/api/v2/organizations/{org}/groups/ai/spend?group_ids=...` with the IDs from step 1.

The groups endpoint from 1) is currently not paginated, but if pagination is added later, this design keeps the two responses in sync. This spend endpoint intentionally takes `group_ids` rather than paginating on its own, since it depends on the group set from step 1. Pagination could be added in the future, especially for Cost Control-focused pages.

<img width="2880" height="1460" alt="image" src="https://github.com/user-attachments/assets/ea83b74d-6a4f-45a6-af2f-1024e019da07" />

## Changes

- Add `codersdk.OrganizationGroupsAISpend` and `OrganizationGroupAISpend` types, plus a shared `AISpendPeriodWindow` embedded in the spend response.
- Add `GetOrganizationGroupsAISpend` SQL query with a dbauthz per-row filter that mirrors `GET /organizations/{org}/groups`.
- Add handler and route under `/organizations/{organization}/groups/ai/spend` with a required `group_ids` query param (cap 100). Callers with more than 100 groups are expected to batch across multiple requests.
- Add codersdk client method.
- Tests: dbauthz, raw SQL, endpoint, and role-access.

Closes https://linear.app/codercom/issue/AIGOV-466/backend-organization-groups-endpoint-with-groups-spend

> [!NOTE]
> Initially generated by Claude Opus 4.7, modified and reviewed by @ssncferreira
This commit is contained in:
Susana Ferreira
2026-07-20 12:54:52 +01:00
committed by GitHub
parent b511a68ab0
commit 2adc8f5272
19 changed files with 1216 additions and 23 deletions
+81
View File
@@ -4752,6 +4752,49 @@ const docTemplate = `{
]
}
},
"/api/v2/organizations/{organization}/groups/ai/spend": {
"get": {
"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",
"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 (maximum 100)",
"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 +21667,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": {
+77
View File
@@ -4193,6 +4193,45 @@
]
}
},
"/api/v2/organizations/{organization}/groups/ai/spend": {
"get": {
"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",
"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 (maximum 100)",
"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 +19770,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": {
+11
View File
@@ -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 {
+4
View File
@@ -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.
+16
View File
@@ -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})
+8
View File
@@ -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)
+15
View File
@@ -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()
+4
View File
@@ -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())
}
+5
View File
@@ -651,6 +651,11 @@ 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.
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)
+306
View File
@@ -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() {
+62
View File
@@ -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
ORDER BY groups.id
`
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.
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
+20
View File
@@ -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
WHERE groups.organization_id = @organization_id
AND groups.id = ANY(@group_ids::uuid[])
GROUP BY groups.id, budget.spend_limit_micros
ORDER BY groups.id;
+61 -7
View File
@@ -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
@@ -40,13 +42,9 @@ type UserAIBudgetSummary struct {
LimitSource *AIBudgetLimitSource `json:"limit_source"`
}
// UserAISpendStatus is the current AI spend snapshot for a user within
// the active budget period.
type UserAISpendStatus struct {
UserAIBudgetSummary
// CurrentSpendMicros is the user's spend on their effective group over
// the current budget period.
CurrentSpendMicros int64 `json:"current_spend_micros"`
// 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"`
@@ -55,6 +53,35 @@ type UserAISpendStatus struct {
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"`
}
// 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.
CurrentSpendMicros int64 `json:"current_spend_micros"`
}
type AIBridgeSession struct {
ID string `json:"id"`
Initiator MinimalUser `json:"initiator"`
@@ -444,3 +471,30 @@ 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. 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 := 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,
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)
}
+50
View File
@@ -1839,6 +1839,56 @@ 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
```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' \
-H 'Coder-Session-Token: API_KEY'
```
`GET /api/v2/organizations/{organization}/groups/ai/spend`
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
| 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
> 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
+42
View File
@@ -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
+88 -5
View File
@@ -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
@@ -877,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
@@ -890,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)
@@ -914,8 +921,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 +948,77 @@ func (api *API) userAISpendStatus(rw http.ResponseWriter, r *http.Request) {
httpapi.Write(ctx, rw, http.StatusOK, resp)
}
// @Summary Get organization groups AI spend
// @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
// @Tags Enterprise
// @Param organization path string true "Organization ID" format(uuid)
// @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) {
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
}
periodWindow, err := api.currentAIBudgetWindow()
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 {
resp.Groups = append(resp.Groups, db2sdk.OrganizationGroupAISpend(row))
}
httpapi.Write(ctx, rw, http.StatusOK, resp)
}
+308
View File
@@ -3206,6 +3206,314 @@ 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{
codersdk.FeatureTemplateRBAC: 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())
require.Contains(t, sdkErr.Message, "AI Gateway is a Premium feature")
})
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.FeatureTemplateRBAC: 1,
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())
require.Contains(t, sdkErr.Message, "ai-gateway-cost-control")
})
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("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()
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).
+9
View File
@@ -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),
+49 -11
View File
@@ -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