From 60ee59da05920041e2b82a5ea49968ea0abea5ee Mon Sep 17 00:00:00 2001 From: Susana Cardoso Ferreira Date: Thu, 9 Jul 2026 16:43:56 +0000 Subject: [PATCH 1/8] feat: add GET /groups/{group}/members/ai/spend --- coderd/apidoc/docs.go | 144 ++++++ coderd/apidoc/swagger.json | 136 ++++++ 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 | 11 + coderd/database/querier_test.go | 521 ++++++++++++++++++++++ coderd/database/queries.sql.go | 130 ++++++ coderd/database/queries/aicostcontrol.sql | 84 ++++ codersdk/aibridge.go | 56 +++ docs/reference/api/enterprise.md | 101 +++++ docs/reference/api/schemas.md | 48 ++ enterprise/coderd/aibridge.go | 103 +++++ enterprise/coderd/aibridge_test.go | 466 +++++++++++++++++++ enterprise/coderd/coderd.go | 18 + site/src/api/typesGenerated.ts | 41 ++ 18 files changed, 1906 insertions(+) diff --git a/coderd/apidoc/docs.go b/coderd/apidoc/docs.go index d3a9660ce60c2..9ccbaf280f2a6 100644 --- a/coderd/apidoc/docs.go +++ b/coderd/apidoc/docs.go @@ -3384,6 +3384,49 @@ const docTemplate = `{ ] } }, + "/api/v2/groups/{group}/members/ai/spend": { + "get": { + "description": "Returns aggregate AI spend attributed to the group per requested user. User IDs that are not members of the group, or that the caller has no read access to, are silently omitted.", + "produces": [ + "application/json" + ], + "tags": [ + "Enterprise" + ], + "summary": "Get group members AI spend", + "operationId": "get-group-members-ai-spend", + "parameters": [ + { + "type": "string", + "format": "uuid", + "description": "Group ID", + "name": "group", + "in": "path", + "required": true + }, + { + "type": "string", + "description": "Comma-separated list of user IDs (maximum 100)", + "name": "user_ids", + "in": "query", + "required": true + } + ], + "responses": { + "200": { + "description": "OK", + "schema": { + "$ref": "#/definitions/codersdk.GroupMembersAISpend" + } + } + }, + "security": [ + { + "CoderSessionToken": [] + } + ] + } + }, "/api/v2/init-script/{os}/{arch}": { "get": { "produces": [ @@ -4904,6 +4947,56 @@ const docTemplate = `{ ] } }, + "/api/v2/organizations/{organization}/groups/{groupName}/members/ai/spend": { + "get": { + "description": "Returns aggregate AI spend attributed to the group per requested user. User IDs that are not members of the group, or that the caller has no read access to, are silently omitted.", + "produces": [ + "application/json" + ], + "tags": [ + "Enterprise" + ], + "summary": "Get group members AI spend by organization", + "operationId": "get-group-members-ai-spend-by-organization", + "parameters": [ + { + "type": "string", + "format": "uuid", + "description": "Organization ID", + "name": "organization", + "in": "path", + "required": true + }, + { + "type": "string", + "description": "Group name", + "name": "groupName", + "in": "path", + "required": true + }, + { + "type": "string", + "description": "Comma-separated list of user IDs (maximum 100)", + "name": "user_ids", + "in": "query", + "required": true + } + ], + "responses": { + "200": { + "description": "OK", + "schema": { + "$ref": "#/definitions/codersdk.GroupMembersAISpend" + } + } + }, + "security": [ + { + "CoderSessionToken": [] + } + ] + } + }, "/api/v2/organizations/{organization}/members": { "get": { "produces": [ @@ -20258,6 +20351,57 @@ const docTemplate = `{ } } }, + "codersdk.GroupMemberAISpend": { + "type": "object", + "properties": { + "effective_group_id": { + "description": "EffectiveGroupID is the user's effective budget group within the queried\ngroup's organization. Null when no effective budget group is visible in\nthis organization, including when the user's budget resolves to a group\nin another organization.", + "type": "string", + "format": "uuid" + }, + "group_spend_micros": { + "description": "GroupSpendMicros is the user's spend attributed to the queried group\nover the current budget period.", + "type": "integer" + }, + "limit_source": { + "description": "LimitSource identifies the tier that produced the limit. Null when the\nuser's budget resolves to another group.", + "allOf": [ + { + "$ref": "#/definitions/codersdk.AIBudgetLimitSource" + } + ] + }, + "spend_limit_micros": { + "description": "SpendLimitMicros is the spend limit when the queried group is this\nuser's effective budget source. Null when the user's budget resolves to\nanother group.", + "type": "integer" + }, + "user_id": { + "type": "string", + "format": "uuid" + } + } + }, + "codersdk.GroupMembersAISpend": { + "type": "object", + "properties": { + "members": { + "type": "array", + "items": { + "$ref": "#/definitions/codersdk.GroupMemberAISpend" + } + }, + "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.GroupMembersResponse": { "type": "object", "properties": { diff --git a/coderd/apidoc/swagger.json b/coderd/apidoc/swagger.json index 30881bdc5bfc7..b984e43951ed3 100644 --- a/coderd/apidoc/swagger.json +++ b/coderd/apidoc/swagger.json @@ -2987,6 +2987,45 @@ ] } }, + "/api/v2/groups/{group}/members/ai/spend": { + "get": { + "description": "Returns aggregate AI spend attributed to the group per requested user. User IDs that are not members of the group, or that the caller has no read access to, are silently omitted.", + "produces": ["application/json"], + "tags": ["Enterprise"], + "summary": "Get group members AI spend", + "operationId": "get-group-members-ai-spend", + "parameters": [ + { + "type": "string", + "format": "uuid", + "description": "Group ID", + "name": "group", + "in": "path", + "required": true + }, + { + "type": "string", + "description": "Comma-separated list of user IDs (maximum 100)", + "name": "user_ids", + "in": "query", + "required": true + } + ], + "responses": { + "200": { + "description": "OK", + "schema": { + "$ref": "#/definitions/codersdk.GroupMembersAISpend" + } + } + }, + "security": [ + { + "CoderSessionToken": [] + } + ] + } + }, "/api/v2/init-script/{os}/{arch}": { "get": { "produces": ["text/plain"], @@ -4333,6 +4372,52 @@ ] } }, + "/api/v2/organizations/{organization}/groups/{groupName}/members/ai/spend": { + "get": { + "description": "Returns aggregate AI spend attributed to the group per requested user. User IDs that are not members of the group, or that the caller has no read access to, are silently omitted.", + "produces": ["application/json"], + "tags": ["Enterprise"], + "summary": "Get group members AI spend by organization", + "operationId": "get-group-members-ai-spend-by-organization", + "parameters": [ + { + "type": "string", + "format": "uuid", + "description": "Organization ID", + "name": "organization", + "in": "path", + "required": true + }, + { + "type": "string", + "description": "Group name", + "name": "groupName", + "in": "path", + "required": true + }, + { + "type": "string", + "description": "Comma-separated list of user IDs (maximum 100)", + "name": "user_ids", + "in": "query", + "required": true + } + ], + "responses": { + "200": { + "description": "OK", + "schema": { + "$ref": "#/definitions/codersdk.GroupMembersAISpend" + } + } + }, + "security": [ + { + "CoderSessionToken": [] + } + ] + } + }, "/api/v2/organizations/{organization}/members": { "get": { "produces": ["application/json"], @@ -18436,6 +18521,57 @@ } } }, + "codersdk.GroupMemberAISpend": { + "type": "object", + "properties": { + "effective_group_id": { + "description": "EffectiveGroupID is the user's effective budget group within the queried\ngroup's organization. Null when no effective budget group is visible in\nthis organization, including when the user's budget resolves to a group\nin another organization.", + "type": "string", + "format": "uuid" + }, + "group_spend_micros": { + "description": "GroupSpendMicros is the user's spend attributed to the queried group\nover the current budget period.", + "type": "integer" + }, + "limit_source": { + "description": "LimitSource identifies the tier that produced the limit. Null when the\nuser's budget resolves to another group.", + "allOf": [ + { + "$ref": "#/definitions/codersdk.AIBudgetLimitSource" + } + ] + }, + "spend_limit_micros": { + "description": "SpendLimitMicros is the spend limit when the queried group is this\nuser's effective budget source. Null when the user's budget resolves to\nanother group.", + "type": "integer" + }, + "user_id": { + "type": "string", + "format": "uuid" + } + } + }, + "codersdk.GroupMembersAISpend": { + "type": "object", + "properties": { + "members": { + "type": "array", + "items": { + "$ref": "#/definitions/codersdk.GroupMemberAISpend" + } + }, + "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.GroupMembersResponse": { "type": "object", "properties": { diff --git a/coderd/database/dbauthz/dbauthz.go b/coderd/database/dbauthz/dbauthz.go index 48c26244c17e7..3161f3c3e9db0 100644 --- a/coderd/database/dbauthz/dbauthz.go +++ b/coderd/database/dbauthz/dbauthz.go @@ -3883,6 +3883,10 @@ func (q *querier) GetGroupMembers(ctx context.Context, includeSystem bool) ([]da return q.db.GetGroupMembers(ctx, includeSystem) } +func (q *querier) GetGroupMembersAISpend(ctx context.Context, arg database.GetGroupMembersAISpendParams) ([]database.GetGroupMembersAISpendRow, error) { + return fetchWithPostFilter(q.auth, policy.ActionRead, q.db.GetGroupMembersAISpend)(ctx, arg) +} + func (q *querier) GetGroupMembersByGroupID(ctx context.Context, arg database.GetGroupMembersByGroupIDParams) ([]database.GroupMember, error) { return fetchWithPostFilter(q.auth, policy.ActionRead, q.db.GetGroupMembersByGroupID)(ctx, arg) } diff --git a/coderd/database/dbauthz/dbauthz_test.go b/coderd/database/dbauthz/dbauthz_test.go index dd62a2fb16bb9..c029cd3761146 100644 --- a/coderd/database/dbauthz/dbauthz_test.go +++ b/coderd/database/dbauthz/dbauthz_test.go @@ -6954,6 +6954,22 @@ func (s *MethodTestSuite) TestAIBridge() { Returns([]database.GetOrganizationGroupsAISpendRow{row1, row2}) })) + s.Run("GetGroupMembersAISpend", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { + group := testutil.Fake(s.T(), faker, database.Group{}) + row1 := testutil.Fake(s.T(), faker, database.GetGroupMembersAISpendRow{OrganizationID: group.OrganizationID}) + row2 := testutil.Fake(s.T(), faker, database.GetGroupMembersAISpendRow{OrganizationID: group.OrganizationID}) + arg := database.GetGroupMembersAISpendParams{ + GroupID: group.ID, + UserIds: []uuid.UUID{row1.UserID, row2.UserID}, + PeriodStart: time.Now().UTC().Truncate(24 * time.Hour), + } + dbm.EXPECT().GetGroupMembersAISpend(gomock.Any(), arg). + Return([]database.GetGroupMembersAISpendRow{row1, row2}, nil).AnyTimes() + check.Args(arg). + Asserts(row1, policy.ActionRead, row2, policy.ActionRead). + Returns([]database.GetGroupMembersAISpendRow{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 6225be9058d87..57c742457f0e4 100644 --- a/coderd/database/dbmetrics/querymetrics.go +++ b/coderd/database/dbmetrics/querymetrics.go @@ -2161,6 +2161,14 @@ func (m queryMetricsStore) GetGroupMembers(ctx context.Context, includeSystem bo return r0, r1 } +func (m queryMetricsStore) GetGroupMembersAISpend(ctx context.Context, arg database.GetGroupMembersAISpendParams) ([]database.GetGroupMembersAISpendRow, error) { + start := time.Now() + r0, r1 := m.s.GetGroupMembersAISpend(ctx, arg) + m.queryLatencies.WithLabelValues("GetGroupMembersAISpend").Observe(time.Since(start).Seconds()) + m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "GetGroupMembersAISpend").Inc() + return r0, r1 +} + func (m queryMetricsStore) GetGroupMembersByGroupID(ctx context.Context, arg database.GetGroupMembersByGroupIDParams) ([]database.GroupMember, error) { start := time.Now() r0, r1 := m.s.GetGroupMembersByGroupID(ctx, arg) diff --git a/coderd/database/dbmock/dbmock.go b/coderd/database/dbmock/dbmock.go index 1c403c54f40f7..9103654b6d3e2 100644 --- a/coderd/database/dbmock/dbmock.go +++ b/coderd/database/dbmock/dbmock.go @@ -4003,6 +4003,21 @@ func (mr *MockStoreMockRecorder) GetGroupMembers(ctx, includeSystem any) *gomock return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetGroupMembers", reflect.TypeOf((*MockStore)(nil).GetGroupMembers), ctx, includeSystem) } +// GetGroupMembersAISpend mocks base method. +func (m *MockStore) GetGroupMembersAISpend(ctx context.Context, arg database.GetGroupMembersAISpendParams) ([]database.GetGroupMembersAISpendRow, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "GetGroupMembersAISpend", ctx, arg) + ret0, _ := ret[0].([]database.GetGroupMembersAISpendRow) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// GetGroupMembersAISpend indicates an expected call of GetGroupMembersAISpend. +func (mr *MockStoreMockRecorder) GetGroupMembersAISpend(ctx, arg any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetGroupMembersAISpend", reflect.TypeOf((*MockStore)(nil).GetGroupMembersAISpend), ctx, arg) +} + // GetGroupMembersByGroupID mocks base method. func (m *MockStore) GetGroupMembersByGroupID(ctx context.Context, arg database.GetGroupMembersByGroupIDParams) ([]database.GroupMember, error) { m.ctrl.T.Helper() diff --git a/coderd/database/modelmethods.go b/coderd/database/modelmethods.go index 1fba7b4d99ee4..76ca27166bad0 100644 --- a/coderd/database/modelmethods.go +++ b/coderd/database/modelmethods.go @@ -470,6 +470,10 @@ func (gm GetGroupMembersByGroupIDPaginatedRow) RBACObject() rbac.Object { return rbac.ResourceGroupMember.WithID(gm.UserID).InOrg(gm.OrganizationID).WithOwner(gm.UserID.String()) } +func (r GetGroupMembersAISpendRow) RBACObject() rbac.Object { + return rbac.ResourceGroupMember.WithID(r.UserID).InOrg(r.OrganizationID).WithOwner(r.UserID.String()) +} + // PrebuiltWorkspaceResource defines the interface for types that can be identified as prebuilt workspaces // and converted to their corresponding prebuilt workspace RBAC object. type PrebuiltWorkspaceResource interface { diff --git a/coderd/database/querier.go b/coderd/database/querier.go index b9bb106391310..c311608456c51 100644 --- a/coderd/database/querier.go +++ b/coderd/database/querier.go @@ -577,6 +577,17 @@ type sqlcQuerier interface { GetGroupByID(ctx context.Context, id uuid.UUID) (Group, error) GetGroupByOrgAndName(ctx context.Context, arg GetGroupByOrgAndNameParams) (Group, error) GetGroupMembers(ctx context.Context, includeSystem bool) ([]GroupMember, error) + // Returns each user's AI spend attributed to the queried group, on or after + // period_start until NOW. Only current members of the queried group are + // returned. spend_limit_micros and limit_source are populated only when the + // queried group is the user's effective budget source. The effective_group_id + // is null when the user has no configured budget or when the effective group + // belongs to a different organization than the queried group. + // The period_start parameter is normalized to its UTC calendar day. + // Spend is aggregated for the queried group, not the user's effective group. + // A LEFT JOIN leaves spend_limit_micros and limit_source null for users + // whose effective budget source is not the queried group. + GetGroupMembersAISpend(ctx context.Context, arg GetGroupMembersAISpendParams) ([]GetGroupMembersAISpendRow, error) GetGroupMembersByGroupID(ctx context.Context, arg GetGroupMembersByGroupIDParams) ([]GroupMember, error) GetGroupMembersByGroupIDPaginated(ctx context.Context, arg GetGroupMembersByGroupIDPaginatedParams) ([]GetGroupMembersByGroupIDPaginatedRow, error) // Returns the total count of members in a group. Shows the total diff --git a/coderd/database/querier_test.go b/coderd/database/querier_test.go index 38c99dd7e0e82..210445999177d 100644 --- a/coderd/database/querier_test.go +++ b/coderd/database/querier_test.go @@ -13106,6 +13106,527 @@ func TestGetOrganizationGroupsAISpend(t *testing.T) { }) } +func TestGetGroupMembersAISpend(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 + + t.Run("NoBudgetNoSpend", func(t *testing.T) { + t.Parallel() + db, _ := dbtestutil.NewDB(t) + ctx := testutil.Context(t, testutil.WaitShort) + + // Given: a group member with no budget and no 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}) + dbgen.OrganizationMember(t, db, database.OrganizationMember{UserID: user.ID, OrganizationID: org.ID}) + dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: group.ID, UserID: user.ID}) + + // When: querying spend for the user. + got, err := db.GetGroupMembersAISpend(ctx, database.GetGroupMembersAISpendParams{ + GroupID: group.ID, + UserIds: []uuid.UUID{user.ID}, + PeriodStart: monthStart, + }) + require.NoError(t, err) + + // Then: one row with no effective group and zero spend. + require.Len(t, got, 1) + require.Equal(t, user.ID, got[0].UserID) + require.Equal(t, org.ID, got[0].OrganizationID) + require.False(t, got[0].EffectiveGroupID.Valid, "expected no effective group") + require.Equal(t, int64(0), got[0].GroupSpendMicros) + }) + + t.Run("GroupBudget", func(t *testing.T) { + t.Parallel() + db, _ := dbtestutil.NewDB(t) + ctx := testutil.Context(t, testutil.WaitShort) + + // Given: a user in the queried group and the queried group has a budget. + user := dbgen.User(t, db, database.User{}) + org := dbgen.Organization(t, db, database.Organization{}) + group := dbgen.Group(t, db, database.Group{OrganizationID: org.ID}) + dbgen.OrganizationMember(t, db, database.OrganizationMember{UserID: user.ID, OrganizationID: org.ID}) + dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: group.ID, UserID: user.ID}) + _, err := db.UpsertGroupAIBudget(ctx, database.UpsertGroupAIBudgetParams{ + GroupID: group.ID, + SpendLimitMicros: 1_000_000, + }) + require.NoError(t, err) + + // When: querying spend for the user. + got, err := db.GetGroupMembersAISpend(ctx, database.GetGroupMembersAISpendParams{ + GroupID: group.ID, + UserIds: []uuid.UUID{user.ID}, + PeriodStart: monthStart, + }) + require.NoError(t, err) + + // Then: the effective group is the queried group, and the limit and + // source reflect the group budget. + require.Len(t, got, 1) + require.Equal(t, uuid.NullUUID{UUID: group.ID, Valid: true}, got[0].EffectiveGroupID) + require.Equal(t, sql.NullInt64{Int64: 1_000_000, Valid: true}, got[0].SpendLimitMicros) + require.Equal(t, sql.NullString{String: "group", Valid: true}, got[0].LimitSource) + }) + + t.Run("OverrideOnQueriedGroup", func(t *testing.T) { + t.Parallel() + db, _ := dbtestutil.NewDB(t) + ctx := testutil.Context(t, testutil.WaitShort) + + // Given: a group member with a user override targeting the queried group. + user := dbgen.User(t, db, database.User{}) + org := dbgen.Organization(t, db, database.Organization{}) + group := dbgen.Group(t, db, database.Group{OrganizationID: org.ID}) + dbgen.OrganizationMember(t, db, database.OrganizationMember{UserID: user.ID, OrganizationID: org.ID}) + dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: group.ID, UserID: user.ID}) + _, err := db.UpsertUserAIBudgetOverride(ctx, database.UpsertUserAIBudgetOverrideParams{ + UserID: user.ID, + GroupID: group.ID, + SpendLimitMicros: 500_000, + }) + require.NoError(t, err) + + // When: querying spend for the user. + got, err := db.GetGroupMembersAISpend(ctx, database.GetGroupMembersAISpendParams{ + GroupID: group.ID, + UserIds: []uuid.UUID{user.ID}, + PeriodStart: monthStart, + }) + require.NoError(t, err) + + // Then: the limit comes from the override, with the user_override source. + require.Len(t, got, 1) + require.Equal(t, uuid.NullUUID{UUID: group.ID, Valid: true}, got[0].EffectiveGroupID) + require.Equal(t, sql.NullInt64{Int64: 500_000, Valid: true}, got[0].SpendLimitMicros) + require.Equal(t, sql.NullString{String: "user_override", Valid: true}, got[0].LimitSource) + }) + + t.Run("Spend", func(t *testing.T) { + t.Parallel() + db, _ := dbtestutil.NewDB(t) + ctx := testutil.Context(t, testutil.WaitShort) + + // Given: a group member with spend attributed to the queried group. + user := dbgen.User(t, db, database.User{}) + org := dbgen.Organization(t, db, database.Organization{}) + group := dbgen.Group(t, db, database.Group{OrganizationID: org.ID}) + dbgen.OrganizationMember(t, db, database.OrganizationMember{UserID: user.ID, OrganizationID: org.ID}) + dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: group.ID, UserID: user.ID}) + _, err := db.IncrementUserAIDailySpend(ctx, database.IncrementUserAIDailySpendParams{ + UserID: user.ID, EffectiveGroupID: group.ID, Day: now, CostMicros: 250, + }) + require.NoError(t, err) + + // When: querying spend for the user. + got, err := db.GetGroupMembersAISpend(ctx, database.GetGroupMembersAISpendParams{ + GroupID: group.ID, + UserIds: []uuid.UUID{user.ID}, + PeriodStart: monthStart, + }) + require.NoError(t, err) + + // Then: the queried-group spend is aggregated. + require.Len(t, got, 1) + require.Equal(t, int64(250), got[0].GroupSpendMicros) + }) + + t.Run("SpendWithDifferentEffectiveGroup", func(t *testing.T) { + t.Parallel() + db, _ := dbtestutil.NewDB(t) + ctx := testutil.Context(t, testutil.WaitShort) + + // Given: a group member with spend attributed to the queried group, + // whose current effective group is a different same-org group. + user := dbgen.User(t, db, database.User{}) + org := dbgen.Organization(t, db, database.Organization{}) + queried := dbgen.Group(t, db, database.Group{OrganizationID: org.ID}) + other := dbgen.Group(t, db, database.Group{OrganizationID: org.ID}) + dbgen.OrganizationMember(t, db, database.OrganizationMember{UserID: user.ID, OrganizationID: org.ID}) + dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: queried.ID, UserID: user.ID}) + dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: other.ID, UserID: user.ID}) + _, err := db.UpsertGroupAIBudget(ctx, database.UpsertGroupAIBudgetParams{ + GroupID: other.ID, + SpendLimitMicros: 1_000_000, + }) + require.NoError(t, err) + _, err = db.IncrementUserAIDailySpend(ctx, database.IncrementUserAIDailySpendParams{ + UserID: user.ID, EffectiveGroupID: queried.ID, Day: now, CostMicros: 250, + }) + require.NoError(t, err) + + // When: querying the queried group's spend for the user. + got, err := db.GetGroupMembersAISpend(ctx, database.GetGroupMembersAISpendParams{ + GroupID: queried.ID, + UserIds: []uuid.UUID{user.ID}, + PeriodStart: monthStart, + }) + require.NoError(t, err) + + // Then: queried-group spend is returned, effective_group_id is the + // other group, and the limit and source are null because the queried + // group is not the effective source. + require.Len(t, got, 1) + require.Equal(t, uuid.NullUUID{UUID: other.ID, Valid: true}, got[0].EffectiveGroupID) + require.False(t, got[0].SpendLimitMicros.Valid) + require.False(t, got[0].LimitSource.Valid) + require.Equal(t, int64(250), got[0].GroupSpendMicros) + }) + + t.Run("HidesFormerMember", func(t *testing.T) { + t.Parallel() + db, _ := dbtestutil.NewDB(t) + ctx := testutil.Context(t, testutil.WaitShort) + + // Given: a user who has historical spend attributed to the queried + // group but is not currently a member of it. + user := dbgen.User(t, db, database.User{}) + org := dbgen.Organization(t, db, database.Organization{}) + queried := dbgen.Group(t, db, database.Group{OrganizationID: org.ID}) + dbgen.OrganizationMember(t, db, database.OrganizationMember{UserID: user.ID, OrganizationID: org.ID}) + _, err := db.IncrementUserAIDailySpend(ctx, database.IncrementUserAIDailySpendParams{ + UserID: user.ID, EffectiveGroupID: queried.ID, Day: now, CostMicros: 500, + }) + require.NoError(t, err) + + // When: querying the queried group's spend for the user. + got, err := db.GetGroupMembersAISpend(ctx, database.GetGroupMembersAISpendParams{ + GroupID: queried.ID, + UserIds: []uuid.UUID{user.ID}, + PeriodStart: monthStart, + }) + require.NoError(t, err) + + // Then: the user is filtered out and their historical spend is not returned. + require.Empty(t, got) + }) + + t.Run("ExcludesOtherGroupSpend", func(t *testing.T) { + t.Parallel() + db, _ := dbtestutil.NewDB(t) + ctx := testutil.Context(t, testutil.WaitShort) + + // Given: a group member with spend attributed to a different group in the same org. + user := dbgen.User(t, db, database.User{}) + org := dbgen.Organization(t, db, database.Organization{}) + group := dbgen.Group(t, db, database.Group{OrganizationID: org.ID}) + otherGroup := dbgen.Group(t, db, database.Group{OrganizationID: org.ID}) + dbgen.OrganizationMember(t, db, database.OrganizationMember{UserID: user.ID, OrganizationID: org.ID}) + dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: group.ID, UserID: user.ID}) + _, err := db.IncrementUserAIDailySpend(ctx, database.IncrementUserAIDailySpendParams{ + UserID: user.ID, EffectiveGroupID: otherGroup.ID, Day: now, CostMicros: 500, + }) + require.NoError(t, err) + + // When: querying spend for the queried group. + got, err := db.GetGroupMembersAISpend(ctx, database.GetGroupMembersAISpendParams{ + GroupID: group.ID, + UserIds: []uuid.UUID{user.ID}, + PeriodStart: monthStart, + }) + require.NoError(t, err) + + // Then: spend attributed to the other group is not counted. + require.Len(t, got, 1) + require.Equal(t, int64(0), got[0].GroupSpendMicros) + }) + + t.Run("MultipleMembers", func(t *testing.T) { + t.Parallel() + db, _ := dbtestutil.NewDB(t) + ctx := testutil.Context(t, testutil.WaitShort) + + // Given: two group members with different spend attributed to the 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}) + dbgen.OrganizationMember(t, db, database.OrganizationMember{UserID: userA.ID, OrganizationID: org.ID}) + dbgen.OrganizationMember(t, db, database.OrganizationMember{UserID: userB.ID, OrganizationID: org.ID}) + dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: group.ID, UserID: userA.ID}) + dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: group.ID, UserID: userB.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: 250, + }) + require.NoError(t, err) + + // When: querying spend for both users. + got, err := db.GetGroupMembersAISpend(ctx, database.GetGroupMembersAISpendParams{ + GroupID: group.ID, + UserIds: []uuid.UUID{userA.ID, userB.ID}, + PeriodStart: monthStart, + }) + require.NoError(t, err) + + // Then: both users are returned with their own aggregate spend. + require.Len(t, got, 2) + byID := make(map[uuid.UUID]database.GetGroupMembersAISpendRow, len(got)) + for _, r := range got { + byID[r.UserID] = r + } + require.Equal(t, int64(100), byID[userA.ID].GroupSpendMicros) + require.Equal(t, int64(250), byID[userB.ID].GroupSpendMicros) + }) + + t.Run("ExcludesNonMembers", func(t *testing.T) { + t.Parallel() + db, _ := dbtestutil.NewDB(t) + ctx := testutil.Context(t, testutil.WaitShort) + + // Given: a member of the queried group and a non-member in the same org. + member := dbgen.User(t, db, database.User{}) + nonMember := dbgen.User(t, db, database.User{}) + org := dbgen.Organization(t, db, database.Organization{}) + group := dbgen.Group(t, db, database.Group{OrganizationID: org.ID}) + dbgen.OrganizationMember(t, db, database.OrganizationMember{UserID: member.ID, OrganizationID: org.ID}) + dbgen.OrganizationMember(t, db, database.OrganizationMember{UserID: nonMember.ID, OrganizationID: org.ID}) + dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: group.ID, UserID: member.ID}) + + // When: querying with both user IDs. + got, err := db.GetGroupMembersAISpend(ctx, database.GetGroupMembersAISpendParams{ + GroupID: group.ID, + UserIds: []uuid.UUID{member.ID, nonMember.ID}, + PeriodStart: monthStart, + }) + require.NoError(t, err) + + // Then: only the group member is returned. + require.Len(t, got, 1) + require.Equal(t, member.ID, got[0].UserID) + }) + + t.Run("CrossOrgEffectiveGroupMasked", func(t *testing.T) { + t.Parallel() + db, _ := dbtestutil.NewDB(t) + ctx := testutil.Context(t, testutil.WaitShort) + + // Given: a member of the queried group whose highest-limit budget + // group is in a different org (via cross-org group membership). + 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}) + dbgen.OrganizationMember(t, db, database.OrganizationMember{UserID: user.ID, OrganizationID: org.ID}) + dbgen.OrganizationMember(t, db, database.OrganizationMember{UserID: user.ID, OrganizationID: otherOrg.ID}) + dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: group.ID, UserID: user.ID}) + dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: otherOrgGroup.ID, UserID: user.ID}) + _, err := db.UpsertGroupAIBudget(ctx, database.UpsertGroupAIBudgetParams{ + GroupID: otherOrgGroup.ID, + SpendLimitMicros: 9_999_999, + }) + require.NoError(t, err) + + // When: querying spend for the user in the queried group's org. + got, err := db.GetGroupMembersAISpend(ctx, database.GetGroupMembersAISpendParams{ + GroupID: group.ID, + UserIds: []uuid.UUID{user.ID}, + PeriodStart: monthStart, + }) + require.NoError(t, err) + + // Then: effective_group_id is masked to NULL, the highest-limit group is cross-org. + require.Len(t, got, 1) + require.False(t, got[0].EffectiveGroupID.Valid, "cross-org effective group must be masked") + require.False(t, got[0].SpendLimitMicros.Valid) + require.False(t, got[0].LimitSource.Valid) + }) + + t.Run("OverrideWins", func(t *testing.T) { + t.Parallel() + db, _ := dbtestutil.NewDB(t) + ctx := testutil.Context(t, testutil.WaitShort) + + // Given: a member of the queried group who is also in two other + // budgeted groups, one being the natural highest-limit and the other + // set via a user override. + user := dbgen.User(t, db, database.User{}) + org := dbgen.Organization(t, db, database.Organization{}) + queried := dbgen.Group(t, db, database.Group{OrganizationID: org.ID}) + overrideTarget := dbgen.Group(t, db, database.Group{OrganizationID: org.ID}) + highestLimit := dbgen.Group(t, db, database.Group{OrganizationID: org.ID}) + dbgen.OrganizationMember(t, db, database.OrganizationMember{UserID: user.ID, OrganizationID: org.ID}) + dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: queried.ID, UserID: user.ID}) + dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: overrideTarget.ID, UserID: user.ID}) + dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: highestLimit.ID, UserID: user.ID}) + _, err := db.UpsertGroupAIBudget(ctx, database.UpsertGroupAIBudgetParams{ + GroupID: overrideTarget.ID, + SpendLimitMicros: 1_000_000, + }) + require.NoError(t, err) + _, err = db.UpsertGroupAIBudget(ctx, database.UpsertGroupAIBudgetParams{ + GroupID: highestLimit.ID, + SpendLimitMicros: 5_000_000, + }) + require.NoError(t, err) + _, err = db.UpsertUserAIBudgetOverride(ctx, database.UpsertUserAIBudgetOverrideParams{ + UserID: user.ID, + GroupID: overrideTarget.ID, + SpendLimitMicros: 500_000, + }) + require.NoError(t, err) + + // When: querying spend for the user. + got, err := db.GetGroupMembersAISpend(ctx, database.GetGroupMembersAISpendParams{ + GroupID: queried.ID, + UserIds: []uuid.UUID{user.ID}, + PeriodStart: monthStart, + }) + require.NoError(t, err) + + // Then: the override target wins over the highest-limit group. + require.Len(t, got, 1) + require.Equal(t, uuid.NullUUID{UUID: overrideTarget.ID, Valid: true}, got[0].EffectiveGroupID) + }) + + t.Run("EveryoneGroupCounts", func(t *testing.T) { + t.Parallel() + db, _ := dbtestutil.NewDB(t) + ctx := testutil.Context(t, testutil.WaitShort) + + // Given: a member of the queried group whose only budgeted group is + // the org's implicit Everyone group. + user := dbgen.User(t, db, database.User{}) + org := dbgen.Organization(t, db, database.Organization{}) + queried := dbgen.Group(t, db, database.Group{OrganizationID: org.ID}) + dbgen.OrganizationMember(t, db, database.OrganizationMember{UserID: user.ID, OrganizationID: org.ID}) + dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: queried.ID, UserID: user.ID}) + // The Everyone group has ID equal to the organization ID and must be + // inserted explicitly for this test's FK constraint on group_ai_budgets. + //nolint:gocritic // Requires system context. + _, err := db.InsertAllUsersGroup(dbauthz.AsSystemRestricted(ctx), org.ID) + require.NoError(t, err) + _, err = db.UpsertGroupAIBudget(ctx, database.UpsertGroupAIBudgetParams{ + GroupID: org.ID, + SpendLimitMicros: 1_000_000, + }) + require.NoError(t, err) + + // When: querying spend for the user. + got, err := db.GetGroupMembersAISpend(ctx, database.GetGroupMembersAISpendParams{ + GroupID: queried.ID, + UserIds: []uuid.UUID{user.ID}, + PeriodStart: monthStart, + }) + require.NoError(t, err) + + // Then: effective_group_id resolves to the Everyone group (org.ID). + require.Len(t, got, 1) + require.Equal(t, uuid.NullUUID{UUID: org.ID, Valid: true}, got[0].EffectiveGroupID) + }) + + t.Run("ExcludesSpendBeforePeriodStart", func(t *testing.T) { + t.Parallel() + db, _ := dbtestutil.NewDB(t) + ctx := testutil.Context(t, testutil.WaitShort) + + // Given: a group member with 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}) + dbgen.OrganizationMember(t, db, database.OrganizationMember{UserID: user.ID, OrganizationID: org.ID}) + dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: group.ID, UserID: user.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.GetGroupMembersAISpend(ctx, database.GetGroupMembersAISpendParams{ + GroupID: group.ID, + UserIds: []uuid.UUID{user.ID}, + PeriodStart: monthStart, + }) + require.NoError(t, err) + + // Then: only current-period spend is aggregated. + require.Len(t, got, 1) + require.Equal(t, int64(25), got[0].GroupSpendMicros) + }) + + t.Run("AggregatesSpendAcrossDays", func(t *testing.T) { + t.Parallel() + db, _ := dbtestutil.NewDB(t) + ctx := testutil.Context(t, testutil.WaitShort) + + // Given: a group member with spend on multiple days in the 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}) + dbgen.OrganizationMember(t, db, database.OrganizationMember{UserID: user.ID, OrganizationID: org.ID}) + dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: group.ID, UserID: user.ID}) + days := []time.Time{monthStart, monthStart.AddDate(0, 0, 7), monthStart.AddDate(0, 0, 14)} + for i, day := range days { + _, err := db.IncrementUserAIDailySpend(ctx, database.IncrementUserAIDailySpendParams{ + UserID: user.ID, EffectiveGroupID: group.ID, Day: day, CostMicros: int64((i + 1) * 100), + }) + require.NoError(t, err) + } + + // When: querying the group's member spend. + got, err := db.GetGroupMembersAISpend(ctx, database.GetGroupMembersAISpendParams{ + GroupID: group.ID, + UserIds: []uuid.UUID{user.ID}, + PeriodStart: monthStart, + }) + require.NoError(t, err) + + // Then: per-user spend is summed across all days in the period. + require.Len(t, got, 1) + require.Equal(t, int64(600), got[0].GroupSpendMicros) + }) + + t.Run("NormalizesNonUTCPeriodStart", func(t *testing.T) { + t.Parallel() + db, _ := dbtestutil.NewDB(t) + ctx := testutil.Context(t, testutil.WaitShort) + + // Given: a group member with spend on the prior UTC day and on 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}) + dbgen.OrganizationMember(t, db, database.OrganizationMember{UserID: user.ID, OrganizationID: org.ID}) + dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: group.ID, UserID: user.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.GetGroupMembersAISpend(ctx, database.GetGroupMembersAISpendParams{ + GroupID: group.ID, + UserIds: []uuid.UUID{user.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].GroupSpendMicros, + "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 ac81aac2075e1..ca628587591f2 100644 --- a/coderd/database/queries.sql.go +++ b/coderd/database/queries.sql.go @@ -2533,6 +2533,136 @@ func (q *sqlQuerier) GetGroupAIBudget(ctx context.Context, groupID uuid.UUID) (G return i, err } +const getGroupMembersAISpend = `-- name: GetGroupMembersAISpend :many +WITH queried_group AS ( + -- The queried group's org, used to detect cross-org effective groups. + SELECT id, organization_id + FROM groups + WHERE id = $1 +), +filtered_users AS ( + -- Users from @user_ids that are members of the queried group. Uses + -- group_members_expanded so the implicit Everyone group counts. + SELECT DISTINCT user_id + FROM group_members_expanded + WHERE group_id = $1 + AND user_id = ANY($3::uuid[]) +), +user_highest_group AS ( + -- Per user, the highest-limit group they belong to. Uses + -- group_members_expanded so the implicit Everyone group counts. + SELECT DISTINCT ON (member.user_id) + member.user_id, + budget.group_id, + budget.spend_limit_micros + FROM group_ai_budgets budget + JOIN group_members_expanded member ON member.group_id = budget.group_id + WHERE member.user_id IN (SELECT user_id FROM filtered_users) + ORDER BY member.user_id, budget.spend_limit_micros DESC, member.group_name ASC, budget.group_id ASC +), +effective AS ( + -- Effective budget per user: a per-user override wins over the + -- highest-limit group. + SELECT + filtered_users.user_id, + COALESCE(override.group_id, user_highest_group.group_id) AS effective_group_id, + COALESCE(override.spend_limit_micros, user_highest_group.spend_limit_micros) AS spend_limit_micros, + (CASE + WHEN override.group_id IS NOT NULL THEN 'user_override' + WHEN user_highest_group.group_id IS NOT NULL THEN 'group' + END)::text AS limit_source + FROM filtered_users + LEFT JOIN user_ai_budget_overrides override ON override.user_id = filtered_users.user_id + LEFT JOIN user_highest_group ON user_highest_group.user_id = filtered_users.user_id +), +applied_budget AS ( + -- The limit and source only for users whose effective budget source is the + -- queried group. + SELECT user_id, spend_limit_micros, limit_source + FROM effective + WHERE effective_group_id = $1 +) +SELECT + effective.user_id, + queried_group.organization_id, + effective_group.id AS effective_group_id, + applied_budget.spend_limit_micros, + applied_budget.limit_source, + COALESCE(SUM(spend.spend_micros), 0)::BIGINT AS group_spend_micros +FROM effective +CROSS JOIN queried_group +LEFT JOIN groups effective_group + ON effective_group.id = effective.effective_group_id + AND effective_group.organization_id = queried_group.organization_id +LEFT JOIN applied_budget ON applied_budget.user_id = effective.user_id +LEFT JOIN ai_user_daily_spend spend + ON spend.user_id = effective.user_id + AND spend.effective_group_id = $1 + AND spend.day >= (($2::timestamptz) AT TIME ZONE 'UTC')::date +GROUP BY + effective.user_id, + queried_group.organization_id, + effective_group.id, + applied_budget.spend_limit_micros, + applied_budget.limit_source +ORDER BY effective.user_id +` + +type GetGroupMembersAISpendParams struct { + GroupID uuid.UUID `db:"group_id" json:"group_id"` + PeriodStart time.Time `db:"period_start" json:"period_start"` + UserIds []uuid.UUID `db:"user_ids" json:"user_ids"` +} + +type GetGroupMembersAISpendRow struct { + UserID uuid.UUID `db:"user_id" json:"user_id"` + OrganizationID uuid.UUID `db:"organization_id" json:"organization_id"` + EffectiveGroupID uuid.NullUUID `db:"effective_group_id" json:"effective_group_id"` + SpendLimitMicros sql.NullInt64 `db:"spend_limit_micros" json:"spend_limit_micros"` + LimitSource sql.NullString `db:"limit_source" json:"limit_source"` + GroupSpendMicros int64 `db:"group_spend_micros" json:"group_spend_micros"` +} + +// Returns each user's AI spend attributed to the queried group, on or after +// period_start until NOW. Only current members of the queried group are +// returned. spend_limit_micros and limit_source are populated only when the +// queried group is the user's effective budget source. The effective_group_id +// is null when the user has no configured budget or when the effective group +// belongs to a different organization than the queried group. +// The period_start parameter is normalized to its UTC calendar day. +// Spend is aggregated for the queried group, not the user's effective group. +// A LEFT JOIN leaves spend_limit_micros and limit_source null for users +// whose effective budget source is not the queried group. +func (q *sqlQuerier) GetGroupMembersAISpend(ctx context.Context, arg GetGroupMembersAISpendParams) ([]GetGroupMembersAISpendRow, error) { + rows, err := q.db.QueryContext(ctx, getGroupMembersAISpend, arg.GroupID, arg.PeriodStart, pq.Array(arg.UserIds)) + if err != nil { + return nil, err + } + defer rows.Close() + var items []GetGroupMembersAISpendRow + for rows.Next() { + var i GetGroupMembersAISpendRow + if err := rows.Scan( + &i.UserID, + &i.OrganizationID, + &i.EffectiveGroupID, + &i.SpendLimitMicros, + &i.LimitSource, + &i.GroupSpendMicros, + ); 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 getHighestGroupAIBudgetByUser = `-- name: GetHighestGroupAIBudgetByUser :one SELECT gaib.group_id, diff --git a/coderd/database/queries/aicostcontrol.sql b/coderd/database/queries/aicostcontrol.sql index 7738e3e09e073..fe8de3173932a 100644 --- a/coderd/database/queries/aicostcontrol.sql +++ b/coderd/database/queries/aicostcontrol.sql @@ -121,3 +121,87 @@ 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; + +-- name: GetGroupMembersAISpend :many +-- Returns each user's AI spend attributed to the queried group, on or after +-- period_start until NOW. Only current members of the queried group are +-- returned. spend_limit_micros and limit_source are populated only when the +-- queried group is the user's effective budget source. The effective_group_id +-- is null when the user has no configured budget or when the effective group +-- belongs to a different organization than the queried group. +-- The period_start parameter is normalized to its UTC calendar day. +WITH queried_group AS ( + -- The queried group's org, used to detect cross-org effective groups. + SELECT id, organization_id + FROM groups + WHERE id = @group_id +), +filtered_users AS ( + -- Users from @user_ids that are members of the queried group. Uses + -- group_members_expanded so the implicit Everyone group counts. + SELECT DISTINCT user_id + FROM group_members_expanded + WHERE group_id = @group_id + AND user_id = ANY(@user_ids::uuid[]) +), +user_highest_group AS ( + -- Per user, the highest-limit group they belong to. Uses + -- group_members_expanded so the implicit Everyone group counts. + SELECT DISTINCT ON (member.user_id) + member.user_id, + budget.group_id, + budget.spend_limit_micros + FROM group_ai_budgets budget + JOIN group_members_expanded member ON member.group_id = budget.group_id + WHERE member.user_id IN (SELECT user_id FROM filtered_users) + ORDER BY member.user_id, budget.spend_limit_micros DESC, member.group_name ASC, budget.group_id ASC +), +effective AS ( + -- Effective budget per user: a per-user override wins over the + -- highest-limit group. + SELECT + filtered_users.user_id, + COALESCE(override.group_id, user_highest_group.group_id) AS effective_group_id, + COALESCE(override.spend_limit_micros, user_highest_group.spend_limit_micros) AS spend_limit_micros, + (CASE + WHEN override.group_id IS NOT NULL THEN 'user_override' + WHEN user_highest_group.group_id IS NOT NULL THEN 'group' + END)::text AS limit_source + FROM filtered_users + LEFT JOIN user_ai_budget_overrides override ON override.user_id = filtered_users.user_id + LEFT JOIN user_highest_group ON user_highest_group.user_id = filtered_users.user_id +), +applied_budget AS ( + -- The limit and source only for users whose effective budget source is the + -- queried group. + SELECT user_id, spend_limit_micros, limit_source + FROM effective + WHERE effective_group_id = @group_id +) +-- Spend is aggregated for the queried group, not the user's effective group. +SELECT + effective.user_id, + queried_group.organization_id, + effective_group.id AS effective_group_id, + applied_budget.spend_limit_micros, + applied_budget.limit_source, + COALESCE(SUM(spend.spend_micros), 0)::BIGINT AS group_spend_micros +FROM effective +CROSS JOIN queried_group +LEFT JOIN groups effective_group + ON effective_group.id = effective.effective_group_id + AND effective_group.organization_id = queried_group.organization_id +-- A LEFT JOIN leaves spend_limit_micros and limit_source null for users +-- whose effective budget source is not the queried group. +LEFT JOIN applied_budget ON applied_budget.user_id = effective.user_id +LEFT JOIN ai_user_daily_spend spend + ON spend.user_id = effective.user_id + AND spend.effective_group_id = @group_id + AND spend.day >= ((@period_start::timestamptz) AT TIME ZONE 'UTC')::date +GROUP BY + effective.user_id, + queried_group.organization_id, + effective_group.id, + applied_budget.spend_limit_micros, + applied_budget.limit_source +ORDER BY effective.user_id; diff --git a/codersdk/aibridge.go b/codersdk/aibridge.go index 7cec7fbfa3116..9542e44621b92 100644 --- a/codersdk/aibridge.go +++ b/codersdk/aibridge.go @@ -82,6 +82,34 @@ type OrganizationGroupAISpend struct { CurrentSpendMicros int64 `json:"current_spend_micros"` } +// GroupMembersAISpend reports per-member AI spend attributed to a specific +// group in the active budget period. +type GroupMembersAISpend struct { + AISpendPeriodWindow + Members []GroupMemberAISpend `json:"members"` +} + +// GroupMemberAISpend is a single member's AI spend attributed to the queried +// group in the current budget period. +type GroupMemberAISpend struct { + UserID uuid.UUID `json:"user_id" format:"uuid"` + // EffectiveGroupID is the user's effective budget group within the queried + // group's organization. Null when no effective budget group is visible in + // this organization, including when the user's budget resolves to a group + // in another organization. + EffectiveGroupID *uuid.UUID `json:"effective_group_id" format:"uuid"` + // SpendLimitMicros is the spend limit when the queried group is this + // user's effective budget source. Null when the user's budget resolves to + // another group. + SpendLimitMicros *int64 `json:"spend_limit_micros"` + // LimitSource identifies the tier that produced the limit. Null when the + // user's budget resolves to another group. + LimitSource *AIBudgetLimitSource `json:"limit_source"` + // GroupSpendMicros is the user's spend attributed to the queried group + // over the current budget period. + GroupSpendMicros int64 `json:"group_spend_micros"` +} + type AIBridgeSession struct { ID string `json:"id"` Initiator MinimalUser `json:"initiator"` @@ -498,3 +526,31 @@ func (c *Client) OrganizationGroupsAISpend(ctx context.Context, organization uui var resp OrganizationGroupsAISpend return resp, json.NewDecoder(res.Body).Decode(&resp) } + +// GroupMembersAISpend returns AI spend attributed to the given group for the +// specified users within the active budget period. +func (c *Client) GroupMembersAISpend(ctx context.Context, group uuid.UUID, userIDs []uuid.UUID) (GroupMembersAISpend, error) { + ids := make([]string, len(userIDs)) + for i, id := range userIDs { + ids[i] = id.String() + } + res, err := c.Request(ctx, http.MethodGet, + fmt.Sprintf("/api/v2/groups/%s/members/ai/spend", group.String()), + nil, + func(r *http.Request) { + q := r.URL.Query() + q.Set("user_ids", strings.Join(ids, ",")) + r.URL.RawQuery = q.Encode() + }, + ) + if err != nil { + return GroupMembersAISpend{}, xerrors.Errorf("make request: %w", err) + } + defer res.Body.Close() + + if res.StatusCode != http.StatusOK { + return GroupMembersAISpend{}, ReadBodyAsError(res) + } + var resp GroupMembersAISpend + return resp, json.NewDecoder(res.Body).Decode(&resp) +} diff --git a/docs/reference/api/enterprise.md b/docs/reference/api/enterprise.md index 5497703829c8f..b821d066176e5 100644 --- a/docs/reference/api/enterprise.md +++ b/docs/reference/api/enterprise.md @@ -1092,6 +1092,56 @@ curl -X GET http://coder-server:8080/api/v2/groups/{group}/members \ To perform this operation, you must be authenticated. [Learn more](authentication.md). +## Get group members AI spend + +### Code samples + +```shell +# Example request using curl +curl -X GET http://coder-server:8080/api/v2/groups/{group}/members/ai/spend?user_ids=string \ + -H 'Accept: application/json' \ + -H 'Coder-Session-Token: API_KEY' +``` + +`GET /api/v2/groups/{group}/members/ai/spend` + +Returns aggregate AI spend attributed to the group per requested user. User IDs that are not members of the group, or that the caller has no read access to, are silently omitted. + +### Parameters + +| Name | In | Type | Required | Description | +|------------|-------|--------------|----------|------------------------------------------------| +| `group` | path | string(uuid) | true | Group ID | +| `user_ids` | query | string | true | Comma-separated list of user IDs (maximum 100) | + +### Example responses + +> 200 Response + +```json +{ + "members": [ + { + "effective_group_id": "85e2b926-ddfb-4c66-b68e-b66e5acec6c0", + "group_spend_micros": 0, + "limit_source": "user_override", + "spend_limit_micros": 0, + "user_id": "a169451c-8525-4352-b8ca-070dd449a1a5" + } + ], + "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.GroupMembersAISpend](schemas.md#codersdkgroupmembersaispend) | + +To perform this operation, you must be authenticated. [Learn more](authentication.md). + ## Get licenses ### Code samples @@ -2010,6 +2060,57 @@ curl -X GET http://coder-server:8080/api/v2/organizations/{organization}/groups/ To perform this operation, you must be authenticated. [Learn more](authentication.md). +## Get group members AI spend by organization + +### Code samples + +```shell +# Example request using curl +curl -X GET http://coder-server:8080/api/v2/organizations/{organization}/groups/{groupName}/members/ai/spend?user_ids=string \ + -H 'Accept: application/json' \ + -H 'Coder-Session-Token: API_KEY' +``` + +`GET /api/v2/organizations/{organization}/groups/{groupName}/members/ai/spend` + +Returns aggregate AI spend attributed to the group per requested user. User IDs that are not members of the group, or that the caller has no read access to, are silently omitted. + +### Parameters + +| Name | In | Type | Required | Description | +|----------------|-------|--------------|----------|------------------------------------------------| +| `organization` | path | string(uuid) | true | Organization ID | +| `groupName` | path | string | true | Group name | +| `user_ids` | query | string | true | Comma-separated list of user IDs (maximum 100) | + +### Example responses + +> 200 Response + +```json +{ + "members": [ + { + "effective_group_id": "85e2b926-ddfb-4c66-b68e-b66e5acec6c0", + "group_spend_micros": 0, + "limit_source": "user_override", + "spend_limit_micros": 0, + "user_id": "a169451c-8525-4352-b8ca-070dd449a1a5" + } + ], + "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.GroupMembersAISpend](schemas.md#codersdkgroupmembersaispend) | + +To perform this operation, you must be authenticated. [Learn more](authentication.md). + ## Get workspace quota by user ### Code samples diff --git a/docs/reference/api/schemas.md b/docs/reference/api/schemas.md index e938a624b7f8e..9d5cc9a4d7113 100644 --- a/docs/reference/api/schemas.md +++ b/docs/reference/api/schemas.md @@ -7685,6 +7685,54 @@ Only certain features set these fields: - FeatureManagedAgentLimit| | `spend_limit_micros` | integer | false | | | | `updated_at` | string | false | | | +## codersdk.GroupMemberAISpend + +```json +{ + "effective_group_id": "85e2b926-ddfb-4c66-b68e-b66e5acec6c0", + "group_spend_micros": 0, + "limit_source": "user_override", + "spend_limit_micros": 0, + "user_id": "a169451c-8525-4352-b8ca-070dd449a1a5" +} +``` + +### Properties + +| Name | Type | Required | Restrictions | Description | +|----------------------|--------------------------------------------------------------|----------|--------------|-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------| +| `effective_group_id` | string | false | | Effective group ID is the user's effective budget group within the queried group's organization. Null when no effective budget group is visible in this organization, including when the user's budget resolves to a group in another organization. | +| `group_spend_micros` | integer | false | | Group spend micros is the user's spend attributed to the queried group over the current budget period. | +| `limit_source` | [codersdk.AIBudgetLimitSource](#codersdkaibudgetlimitsource) | false | | Limit source identifies the tier that produced the limit. Null when the user's budget resolves to another group. | +| `spend_limit_micros` | integer | false | | Spend limit micros is the spend limit when the queried group is this user's effective budget source. Null when the user's budget resolves to another group. | +| `user_id` | string | false | | | + +## codersdk.GroupMembersAISpend + +```json +{ + "members": [ + { + "effective_group_id": "85e2b926-ddfb-4c66-b68e-b66e5acec6c0", + "group_spend_micros": 0, + "limit_source": "user_override", + "spend_limit_micros": 0, + "user_id": "a169451c-8525-4352-b8ca-070dd449a1a5" + } + ], + "period_end": "2019-08-24T14:15:22Z", + "period_start": "2019-08-24T14:15:22Z" +} +``` + +### Properties + +| Name | Type | Required | Restrictions | Description | +|----------------|---------------------------------------------------------------------|----------|--------------|-------------------------------------------------------------------------| +| `members` | array of [codersdk.GroupMemberAISpend](#codersdkgroupmemberaispend) | 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.GroupMembersResponse ```json diff --git a/enterprise/coderd/aibridge.go b/enterprise/coderd/aibridge.go index d128a82aa870a..9fea257703ff3 100644 --- a/enterprise/coderd/aibridge.go +++ b/enterprise/coderd/aibridge.go @@ -38,6 +38,7 @@ const ( // requests. This is hardcoded to keep configuration simple. aiBridgeRateLimitWindow = time.Second maxOrganizationGroupsAISpendGroupIDs = 100 + maxGroupMembersAISpendUserIDs = 100 ) // errInvalidCursor is returned when a pagination cursor does not @@ -1022,3 +1023,105 @@ func (api *API) organizationGroupsAISpend(rw http.ResponseWriter, r *http.Reques httpapi.Write(ctx, rw, http.StatusOK, resp) } + +// @Summary Get group members AI spend by organization +// @Description Returns aggregate AI spend attributed to the group per requested user. User IDs that are not members of the group, or that the caller has no read access to, are silently omitted. +// @ID get-group-members-ai-spend-by-organization +// @Security CoderSessionToken +// @Produce json +// @Tags Enterprise +// @Param organization path string true "Organization ID" format(uuid) +// @Param groupName path string true "Group name" +// @Param user_ids query string true "Comma-separated list of user IDs (maximum 100)" +// @Success 200 {object} codersdk.GroupMembersAISpend +// @Router /api/v2/organizations/{organization}/groups/{groupName}/members/ai/spend [get] +func (api *API) groupMembersAISpendByOrganization(rw http.ResponseWriter, r *http.Request) { + api.groupMembersAISpend(rw, r) +} + +// @Summary Get group members AI spend +// @Description Returns aggregate AI spend attributed to the group per requested user. User IDs that are not members of the group, or that the caller has no read access to, are silently omitted. +// @ID get-group-members-ai-spend +// @Security CoderSessionToken +// @Produce json +// @Tags Enterprise +// @Param group path string true "Group ID" format(uuid) +// @Param user_ids query string true "Comma-separated list of user IDs (maximum 100)" +// @Success 200 {object} codersdk.GroupMembersAISpend +// @Router /api/v2/groups/{group}/members/ai/spend [get] +func (api *API) groupMembersAISpend(rw http.ResponseWriter, r *http.Request) { + ctx := r.Context() + group := httpmw.GroupParam(r) + logger := api.Logger.With(slog.F("group_id", group.ID)) + + parser := httpapi.NewQueryParamParser() + parser.RequiredNotEmpty("user_ids") + userIDs := parser.UUIDs(r.URL.Query(), nil, "user_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(userIDs) > maxGroupMembersAISpendUserIDs { + httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{ + Message: fmt.Sprintf( + "user_ids has %d entries, maximum is %d.", + len(userIDs), maxGroupMembersAISpendUserIDs, + ), + }) + 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.GetGroupMembersAISpend(ctx, database.GetGroupMembersAISpendParams{ + GroupID: group.ID, + UserIds: userIDs, + PeriodStart: periodWindow.Start, + }) + if err != nil { + logger.Error(ctx, "failed to get group members AI spend", slog.Error(err)) + httpapi.InternalServerError(rw, err) + return + } + + resp := codersdk.GroupMembersAISpend{ + AISpendPeriodWindow: codersdk.AISpendPeriodWindow{ + PeriodStart: periodWindow.Start, + PeriodEnd: periodWindow.End, + }, + Members: make([]codersdk.GroupMemberAISpend, 0, len(rows)), + } + for _, row := range rows { + entry := codersdk.GroupMemberAISpend{ + UserID: row.UserID, + GroupSpendMicros: row.GroupSpendMicros, + } + if row.EffectiveGroupID.Valid { + uid := row.EffectiveGroupID.UUID + entry.EffectiveGroupID = &uid + } + if row.SpendLimitMicros.Valid { + entry.SpendLimitMicros = &row.SpendLimitMicros.Int64 + } + if row.LimitSource.Valid { + source := codersdk.AIBudgetLimitSource(row.LimitSource.String) + entry.LimitSource = &source + } + resp.Members = append(resp.Members, entry) + } + + httpapi.Write(ctx, rw, http.StatusOK, resp) +} diff --git a/enterprise/coderd/aibridge_test.go b/enterprise/coderd/aibridge_test.go index e3d2844769d7b..a07e849386899 100644 --- a/enterprise/coderd/aibridge_test.go +++ b/enterprise/coderd/aibridge_test.go @@ -3514,6 +3514,472 @@ func TestOrganizationGroupsAISpendRoleAccess(t *testing.T) { } } +func TestGroupMembersAISpend(t *testing.T) { + t.Parallel() + + t.Run("RequiresLicenseFeature", func(t *testing.T) { + t.Parallel() + + dv := coderdtest.DeploymentValues(t) + 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, + }, + }, + }) + adminClient, _ := coderdtest.CreateAnotherUser(t, ownerClient, owner.OrganizationID, rbac.RoleUserAdmin()) + ctx := testutil.Context(t, testutil.WaitLong) + group, err := adminClient.CreateGroup(ctx, owner.OrganizationID, codersdk.CreateGroupRequest{ + Name: "req-license-feature-members-group", + }) + require.NoError(t, err) + + _, err = adminClient.GroupMembersAISpend(ctx, group.ID, []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) + ownerClient, owner := coderdenttest.New(t, &coderdenttest.Options{ + Options: &coderdtest.Options{DeploymentValues: dv}, + LicenseOptions: &coderdenttest.LicenseOptions{ + Features: license.Features{ + codersdk.FeatureTemplateRBAC: 1, + codersdk.FeatureAIBridge: 1, + }, + }, + }) + adminClient, _ := coderdtest.CreateAnotherUser(t, ownerClient, owner.OrganizationID, rbac.RoleUserAdmin()) + ctx := testutil.Context(t, testutil.WaitLong) + group, err := adminClient.CreateGroup(ctx, owner.OrganizationID, codersdk.CreateGroupRequest{ + Name: "req-experiment-members-group", + }) + require.NoError(t, err) + + _, err = adminClient.GroupMembersAISpend(ctx, group.ID, []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("MissingUserIDs", func(t *testing.T) { + t.Parallel() + + adminClient, _, group := setupAICostControlTest(t, aiCostControlTestOptions{GroupName: "missing-ids-members-group"}) + ctx := testutil.Context(t, testutil.WaitLong) + + // When: querying with no user_ids. + _, err := adminClient.GroupMembersAISpend(ctx, group.ID, nil) + + // Then: request fails with 400. + var sdkErr *codersdk.Error + require.ErrorAs(t, err, &sdkErr) + require.Equal(t, http.StatusBadRequest, sdkErr.StatusCode()) + }) + + t.Run("InclusiveMaxUserIDs", func(t *testing.T) { + t.Parallel() + + adminClient, _, group := setupAICostControlTest(t, aiCostControlTestOptions{GroupName: "inclusive-max-user-ids-group"}) + ctx := testutil.Context(t, testutil.WaitLong) + + // Given: 100 user_ids, exactly at the cap. + ids := make([]uuid.UUID, 100) + for i := range ids { + ids[i] = uuid.New() + } + + // When: querying spend. + _, err := adminClient.GroupMembersAISpend(ctx, group.ID, ids) + + // Then: request succeeds. + require.NoError(t, err) + }) + + t.Run("TooManyUserIDs", func(t *testing.T) { + t.Parallel() + + adminClient, _, group := setupAICostControlTest(t, aiCostControlTestOptions{GroupName: "too-many-user-ids-group"}) + ctx := testutil.Context(t, testutil.WaitLong) + + // Given: 101 user_ids, above the cap of 100. + ids := make([]uuid.UUID, 101) + for i := range ids { + ids[i] = uuid.New() + } + + // When: querying spend. + _, err := adminClient.GroupMembersAISpend(ctx, group.ID, ids) + + // Then: request fails with 400. + var sdkErr *codersdk.Error + require.ErrorAs(t, err, &sdkErr) + require.Equal(t, http.StatusBadRequest, sdkErr.StatusCode()) + }) + + t.Run("MalformedUserID", func(t *testing.T) { + t.Parallel() + + adminClient, _, group := setupAICostControlTest(t, aiCostControlTestOptions{GroupName: "malformed-user-id-group"}) + ctx := testutil.Context(t, testutil.WaitLong) + + // Given: a malformed UUID passed via raw HTTP. + res, err := adminClient.Request(ctx, http.MethodGet, + "/api/v2/groups/"+group.ID.String()+"/members/ai/spend", + nil, + func(r *http.Request) { + q := r.URL.Query() + q.Set("user_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("UserInOtherOrgExcluded", func(t *testing.T) { + t.Parallel() + + db, ps := dbtestutil.NewDB(t) + adminClient, targetUser, group := setupAICostControlTest(t, aiCostControlTestOptions{ + GroupName: "primary-org-members-group", + Database: db, + Pubsub: ps, + }) + otherOrg := dbgen.Organization(t, db, database.Organization{}) + otherOrgUser := dbgen.User(t, db, database.User{}) + dbgen.OrganizationMember(t, db, database.OrganizationMember{UserID: otherOrgUser.ID, OrganizationID: otherOrg.ID}) + ctx := testutil.Context(t, testutil.WaitLong) + + // When: querying the group with a user from the primary org and one from another org. + resp, err := adminClient.GroupMembersAISpend(ctx, group.ID, []uuid.UUID{targetUser.ID, otherOrgUser.ID}) + require.NoError(t, err) + + // Then: only the primary-org user is returned. + require.Len(t, resp.Members, 1) + require.Equal(t, targetUser.ID, resp.Members[0].UserID) + }) + + t.Run("CrossOrgEffectiveGroupMasked", func(t *testing.T) { + t.Parallel() + + dv := coderdtest.DeploymentValues(t) + dv.AI.BridgeConfig.Enabled = serpent.Bool(true) + dv.Experiments = []string{string(codersdk.ExperimentAIGatewayCostControl)} + db, ps := dbtestutil.NewDB(t) + ownerClient, owner := coderdenttest.New(t, &coderdenttest.Options{ + Options: &coderdtest.Options{DeploymentValues: dv, Database: db, Pubsub: ps}, + LicenseOptions: &coderdenttest.LicenseOptions{ + Features: license.Features{ + codersdk.FeatureTemplateRBAC: 1, + codersdk.FeatureAIBridge: 1, + codersdk.FeatureMultipleOrganizations: 1, + }, + }, + }) + userAdminClient, _ := coderdtest.CreateAnotherUser(t, ownerClient, owner.OrganizationID, rbac.RoleUserAdmin()) + ctx := testutil.Context(t, testutil.WaitLong) + + // Given: a member of the queried group whose highest-limit budget + // group lives in a different org. + otherOrg := coderdenttest.CreateOrganization(t, ownerClient, coderdenttest.CreateOrganizationOptions{}) + _, targetUser := coderdtest.CreateAnotherUser(t, ownerClient, owner.OrganizationID) + dbgen.OrganizationMember(t, db, database.OrganizationMember{UserID: targetUser.ID, OrganizationID: otherOrg.ID}) + queried, err := userAdminClient.CreateGroup(ctx, owner.OrganizationID, codersdk.CreateGroupRequest{ + Name: "queried-cross-org-mask-group", + }) + require.NoError(t, err) + _, err = userAdminClient.PatchGroup(ctx, queried.ID, codersdk.PatchGroupRequest{ + AddUsers: []string{targetUser.ID.String()}, + }) + require.NoError(t, err) + crossOrgGroup, err := userAdminClient.CreateGroup(ctx, otherOrg.ID, codersdk.CreateGroupRequest{ + Name: "cross-org-budget-group", + }) + require.NoError(t, err) + _, err = userAdminClient.PatchGroup(ctx, crossOrgGroup.ID, codersdk.PatchGroupRequest{ + AddUsers: []string{targetUser.ID.String()}, + }) + require.NoError(t, err) + _, err = userAdminClient.UpsertGroupAIBudget(ctx, crossOrgGroup.ID, codersdk.UpsertGroupAIBudgetRequest{ + SpendLimitMicros: 9_999_999, + }) + require.NoError(t, err) + + // When: the owner, who can read both orgs, queries the primary group. + //nolint:gocritic // The test asserts that even an owner sees the mask. + resp, err := ownerClient.GroupMembersAISpend(ctx, queried.ID, []uuid.UUID{targetUser.ID}) + require.NoError(t, err) + + // Then: effective_group_id is nil, even though the caller can read both orgs. + require.Len(t, resp.Members, 1) + require.Equal(t, targetUser.ID, resp.Members[0].UserID) + require.Nil(t, resp.Members[0].EffectiveGroupID, "cross-org effective group must be masked even for the owner") + }) + + t.Run("UserOverrideLimitSource", func(t *testing.T) { + t.Parallel() + + adminClient, targetUser, group := setupAICostControlTest(t, aiCostControlTestOptions{GroupName: "override-source-group"}) + ctx := testutil.Context(t, testutil.WaitLong) + + // Given: a user override targeting the queried group. + _, err := adminClient.UpsertUserAIBudgetOverride(ctx, targetUser.ID, codersdk.UpsertUserAIBudgetOverrideRequest{ + GroupID: group.ID, + SpendLimitMicros: 500_000_000, + }) + require.NoError(t, err) + + // When: querying the group's member spend. + resp, err := adminClient.GroupMembersAISpend(ctx, group.ID, []uuid.UUID{targetUser.ID}) + require.NoError(t, err) + + // Then: the limit and source reflect the override. + require.Len(t, resp.Members, 1) + require.NotNil(t, resp.Members[0].SpendLimitMicros) + require.Equal(t, int64(500_000_000), *resp.Members[0].SpendLimitMicros) + require.NotNil(t, resp.Members[0].LimitSource) + require.Equal(t, codersdk.AIBudgetLimitSourceUserOverride, *resp.Members[0].LimitSource) + }) + + t.Run("MemberCanOnlyReadOwnRow", func(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, + }, + }, + }) + userAdminClient, _ := coderdtest.CreateAnotherUser(t, ownerClient, owner.OrganizationID, rbac.RoleUserAdmin()) + memberClient, member := coderdtest.CreateAnotherUser(t, ownerClient, owner.OrganizationID) + _, otherMember := coderdtest.CreateAnotherUser(t, ownerClient, owner.OrganizationID) + + ctx := testutil.Context(t, testutil.WaitLong) + group, err := userAdminClient.CreateGroup(ctx, owner.OrganizationID, codersdk.CreateGroupRequest{ + Name: "member-privacy-group", + }) + require.NoError(t, err) + _, err = userAdminClient.PatchGroup(ctx, group.ID, codersdk.PatchGroupRequest{ + AddUsers: []string{member.ID.String(), otherMember.ID.String()}, + }) + require.NoError(t, err) + + // When: the member queries spend for themselves and another member. + resp, err := memberClient.GroupMembersAISpend(ctx, group.ID, []uuid.UUID{member.ID, otherMember.ID}) + require.NoError(t, err) + + // Then: only the caller's own row is returned. + require.Len(t, resp.Members, 1) + require.Equal(t, member.ID, resp.Members[0].UserID) + }) + + t.Run("OrgScopedRoute", func(t *testing.T) { + t.Parallel() + + adminClient, targetUser, group := setupAICostControlTest(t, aiCostControlTestOptions{GroupName: "org-scoped-route-group"}) + ctx := testutil.Context(t, testutil.WaitLong) + + // When: hitting the org-scoped alias route via raw HTTP. + res, err := adminClient.Request(ctx, http.MethodGet, + "/api/v2/organizations/"+group.OrganizationID.String()+"/groups/"+group.Name+"/members/ai/spend", + nil, + func(r *http.Request) { + q := r.URL.Query() + q.Set("user_ids", targetUser.ID.String()) + r.URL.RawQuery = q.Encode() + }, + ) + require.NoError(t, err) + defer res.Body.Close() + + // Then: 200 with the target user in the response. + require.Equal(t, http.StatusOK, res.StatusCode) + var got codersdk.GroupMembersAISpend + require.NoError(t, json.NewDecoder(res.Body).Decode(&got)) + require.Len(t, got.Members, 1) + require.Equal(t, targetUser.ID, got.Members[0].UserID) + }) + + groupSource := codersdk.AIBudgetLimitSourceGroup + tests := []struct { + name string + setBudget bool + spent int64 + wantEffectiveSet bool + wantSpendLimit *int64 + wantLimitSource *codersdk.AIBudgetLimitSource + wantSpendMicros int64 + }{ + { + name: "NoBudgetNoSpend", + }, + { + name: "BudgetZeroSpend", + setBudget: true, + wantEffectiveSet: true, + wantSpendLimit: ptr.Ref(int64(1_000_000_000)), + wantLimitSource: &groupSource, + }, + { + name: "BudgetWithSpend", + setBudget: true, + spent: 250_000_000, + wantEffectiveSet: true, + wantSpendLimit: ptr.Ref(int64(1_000_000_000)), + wantLimitSource: &groupSource, + wantSpendMicros: 250_000_000, + }, + { + name: "NoBudgetWithSpend", + spent: 100_000_000, + wantSpendMicros: 100_000_000, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + // Given: an admin, a group with a member, optionally a budget on the + // group and seeded spend attributed to the group. + clock := quartz.NewMock(t) + db, ps := dbtestutil.NewDB(t) + adminClient, targetUser, group := setupAICostControlTest(t, aiCostControlTestOptions{ + GroupName: "members-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: 1_000_000_000, + }) + 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 member spend. + got, err := adminClient.GroupMembersAISpend(ctx, group.ID, []uuid.UUID{targetUser.ID}) + require.NoError(t, err) + + // Then: one row is returned with the expected fields. + require.Equal(t, wantPeriodStart, got.PeriodStart) + require.Equal(t, wantPeriodEnd, got.PeriodEnd) + require.Len(t, got.Members, 1) + require.Equal(t, targetUser.ID, got.Members[0].UserID) + if tt.wantEffectiveSet { + require.NotNil(t, got.Members[0].EffectiveGroupID) + require.Equal(t, group.ID, *got.Members[0].EffectiveGroupID) + } else { + require.Nil(t, got.Members[0].EffectiveGroupID) + } + require.Equal(t, tt.wantSpendLimit, got.Members[0].SpendLimitMicros) + require.Equal(t, tt.wantLimitSource, got.Members[0].LimitSource) + require.Equal(t, tt.wantSpendMicros, got.Members[0].GroupSpendMicros) + }) + } +} + +func TestGroupMembersAISpendRoleAccess(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, member := 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-members-group", + }) + require.NoError(t, err) + _, err = userAdminClient.PatchGroup(ctx, group.ID, codersdk.PatchGroupRequest{ + AddUsers: []string{member.ID.String()}, + }) + require.NoError(t, err) + + cases := []struct { + name string + client *codersdk.Client + wantMember bool + }{ + {name: "Owner", client: ownerClient, wantMember: true}, + {name: "UserAdmin", client: userAdminClient, wantMember: true}, + {name: "OrgAdmin", client: orgAdminClient, wantMember: true}, + {name: "OrgUserAdmin", client: orgUserAdminClient, wantMember: true}, + {name: "Member", client: memberClient, wantMember: true}, + {name: "OtherOrgMember", client: otherOrgMemberClient, wantMember: 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.GroupMembersAISpend(ctx, group.ID, []uuid.UUID{member.ID}) + if !tc.wantMember { + 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.Members, 1) + require.Equal(t, member.ID, resp.Members[0].UserID) + }) + } +} + // 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 c08680c6fc35b..764c03eda9b56 100644 --- a/enterprise/coderd/coderd.go +++ b/enterprise/coderd/coderd.go @@ -519,6 +519,15 @@ func New(ctx context.Context, options *Options) (_ *API, err error) { r.Get("/", api.groupByOrganization) r.Get("/members", api.groupMembersByOrganization) + r.Route("/members/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.groupMembersAISpendByOrganization) + }) }) }) r.Route("/provisionerkeys", func(r chi.Router) { @@ -604,6 +613,15 @@ func New(ctx context.Context, options *Options) (_ *API, err error) { r.Patch("/", api.patchGroup) r.Delete("/", api.deleteGroup) r.Get("/members", api.groupMembers) + r.Route("/members/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.groupMembersAISpend) + }) r.Route("/ai/budget", func(r chi.Router) { // AI cost controls are a paid feature (AI Governance add-on). r.Use(api.RequireFeatureMW(codersdk.FeatureAIBridge)) diff --git a/site/src/api/typesGenerated.ts b/site/src/api/typesGenerated.ts index 1f7c26a1a4a6a..132f16d6312b6 100644 --- a/site/src/api/typesGenerated.ts +++ b/site/src/api/typesGenerated.ts @@ -5176,6 +5176,47 @@ export interface GroupArguments { readonly GroupIDs: readonly string[]; } +// From codersdk/aibridge.go +/** + * GroupMemberAISpend is a single member's AI spend attributed to the queried + * group in the current budget period. + */ +export interface GroupMemberAISpend { + readonly user_id: string; + /** + * EffectiveGroupID is the user's effective budget group within the queried + * group's organization. Null when no effective budget group is visible in + * this organization, including when the user's budget resolves to a group + * in another organization. + */ + readonly effective_group_id: string | null; + /** + * SpendLimitMicros is the spend limit when the queried group is this + * user's effective budget source. Null when the user's budget resolves to + * another group. + */ + readonly spend_limit_micros: number | null; + /** + * LimitSource identifies the tier that produced the limit. Null when the + * user's budget resolves to another group. + */ + readonly limit_source: AIBudgetLimitSource | null; + /** + * GroupSpendMicros is the user's spend attributed to the queried group + * over the current budget period. + */ + readonly group_spend_micros: number; +} + +// From codersdk/aibridge.go +/** + * GroupMembersAISpend reports per-member AI spend attributed to a specific + * group in the active budget period. + */ +export interface GroupMembersAISpend extends AISpendPeriodWindow { + readonly members: readonly GroupMemberAISpend[]; +} + // From codersdk/groups.go export interface GroupMembersResponse { readonly users: readonly ReducedUser[]; From 3f56dd4497cb4e1a34c09f46cefddf68fade39d5 Mon Sep 17 00:00:00 2001 From: Susana Cardoso Ferreira Date: Thu, 16 Jul 2026 09:32:29 +0000 Subject: [PATCH 2/8] chore: address review comments --- coderd/apidoc/docs.go | 4 +- coderd/apidoc/swagger.json | 4 +- coderd/database/querier_test.go | 48 +++++++++++++++++++++++ coderd/database/queries.sql.go | 2 +- coderd/database/queries/aicostcontrol.sql | 2 +- codersdk/aibridge.go | 14 +++---- docs/reference/api/schemas.md | 4 +- site/src/api/typesGenerated.ts | 5 ++- 8 files changed, 66 insertions(+), 17 deletions(-) diff --git a/coderd/apidoc/docs.go b/coderd/apidoc/docs.go index 9ccbaf280f2a6..a0195aa10913d 100644 --- a/coderd/apidoc/docs.go +++ b/coderd/apidoc/docs.go @@ -20364,7 +20364,7 @@ const docTemplate = `{ "type": "integer" }, "limit_source": { - "description": "LimitSource identifies the tier that produced the limit. Null when the\nuser's budget resolves to another group.", + "description": "LimitSource identifies the tier that produced the limit. Null when the\nuser's budget resolves to another group or no budget applies to the user\n(unlimited).", "allOf": [ { "$ref": "#/definitions/codersdk.AIBudgetLimitSource" @@ -20372,7 +20372,7 @@ const docTemplate = `{ ] }, "spend_limit_micros": { - "description": "SpendLimitMicros is the spend limit when the queried group is this\nuser's effective budget source. Null when the user's budget resolves to\nanother group.", + "description": "SpendLimitMicros is the spend limit when the queried group is this\nuser's effective budget source. Null when the user's budget resolves to\nanother group or no budget applies to the user (unlimited).", "type": "integer" }, "user_id": { diff --git a/coderd/apidoc/swagger.json b/coderd/apidoc/swagger.json index b984e43951ed3..1dd0cb5c0bd27 100644 --- a/coderd/apidoc/swagger.json +++ b/coderd/apidoc/swagger.json @@ -18534,7 +18534,7 @@ "type": "integer" }, "limit_source": { - "description": "LimitSource identifies the tier that produced the limit. Null when the\nuser's budget resolves to another group.", + "description": "LimitSource identifies the tier that produced the limit. Null when the\nuser's budget resolves to another group or no budget applies to the user\n(unlimited).", "allOf": [ { "$ref": "#/definitions/codersdk.AIBudgetLimitSource" @@ -18542,7 +18542,7 @@ ] }, "spend_limit_micros": { - "description": "SpendLimitMicros is the spend limit when the queried group is this\nuser's effective budget source. Null when the user's budget resolves to\nanother group.", + "description": "SpendLimitMicros is the spend limit when the queried group is this\nuser's effective budget source. Null when the user's budget resolves to\nanother group or no budget applies to the user (unlimited).", "type": "integer" }, "user_id": { diff --git a/coderd/database/querier_test.go b/coderd/database/querier_test.go index 210445999177d..6716f300be58e 100644 --- a/coderd/database/querier_test.go +++ b/coderd/database/querier_test.go @@ -13426,6 +13426,12 @@ func TestGetGroupMembersAISpend(t *testing.T) { SpendLimitMicros: 9_999_999, }) require.NoError(t, err) + // Seed spend attributed to the queried group so we can assert it is + // still returned even when the effective group is masked. + _, err = db.IncrementUserAIDailySpend(ctx, database.IncrementUserAIDailySpendParams{ + UserID: user.ID, EffectiveGroupID: group.ID, Day: now, CostMicros: 250, + }) + require.NoError(t, err) // When: querying spend for the user in the queried group's org. got, err := db.GetGroupMembersAISpend(ctx, database.GetGroupMembersAISpendParams{ @@ -13436,10 +13442,12 @@ func TestGetGroupMembersAISpend(t *testing.T) { require.NoError(t, err) // Then: effective_group_id is masked to NULL, the highest-limit group is cross-org. + // The queried-group spend is still returned. require.Len(t, got, 1) require.False(t, got[0].EffectiveGroupID.Valid, "cross-org effective group must be masked") require.False(t, got[0].SpendLimitMicros.Valid) require.False(t, got[0].LimitSource.Valid) + require.Equal(t, int64(250), got[0].GroupSpendMicros) }) t.Run("OverrideWins", func(t *testing.T) { @@ -13489,6 +13497,46 @@ func TestGetGroupMembersAISpend(t *testing.T) { require.Equal(t, uuid.NullUUID{UUID: overrideTarget.ID, Valid: true}, got[0].EffectiveGroupID) }) + t.Run("TieBreakByGroupName", func(t *testing.T) { + t.Parallel() + db, _ := dbtestutil.NewDB(t) + ctx := testutil.Context(t, testutil.WaitShort) + + // Given: a member of the queried group who is in two same-org groups + // with identical spend limits. + user := dbgen.User(t, db, database.User{}) + org := dbgen.Organization(t, db, database.Organization{}) + queried := dbgen.Group(t, db, database.Group{OrganizationID: org.ID}) + groupA := dbgen.Group(t, db, database.Group{OrganizationID: org.ID, Name: "aaa-tie-group"}) + groupB := dbgen.Group(t, db, database.Group{OrganizationID: org.ID, Name: "bbb-tie-group"}) + dbgen.OrganizationMember(t, db, database.OrganizationMember{UserID: user.ID, OrganizationID: org.ID}) + dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: queried.ID, UserID: user.ID}) + dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: groupA.ID, UserID: user.ID}) + dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: groupB.ID, UserID: user.ID}) + _, err := db.UpsertGroupAIBudget(ctx, database.UpsertGroupAIBudgetParams{ + GroupID: groupA.ID, + SpendLimitMicros: 1_000_000, + }) + require.NoError(t, err) + _, err = db.UpsertGroupAIBudget(ctx, database.UpsertGroupAIBudgetParams{ + GroupID: groupB.ID, + SpendLimitMicros: 1_000_000, + }) + require.NoError(t, err) + + // When: querying spend for the user. + got, err := db.GetGroupMembersAISpend(ctx, database.GetGroupMembersAISpendParams{ + GroupID: queried.ID, + UserIds: []uuid.UUID{user.ID}, + PeriodStart: monthStart, + }) + require.NoError(t, err) + + // Then: the tie is broken by group name ascending, so groupA wins. + require.Len(t, got, 1) + require.Equal(t, uuid.NullUUID{UUID: groupA.ID, Valid: true}, got[0].EffectiveGroupID) + }) + t.Run("EveryoneGroupCounts", func(t *testing.T) { t.Parallel() db, _ := dbtestutil.NewDB(t) diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go index ca628587591f2..df09ba024ff04 100644 --- a/coderd/database/queries.sql.go +++ b/coderd/database/queries.sql.go @@ -2536,7 +2536,7 @@ func (q *sqlQuerier) GetGroupAIBudget(ctx context.Context, groupID uuid.UUID) (G const getGroupMembersAISpend = `-- name: GetGroupMembersAISpend :many WITH queried_group AS ( -- The queried group's org, used to detect cross-org effective groups. - SELECT id, organization_id + SELECT organization_id FROM groups WHERE id = $1 ), diff --git a/coderd/database/queries/aicostcontrol.sql b/coderd/database/queries/aicostcontrol.sql index fe8de3173932a..5d4c5729a3efd 100644 --- a/coderd/database/queries/aicostcontrol.sql +++ b/coderd/database/queries/aicostcontrol.sql @@ -132,7 +132,7 @@ ORDER BY groups.id; -- The period_start parameter is normalized to its UTC calendar day. WITH queried_group AS ( -- The queried group's org, used to detect cross-org effective groups. - SELECT id, organization_id + SELECT organization_id FROM groups WHERE id = @group_id ), diff --git a/codersdk/aibridge.go b/codersdk/aibridge.go index 9542e44621b92..78f8491b29030 100644 --- a/codersdk/aibridge.go +++ b/codersdk/aibridge.go @@ -100,10 +100,11 @@ type GroupMemberAISpend struct { EffectiveGroupID *uuid.UUID `json:"effective_group_id" format:"uuid"` // SpendLimitMicros is the spend limit when the queried group is this // user's effective budget source. Null when the user's budget resolves to - // another group. + // another group or no budget applies to the user (unlimited). SpendLimitMicros *int64 `json:"spend_limit_micros"` // LimitSource identifies the tier that produced the limit. Null when the - // user's budget resolves to another group. + // user's budget resolves to another group or no budget applies to the user + // (unlimited). LimitSource *AIBudgetLimitSource `json:"limit_source"` // GroupSpendMicros is the user's spend attributed to the queried group // over the current budget period. @@ -528,12 +529,11 @@ func (c *Client) OrganizationGroupsAISpend(ctx context.Context, organization uui } // GroupMembersAISpend returns AI spend attributed to the given group for the -// specified users within the active budget period. +// specified users within the active budget period. At most 100 user IDs may be +// requested per call, and callers with more members are expected to batch +// across multiple requests. func (c *Client) GroupMembersAISpend(ctx context.Context, group uuid.UUID, userIDs []uuid.UUID) (GroupMembersAISpend, error) { - ids := make([]string, len(userIDs)) - for i, id := range userIDs { - ids[i] = id.String() - } + ids := slice.List(userIDs, func(id uuid.UUID) string { return id.String() }) res, err := c.Request(ctx, http.MethodGet, fmt.Sprintf("/api/v2/groups/%s/members/ai/spend", group.String()), nil, diff --git a/docs/reference/api/schemas.md b/docs/reference/api/schemas.md index 9d5cc9a4d7113..1daf3bc906cba 100644 --- a/docs/reference/api/schemas.md +++ b/docs/reference/api/schemas.md @@ -7703,8 +7703,8 @@ Only certain features set these fields: - FeatureManagedAgentLimit| |----------------------|--------------------------------------------------------------|----------|--------------|-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------| | `effective_group_id` | string | false | | Effective group ID is the user's effective budget group within the queried group's organization. Null when no effective budget group is visible in this organization, including when the user's budget resolves to a group in another organization. | | `group_spend_micros` | integer | false | | Group spend micros is the user's spend attributed to the queried group over the current budget period. | -| `limit_source` | [codersdk.AIBudgetLimitSource](#codersdkaibudgetlimitsource) | false | | Limit source identifies the tier that produced the limit. Null when the user's budget resolves to another group. | -| `spend_limit_micros` | integer | false | | Spend limit micros is the spend limit when the queried group is this user's effective budget source. Null when the user's budget resolves to another group. | +| `limit_source` | [codersdk.AIBudgetLimitSource](#codersdkaibudgetlimitsource) | false | | Limit source identifies the tier that produced the limit. Null when the user's budget resolves to another group or no budget applies to the user (unlimited). | +| `spend_limit_micros` | integer | false | | Spend limit micros is the spend limit when the queried group is this user's effective budget source. Null when the user's budget resolves to another group or no budget applies to the user (unlimited). | | `user_id` | string | false | | | ## codersdk.GroupMembersAISpend diff --git a/site/src/api/typesGenerated.ts b/site/src/api/typesGenerated.ts index 132f16d6312b6..3be2305222f49 100644 --- a/site/src/api/typesGenerated.ts +++ b/site/src/api/typesGenerated.ts @@ -5193,12 +5193,13 @@ export interface GroupMemberAISpend { /** * SpendLimitMicros is the spend limit when the queried group is this * user's effective budget source. Null when the user's budget resolves to - * another group. + * another group or no budget applies to the user (unlimited). */ readonly spend_limit_micros: number | null; /** * LimitSource identifies the tier that produced the limit. Null when the - * user's budget resolves to another group. + * user's budget resolves to another group or no budget applies to the user + * (unlimited). */ readonly limit_source: AIBudgetLimitSource | null; /** From 0ad3f920ea9f844e3a04c360519795f86eff703f Mon Sep 17 00:00:00 2001 From: Susana Cardoso Ferreira Date: Thu, 16 Jul 2026 09:52:19 +0000 Subject: [PATCH 3/8] chore: minor improvements --- coderd/apidoc/docs.go | 4 ++-- coderd/apidoc/swagger.json | 4 ++-- docs/reference/api/enterprise.md | 8 ++++++-- enterprise/coderd/aibridge.go | 8 ++++++-- 4 files changed, 16 insertions(+), 8 deletions(-) diff --git a/coderd/apidoc/docs.go b/coderd/apidoc/docs.go index a0195aa10913d..3302426bac1ff 100644 --- a/coderd/apidoc/docs.go +++ b/coderd/apidoc/docs.go @@ -3386,7 +3386,7 @@ const docTemplate = `{ }, "/api/v2/groups/{group}/members/ai/spend": { "get": { - "description": "Returns aggregate AI spend attributed to the group per requested user. User IDs that are not members of the group, or that the caller has no read access to, are silently omitted.", + "description": "Returns aggregate AI spend attributed to the group per requested user.\nA maximum of 100 user IDs may be requested per call, and requests with more are rejected, so callers are expected to batch across multiple requests.\nUser IDs that are not members of the group, or that the caller has no read access to, are silently omitted.", "produces": [ "application/json" ], @@ -4949,7 +4949,7 @@ const docTemplate = `{ }, "/api/v2/organizations/{organization}/groups/{groupName}/members/ai/spend": { "get": { - "description": "Returns aggregate AI spend attributed to the group per requested user. User IDs that are not members of the group, or that the caller has no read access to, are silently omitted.", + "description": "Returns aggregate AI spend attributed to the group per requested user.\nA maximum of 100 user IDs may be requested per call, and requests with more are rejected, so callers are expected to batch across multiple requests.\nUser IDs that are not members of the group, or that the caller has no read access to, are silently omitted.", "produces": [ "application/json" ], diff --git a/coderd/apidoc/swagger.json b/coderd/apidoc/swagger.json index 1dd0cb5c0bd27..9d72cb80d1647 100644 --- a/coderd/apidoc/swagger.json +++ b/coderd/apidoc/swagger.json @@ -2989,7 +2989,7 @@ }, "/api/v2/groups/{group}/members/ai/spend": { "get": { - "description": "Returns aggregate AI spend attributed to the group per requested user. User IDs that are not members of the group, or that the caller has no read access to, are silently omitted.", + "description": "Returns aggregate AI spend attributed to the group per requested user.\nA maximum of 100 user IDs may be requested per call, and requests with more are rejected, so callers are expected to batch across multiple requests.\nUser IDs that are not members of the group, or that the caller has no read access to, are silently omitted.", "produces": ["application/json"], "tags": ["Enterprise"], "summary": "Get group members AI spend", @@ -4374,7 +4374,7 @@ }, "/api/v2/organizations/{organization}/groups/{groupName}/members/ai/spend": { "get": { - "description": "Returns aggregate AI spend attributed to the group per requested user. User IDs that are not members of the group, or that the caller has no read access to, are silently omitted.", + "description": "Returns aggregate AI spend attributed to the group per requested user.\nA maximum of 100 user IDs may be requested per call, and requests with more are rejected, so callers are expected to batch across multiple requests.\nUser IDs that are not members of the group, or that the caller has no read access to, are silently omitted.", "produces": ["application/json"], "tags": ["Enterprise"], "summary": "Get group members AI spend by organization", diff --git a/docs/reference/api/enterprise.md b/docs/reference/api/enterprise.md index b821d066176e5..ea9e937c5b7ca 100644 --- a/docs/reference/api/enterprise.md +++ b/docs/reference/api/enterprise.md @@ -1105,7 +1105,9 @@ curl -X GET http://coder-server:8080/api/v2/groups/{group}/members/ai/spend?user `GET /api/v2/groups/{group}/members/ai/spend` -Returns aggregate AI spend attributed to the group per requested user. User IDs that are not members of the group, or that the caller has no read access to, are silently omitted. +Returns aggregate AI spend attributed to the group per requested user. +A maximum of 100 user IDs may be requested per call, and requests with more are rejected, so callers are expected to batch across multiple requests. +User IDs that are not members of the group, or that the caller has no read access to, are silently omitted. ### Parameters @@ -2073,7 +2075,9 @@ curl -X GET http://coder-server:8080/api/v2/organizations/{organization}/groups/ `GET /api/v2/organizations/{organization}/groups/{groupName}/members/ai/spend` -Returns aggregate AI spend attributed to the group per requested user. User IDs that are not members of the group, or that the caller has no read access to, are silently omitted. +Returns aggregate AI spend attributed to the group per requested user. +A maximum of 100 user IDs may be requested per call, and requests with more are rejected, so callers are expected to batch across multiple requests. +User IDs that are not members of the group, or that the caller has no read access to, are silently omitted. ### Parameters diff --git a/enterprise/coderd/aibridge.go b/enterprise/coderd/aibridge.go index 9fea257703ff3..a658c2fa05321 100644 --- a/enterprise/coderd/aibridge.go +++ b/enterprise/coderd/aibridge.go @@ -1025,7 +1025,9 @@ func (api *API) organizationGroupsAISpend(rw http.ResponseWriter, r *http.Reques } // @Summary Get group members AI spend by organization -// @Description Returns aggregate AI spend attributed to the group per requested user. User IDs that are not members of the group, or that the caller has no read access to, are silently omitted. +// @Description Returns aggregate AI spend attributed to the group per requested user. +// @Description A maximum of 100 user IDs may be requested per call, and requests with more are rejected, so callers are expected to batch across multiple requests. +// @Description User IDs that are not members of the group, or that the caller has no read access to, are silently omitted. // @ID get-group-members-ai-spend-by-organization // @Security CoderSessionToken // @Produce json @@ -1040,7 +1042,9 @@ func (api *API) groupMembersAISpendByOrganization(rw http.ResponseWriter, r *htt } // @Summary Get group members AI spend -// @Description Returns aggregate AI spend attributed to the group per requested user. User IDs that are not members of the group, or that the caller has no read access to, are silently omitted. +// @Description Returns aggregate AI spend attributed to the group per requested user. +// @Description A maximum of 100 user IDs may be requested per call, and requests with more are rejected, so callers are expected to batch across multiple requests. +// @Description User IDs that are not members of the group, or that the caller has no read access to, are silently omitted. // @ID get-group-members-ai-spend // @Security CoderSessionToken // @Produce json From 5aad37cf15d2d0e9703a3ee52f952ecbbad7bae3 Mon Sep 17 00:00:00 2001 From: Susana Cardoso Ferreira Date: Thu, 16 Jul 2026 10:35:54 +0000 Subject: [PATCH 4/8] fix: make gen --- docs/reference/api/enterprise.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/docs/reference/api/enterprise.md b/docs/reference/api/enterprise.md index ea9e937c5b7ca..a48e9d2009892 100644 --- a/docs/reference/api/enterprise.md +++ b/docs/reference/api/enterprise.md @@ -1096,7 +1096,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/groups/{group}/members/ai/spend?user_ids=string \ -H 'Accept: application/json' \ @@ -2066,7 +2066,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/{groupName}/members/ai/spend?user_ids=string \ -H 'Accept: application/json' \ From 1f9df60959ced788d5b51e06c83adf2c7e8d7f45 Mon Sep 17 00:00:00 2001 From: Susana Cardoso Ferreira Date: Thu, 16 Jul 2026 10:47:43 +0000 Subject: [PATCH 5/8] chore: address comments --- coderd/database/queries.sql.go | 6 +++--- coderd/database/queries/aicostcontrol.sql | 6 +++--- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go index df09ba024ff04..5e99aa6d48bd6 100644 --- a/coderd/database/queries.sql.go +++ b/coderd/database/queries.sql.go @@ -2565,7 +2565,7 @@ effective AS ( -- highest-limit group. SELECT filtered_users.user_id, - COALESCE(override.group_id, user_highest_group.group_id) AS effective_group_id, + COALESCE(override.group_id, user_highest_group.group_id) AS raw_effective_group_id, COALESCE(override.spend_limit_micros, user_highest_group.spend_limit_micros) AS spend_limit_micros, (CASE WHEN override.group_id IS NOT NULL THEN 'user_override' @@ -2580,7 +2580,7 @@ applied_budget AS ( -- queried group. SELECT user_id, spend_limit_micros, limit_source FROM effective - WHERE effective_group_id = $1 + WHERE raw_effective_group_id = $1 ) SELECT effective.user_id, @@ -2592,7 +2592,7 @@ SELECT FROM effective CROSS JOIN queried_group LEFT JOIN groups effective_group - ON effective_group.id = effective.effective_group_id + ON effective_group.id = effective.raw_effective_group_id AND effective_group.organization_id = queried_group.organization_id LEFT JOIN applied_budget ON applied_budget.user_id = effective.user_id LEFT JOIN ai_user_daily_spend spend diff --git a/coderd/database/queries/aicostcontrol.sql b/coderd/database/queries/aicostcontrol.sql index 5d4c5729a3efd..d51f65279d5c0 100644 --- a/coderd/database/queries/aicostcontrol.sql +++ b/coderd/database/queries/aicostcontrol.sql @@ -161,7 +161,7 @@ effective AS ( -- highest-limit group. SELECT filtered_users.user_id, - COALESCE(override.group_id, user_highest_group.group_id) AS effective_group_id, + COALESCE(override.group_id, user_highest_group.group_id) AS raw_effective_group_id, COALESCE(override.spend_limit_micros, user_highest_group.spend_limit_micros) AS spend_limit_micros, (CASE WHEN override.group_id IS NOT NULL THEN 'user_override' @@ -176,7 +176,7 @@ applied_budget AS ( -- queried group. SELECT user_id, spend_limit_micros, limit_source FROM effective - WHERE effective_group_id = @group_id + WHERE raw_effective_group_id = @group_id ) -- Spend is aggregated for the queried group, not the user's effective group. SELECT @@ -189,7 +189,7 @@ SELECT FROM effective CROSS JOIN queried_group LEFT JOIN groups effective_group - ON effective_group.id = effective.effective_group_id + ON effective_group.id = effective.raw_effective_group_id AND effective_group.organization_id = queried_group.organization_id -- A LEFT JOIN leaves spend_limit_micros and limit_source null for users -- whose effective budget source is not the queried group. From 2e92724de04e75a666b99363fb670f20c838f285 Mon Sep 17 00:00:00 2001 From: Susana Cardoso Ferreira Date: Thu, 16 Jul 2026 11:56:16 +0000 Subject: [PATCH 6/8] test: minor improvements --- coderd/database/db2sdk/db2sdk.go | 18 + coderd/database/querier_test.go | 547 +++++++++++++++-------------- enterprise/coderd/aibridge.go | 17 +- enterprise/coderd/aibridge_test.go | 289 +++++++-------- 4 files changed, 434 insertions(+), 437 deletions(-) diff --git a/coderd/database/db2sdk/db2sdk.go b/coderd/database/db2sdk/db2sdk.go index 9a510674ea7cc..af41242168213 100644 --- a/coderd/database/db2sdk/db2sdk.go +++ b/coderd/database/db2sdk/db2sdk.go @@ -1469,6 +1469,24 @@ func OrganizationGroupAISpend(row database.GetOrganizationGroupsAISpendRow) code return group } +func GroupMemberAISpend(row database.GetGroupMembersAISpendRow) codersdk.GroupMemberAISpend { + member := codersdk.GroupMemberAISpend{ + UserID: row.UserID, + GroupSpendMicros: row.GroupSpendMicros, + } + if row.EffectiveGroupID.Valid { + member.EffectiveGroupID = &row.EffectiveGroupID.UUID + } + if row.SpendLimitMicros.Valid { + member.SpendLimitMicros = &row.SpendLimitMicros.Int64 + } + if row.LimitSource.Valid { + source := codersdk.AIBudgetLimitSource(row.LimitSource.String) + member.LimitSource = &source + } + return member +} + func InvalidatedPresets(invalidatedPresets []database.UpdatePresetsLastInvalidatedAtRow) []codersdk.InvalidatedPreset { var presets []codersdk.InvalidatedPreset for _, p := range invalidatedPresets { diff --git a/coderd/database/querier_test.go b/coderd/database/querier_test.go index 6716f300be58e..b0330b632e503 100644 --- a/coderd/database/querier_test.go +++ b/coderd/database/querier_test.go @@ -13114,19 +13114,182 @@ func TestGetGroupMembersAISpend(t *testing.T) { now := monthStart.AddDate(0, 0, 14) // 2024-06-15 prevMonthLastDay := monthStart.AddDate(0, 0, -1) // 2024-05-31 - t.Run("NoBudgetNoSpend", func(t *testing.T) { + tests := []struct { + name string + groupLimit int64 + overrideLimit int64 + spend int64 + wantEffectiveGroup bool + wantLimit sql.NullInt64 + wantSource sql.NullString + wantSpend int64 + }{ + { + name: "NoBudgetNoSpend", + wantEffectiveGroup: false, + wantLimit: sql.NullInt64{}, + wantSource: sql.NullString{}, + wantSpend: 0, + }, + { + name: "GroupBudget", + groupLimit: 1_000_000, + wantEffectiveGroup: true, + wantLimit: sql.NullInt64{Int64: 1_000_000, Valid: true}, + wantSource: sql.NullString{String: "group", Valid: true}, + wantSpend: 0, + }, + { + name: "OverrideBudget", + overrideLimit: 500_000, + wantEffectiveGroup: true, + wantLimit: sql.NullInt64{Int64: 500_000, Valid: true}, + wantSource: sql.NullString{String: "user_override", Valid: true}, + wantSpend: 0, + }, + { + name: "NoBudgetWithSpend", + spend: 250, + wantEffectiveGroup: false, + wantLimit: sql.NullInt64{}, + wantSource: sql.NullString{}, + wantSpend: 250, + }, + { + name: "BudgetWithSpend", + groupLimit: 1_000_000, + spend: 250, + wantEffectiveGroup: true, + wantLimit: sql.NullInt64{Int64: 1_000_000, Valid: true}, + wantSource: sql.NullString{String: "group", Valid: true}, + wantSpend: 250, + }, + } + + 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: a member of the queried group, optionally with a group + // budget, a user override, 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}) + dbgen.OrganizationMember(t, db, database.OrganizationMember{UserID: user.ID, OrganizationID: org.ID}) + dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: group.ID, UserID: user.ID}) + if tt.groupLimit > 0 { + _, err := db.UpsertGroupAIBudget(ctx, database.UpsertGroupAIBudgetParams{ + GroupID: group.ID, + SpendLimitMicros: tt.groupLimit, + }) + require.NoError(t, err) + } + if tt.overrideLimit > 0 { + _, err := db.UpsertUserAIBudgetOverride(ctx, database.UpsertUserAIBudgetOverrideParams{ + UserID: user.ID, + GroupID: group.ID, + SpendLimitMicros: tt.overrideLimit, + }) + require.NoError(t, err) + } + if tt.spend > 0 { + _, err := db.IncrementUserAIDailySpend(ctx, database.IncrementUserAIDailySpendParams{ + UserID: user.ID, EffectiveGroupID: group.ID, Day: now, CostMicros: tt.spend, + }) + require.NoError(t, err) + } + + // When: querying spend for the user. + got, err := db.GetGroupMembersAISpend(ctx, database.GetGroupMembersAISpendParams{ + GroupID: group.ID, + UserIds: []uuid.UUID{user.ID}, + PeriodStart: monthStart, + }) + require.NoError(t, err) + + // Then: one row with the expected effective group, limit, source, and spend. + require.Len(t, got, 1) + require.Equal(t, user.ID, got[0].UserID) + require.Equal(t, org.ID, got[0].OrganizationID) + if tt.wantEffectiveGroup { + require.Equal(t, uuid.NullUUID{UUID: group.ID, Valid: true}, got[0].EffectiveGroupID) + } else { + require.False(t, got[0].EffectiveGroupID.Valid, "expected no effective group") + } + require.Equal(t, tt.wantLimit, got[0].SpendLimitMicros) + require.Equal(t, tt.wantSource, got[0].LimitSource) + require.Equal(t, tt.wantSpend, got[0].GroupSpendMicros) + }) + } + + t.Run("MultipleMembers", func(t *testing.T) { t.Parallel() db, _ := dbtestutil.NewDB(t) ctx := testutil.Context(t, testutil.WaitShort) - // Given: a group member with no budget and no spend. + // Given: two group members with different spend attributed to the 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}) + dbgen.OrganizationMember(t, db, database.OrganizationMember{UserID: userA.ID, OrganizationID: org.ID}) + dbgen.OrganizationMember(t, db, database.OrganizationMember{UserID: userB.ID, OrganizationID: org.ID}) + dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: group.ID, UserID: userA.ID}) + dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: group.ID, UserID: userB.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: 250, + }) + require.NoError(t, err) + + // When: querying spend for both users. + got, err := db.GetGroupMembersAISpend(ctx, database.GetGroupMembersAISpendParams{ + GroupID: group.ID, + UserIds: []uuid.UUID{userA.ID, userB.ID}, + PeriodStart: monthStart, + }) + require.NoError(t, err) + + // Then: both users are returned with their own aggregate spend. + require.Len(t, got, 2) + byID := make(map[uuid.UUID]database.GetGroupMembersAISpendRow, len(got)) + for _, row := range got { + byID[row.UserID] = row + } + require.Equal(t, int64(100), byID[userA.ID].GroupSpendMicros) + require.Equal(t, int64(250), byID[userB.ID].GroupSpendMicros) + for _, row := range got { + require.False(t, row.EffectiveGroupID.Valid) + require.False(t, row.SpendLimitMicros.Valid) + require.False(t, row.LimitSource.Valid) + } + }) + + t.Run("AggregatesSpendAcrossDays", func(t *testing.T) { + t.Parallel() + db, _ := dbtestutil.NewDB(t) + ctx := testutil.Context(t, testutil.WaitShort) + + // Given: a group member with spend on multiple days in the 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}) dbgen.OrganizationMember(t, db, database.OrganizationMember{UserID: user.ID, OrganizationID: org.ID}) dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: group.ID, UserID: user.ID}) + days := []time.Time{monthStart, monthStart.AddDate(0, 0, 7), monthStart.AddDate(0, 0, 14)} + for i, day := range days { + _, err := db.IncrementUserAIDailySpend(ctx, database.IncrementUserAIDailySpendParams{ + UserID: user.ID, EffectiveGroupID: group.ID, Day: day, CostMicros: int64((i + 1) * 100), + }) + require.NoError(t, err) + } - // When: querying spend for the user. + // When: querying the group's member spend. got, err := db.GetGroupMembersAISpend(ctx, database.GetGroupMembersAISpendParams{ GroupID: group.ID, UserIds: []uuid.UUID{user.ID}, @@ -13134,107 +13297,144 @@ func TestGetGroupMembersAISpend(t *testing.T) { }) require.NoError(t, err) - // Then: one row with no effective group and zero spend. + // Then: per-user spend is summed across all days in the period. require.Len(t, got, 1) - require.Equal(t, user.ID, got[0].UserID) - require.Equal(t, org.ID, got[0].OrganizationID) - require.False(t, got[0].EffectiveGroupID.Valid, "expected no effective group") - require.Equal(t, int64(0), got[0].GroupSpendMicros) + require.False(t, got[0].EffectiveGroupID.Valid) + require.False(t, got[0].SpendLimitMicros.Valid) + require.False(t, got[0].LimitSource.Valid) + require.Equal(t, int64(600), got[0].GroupSpendMicros) }) - t.Run("GroupBudget", func(t *testing.T) { + t.Run("OverrideWins", func(t *testing.T) { t.Parallel() db, _ := dbtestutil.NewDB(t) ctx := testutil.Context(t, testutil.WaitShort) - // Given: a user in the queried group and the queried group has a budget. + // Given: a member of the queried group who is also in two other + // budgeted groups, one being the natural highest-limit and the other + // set via a user override. user := dbgen.User(t, db, database.User{}) org := dbgen.Organization(t, db, database.Organization{}) - group := dbgen.Group(t, db, database.Group{OrganizationID: org.ID}) + queried := dbgen.Group(t, db, database.Group{OrganizationID: org.ID}) + overrideTarget := dbgen.Group(t, db, database.Group{OrganizationID: org.ID}) + highestLimit := dbgen.Group(t, db, database.Group{OrganizationID: org.ID}) dbgen.OrganizationMember(t, db, database.OrganizationMember{UserID: user.ID, OrganizationID: org.ID}) - dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: group.ID, UserID: user.ID}) + dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: queried.ID, UserID: user.ID}) + dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: overrideTarget.ID, UserID: user.ID}) + dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: highestLimit.ID, UserID: user.ID}) _, err := db.UpsertGroupAIBudget(ctx, database.UpsertGroupAIBudgetParams{ - GroupID: group.ID, + GroupID: overrideTarget.ID, SpendLimitMicros: 1_000_000, }) require.NoError(t, err) + _, err = db.UpsertGroupAIBudget(ctx, database.UpsertGroupAIBudgetParams{ + GroupID: highestLimit.ID, + SpendLimitMicros: 5_000_000, + }) + require.NoError(t, err) + _, err = db.UpsertUserAIBudgetOverride(ctx, database.UpsertUserAIBudgetOverrideParams{ + UserID: user.ID, + GroupID: overrideTarget.ID, + SpendLimitMicros: 500_000, + }) + require.NoError(t, err) // When: querying spend for the user. got, err := db.GetGroupMembersAISpend(ctx, database.GetGroupMembersAISpendParams{ - GroupID: group.ID, + GroupID: queried.ID, UserIds: []uuid.UUID{user.ID}, PeriodStart: monthStart, }) require.NoError(t, err) - // Then: the effective group is the queried group, and the limit and - // source reflect the group budget. + // Then: the override target wins over the highest-limit group. require.Len(t, got, 1) - require.Equal(t, uuid.NullUUID{UUID: group.ID, Valid: true}, got[0].EffectiveGroupID) - require.Equal(t, sql.NullInt64{Int64: 1_000_000, Valid: true}, got[0].SpendLimitMicros) - require.Equal(t, sql.NullString{String: "group", Valid: true}, got[0].LimitSource) + require.Equal(t, uuid.NullUUID{UUID: overrideTarget.ID, Valid: true}, got[0].EffectiveGroupID) + require.False(t, got[0].SpendLimitMicros.Valid) + require.False(t, got[0].LimitSource.Valid) + require.Equal(t, int64(0), got[0].GroupSpendMicros) }) - t.Run("OverrideOnQueriedGroup", func(t *testing.T) { + t.Run("EqualBudgetTieBreak", func(t *testing.T) { t.Parallel() db, _ := dbtestutil.NewDB(t) ctx := testutil.Context(t, testutil.WaitShort) - // Given: a group member with a user override targeting the queried group. + // Given: a member of the queried group who is in two same-org groups + // with identical spend limits. user := dbgen.User(t, db, database.User{}) org := dbgen.Organization(t, db, database.Organization{}) - group := dbgen.Group(t, db, database.Group{OrganizationID: org.ID}) + queried := dbgen.Group(t, db, database.Group{OrganizationID: org.ID}) + groupA := dbgen.Group(t, db, database.Group{OrganizationID: org.ID, Name: "aaa-tie-group"}) + groupB := dbgen.Group(t, db, database.Group{OrganizationID: org.ID, Name: "bbb-tie-group"}) dbgen.OrganizationMember(t, db, database.OrganizationMember{UserID: user.ID, OrganizationID: org.ID}) - dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: group.ID, UserID: user.ID}) - _, err := db.UpsertUserAIBudgetOverride(ctx, database.UpsertUserAIBudgetOverrideParams{ - UserID: user.ID, - GroupID: group.ID, - SpendLimitMicros: 500_000, + dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: queried.ID, UserID: user.ID}) + dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: groupA.ID, UserID: user.ID}) + dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: groupB.ID, UserID: user.ID}) + _, err := db.UpsertGroupAIBudget(ctx, database.UpsertGroupAIBudgetParams{ + GroupID: groupA.ID, + SpendLimitMicros: 1_000_000, + }) + require.NoError(t, err) + _, err = db.UpsertGroupAIBudget(ctx, database.UpsertGroupAIBudgetParams{ + GroupID: groupB.ID, + SpendLimitMicros: 1_000_000, }) require.NoError(t, err) // When: querying spend for the user. got, err := db.GetGroupMembersAISpend(ctx, database.GetGroupMembersAISpendParams{ - GroupID: group.ID, + GroupID: queried.ID, UserIds: []uuid.UUID{user.ID}, PeriodStart: monthStart, }) require.NoError(t, err) - // Then: the limit comes from the override, with the user_override source. + // Then: the tie is broken by group name ascending, so groupA wins. require.Len(t, got, 1) - require.Equal(t, uuid.NullUUID{UUID: group.ID, Valid: true}, got[0].EffectiveGroupID) - require.Equal(t, sql.NullInt64{Int64: 500_000, Valid: true}, got[0].SpendLimitMicros) - require.Equal(t, sql.NullString{String: "user_override", Valid: true}, got[0].LimitSource) + require.Equal(t, uuid.NullUUID{UUID: groupA.ID, Valid: true}, got[0].EffectiveGroupID) + require.False(t, got[0].SpendLimitMicros.Valid) + require.False(t, got[0].LimitSource.Valid) + require.Equal(t, int64(0), got[0].GroupSpendMicros) }) - t.Run("Spend", func(t *testing.T) { + t.Run("EveryoneGroupCounts", func(t *testing.T) { t.Parallel() db, _ := dbtestutil.NewDB(t) ctx := testutil.Context(t, testutil.WaitShort) - // Given: a group member with spend attributed to the queried group. + // Given: a member of the queried group whose only budgeted group is + // the org's implicit Everyone group. user := dbgen.User(t, db, database.User{}) org := dbgen.Organization(t, db, database.Organization{}) - group := dbgen.Group(t, db, database.Group{OrganizationID: org.ID}) + queried := dbgen.Group(t, db, database.Group{OrganizationID: org.ID}) dbgen.OrganizationMember(t, db, database.OrganizationMember{UserID: user.ID, OrganizationID: org.ID}) - dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: group.ID, UserID: user.ID}) - _, err := db.IncrementUserAIDailySpend(ctx, database.IncrementUserAIDailySpendParams{ - UserID: user.ID, EffectiveGroupID: group.ID, Day: now, CostMicros: 250, + dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: queried.ID, UserID: user.ID}) + // The Everyone group has ID equal to the organization ID and must be + // inserted explicitly for this test's FK constraint on group_ai_budgets. + //nolint:gocritic // Requires system context. + _, err := db.InsertAllUsersGroup(dbauthz.AsSystemRestricted(ctx), org.ID) + require.NoError(t, err) + _, err = db.UpsertGroupAIBudget(ctx, database.UpsertGroupAIBudgetParams{ + GroupID: org.ID, + SpendLimitMicros: 1_000_000, }) require.NoError(t, err) // When: querying spend for the user. got, err := db.GetGroupMembersAISpend(ctx, database.GetGroupMembersAISpendParams{ - GroupID: group.ID, + GroupID: queried.ID, UserIds: []uuid.UUID{user.ID}, PeriodStart: monthStart, }) require.NoError(t, err) - // Then: the queried-group spend is aggregated. + // Then: effective_group_id resolves to the Everyone group. require.Len(t, got, 1) - require.Equal(t, int64(250), got[0].GroupSpendMicros) + require.Equal(t, uuid.NullUUID{UUID: org.ID, Valid: true}, got[0].EffectiveGroupID) + require.False(t, got[0].SpendLimitMicros.Valid) + require.False(t, got[0].LimitSource.Valid) + require.Equal(t, int64(0), got[0].GroupSpendMicros) }) t.Run("SpendWithDifferentEffectiveGroup", func(t *testing.T) { @@ -13279,34 +13479,6 @@ func TestGetGroupMembersAISpend(t *testing.T) { require.Equal(t, int64(250), got[0].GroupSpendMicros) }) - t.Run("HidesFormerMember", func(t *testing.T) { - t.Parallel() - db, _ := dbtestutil.NewDB(t) - ctx := testutil.Context(t, testutil.WaitShort) - - // Given: a user who has historical spend attributed to the queried - // group but is not currently a member of it. - user := dbgen.User(t, db, database.User{}) - org := dbgen.Organization(t, db, database.Organization{}) - queried := dbgen.Group(t, db, database.Group{OrganizationID: org.ID}) - dbgen.OrganizationMember(t, db, database.OrganizationMember{UserID: user.ID, OrganizationID: org.ID}) - _, err := db.IncrementUserAIDailySpend(ctx, database.IncrementUserAIDailySpendParams{ - UserID: user.ID, EffectiveGroupID: queried.ID, Day: now, CostMicros: 500, - }) - require.NoError(t, err) - - // When: querying the queried group's spend for the user. - got, err := db.GetGroupMembersAISpend(ctx, database.GetGroupMembersAISpendParams{ - GroupID: queried.ID, - UserIds: []uuid.UUID{user.ID}, - PeriodStart: monthStart, - }) - require.NoError(t, err) - - // Then: the user is filtered out and their historical spend is not returned. - require.Empty(t, got) - }) - t.Run("ExcludesOtherGroupSpend", func(t *testing.T) { t.Parallel() db, _ := dbtestutil.NewDB(t) @@ -13334,75 +13506,69 @@ func TestGetGroupMembersAISpend(t *testing.T) { // Then: spend attributed to the other group is not counted. require.Len(t, got, 1) + require.False(t, got[0].EffectiveGroupID.Valid) + require.False(t, got[0].SpendLimitMicros.Valid) + require.False(t, got[0].LimitSource.Valid) require.Equal(t, int64(0), got[0].GroupSpendMicros) }) - t.Run("MultipleMembers", func(t *testing.T) { + t.Run("ExcludesNonMembers", func(t *testing.T) { t.Parallel() db, _ := dbtestutil.NewDB(t) ctx := testutil.Context(t, testutil.WaitShort) - // Given: two group members with different spend attributed to the group. - userA := dbgen.User(t, db, database.User{}) - userB := dbgen.User(t, db, database.User{}) + // Given: a member of the queried group and a non-member in the same org. + member := dbgen.User(t, db, database.User{}) + nonMember := dbgen.User(t, db, database.User{}) org := dbgen.Organization(t, db, database.Organization{}) group := dbgen.Group(t, db, database.Group{OrganizationID: org.ID}) - dbgen.OrganizationMember(t, db, database.OrganizationMember{UserID: userA.ID, OrganizationID: org.ID}) - dbgen.OrganizationMember(t, db, database.OrganizationMember{UserID: userB.ID, OrganizationID: org.ID}) - dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: group.ID, UserID: userA.ID}) - dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: group.ID, UserID: userB.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: 250, - }) - require.NoError(t, err) + dbgen.OrganizationMember(t, db, database.OrganizationMember{UserID: member.ID, OrganizationID: org.ID}) + dbgen.OrganizationMember(t, db, database.OrganizationMember{UserID: nonMember.ID, OrganizationID: org.ID}) + dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: group.ID, UserID: member.ID}) - // When: querying spend for both users. + // When: querying with both user IDs. got, err := db.GetGroupMembersAISpend(ctx, database.GetGroupMembersAISpendParams{ GroupID: group.ID, - UserIds: []uuid.UUID{userA.ID, userB.ID}, + UserIds: []uuid.UUID{member.ID, nonMember.ID}, PeriodStart: monthStart, }) require.NoError(t, err) - // Then: both users are returned with their own aggregate spend. - require.Len(t, got, 2) - byID := make(map[uuid.UUID]database.GetGroupMembersAISpendRow, len(got)) - for _, r := range got { - byID[r.UserID] = r - } - require.Equal(t, int64(100), byID[userA.ID].GroupSpendMicros) - require.Equal(t, int64(250), byID[userB.ID].GroupSpendMicros) + // Then: only the group member is returned. + require.Len(t, got, 1) + require.Equal(t, member.ID, got[0].UserID) + require.False(t, got[0].EffectiveGroupID.Valid) + require.False(t, got[0].SpendLimitMicros.Valid) + require.False(t, got[0].LimitSource.Valid) + require.Equal(t, int64(0), got[0].GroupSpendMicros) }) - t.Run("ExcludesNonMembers", func(t *testing.T) { + t.Run("HidesFormerMember", func(t *testing.T) { t.Parallel() db, _ := dbtestutil.NewDB(t) ctx := testutil.Context(t, testutil.WaitShort) - // Given: a member of the queried group and a non-member in the same org. - member := dbgen.User(t, db, database.User{}) - nonMember := dbgen.User(t, db, database.User{}) + // Given: a user who has historical spend attributed to the queried + // group but is not currently a member of it. + user := dbgen.User(t, db, database.User{}) org := dbgen.Organization(t, db, database.Organization{}) - group := dbgen.Group(t, db, database.Group{OrganizationID: org.ID}) - dbgen.OrganizationMember(t, db, database.OrganizationMember{UserID: member.ID, OrganizationID: org.ID}) - dbgen.OrganizationMember(t, db, database.OrganizationMember{UserID: nonMember.ID, OrganizationID: org.ID}) - dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: group.ID, UserID: member.ID}) + queried := dbgen.Group(t, db, database.Group{OrganizationID: org.ID}) + dbgen.OrganizationMember(t, db, database.OrganizationMember{UserID: user.ID, OrganizationID: org.ID}) + _, err := db.IncrementUserAIDailySpend(ctx, database.IncrementUserAIDailySpendParams{ + UserID: user.ID, EffectiveGroupID: queried.ID, Day: now, CostMicros: 500, + }) + require.NoError(t, err) - // When: querying with both user IDs. + // When: querying the queried group's spend for the user. got, err := db.GetGroupMembersAISpend(ctx, database.GetGroupMembersAISpendParams{ - GroupID: group.ID, - UserIds: []uuid.UUID{member.ID, nonMember.ID}, + GroupID: queried.ID, + UserIds: []uuid.UUID{user.ID}, PeriodStart: monthStart, }) require.NoError(t, err) - // Then: only the group member is returned. - require.Len(t, got, 1) - require.Equal(t, member.ID, got[0].UserID) + // Then: the user is filtered out and their historical spend is not returned. + require.Empty(t, got) }) t.Run("CrossOrgEffectiveGroupMasked", func(t *testing.T) { @@ -13411,7 +13577,7 @@ func TestGetGroupMembersAISpend(t *testing.T) { ctx := testutil.Context(t, testutil.WaitShort) // Given: a member of the queried group whose highest-limit budget - // group is in a different org (via cross-org group membership). + // group is in a different org. user := dbgen.User(t, db, database.User{}) org := dbgen.Organization(t, db, database.Organization{}) otherOrg := dbgen.Organization(t, db, database.Organization{}) @@ -13450,129 +13616,6 @@ func TestGetGroupMembersAISpend(t *testing.T) { require.Equal(t, int64(250), got[0].GroupSpendMicros) }) - t.Run("OverrideWins", func(t *testing.T) { - t.Parallel() - db, _ := dbtestutil.NewDB(t) - ctx := testutil.Context(t, testutil.WaitShort) - - // Given: a member of the queried group who is also in two other - // budgeted groups, one being the natural highest-limit and the other - // set via a user override. - user := dbgen.User(t, db, database.User{}) - org := dbgen.Organization(t, db, database.Organization{}) - queried := dbgen.Group(t, db, database.Group{OrganizationID: org.ID}) - overrideTarget := dbgen.Group(t, db, database.Group{OrganizationID: org.ID}) - highestLimit := dbgen.Group(t, db, database.Group{OrganizationID: org.ID}) - dbgen.OrganizationMember(t, db, database.OrganizationMember{UserID: user.ID, OrganizationID: org.ID}) - dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: queried.ID, UserID: user.ID}) - dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: overrideTarget.ID, UserID: user.ID}) - dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: highestLimit.ID, UserID: user.ID}) - _, err := db.UpsertGroupAIBudget(ctx, database.UpsertGroupAIBudgetParams{ - GroupID: overrideTarget.ID, - SpendLimitMicros: 1_000_000, - }) - require.NoError(t, err) - _, err = db.UpsertGroupAIBudget(ctx, database.UpsertGroupAIBudgetParams{ - GroupID: highestLimit.ID, - SpendLimitMicros: 5_000_000, - }) - require.NoError(t, err) - _, err = db.UpsertUserAIBudgetOverride(ctx, database.UpsertUserAIBudgetOverrideParams{ - UserID: user.ID, - GroupID: overrideTarget.ID, - SpendLimitMicros: 500_000, - }) - require.NoError(t, err) - - // When: querying spend for the user. - got, err := db.GetGroupMembersAISpend(ctx, database.GetGroupMembersAISpendParams{ - GroupID: queried.ID, - UserIds: []uuid.UUID{user.ID}, - PeriodStart: monthStart, - }) - require.NoError(t, err) - - // Then: the override target wins over the highest-limit group. - require.Len(t, got, 1) - require.Equal(t, uuid.NullUUID{UUID: overrideTarget.ID, Valid: true}, got[0].EffectiveGroupID) - }) - - t.Run("TieBreakByGroupName", func(t *testing.T) { - t.Parallel() - db, _ := dbtestutil.NewDB(t) - ctx := testutil.Context(t, testutil.WaitShort) - - // Given: a member of the queried group who is in two same-org groups - // with identical spend limits. - user := dbgen.User(t, db, database.User{}) - org := dbgen.Organization(t, db, database.Organization{}) - queried := dbgen.Group(t, db, database.Group{OrganizationID: org.ID}) - groupA := dbgen.Group(t, db, database.Group{OrganizationID: org.ID, Name: "aaa-tie-group"}) - groupB := dbgen.Group(t, db, database.Group{OrganizationID: org.ID, Name: "bbb-tie-group"}) - dbgen.OrganizationMember(t, db, database.OrganizationMember{UserID: user.ID, OrganizationID: org.ID}) - dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: queried.ID, UserID: user.ID}) - dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: groupA.ID, UserID: user.ID}) - dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: groupB.ID, UserID: user.ID}) - _, err := db.UpsertGroupAIBudget(ctx, database.UpsertGroupAIBudgetParams{ - GroupID: groupA.ID, - SpendLimitMicros: 1_000_000, - }) - require.NoError(t, err) - _, err = db.UpsertGroupAIBudget(ctx, database.UpsertGroupAIBudgetParams{ - GroupID: groupB.ID, - SpendLimitMicros: 1_000_000, - }) - require.NoError(t, err) - - // When: querying spend for the user. - got, err := db.GetGroupMembersAISpend(ctx, database.GetGroupMembersAISpendParams{ - GroupID: queried.ID, - UserIds: []uuid.UUID{user.ID}, - PeriodStart: monthStart, - }) - require.NoError(t, err) - - // Then: the tie is broken by group name ascending, so groupA wins. - require.Len(t, got, 1) - require.Equal(t, uuid.NullUUID{UUID: groupA.ID, Valid: true}, got[0].EffectiveGroupID) - }) - - t.Run("EveryoneGroupCounts", func(t *testing.T) { - t.Parallel() - db, _ := dbtestutil.NewDB(t) - ctx := testutil.Context(t, testutil.WaitShort) - - // Given: a member of the queried group whose only budgeted group is - // the org's implicit Everyone group. - user := dbgen.User(t, db, database.User{}) - org := dbgen.Organization(t, db, database.Organization{}) - queried := dbgen.Group(t, db, database.Group{OrganizationID: org.ID}) - dbgen.OrganizationMember(t, db, database.OrganizationMember{UserID: user.ID, OrganizationID: org.ID}) - dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: queried.ID, UserID: user.ID}) - // The Everyone group has ID equal to the organization ID and must be - // inserted explicitly for this test's FK constraint on group_ai_budgets. - //nolint:gocritic // Requires system context. - _, err := db.InsertAllUsersGroup(dbauthz.AsSystemRestricted(ctx), org.ID) - require.NoError(t, err) - _, err = db.UpsertGroupAIBudget(ctx, database.UpsertGroupAIBudgetParams{ - GroupID: org.ID, - SpendLimitMicros: 1_000_000, - }) - require.NoError(t, err) - - // When: querying spend for the user. - got, err := db.GetGroupMembersAISpend(ctx, database.GetGroupMembersAISpendParams{ - GroupID: queried.ID, - UserIds: []uuid.UUID{user.ID}, - PeriodStart: monthStart, - }) - require.NoError(t, err) - - // Then: effective_group_id resolves to the Everyone group (org.ID). - require.Len(t, got, 1) - require.Equal(t, uuid.NullUUID{UUID: org.ID, Valid: true}, got[0].EffectiveGroupID) - }) - t.Run("ExcludesSpendBeforePeriodStart", func(t *testing.T) { t.Parallel() db, _ := dbtestutil.NewDB(t) @@ -13603,41 +13646,12 @@ func TestGetGroupMembersAISpend(t *testing.T) { // Then: only current-period spend is aggregated. require.Len(t, got, 1) + require.False(t, got[0].EffectiveGroupID.Valid) + require.False(t, got[0].SpendLimitMicros.Valid) + require.False(t, got[0].LimitSource.Valid) require.Equal(t, int64(25), got[0].GroupSpendMicros) }) - t.Run("AggregatesSpendAcrossDays", func(t *testing.T) { - t.Parallel() - db, _ := dbtestutil.NewDB(t) - ctx := testutil.Context(t, testutil.WaitShort) - - // Given: a group member with spend on multiple days in the 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}) - dbgen.OrganizationMember(t, db, database.OrganizationMember{UserID: user.ID, OrganizationID: org.ID}) - dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: group.ID, UserID: user.ID}) - days := []time.Time{monthStart, monthStart.AddDate(0, 0, 7), monthStart.AddDate(0, 0, 14)} - for i, day := range days { - _, err := db.IncrementUserAIDailySpend(ctx, database.IncrementUserAIDailySpendParams{ - UserID: user.ID, EffectiveGroupID: group.ID, Day: day, CostMicros: int64((i + 1) * 100), - }) - require.NoError(t, err) - } - - // When: querying the group's member spend. - got, err := db.GetGroupMembersAISpend(ctx, database.GetGroupMembersAISpendParams{ - GroupID: group.ID, - UserIds: []uuid.UUID{user.ID}, - PeriodStart: monthStart, - }) - require.NoError(t, err) - - // Then: per-user spend is summed across all days in the period. - require.Len(t, got, 1) - require.Equal(t, int64(600), got[0].GroupSpendMicros) - }) - t.Run("NormalizesNonUTCPeriodStart", func(t *testing.T) { t.Parallel() db, _ := dbtestutil.NewDB(t) @@ -13670,6 +13684,9 @@ func TestGetGroupMembersAISpend(t *testing.T) { // Then: the prior UTC day's spend is excluded from the aggregate. require.Len(t, got, 1) + require.False(t, got[0].EffectiveGroupID.Valid) + require.False(t, got[0].SpendLimitMicros.Valid) + require.False(t, got[0].LimitSource.Valid) require.Equal(t, int64(25), got[0].GroupSpendMicros, "sum must exclude prevMonthLastDay row after normalization") }) diff --git a/enterprise/coderd/aibridge.go b/enterprise/coderd/aibridge.go index a658c2fa05321..49ba9764b5d95 100644 --- a/enterprise/coderd/aibridge.go +++ b/enterprise/coderd/aibridge.go @@ -1109,22 +1109,7 @@ func (api *API) groupMembersAISpend(rw http.ResponseWriter, r *http.Request) { Members: make([]codersdk.GroupMemberAISpend, 0, len(rows)), } for _, row := range rows { - entry := codersdk.GroupMemberAISpend{ - UserID: row.UserID, - GroupSpendMicros: row.GroupSpendMicros, - } - if row.EffectiveGroupID.Valid { - uid := row.EffectiveGroupID.UUID - entry.EffectiveGroupID = &uid - } - if row.SpendLimitMicros.Valid { - entry.SpendLimitMicros = &row.SpendLimitMicros.Int64 - } - if row.LimitSource.Valid { - source := codersdk.AIBudgetLimitSource(row.LimitSource.String) - entry.LimitSource = &source - } - resp.Members = append(resp.Members, entry) + resp.Members = append(resp.Members, db2sdk.GroupMemberAISpend(row)) } httpapi.Write(ctx, rw, http.StatusOK, resp) diff --git a/enterprise/coderd/aibridge_test.go b/enterprise/coderd/aibridge_test.go index a07e849386899..092501294cb65 100644 --- a/enterprise/coderd/aibridge_test.go +++ b/enterprise/coderd/aibridge_test.go @@ -3671,8 +3671,118 @@ func TestGroupMembersAISpend(t *testing.T) { // Then: only the primary-org user is returned. require.Len(t, resp.Members, 1) require.Equal(t, targetUser.ID, resp.Members[0].UserID) + require.Nil(t, resp.Members[0].EffectiveGroupID) + require.Nil(t, resp.Members[0].SpendLimitMicros) + require.Nil(t, resp.Members[0].LimitSource) + require.Equal(t, int64(0), resp.Members[0].GroupSpendMicros) }) + tests := []struct { + name string + groupLimit int64 + overrideLimit int64 + spent int64 + wantEffectiveGroup bool + wantSpendLimit *int64 + wantLimitSource *codersdk.AIBudgetLimitSource + wantSpendMicros int64 + }{ + { + name: "NoBudgetNoSpend", + }, + { + name: "BudgetZeroSpend", + groupLimit: 1_000_000_000, + wantEffectiveGroup: true, + wantSpendLimit: ptr.Ref(int64(1_000_000_000)), + wantLimitSource: ptr.Ref(codersdk.AIBudgetLimitSourceGroup), + }, + { + name: "BudgetWithSpend", + groupLimit: 1_000_000_000, + spent: 250_000_000, + wantEffectiveGroup: true, + wantSpendLimit: ptr.Ref(int64(1_000_000_000)), + wantLimitSource: ptr.Ref(codersdk.AIBudgetLimitSourceGroup), + wantSpendMicros: 250_000_000, + }, + { + name: "NoBudgetWithSpend", + spent: 100_000_000, + wantSpendMicros: 100_000_000, + }, + { + name: "OverrideBudget", + overrideLimit: 500_000_000, + wantEffectiveGroup: true, + wantSpendLimit: ptr.Ref(int64(500_000_000)), + wantLimitSource: ptr.Ref(codersdk.AIBudgetLimitSourceUserOverride), + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + // Given: an admin, a group with a member, optionally a group budget, + // a user override, and seeded spend attributed to the group. + clock := quartz.NewMock(t) + db, ps := dbtestutil.NewDB(t) + adminClient, targetUser, group := setupAICostControlTest(t, aiCostControlTestOptions{ + GroupName: "members-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.groupLimit > 0 { + _, err := adminClient.UpsertGroupAIBudget(ctx, group.ID, codersdk.UpsertGroupAIBudgetRequest{ + SpendLimitMicros: tt.groupLimit, + }) + require.NoError(t, err) + } + if tt.overrideLimit > 0 { + _, err := adminClient.UpsertUserAIBudgetOverride(ctx, targetUser.ID, codersdk.UpsertUserAIBudgetOverrideRequest{ + GroupID: group.ID, + SpendLimitMicros: tt.overrideLimit, + }) + 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 member spend. + got, err := adminClient.GroupMembersAISpend(ctx, group.ID, []uuid.UUID{targetUser.ID}) + require.NoError(t, err) + + // Then: one row is returned with the expected fields. + require.Equal(t, wantPeriodStart, got.PeriodStart) + require.Equal(t, wantPeriodEnd, got.PeriodEnd) + require.Len(t, got.Members, 1) + require.Equal(t, targetUser.ID, got.Members[0].UserID) + if tt.wantEffectiveGroup { + require.NotNil(t, got.Members[0].EffectiveGroupID) + require.Equal(t, group.ID, *got.Members[0].EffectiveGroupID) + } else { + require.Nil(t, got.Members[0].EffectiveGroupID) + } + require.Equal(t, tt.wantSpendLimit, got.Members[0].SpendLimitMicros) + require.Equal(t, tt.wantLimitSource, got.Members[0].LimitSource) + require.Equal(t, tt.wantSpendMicros, got.Members[0].GroupSpendMicros) + }) + } + t.Run("CrossOrgEffectiveGroupMasked", func(t *testing.T) { t.Parallel() @@ -3728,69 +3838,9 @@ func TestGroupMembersAISpend(t *testing.T) { require.Len(t, resp.Members, 1) require.Equal(t, targetUser.ID, resp.Members[0].UserID) require.Nil(t, resp.Members[0].EffectiveGroupID, "cross-org effective group must be masked even for the owner") - }) - - t.Run("UserOverrideLimitSource", func(t *testing.T) { - t.Parallel() - - adminClient, targetUser, group := setupAICostControlTest(t, aiCostControlTestOptions{GroupName: "override-source-group"}) - ctx := testutil.Context(t, testutil.WaitLong) - - // Given: a user override targeting the queried group. - _, err := adminClient.UpsertUserAIBudgetOverride(ctx, targetUser.ID, codersdk.UpsertUserAIBudgetOverrideRequest{ - GroupID: group.ID, - SpendLimitMicros: 500_000_000, - }) - require.NoError(t, err) - - // When: querying the group's member spend. - resp, err := adminClient.GroupMembersAISpend(ctx, group.ID, []uuid.UUID{targetUser.ID}) - require.NoError(t, err) - - // Then: the limit and source reflect the override. - require.Len(t, resp.Members, 1) - require.NotNil(t, resp.Members[0].SpendLimitMicros) - require.Equal(t, int64(500_000_000), *resp.Members[0].SpendLimitMicros) - require.NotNil(t, resp.Members[0].LimitSource) - require.Equal(t, codersdk.AIBudgetLimitSourceUserOverride, *resp.Members[0].LimitSource) - }) - - t.Run("MemberCanOnlyReadOwnRow", func(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, - }, - }, - }) - userAdminClient, _ := coderdtest.CreateAnotherUser(t, ownerClient, owner.OrganizationID, rbac.RoleUserAdmin()) - memberClient, member := coderdtest.CreateAnotherUser(t, ownerClient, owner.OrganizationID) - _, otherMember := coderdtest.CreateAnotherUser(t, ownerClient, owner.OrganizationID) - - ctx := testutil.Context(t, testutil.WaitLong) - group, err := userAdminClient.CreateGroup(ctx, owner.OrganizationID, codersdk.CreateGroupRequest{ - Name: "member-privacy-group", - }) - require.NoError(t, err) - _, err = userAdminClient.PatchGroup(ctx, group.ID, codersdk.PatchGroupRequest{ - AddUsers: []string{member.ID.String(), otherMember.ID.String()}, - }) - require.NoError(t, err) - - // When: the member queries spend for themselves and another member. - resp, err := memberClient.GroupMembersAISpend(ctx, group.ID, []uuid.UUID{member.ID, otherMember.ID}) - require.NoError(t, err) - - // Then: only the caller's own row is returned. - require.Len(t, resp.Members, 1) - require.Equal(t, member.ID, resp.Members[0].UserID) + require.Nil(t, resp.Members[0].SpendLimitMicros) + require.Nil(t, resp.Members[0].LimitSource) + require.Equal(t, int64(0), resp.Members[0].GroupSpendMicros) }) t.Run("OrgScopedRoute", func(t *testing.T) { @@ -3818,99 +3868,11 @@ func TestGroupMembersAISpend(t *testing.T) { require.NoError(t, json.NewDecoder(res.Body).Decode(&got)) require.Len(t, got.Members, 1) require.Equal(t, targetUser.ID, got.Members[0].UserID) + require.Nil(t, got.Members[0].EffectiveGroupID) + require.Nil(t, got.Members[0].SpendLimitMicros) + require.Nil(t, got.Members[0].LimitSource) + require.Equal(t, int64(0), got.Members[0].GroupSpendMicros) }) - - groupSource := codersdk.AIBudgetLimitSourceGroup - tests := []struct { - name string - setBudget bool - spent int64 - wantEffectiveSet bool - wantSpendLimit *int64 - wantLimitSource *codersdk.AIBudgetLimitSource - wantSpendMicros int64 - }{ - { - name: "NoBudgetNoSpend", - }, - { - name: "BudgetZeroSpend", - setBudget: true, - wantEffectiveSet: true, - wantSpendLimit: ptr.Ref(int64(1_000_000_000)), - wantLimitSource: &groupSource, - }, - { - name: "BudgetWithSpend", - setBudget: true, - spent: 250_000_000, - wantEffectiveSet: true, - wantSpendLimit: ptr.Ref(int64(1_000_000_000)), - wantLimitSource: &groupSource, - wantSpendMicros: 250_000_000, - }, - { - name: "NoBudgetWithSpend", - spent: 100_000_000, - wantSpendMicros: 100_000_000, - }, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - t.Parallel() - - // Given: an admin, a group with a member, optionally a budget on the - // group and seeded spend attributed to the group. - clock := quartz.NewMock(t) - db, ps := dbtestutil.NewDB(t) - adminClient, targetUser, group := setupAICostControlTest(t, aiCostControlTestOptions{ - GroupName: "members-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: 1_000_000_000, - }) - 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 member spend. - got, err := adminClient.GroupMembersAISpend(ctx, group.ID, []uuid.UUID{targetUser.ID}) - require.NoError(t, err) - - // Then: one row is returned with the expected fields. - require.Equal(t, wantPeriodStart, got.PeriodStart) - require.Equal(t, wantPeriodEnd, got.PeriodEnd) - require.Len(t, got.Members, 1) - require.Equal(t, targetUser.ID, got.Members[0].UserID) - if tt.wantEffectiveSet { - require.NotNil(t, got.Members[0].EffectiveGroupID) - require.Equal(t, group.ID, *got.Members[0].EffectiveGroupID) - } else { - require.Nil(t, got.Members[0].EffectiveGroupID) - } - require.Equal(t, tt.wantSpendLimit, got.Members[0].SpendLimitMicros) - require.Equal(t, tt.wantLimitSource, got.Members[0].LimitSource) - require.Equal(t, tt.wantSpendMicros, got.Members[0].GroupSpendMicros) - }) - } } func TestGroupMembersAISpendRoleAccess(t *testing.T) { @@ -3933,6 +3895,7 @@ func TestGroupMembersAISpendRoleAccess(t *testing.T) { orgAdminClient, _ := coderdtest.CreateAnotherUser(t, ownerClient, owner.OrganizationID, rbac.ScopedRoleOrgAdmin(owner.OrganizationID)) orgUserAdminClient, _ := coderdtest.CreateAnotherUser(t, ownerClient, owner.OrganizationID, rbac.ScopedRoleOrgUserAdmin(owner.OrganizationID)) memberClient, member := coderdtest.CreateAnotherUser(t, ownerClient, owner.OrganizationID) + _, otherMember := coderdtest.CreateAnotherUser(t, ownerClient, owner.OrganizationID) otherOrg := coderdenttest.CreateOrganization(t, ownerClient, coderdenttest.CreateOrganizationOptions{}) otherOrgMemberClient, _ := coderdtest.CreateAnotherUser(t, ownerClient, otherOrg.ID) @@ -3943,7 +3906,7 @@ func TestGroupMembersAISpendRoleAccess(t *testing.T) { }) require.NoError(t, err) _, err = userAdminClient.PatchGroup(ctx, group.ID, codersdk.PatchGroupRequest{ - AddUsers: []string{member.ID.String()}, + AddUsers: []string{member.ID.String(), otherMember.ID.String()}, }) require.NoError(t, err) @@ -3978,6 +3941,20 @@ func TestGroupMembersAISpendRoleAccess(t *testing.T) { require.Equal(t, member.ID, resp.Members[0].UserID) }) } + + t.Run("MemberCanOnlyReadOwnRow", func(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + + // When: a member queries spend for themselves and another member. + resp, err := memberClient.GroupMembersAISpend(ctx, group.ID, []uuid.UUID{member.ID, otherMember.ID}) + require.NoError(t, err) + + // Then: only the caller's own row is returned. + require.Len(t, resp.Members, 1) + require.Equal(t, member.ID, resp.Members[0].UserID) + }) } // aiCostControlTestOptions configures the setup of an AI cost control test From d7f636998ef55104b2dfb2f3ae64aad78f321fc9 Mon Sep 17 00:00:00 2001 From: Susana Cardoso Ferreira Date: Thu, 16 Jul 2026 16:00:32 +0000 Subject: [PATCH 7/8] refactor: nest group member spend limit and source under group_budget --- coderd/apidoc/docs.go | 25 +++++++++++------- coderd/apidoc/swagger.json | 25 +++++++++++------- coderd/database/db2sdk/db2sdk.go | 9 +++---- codersdk/aibridge.go | 19 ++++++++------ docs/reference/api/enterprise.md | 12 ++++++--- docs/reference/api/schemas.md | 41 ++++++++++++++++++++++-------- enterprise/coderd/aibridge_test.go | 35 ++++++++++++------------- site/src/api/typesGenerated.ts | 24 +++++++++-------- 8 files changed, 117 insertions(+), 73 deletions(-) diff --git a/coderd/apidoc/docs.go b/coderd/apidoc/docs.go index 3302426bac1ff..d5a35e6601a0c 100644 --- a/coderd/apidoc/docs.go +++ b/coderd/apidoc/docs.go @@ -15486,6 +15486,17 @@ const docTemplate = `{ } } }, + "codersdk.AIGroupBudget": { + "type": "object", + "properties": { + "limit_source": { + "$ref": "#/definitions/codersdk.AIBudgetLimitSource" + }, + "spend_limit_micros": { + "type": "integer" + } + } + }, "codersdk.AIProvider": { "type": "object", "properties": { @@ -20359,20 +20370,16 @@ const docTemplate = `{ "type": "string", "format": "uuid" }, - "group_spend_micros": { - "description": "GroupSpendMicros is the user's spend attributed to the queried group\nover the current budget period.", - "type": "integer" - }, - "limit_source": { - "description": "LimitSource identifies the tier that produced the limit. Null when the\nuser's budget resolves to another group or no budget applies to the user\n(unlimited).", + "group_budget": { + "description": "GroupBudget is the budget when the queried group is this user's\neffective budget source. Null when the user's budget resolves to another\ngroup or no budget applies to the user.", "allOf": [ { - "$ref": "#/definitions/codersdk.AIBudgetLimitSource" + "$ref": "#/definitions/codersdk.AIGroupBudget" } ] }, - "spend_limit_micros": { - "description": "SpendLimitMicros is the spend limit when the queried group is this\nuser's effective budget source. Null when the user's budget resolves to\nanother group or no budget applies to the user (unlimited).", + "group_spend_micros": { + "description": "GroupSpendMicros is the user's spend attributed to the queried group\nover the current budget period.", "type": "integer" }, "user_id": { diff --git a/coderd/apidoc/swagger.json b/coderd/apidoc/swagger.json index 9d72cb80d1647..c1d1316b51358 100644 --- a/coderd/apidoc/swagger.json +++ b/coderd/apidoc/swagger.json @@ -13813,6 +13813,17 @@ } } }, + "codersdk.AIGroupBudget": { + "type": "object", + "properties": { + "limit_source": { + "$ref": "#/definitions/codersdk.AIBudgetLimitSource" + }, + "spend_limit_micros": { + "type": "integer" + } + } + }, "codersdk.AIProvider": { "type": "object", "properties": { @@ -18529,20 +18540,16 @@ "type": "string", "format": "uuid" }, - "group_spend_micros": { - "description": "GroupSpendMicros is the user's spend attributed to the queried group\nover the current budget period.", - "type": "integer" - }, - "limit_source": { - "description": "LimitSource identifies the tier that produced the limit. Null when the\nuser's budget resolves to another group or no budget applies to the user\n(unlimited).", + "group_budget": { + "description": "GroupBudget is the budget when the queried group is this user's\neffective budget source. Null when the user's budget resolves to another\ngroup or no budget applies to the user.", "allOf": [ { - "$ref": "#/definitions/codersdk.AIBudgetLimitSource" + "$ref": "#/definitions/codersdk.AIGroupBudget" } ] }, - "spend_limit_micros": { - "description": "SpendLimitMicros is the spend limit when the queried group is this\nuser's effective budget source. Null when the user's budget resolves to\nanother group or no budget applies to the user (unlimited).", + "group_spend_micros": { + "description": "GroupSpendMicros is the user's spend attributed to the queried group\nover the current budget period.", "type": "integer" }, "user_id": { diff --git a/coderd/database/db2sdk/db2sdk.go b/coderd/database/db2sdk/db2sdk.go index af41242168213..21aeaad8b5014 100644 --- a/coderd/database/db2sdk/db2sdk.go +++ b/coderd/database/db2sdk/db2sdk.go @@ -1478,11 +1478,10 @@ func GroupMemberAISpend(row database.GetGroupMembersAISpendRow) codersdk.GroupMe member.EffectiveGroupID = &row.EffectiveGroupID.UUID } if row.SpendLimitMicros.Valid { - member.SpendLimitMicros = &row.SpendLimitMicros.Int64 - } - if row.LimitSource.Valid { - source := codersdk.AIBudgetLimitSource(row.LimitSource.String) - member.LimitSource = &source + member.GroupBudget = &codersdk.AIGroupBudget{ + SpendLimitMicros: row.SpendLimitMicros.Int64, + LimitSource: codersdk.AIBudgetLimitSource(row.LimitSource.String), + } } return member } diff --git a/codersdk/aibridge.go b/codersdk/aibridge.go index 78f8491b29030..13253e351ae86 100644 --- a/codersdk/aibridge.go +++ b/codersdk/aibridge.go @@ -27,6 +27,13 @@ const ( AIBudgetLimitSourceGroup AIBudgetLimitSource = "group" ) +// AIGroupBudget is an AI spend limit and the tier that produced it. Both +// fields are always populated together. +type AIGroupBudget struct { + SpendLimitMicros int64 `json:"spend_limit_micros"` + LimitSource AIBudgetLimitSource `json:"limit_source"` +} + // UserAIBudgetSummary is the effective AI budget for a user. When no // budget applies, all fields except UserID are null. type UserAIBudgetSummary struct { @@ -98,14 +105,10 @@ type GroupMemberAISpend struct { // this organization, including when the user's budget resolves to a group // in another organization. EffectiveGroupID *uuid.UUID `json:"effective_group_id" format:"uuid"` - // SpendLimitMicros is the spend limit when the queried group is this - // user's effective budget source. Null when the user's budget resolves to - // another group or no budget applies to the user (unlimited). - SpendLimitMicros *int64 `json:"spend_limit_micros"` - // LimitSource identifies the tier that produced the limit. Null when the - // user's budget resolves to another group or no budget applies to the user - // (unlimited). - LimitSource *AIBudgetLimitSource `json:"limit_source"` + // GroupBudget is the budget when the queried group is this user's + // effective budget source. Null when the user's budget resolves to another + // group or no budget applies to the user. + GroupBudget *AIGroupBudget `json:"group_budget"` // GroupSpendMicros is the user's spend attributed to the queried group // over the current budget period. GroupSpendMicros int64 `json:"group_spend_micros"` diff --git a/docs/reference/api/enterprise.md b/docs/reference/api/enterprise.md index a48e9d2009892..eeaf8032ec0a3 100644 --- a/docs/reference/api/enterprise.md +++ b/docs/reference/api/enterprise.md @@ -1125,9 +1125,11 @@ User IDs that are not members of the group, or that the caller has no read acces "members": [ { "effective_group_id": "85e2b926-ddfb-4c66-b68e-b66e5acec6c0", + "group_budget": { + "limit_source": "user_override", + "spend_limit_micros": 0 + }, "group_spend_micros": 0, - "limit_source": "user_override", - "spend_limit_micros": 0, "user_id": "a169451c-8525-4352-b8ca-070dd449a1a5" } ], @@ -2096,9 +2098,11 @@ User IDs that are not members of the group, or that the caller has no read acces "members": [ { "effective_group_id": "85e2b926-ddfb-4c66-b68e-b66e5acec6c0", + "group_budget": { + "limit_source": "user_override", + "spend_limit_micros": 0 + }, "group_spend_micros": 0, - "limit_source": "user_override", - "spend_limit_micros": 0, "user_id": "a169451c-8525-4352-b8ca-070dd449a1a5" } ], diff --git a/docs/reference/api/schemas.md b/docs/reference/api/schemas.md index 1daf3bc906cba..e7f4ac8151d36 100644 --- a/docs/reference/api/schemas.md +++ b/docs/reference/api/schemas.md @@ -1030,6 +1030,22 @@ | `last_heartbeat_at` | string | false | | | | `name` | string | false | | | +## codersdk.AIGroupBudget + +```json +{ + "limit_source": "user_override", + "spend_limit_micros": 0 +} +``` + +### Properties + +| Name | Type | Required | Restrictions | Description | +|----------------------|--------------------------------------------------------------|----------|--------------|-------------| +| `limit_source` | [codersdk.AIBudgetLimitSource](#codersdkaibudgetlimitsource) | false | | | +| `spend_limit_micros` | integer | false | | | + ## codersdk.AIProvider ```json @@ -7690,22 +7706,23 @@ Only certain features set these fields: - FeatureManagedAgentLimit| ```json { "effective_group_id": "85e2b926-ddfb-4c66-b68e-b66e5acec6c0", + "group_budget": { + "limit_source": "user_override", + "spend_limit_micros": 0 + }, "group_spend_micros": 0, - "limit_source": "user_override", - "spend_limit_micros": 0, "user_id": "a169451c-8525-4352-b8ca-070dd449a1a5" } ``` ### Properties -| Name | Type | Required | Restrictions | Description | -|----------------------|--------------------------------------------------------------|----------|--------------|-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------| -| `effective_group_id` | string | false | | Effective group ID is the user's effective budget group within the queried group's organization. Null when no effective budget group is visible in this organization, including when the user's budget resolves to a group in another organization. | -| `group_spend_micros` | integer | false | | Group spend micros is the user's spend attributed to the queried group over the current budget period. | -| `limit_source` | [codersdk.AIBudgetLimitSource](#codersdkaibudgetlimitsource) | false | | Limit source identifies the tier that produced the limit. Null when the user's budget resolves to another group or no budget applies to the user (unlimited). | -| `spend_limit_micros` | integer | false | | Spend limit micros is the spend limit when the queried group is this user's effective budget source. Null when the user's budget resolves to another group or no budget applies to the user (unlimited). | -| `user_id` | string | false | | | +| Name | Type | Required | Restrictions | Description | +|----------------------|--------------------------------------------------|----------|--------------|-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------| +| `effective_group_id` | string | false | | Effective group ID is the user's effective budget group within the queried group's organization. Null when no effective budget group is visible in this organization, including when the user's budget resolves to a group in another organization. | +| `group_budget` | [codersdk.AIGroupBudget](#codersdkaigroupbudget) | false | | Group budget is the budget when the queried group is this user's effective budget source. Null when the user's budget resolves to another group or no budget applies to the user. | +| `group_spend_micros` | integer | false | | Group spend micros is the user's spend attributed to the queried group over the current budget period. | +| `user_id` | string | false | | | ## codersdk.GroupMembersAISpend @@ -7714,9 +7731,11 @@ Only certain features set these fields: - FeatureManagedAgentLimit| "members": [ { "effective_group_id": "85e2b926-ddfb-4c66-b68e-b66e5acec6c0", + "group_budget": { + "limit_source": "user_override", + "spend_limit_micros": 0 + }, "group_spend_micros": 0, - "limit_source": "user_override", - "spend_limit_micros": 0, "user_id": "a169451c-8525-4352-b8ca-070dd449a1a5" } ], diff --git a/enterprise/coderd/aibridge_test.go b/enterprise/coderd/aibridge_test.go index 092501294cb65..9f3da140997c9 100644 --- a/enterprise/coderd/aibridge_test.go +++ b/enterprise/coderd/aibridge_test.go @@ -3672,8 +3672,7 @@ func TestGroupMembersAISpend(t *testing.T) { require.Len(t, resp.Members, 1) require.Equal(t, targetUser.ID, resp.Members[0].UserID) require.Nil(t, resp.Members[0].EffectiveGroupID) - require.Nil(t, resp.Members[0].SpendLimitMicros) - require.Nil(t, resp.Members[0].LimitSource) + require.Nil(t, resp.Members[0].GroupBudget) require.Equal(t, int64(0), resp.Members[0].GroupSpendMicros) }) @@ -3683,8 +3682,7 @@ func TestGroupMembersAISpend(t *testing.T) { overrideLimit int64 spent int64 wantEffectiveGroup bool - wantSpendLimit *int64 - wantLimitSource *codersdk.AIBudgetLimitSource + wantGroupBudget *codersdk.AIGroupBudget wantSpendMicros int64 }{ { @@ -3694,17 +3692,21 @@ func TestGroupMembersAISpend(t *testing.T) { name: "BudgetZeroSpend", groupLimit: 1_000_000_000, wantEffectiveGroup: true, - wantSpendLimit: ptr.Ref(int64(1_000_000_000)), - wantLimitSource: ptr.Ref(codersdk.AIBudgetLimitSourceGroup), + wantGroupBudget: &codersdk.AIGroupBudget{ + SpendLimitMicros: 1_000_000_000, + LimitSource: codersdk.AIBudgetLimitSourceGroup, + }, }, { name: "BudgetWithSpend", groupLimit: 1_000_000_000, spent: 250_000_000, wantEffectiveGroup: true, - wantSpendLimit: ptr.Ref(int64(1_000_000_000)), - wantLimitSource: ptr.Ref(codersdk.AIBudgetLimitSourceGroup), - wantSpendMicros: 250_000_000, + wantGroupBudget: &codersdk.AIGroupBudget{ + SpendLimitMicros: 1_000_000_000, + LimitSource: codersdk.AIBudgetLimitSourceGroup, + }, + wantSpendMicros: 250_000_000, }, { name: "NoBudgetWithSpend", @@ -3715,8 +3717,10 @@ func TestGroupMembersAISpend(t *testing.T) { name: "OverrideBudget", overrideLimit: 500_000_000, wantEffectiveGroup: true, - wantSpendLimit: ptr.Ref(int64(500_000_000)), - wantLimitSource: ptr.Ref(codersdk.AIBudgetLimitSourceUserOverride), + wantGroupBudget: &codersdk.AIGroupBudget{ + SpendLimitMicros: 500_000_000, + LimitSource: codersdk.AIBudgetLimitSourceUserOverride, + }, }, } @@ -3777,8 +3781,7 @@ func TestGroupMembersAISpend(t *testing.T) { } else { require.Nil(t, got.Members[0].EffectiveGroupID) } - require.Equal(t, tt.wantSpendLimit, got.Members[0].SpendLimitMicros) - require.Equal(t, tt.wantLimitSource, got.Members[0].LimitSource) + require.Equal(t, tt.wantGroupBudget, got.Members[0].GroupBudget) require.Equal(t, tt.wantSpendMicros, got.Members[0].GroupSpendMicros) }) } @@ -3838,8 +3841,7 @@ func TestGroupMembersAISpend(t *testing.T) { require.Len(t, resp.Members, 1) require.Equal(t, targetUser.ID, resp.Members[0].UserID) require.Nil(t, resp.Members[0].EffectiveGroupID, "cross-org effective group must be masked even for the owner") - require.Nil(t, resp.Members[0].SpendLimitMicros) - require.Nil(t, resp.Members[0].LimitSource) + require.Nil(t, resp.Members[0].GroupBudget) require.Equal(t, int64(0), resp.Members[0].GroupSpendMicros) }) @@ -3869,8 +3871,7 @@ func TestGroupMembersAISpend(t *testing.T) { require.Len(t, got.Members, 1) require.Equal(t, targetUser.ID, got.Members[0].UserID) require.Nil(t, got.Members[0].EffectiveGroupID) - require.Nil(t, got.Members[0].SpendLimitMicros) - require.Nil(t, got.Members[0].LimitSource) + require.Nil(t, got.Members[0].GroupBudget) require.Equal(t, int64(0), got.Members[0].GroupSpendMicros) }) } diff --git a/site/src/api/typesGenerated.ts b/site/src/api/typesGenerated.ts index 3be2305222f49..290cc6a3f2af1 100644 --- a/site/src/api/typesGenerated.ts +++ b/site/src/api/typesGenerated.ts @@ -292,6 +292,16 @@ export interface AIGatewayKey { */ export const AIGatewayKeyHeader = "X-Coder-AI-Governance-Gateway-Key"; +// From codersdk/aibridge.go +/** + * AIGroupBudget is an AI spend limit and the tier that produced it. Both + * fields are always populated together. + */ +export interface AIGroupBudget { + readonly spend_limit_micros: number; + readonly limit_source: AIBudgetLimitSource; +} + // From codersdk/aiproviders.go /** * AIProvider represents an AI provider configuration row as returned @@ -5191,17 +5201,11 @@ export interface GroupMemberAISpend { */ readonly effective_group_id: string | null; /** - * SpendLimitMicros is the spend limit when the queried group is this - * user's effective budget source. Null when the user's budget resolves to - * another group or no budget applies to the user (unlimited). + * GroupBudget is the budget when the queried group is this user's + * effective budget source. Null when the user's budget resolves to another + * group or no budget applies to the user. */ - readonly spend_limit_micros: number | null; - /** - * LimitSource identifies the tier that produced the limit. Null when the - * user's budget resolves to another group or no budget applies to the user - * (unlimited). - */ - readonly limit_source: AIBudgetLimitSource | null; + readonly group_budget: AIGroupBudget | null; /** * GroupSpendMicros is the user's spend attributed to the queried group * over the current budget period. From eca5f763e3a8327b088beef1a56a2699ea2ec659 Mon Sep 17 00:00:00 2001 From: Susana Cardoso Ferreira Date: Mon, 20 Jul 2026 11:40:42 +0000 Subject: [PATCH 8/8] chore: add TODO comments for AIGOV-527 --- coderd/aibridge/budget/budget.go | 2 ++ coderd/database/querier.go | 1 + coderd/database/queries.sql.go | 1 + coderd/database/queries/aicostcontrol.sql | 1 + 4 files changed, 5 insertions(+) diff --git a/coderd/aibridge/budget/budget.go b/coderd/aibridge/budget/budget.go index b1ef9742db99d..5b86cdbdd6813 100644 --- a/coderd/aibridge/budget/budget.go +++ b/coderd/aibridge/budget/budget.go @@ -37,6 +37,8 @@ type EffectiveBudget struct { // return value is false when no budget is configured for the user. A per-user // override wins unconditionally; otherwise the budget is selected from the // user's groups according to policy. +// +// TODO(AIGOV-527): unify effective group resolution in a single place. func ResolveUserAIBudget(ctx context.Context, db Store, userID uuid.UUID, policy codersdk.AIBudgetPolicy) (EffectiveBudget, bool, error) { // A per-user override always wins. override, err := db.GetUserAIBudgetOverride(ctx, userID) diff --git a/coderd/database/querier.go b/coderd/database/querier.go index c311608456c51..bbec10b230cbc 100644 --- a/coderd/database/querier.go +++ b/coderd/database/querier.go @@ -584,6 +584,7 @@ type sqlcQuerier interface { // is null when the user has no configured budget or when the effective group // belongs to a different organization than the queried group. // The period_start parameter is normalized to its UTC calendar day. + // TODO(AIGOV-527): unify effective group resolution in a single place. // Spend is aggregated for the queried group, not the user's effective group. // A LEFT JOIN leaves spend_limit_micros and limit_source null for users // whose effective budget source is not the queried group. diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go index 5e99aa6d48bd6..4ba7989c53f0a 100644 --- a/coderd/database/queries.sql.go +++ b/coderd/database/queries.sql.go @@ -2630,6 +2630,7 @@ type GetGroupMembersAISpendRow struct { // is null when the user has no configured budget or when the effective group // belongs to a different organization than the queried group. // The period_start parameter is normalized to its UTC calendar day. +// TODO(AIGOV-527): unify effective group resolution in a single place. // Spend is aggregated for the queried group, not the user's effective group. // A LEFT JOIN leaves spend_limit_micros and limit_source null for users // whose effective budget source is not the queried group. diff --git a/coderd/database/queries/aicostcontrol.sql b/coderd/database/queries/aicostcontrol.sql index d51f65279d5c0..c15a63d55fa32 100644 --- a/coderd/database/queries/aicostcontrol.sql +++ b/coderd/database/queries/aicostcontrol.sql @@ -130,6 +130,7 @@ ORDER BY groups.id; -- is null when the user has no configured budget or when the effective group -- belongs to a different organization than the queried group. -- The period_start parameter is normalized to its UTC calendar day. +-- TODO(AIGOV-527): unify effective group resolution in a single place. WITH queried_group AS ( -- The queried group's org, used to detect cross-org effective groups. SELECT organization_id