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/apidoc/docs.go b/coderd/apidoc/docs.go index d3a9660ce60c2..d5a35e6601a0c 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.\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", + "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.\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", + "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": [ @@ -15393,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": { @@ -20258,6 +20362,53 @@ 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_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.AIGroupBudget" + } + ] + }, + "group_spend_micros": { + "description": "GroupSpendMicros is the user's spend attributed to the queried group\nover the current budget period.", + "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..c1d1316b51358 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.\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", + "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.\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", + "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"], @@ -13728,6 +13813,17 @@ } } }, + "codersdk.AIGroupBudget": { + "type": "object", + "properties": { + "limit_source": { + "$ref": "#/definitions/codersdk.AIBudgetLimitSource" + }, + "spend_limit_micros": { + "type": "integer" + } + } + }, "codersdk.AIProvider": { "type": "object", "properties": { @@ -18436,6 +18532,53 @@ } } }, + "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_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.AIGroupBudget" + } + ] + }, + "group_spend_micros": { + "description": "GroupSpendMicros is the user's spend attributed to the queried group\nover the current budget period.", + "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/db2sdk/db2sdk.go b/coderd/database/db2sdk/db2sdk.go index 9a510674ea7cc..21aeaad8b5014 100644 --- a/coderd/database/db2sdk/db2sdk.go +++ b/coderd/database/db2sdk/db2sdk.go @@ -1469,6 +1469,23 @@ 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.GroupBudget = &codersdk.AIGroupBudget{ + SpendLimitMicros: row.SpendLimitMicros.Int64, + LimitSource: codersdk.AIBudgetLimitSource(row.LimitSource.String), + } + } + return member +} + func InvalidatedPresets(invalidatedPresets []database.UpdatePresetsLastInvalidatedAtRow) []codersdk.InvalidatedPreset { var presets []codersdk.InvalidatedPreset for _, p := range invalidatedPresets { diff --git a/coderd/database/dbauthz/dbauthz.go b/coderd/database/dbauthz/dbauthz.go index 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..bbec10b230cbc 100644 --- a/coderd/database/querier.go +++ b/coderd/database/querier.go @@ -577,6 +577,18 @@ 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. + // 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. + 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..b0330b632e503 100644 --- a/coderd/database/querier_test.go +++ b/coderd/database/querier_test.go @@ -13106,6 +13106,592 @@ 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 + + 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: 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 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.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("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) + require.False(t, got[0].SpendLimitMicros.Valid) + require.False(t, got[0].LimitSource.Valid) + require.Equal(t, int64(0), got[0].GroupSpendMicros) + }) + + t.Run("EqualBudgetTieBreak", 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) + require.False(t, got[0].SpendLimitMicros.Valid) + require.False(t, got[0].LimitSource.Valid) + require.Equal(t, int64(0), got[0].GroupSpendMicros) + }) + + 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. + require.Len(t, got, 1) + 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) { + 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("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.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.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) + 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("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("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. + 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) + // 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{ + 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. + // 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("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.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("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.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") + }) +} + 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..4ba7989c53f0a 100644 --- a/coderd/database/queries.sql.go +++ b/coderd/database/queries.sql.go @@ -2533,6 +2533,137 @@ 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 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 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' + 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 raw_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.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 + 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. +// 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. +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..c15a63d55fa32 100644 --- a/coderd/database/queries/aicostcontrol.sql +++ b/coderd/database/queries/aicostcontrol.sql @@ -121,3 +121,88 @@ 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. +-- 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 + 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 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' + 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 raw_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.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. +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..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 { @@ -82,6 +89,31 @@ 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"` + // 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"` +} + type AIBridgeSession struct { ID string `json:"id"` Initiator MinimalUser `json:"initiator"` @@ -498,3 +530,30 @@ 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. 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 := 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, + 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..eeaf8032ec0a3 100644 --- a/docs/reference/api/enterprise.md +++ b/docs/reference/api/enterprise.md @@ -1092,6 +1092,60 @@ 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 + +```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' \ + -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. +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 + +| 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_budget": { + "limit_source": "user_override", + "spend_limit_micros": 0 + }, + "group_spend_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 +2064,61 @@ 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 + +```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' \ + -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. +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 + +| 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_budget": { + "limit_source": "user_override", + "spend_limit_micros": 0 + }, + "group_spend_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..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 @@ -7685,6 +7701,57 @@ 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_budget": { + "limit_source": "user_override", + "spend_limit_micros": 0 + }, + "group_spend_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_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 + +```json +{ + "members": [ + { + "effective_group_id": "85e2b926-ddfb-4c66-b68e-b66e5acec6c0", + "group_budget": { + "limit_source": "user_override", + "spend_limit_micros": 0 + }, + "group_spend_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..49ba9764b5d95 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,94 @@ 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. +// @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 +// @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. +// @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 +// @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 { + 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 e3d2844769d7b..9f3da140997c9 100644 --- a/enterprise/coderd/aibridge_test.go +++ b/enterprise/coderd/aibridge_test.go @@ -3514,6 +3514,450 @@ 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) + require.Nil(t, resp.Members[0].EffectiveGroupID) + require.Nil(t, resp.Members[0].GroupBudget) + require.Equal(t, int64(0), resp.Members[0].GroupSpendMicros) + }) + + tests := []struct { + name string + groupLimit int64 + overrideLimit int64 + spent int64 + wantEffectiveGroup bool + wantGroupBudget *codersdk.AIGroupBudget + wantSpendMicros int64 + }{ + { + name: "NoBudgetNoSpend", + }, + { + name: "BudgetZeroSpend", + groupLimit: 1_000_000_000, + wantEffectiveGroup: true, + wantGroupBudget: &codersdk.AIGroupBudget{ + SpendLimitMicros: 1_000_000_000, + LimitSource: codersdk.AIBudgetLimitSourceGroup, + }, + }, + { + name: "BudgetWithSpend", + groupLimit: 1_000_000_000, + spent: 250_000_000, + wantEffectiveGroup: true, + wantGroupBudget: &codersdk.AIGroupBudget{ + SpendLimitMicros: 1_000_000_000, + LimitSource: codersdk.AIBudgetLimitSourceGroup, + }, + wantSpendMicros: 250_000_000, + }, + { + name: "NoBudgetWithSpend", + spent: 100_000_000, + wantSpendMicros: 100_000_000, + }, + { + name: "OverrideBudget", + overrideLimit: 500_000_000, + wantEffectiveGroup: true, + wantGroupBudget: &codersdk.AIGroupBudget{ + SpendLimitMicros: 500_000_000, + LimitSource: 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.wantGroupBudget, got.Members[0].GroupBudget) + require.Equal(t, tt.wantSpendMicros, got.Members[0].GroupSpendMicros) + }) + } + + 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") + require.Nil(t, resp.Members[0].GroupBudget) + require.Equal(t, int64(0), resp.Members[0].GroupSpendMicros) + }) + + 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) + require.Nil(t, got.Members[0].EffectiveGroupID) + require.Nil(t, got.Members[0].GroupBudget) + require.Equal(t, int64(0), 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) + _, otherMember := 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(), otherMember.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) + }) + } + + 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 // 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..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 @@ -5176,6 +5186,42 @@ 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; + /** + * 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 group_budget: AIGroupBudget | 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[];