diff --git a/coderd/apidoc/docs.go b/coderd/apidoc/docs.go index f10a3479dc..1253a8a006 100644 --- a/coderd/apidoc/docs.go +++ b/coderd/apidoc/docs.go @@ -20754,7 +20754,7 @@ const docTemplate = `{ "type": "object", "properties": { "current_spend_micros": { - "description": "CurrentSpendMicros is the group's spend over the current budget\nperiod.", + "description": "CurrentSpendMicros is the group's spend over the current budget period.", "type": "integer" }, "group_id": { @@ -20772,7 +20772,11 @@ const docTemplate = `{ "format": "date-time" }, "spend_limit_micros": { - "description": "SpendLimitMicros is the group's configured AI spend limit. Null when\nthe group has no configured budget.", + "description": "SpendLimitMicros is the group's configured AI spend budget per member.\nNull when the group has no configured budget.", + "type": "integer" + }, + "total_spend_limit_micros": { + "description": "TotalSpendLimitMicros is the currently configured combined budget of the\nmembers attributed to this group, with each member's override replacing\ntheir share. Null when the group has no budget, and zero when no members\nare attributed to it.", "type": "integer" } } @@ -22260,7 +22264,7 @@ const docTemplate = `{ "type": "object", "properties": { "current_spend_micros": { - "description": "CurrentSpendMicros is the group's spend over the current budget\nperiod.", + "description": "CurrentSpendMicros is the group's spend over the current budget period.", "type": "integer" }, "group_id": { @@ -22268,7 +22272,11 @@ const docTemplate = `{ "format": "uuid" }, "spend_limit_micros": { - "description": "SpendLimitMicros is the group's configured AI spend limit. Null when\nthe group has no configured budget.", + "description": "SpendLimitMicros is the group's configured AI spend budget per member.\nNull when the group has no configured budget.", + "type": "integer" + }, + "total_spend_limit_micros": { + "description": "TotalSpendLimitMicros is the currently configured combined budget of the\nmembers attributed to this group, with each member's override replacing\ntheir share. Null when the group has no budget, and zero when no members\nare attributed to it.", "type": "integer" } } diff --git a/coderd/apidoc/swagger.json b/coderd/apidoc/swagger.json index fb70c9cf9d..750a556b73 100644 --- a/coderd/apidoc/swagger.json +++ b/coderd/apidoc/swagger.json @@ -18888,7 +18888,7 @@ "type": "object", "properties": { "current_spend_micros": { - "description": "CurrentSpendMicros is the group's spend over the current budget\nperiod.", + "description": "CurrentSpendMicros is the group's spend over the current budget period.", "type": "integer" }, "group_id": { @@ -18906,7 +18906,11 @@ "format": "date-time" }, "spend_limit_micros": { - "description": "SpendLimitMicros is the group's configured AI spend limit. Null when\nthe group has no configured budget.", + "description": "SpendLimitMicros is the group's configured AI spend budget per member.\nNull when the group has no configured budget.", + "type": "integer" + }, + "total_spend_limit_micros": { + "description": "TotalSpendLimitMicros is the currently configured combined budget of the\nmembers attributed to this group, with each member's override replacing\ntheir share. Null when the group has no budget, and zero when no members\nare attributed to it.", "type": "integer" } } @@ -20316,7 +20320,7 @@ "type": "object", "properties": { "current_spend_micros": { - "description": "CurrentSpendMicros is the group's spend over the current budget\nperiod.", + "description": "CurrentSpendMicros is the group's spend over the current budget period.", "type": "integer" }, "group_id": { @@ -20324,7 +20328,11 @@ "format": "uuid" }, "spend_limit_micros": { - "description": "SpendLimitMicros is the group's configured AI spend limit. Null when\nthe group has no configured budget.", + "description": "SpendLimitMicros is the group's configured AI spend budget per member.\nNull when the group has no configured budget.", + "type": "integer" + }, + "total_spend_limit_micros": { + "description": "TotalSpendLimitMicros is the currently configured combined budget of the\nmembers attributed to this group, with each member's override replacing\ntheir share. Null when the group has no budget, and zero when no members\nare attributed to it.", "type": "integer" } } diff --git a/coderd/database/db2sdk/db2sdk.go b/coderd/database/db2sdk/db2sdk.go index c6e49c4dd9..47dff8ac87 100644 --- a/coderd/database/db2sdk/db2sdk.go +++ b/coderd/database/db2sdk/db2sdk.go @@ -1475,6 +1475,9 @@ func OrganizationGroupAISpend(row database.GetOrganizationGroupsAISpendRow) code if row.SpendLimitMicros.Valid { group.SpendLimitMicros = &row.SpendLimitMicros.Int64 } + if row.TotalSpendLimitMicros.Valid { + group.TotalSpendLimitMicros = &row.TotalSpendLimitMicros.Int64 + } return group } diff --git a/coderd/database/querier.go b/coderd/database/querier.go index 0eb2f1d78d..d7dd565b65 100644 --- a/coderd/database/querier.go +++ b/coderd/database/querier.go @@ -691,9 +691,13 @@ type sqlcQuerier interface { GetOrganizationByID(ctx context.Context, id uuid.UUID) (Organization, error) GetOrganizationByName(ctx context.Context, arg GetOrganizationByNameParams) (Organization, error) // Returns AI spend limits and aggregate spend for groups in @group_ids that - // belong to @organization_id, on or after period_start until NOW. The spend - // limit is null when the group has no configured budget. + // belong to @organization_id, on or after period_start until NOW. + // spend_limit_micros is the per-member limit, null when the group has no budget. + // total_spend_limit_micros is the combined budget of the members attributed to + // the group, with each member's override replacing their share. It is null when + // the group has no budget. // The period_start parameter is normalized to its UTC calendar day. + // TODO(AIGOV-527): unify effective group resolution in a single place. GetOrganizationGroupsAISpend(ctx context.Context, arg GetOrganizationGroupsAISpendParams) ([]GetOrganizationGroupsAISpendRow, error) GetOrganizationIDsByMemberIDs(ctx context.Context, ids []uuid.UUID) ([]GetOrganizationIDsByMemberIDsRow, error) GetOrganizationResourceCountByID(ctx context.Context, organizationID uuid.UUID) (GetOrganizationResourceCountByIDRow, error) diff --git a/coderd/database/querier_test.go b/coderd/database/querier_test.go index c5511ebda3..25f839dd10 100644 --- a/coderd/database/querier_test.go +++ b/coderd/database/querier_test.go @@ -13085,45 +13085,174 @@ func TestGetOrganizationGroupsAISpend(t *testing.T) { now := monthStart.AddDate(0, 0, 14) // 2024-06-15 prevMonthLastDay := monthStart.AddDate(0, 0, -1) // 2024-05-31 - type seedRow struct { + // nullInt64 keeps the expectations below readable. + nullInt64 := func(v int64) sql.NullInt64 { + return sql.NullInt64{Int64: v, Valid: true} + } + + type groupBudget struct { + group string + limit int64 + } + type membership struct { + user string + group string + } + type override struct { + user string + group string + limit int64 + } + type spendRow struct { + user string + group string day time.Time spend int64 } + type wantGroup struct { + spendLimit sql.NullInt64 + totalLimit sql.NullInt64 + spend int64 + } + // A case declares its groups, members, overrides and spend by name. Every + // named group is created in one org and queried, and every named user is an + // org member. tests := []struct { - name string - setBudget bool - spendLimit int64 - rows []seedRow - wantCurrentSpend int64 + name string + budgets []groupBudget + members []membership + overrides []override + spend []spendRow + want map[string]wantGroup }{ { - name: "NoBudgetNoSpend", - wantCurrentSpend: 0, + name: "NoBudgetNoSpend", + want: map[string]wantGroup{ + "eng": {spendLimit: sql.NullInt64{}, totalLimit: sql.NullInt64{}}, + }, }, { - name: "ZeroLimitBudget", - setBudget: true, - spendLimit: 0, - wantCurrentSpend: 0, + name: "ZeroLimitBudget", + budgets: []groupBudget{{group: "eng", limit: 0}}, + want: map[string]wantGroup{ + "eng": {spendLimit: nullInt64(0), totalLimit: nullInt64(0)}, + }, }, { - name: "BudgetZeroSpend", - setBudget: true, - spendLimit: 1_000_000, - wantCurrentSpend: 0, + // The group has no members, so nothing is attributed to it and the + // total is zero despite the budget. + name: "BudgetZeroSpend", + budgets: []groupBudget{{group: "eng", limit: 1_000_000}}, + want: map[string]wantGroup{ + "eng": {spendLimit: nullInt64(1_000_000), totalLimit: nullInt64(0)}, + }, }, { - name: "BudgetWithSpend", - setBudget: true, - spendLimit: 1_000_000, - rows: []seedRow{{now, 250}}, - wantCurrentSpend: 250, + // alice has spend attributed to the group but is currently not a member, so the + // total stays zero while the spend counts. + name: "BudgetWithSpend", + budgets: []groupBudget{{group: "eng", limit: 1_000_000}}, + spend: []spendRow{{user: "alice", group: "eng", day: now, spend: 250}}, + want: map[string]wantGroup{ + "eng": {spendLimit: nullInt64(1_000_000), totalLimit: nullInt64(0), spend: 250}, + }, }, { - name: "NoBudgetWithSpend", - rows: []seedRow{{now, 100}}, - wantCurrentSpend: 100, + name: "NoBudgetWithSpend", + spend: []spendRow{{user: "alice", group: "eng", day: now, spend: 100}}, + want: map[string]wantGroup{ + "eng": {spendLimit: sql.NullInt64{}, totalLimit: sql.NullInt64{}, spend: 100}, + }, + }, + { + name: "BudgetedTwoPlainMembers", + budgets: []groupBudget{{group: "eng", limit: 100}}, + members: []membership{{user: "alice", group: "eng"}, {user: "bob", group: "eng"}}, + want: map[string]wantGroup{ + "eng": {spendLimit: nullInt64(100), totalLimit: nullInt64(200)}, + }, + }, + { + name: "BudgetedPlainMemberPlusOverride", + budgets: []groupBudget{{group: "eng", limit: 100}}, + members: []membership{{user: "alice", group: "eng"}, {user: "bob", group: "eng"}}, + overrides: []override{{user: "alice", group: "eng", limit: 1000}}, + // The override replaces its holder's share of the group limit. + want: map[string]wantGroup{ + "eng": {spendLimit: nullInt64(100), totalLimit: nullInt64(1100)}, + }, + }, + { + name: "BudgetedOnlyOverrideMember", + budgets: []groupBudget{{group: "eng", limit: 100}}, + members: []membership{{user: "alice", group: "eng"}}, + overrides: []override{{user: "alice", group: "eng", limit: 1000}}, + want: map[string]wantGroup{ + "eng": {spendLimit: nullInt64(100), totalLimit: nullInt64(1000)}, + }, + }, + { + name: "BudgetedZeroLimitWithMembers", + budgets: []groupBudget{{group: "eng", limit: 0}}, + members: []membership{{user: "alice", group: "eng"}, {user: "bob", group: "eng"}}, + want: map[string]wantGroup{ + "eng": {spendLimit: nullInt64(0), totalLimit: nullInt64(0)}, + }, + }, + { + name: "UnbudgetedWithMembers", + members: []membership{{user: "alice", group: "eng"}, {user: "bob", group: "eng"}}, + want: map[string]wantGroup{ + "eng": {spendLimit: sql.NullInt64{}, totalLimit: sql.NullInt64{}}, + }, + }, + { + name: "UnbudgetedWithOverride", + members: []membership{{user: "alice", group: "eng"}}, + overrides: []override{{user: "alice", group: "eng", limit: 1000}}, + // Members of a group with no budget spend without a cap, so the + // override does not make the group's total finite. + want: map[string]wantGroup{ + "eng": {spendLimit: sql.NullInt64{}, totalLimit: sql.NullInt64{}}, + }, + }, + { + name: "MemberResolvesToHigherBudgetGroup", + budgets: []groupBudget{ + {group: "eng", limit: 100}, + {group: "platform", limit: 500}, + }, + members: []membership{ + {user: "alice", group: "eng"}, + {user: "bob", group: "eng"}, + {user: "alice", group: "platform"}, + }, + // alice is attributed to the higher-limit group, so only bob counts + // toward eng. + want: map[string]wantGroup{ + "eng": {spendLimit: nullInt64(100), totalLimit: nullInt64(100)}, + "platform": {spendLimit: nullInt64(500), totalLimit: nullInt64(500)}, + }, + }, + { + name: "OverrideChargedToOtherGroup", + budgets: []groupBudget{ + {group: "eng", limit: 100}, + {group: "platform", limit: 50}, + }, + members: []membership{ + {user: "alice", group: "eng"}, + {user: "bob", group: "eng"}, + {user: "alice", group: "platform"}, + }, + overrides: []override{{user: "alice", group: "platform", limit: 1000}}, + // The override attributes alice to platform even though eng carries + // the higher group limit. + want: map[string]wantGroup{ + "eng": {spendLimit: nullInt64(100), totalLimit: nullInt64(100)}, + "platform": {spendLimit: nullInt64(50), totalLimit: nullInt64(1000)}, + }, }, } @@ -13133,46 +13262,92 @@ func TestGetOrganizationGroupsAISpend(t *testing.T) { db, _ := dbtestutil.NewDB(t) ctx := testutil.Context(t, testutil.WaitShort) - // Given: an org with a single group, optionally with a budget and seeded spend. - user := dbgen.User(t, db, database.User{}) + // Given: the groups, members, overrides and spend the case declares. org := dbgen.Organization(t, db, database.Organization{}) - group := dbgen.Group(t, db, database.Group{OrganizationID: org.ID}) - if tt.setBudget { + groupIDs := make(map[string]uuid.UUID) + groupID := func(name string) uuid.UUID { + if id, ok := groupIDs[name]; ok { + return id + } + group := dbgen.Group(t, db, database.Group{OrganizationID: org.ID}) + groupIDs[name] = group.ID + return group.ID + } + userIDs := make(map[string]uuid.UUID) + userID := func(name string) uuid.UUID { + if id, ok := userIDs[name]; ok { + return id + } + user := dbgen.User(t, db, database.User{}) + dbgen.OrganizationMember(t, db, database.OrganizationMember{ + UserID: user.ID, + OrganizationID: org.ID, + }) + userIDs[name] = user.ID + return user.ID + } + + for name := range tt.want { + groupID(name) + } + for _, b := range tt.budgets { _, err := db.UpsertGroupAIBudget(ctx, database.UpsertGroupAIBudgetParams{ - GroupID: group.ID, - SpendLimitMicros: tt.spendLimit, + GroupID: groupID(b.group), + SpendLimitMicros: b.limit, }) require.NoError(t, err) } - for _, r := range tt.rows { + for _, m := range tt.members { + dbgen.GroupMember(t, db, database.GroupMemberTable{ + GroupID: groupID(m.group), + UserID: userID(m.user), + }) + } + for _, o := range tt.overrides { + _, err := db.UpsertUserAIBudgetOverride(ctx, database.UpsertUserAIBudgetOverrideParams{ + UserID: userID(o.user), + GroupID: groupID(o.group), + SpendLimitMicros: o.limit, + }) + require.NoError(t, err) + } + for _, s := range tt.spend { _, err := db.IncrementUserAIDailySpend(ctx, database.IncrementUserAIDailySpendParams{ - UserID: user.ID, - EffectiveGroupID: group.ID, - Day: r.day, - CostMicros: r.spend, + UserID: userID(s.user), + EffectiveGroupID: groupID(s.group), + Day: s.day, + CostMicros: s.spend, }) require.NoError(t, err) } - // When: querying spend for the group since monthStart. + // When: querying every group the case declares. + queried := make([]uuid.UUID, 0, len(groupIDs)) + for _, id := range groupIDs { + queried = append(queried, id) + } got, err := db.GetOrganizationGroupsAISpend(ctx, database.GetOrganizationGroupsAISpendParams{ OrganizationID: org.ID, - GroupIds: []uuid.UUID{group.ID}, + GroupIds: queried, PeriodStart: monthStart, }) require.NoError(t, err) - // Then: one row is returned with the group's limit and spend. - require.Len(t, got, 1) - require.Equal(t, group.ID, got[0].GroupID) - require.Equal(t, org.ID, got[0].OrganizationID) - if tt.setBudget { - require.True(t, got[0].SpendLimitMicros.Valid, "expected configured budget") - require.Equal(t, tt.spendLimit, got[0].SpendLimitMicros.Int64, "spend_limit_micros") - } else { - require.False(t, got[0].SpendLimitMicros.Valid, "expected no configured budget") + // Then: each group reports its own limit, spend, and the combined + // limit of the members attributed to it. + require.Len(t, got, len(groupIDs)) + byID := make(map[uuid.UUID]database.GetOrganizationGroupsAISpendRow, len(got)) + for _, r := range got { + byID[r.GroupID] = r + } + for name, want := range tt.want { + row, ok := byID[groupIDs[name]] + require.True(t, ok, "group %q missing from response", name) + require.Equal(t, org.ID, row.OrganizationID) + require.Equal(t, want.spendLimit, row.SpendLimitMicros, "%s spend_limit_micros", name) + require.Equal(t, want.totalLimit, row.TotalSpendLimitMicros, "%s total_spend_limit_micros", name) + require.Equal(t, want.spend, row.CurrentSpendMicros, "%s current_spend_micros", name) } - require.Equal(t, tt.wantCurrentSpend, got[0].CurrentSpendMicros) }) } @@ -13217,10 +13392,14 @@ func TestGetOrganizationGroupsAISpend(t *testing.T) { rowA, ok := byID[groupA.ID] require.True(t, ok, "groupA missing from response") require.Equal(t, sql.NullInt64{Int64: 1_000_000, Valid: true}, rowA.SpendLimitMicros) + // Neither group has members, so the budgeted group totals zero and the + // unbudgeted one is null. + require.Equal(t, sql.NullInt64{Int64: 0, Valid: true}, rowA.TotalSpendLimitMicros) require.Equal(t, int64(250), rowA.CurrentSpendMicros) rowB, ok := byID[groupB.ID] require.True(t, ok, "groupB missing from response") require.Equal(t, sql.NullInt64{}, rowB.SpendLimitMicros) + require.Equal(t, sql.NullInt64{}, rowB.TotalSpendLimitMicros) require.Equal(t, int64(500), rowB.CurrentSpendMicros) }) @@ -13381,6 +13560,81 @@ func TestGetOrganizationGroupsAISpend(t *testing.T) { require.Equal(t, int64(25), got[0].CurrentSpendMicros, "sum must exclude prevMonthLastDay row after normalization") }) + + t.Run("EveryoneGroupCountsOrgMembers", func(t *testing.T) { + t.Parallel() + db, _ := dbtestutil.NewDB(t) + ctx := testutil.Context(t, testutil.WaitShort) + + // Given: an org with two members whose implicit Everyone group carries the + // only budget. + org := dbgen.Organization(t, db, database.Organization{}) + for range 2 { + user := dbgen.User(t, db, database.User{}) + dbgen.OrganizationMember(t, db, database.OrganizationMember{ + UserID: user.ID, + OrganizationID: org.ID, + }) + } + // The Everyone group has ID equal to the organization ID and must be + // inserted explicitly for the group_ai_budgets FK constraint. + //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: 100, + }) + require.NoError(t, err) + + // When: querying the Everyone group. + got, err := db.GetOrganizationGroupsAISpend(ctx, database.GetOrganizationGroupsAISpendParams{ + OrganizationID: org.ID, + GroupIds: []uuid.UUID{org.ID}, + PeriodStart: monthStart, + }) + require.NoError(t, err) + + // Then: every org member counts toward the total. + require.Len(t, got, 1) + require.Equal(t, sql.NullInt64{Int64: 100, Valid: true}, got[0].SpendLimitMicros, "spend_limit_micros") + require.Equal(t, sql.NullInt64{Int64: 200, Valid: true}, got[0].TotalSpendLimitMicros, "total_spend_limit_micros") + require.Equal(t, int64(0), got[0].CurrentSpendMicros, "current_spend_micros") + }) + + t.Run("EveryoneGroupWithoutBudgetIsUnlimited", func(t *testing.T) { + t.Parallel() + db, _ := dbtestutil.NewDB(t) + ctx := testutil.Context(t, testutil.WaitShort) + + // Given: an org with two members and no budget anywhere, which is where + // uncapped users are attributed. + org := dbgen.Organization(t, db, database.Organization{}) + for range 2 { + user := dbgen.User(t, db, database.User{}) + dbgen.OrganizationMember(t, db, database.OrganizationMember{ + UserID: user.ID, + OrganizationID: org.ID, + }) + } + //nolint:gocritic // Requires system context. + _, err := db.InsertAllUsersGroup(dbauthz.AsSystemRestricted(ctx), org.ID) + require.NoError(t, err) + + // When: querying the Everyone group. + got, err := db.GetOrganizationGroupsAISpend(ctx, database.GetOrganizationGroupsAISpendParams{ + OrganizationID: org.ID, + GroupIds: []uuid.UUID{org.ID}, + PeriodStart: monthStart, + }) + require.NoError(t, err) + + // Then: the total is null, so the group reads as unlimited. + require.Len(t, got, 1) + require.Equal(t, sql.NullInt64{}, got[0].SpendLimitMicros, "spend_limit_micros") + require.Equal(t, sql.NullInt64{}, got[0].TotalSpendLimitMicros, "total_spend_limit_micros") + require.Equal(t, int64(0), got[0].CurrentSpendMicros, "current_spend_micros") + }) } func TestGetGroupMembersAISpend(t *testing.T) { diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go index 30bf07a624..d55bbc5a5f 100644 --- a/coderd/database/queries.sql.go +++ b/coderd/database/queries.sql.go @@ -2860,41 +2860,119 @@ func (q *sqlQuerier) GetHighestGroupAIBudgetByUser(ctx context.Context, userID u } const getOrganizationGroupsAISpend = `-- name: GetOrganizationGroupsAISpend :many +WITH queried_groups AS ( + -- The requested groups that belong to the queried organization. + SELECT groups.id, groups.organization_id + FROM groups + WHERE groups.organization_id = $1 + AND groups.id = ANY($2::uuid[]) +), +candidate_users AS ( + -- Members of the queried groups. Uses group_members_expanded so the implicit + -- Everyone group counts. + SELECT DISTINCT member.user_id + FROM group_members_expanded member + WHERE member.group_id IN (SELECT id FROM queried_groups) +), +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 + JOIN organizations ON organizations.id = member.organization_id + JOIN organization_members + ON organization_members.user_id = member.user_id + AND organization_members.organization_id = member.organization_id + WHERE member.user_id IN (SELECT user_id FROM candidate_users) + AND organizations.deleted = false + ORDER BY member.user_id, budget.spend_limit_micros DESC, organization_members.created_at ASC, budget.group_id ASC +), +effective AS ( + -- Effective budget group per user: an override wins over the highest-limit + -- group they belong to. Users with neither are left out, since the group they + -- fall back to has no budget and reports null. + SELECT + candidate_users.user_id, + COALESCE(override.group_id, user_highest_group.group_id) AS effective_group_id, + override.spend_limit_micros AS override_limit_micros + FROM candidate_users + LEFT JOIN user_ai_budget_overrides override ON override.user_id = candidate_users.user_id + LEFT JOIN user_highest_group ON user_highest_group.user_id = candidate_users.user_id +), +group_limits AS ( + -- Per attributed group, how many members take the group's own limit and the + -- combined limit of those carrying an override. + SELECT + effective.effective_group_id AS group_id, + count(*) FILTER (WHERE effective.override_limit_micros IS NULL) AS plain_member_count, + COALESCE(SUM(effective.override_limit_micros), 0)::BIGINT AS override_limit_sum + FROM effective + WHERE effective.effective_group_id IS NOT NULL + GROUP BY effective.effective_group_id +), +group_totals AS ( + -- Combined limit per budgeted group, counting members with no override at the + -- group's own limit and adding the overrides on top. Unbudgeted groups are + -- absent here, so the join below leaves their total null. + SELECT + queried_groups.id AS group_id, + (budget.spend_limit_micros * COALESCE(group_limits.plain_member_count, 0) + + COALESCE(group_limits.override_limit_sum, 0))::BIGINT AS total_spend_limit_micros + FROM queried_groups + JOIN group_ai_budgets budget ON budget.group_id = queried_groups.id + LEFT JOIN group_limits ON group_limits.group_id = queried_groups.id +), +group_spend AS ( + -- Spend per queried group over the period. + SELECT + spend.effective_group_id AS group_id, + COALESCE(SUM(spend.spend_micros), 0)::BIGINT AS current_spend_micros + FROM ai_user_daily_spend spend + WHERE spend.effective_group_id IN (SELECT id FROM queried_groups) + AND spend.day >= (($3::timestamptz) AT TIME ZONE 'UTC')::date + GROUP BY spend.effective_group_id +) SELECT - groups.id AS group_id, - groups.organization_id AS organization_id, + queried_groups.id AS group_id, + queried_groups.organization_id AS organization_id, budget.spend_limit_micros AS spend_limit_micros, - COALESCE(SUM(spend.spend_micros), 0)::BIGINT AS current_spend_micros -FROM groups -LEFT JOIN group_ai_budgets budget ON budget.group_id = groups.id -LEFT JOIN ai_user_daily_spend spend - ON spend.effective_group_id = groups.id - AND spend.day >= (($1::timestamptz) AT TIME ZONE 'UTC')::date -WHERE groups.organization_id = $2 - AND groups.id = ANY($3::uuid[]) -GROUP BY groups.id, budget.spend_limit_micros -ORDER BY groups.id + group_totals.total_spend_limit_micros AS total_spend_limit_micros, + COALESCE(group_spend.current_spend_micros, 0)::BIGINT AS current_spend_micros +FROM queried_groups +LEFT JOIN group_ai_budgets budget ON budget.group_id = queried_groups.id +LEFT JOIN group_totals ON group_totals.group_id = queried_groups.id +LEFT JOIN group_spend ON group_spend.group_id = queried_groups.id +ORDER BY queried_groups.id ` type GetOrganizationGroupsAISpendParams struct { - PeriodStart time.Time `db:"period_start" json:"period_start"` OrganizationID uuid.UUID `db:"organization_id" json:"organization_id"` GroupIds []uuid.UUID `db:"group_ids" json:"group_ids"` + PeriodStart time.Time `db:"period_start" json:"period_start"` } type GetOrganizationGroupsAISpendRow struct { - GroupID uuid.UUID `db:"group_id" json:"group_id"` - OrganizationID uuid.UUID `db:"organization_id" json:"organization_id"` - SpendLimitMicros sql.NullInt64 `db:"spend_limit_micros" json:"spend_limit_micros"` - CurrentSpendMicros int64 `db:"current_spend_micros" json:"current_spend_micros"` + GroupID uuid.UUID `db:"group_id" json:"group_id"` + OrganizationID uuid.UUID `db:"organization_id" json:"organization_id"` + SpendLimitMicros sql.NullInt64 `db:"spend_limit_micros" json:"spend_limit_micros"` + TotalSpendLimitMicros sql.NullInt64 `db:"total_spend_limit_micros" json:"total_spend_limit_micros"` + CurrentSpendMicros int64 `db:"current_spend_micros" json:"current_spend_micros"` } // Returns AI spend limits and aggregate spend for groups in @group_ids that -// belong to @organization_id, on or after period_start until NOW. The spend -// limit is null when the group has no configured budget. +// belong to @organization_id, on or after period_start until NOW. +// spend_limit_micros is the per-member limit, null when the group has no budget. +// total_spend_limit_micros is the combined budget of the members attributed to +// the group, with each member's override replacing their share. It is null when +// the group has no budget. // The period_start parameter is normalized to its UTC calendar day. +// TODO(AIGOV-527): unify effective group resolution in a single place. func (q *sqlQuerier) GetOrganizationGroupsAISpend(ctx context.Context, arg GetOrganizationGroupsAISpendParams) ([]GetOrganizationGroupsAISpendRow, error) { - rows, err := q.db.QueryContext(ctx, getOrganizationGroupsAISpend, arg.PeriodStart, arg.OrganizationID, pq.Array(arg.GroupIds)) + rows, err := q.db.QueryContext(ctx, getOrganizationGroupsAISpend, arg.OrganizationID, pq.Array(arg.GroupIds), arg.PeriodStart) if err != nil { return nil, err } @@ -2906,6 +2984,7 @@ func (q *sqlQuerier) GetOrganizationGroupsAISpend(ctx context.Context, arg GetOr &i.GroupID, &i.OrganizationID, &i.SpendLimitMicros, + &i.TotalSpendLimitMicros, &i.CurrentSpendMicros, ); err != nil { return nil, err diff --git a/coderd/database/queries/aicostcontrol.sql b/coderd/database/queries/aicostcontrol.sql index 6d9753ff6b..43c5970fe2 100644 --- a/coderd/database/queries/aicostcontrol.sql +++ b/coderd/database/queries/aicostcontrol.sql @@ -122,23 +122,100 @@ WHERE user_id = @user_id -- name: GetOrganizationGroupsAISpend :many -- Returns AI spend limits and aggregate spend for groups in @group_ids that --- belong to @organization_id, on or after period_start until NOW. The spend --- limit is null when the group has no configured budget. +-- belong to @organization_id, on or after period_start until NOW. +-- spend_limit_micros is the per-member limit, null when the group has no budget. +-- total_spend_limit_micros is the combined budget of the members attributed to +-- the group, with each member's override replacing their share. It is null when +-- the group has no budget. -- The period_start parameter is normalized to its UTC calendar day. +-- TODO(AIGOV-527): unify effective group resolution in a single place. +WITH queried_groups AS ( + -- The requested groups that belong to the queried organization. + SELECT groups.id, groups.organization_id + FROM groups + WHERE groups.organization_id = @organization_id + AND groups.id = ANY(@group_ids::uuid[]) +), +candidate_users AS ( + -- Members of the queried groups. Uses group_members_expanded so the implicit + -- Everyone group counts. + SELECT DISTINCT member.user_id + FROM group_members_expanded member + WHERE member.group_id IN (SELECT id FROM queried_groups) +), +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 + JOIN organizations ON organizations.id = member.organization_id + JOIN organization_members + ON organization_members.user_id = member.user_id + AND organization_members.organization_id = member.organization_id + WHERE member.user_id IN (SELECT user_id FROM candidate_users) + AND organizations.deleted = false + ORDER BY member.user_id, budget.spend_limit_micros DESC, organization_members.created_at ASC, budget.group_id ASC +), +effective AS ( + -- Effective budget group per user: an override wins over the highest-limit + -- group they belong to. Users with neither are left out, since the group they + -- fall back to has no budget and reports null. + SELECT + candidate_users.user_id, + COALESCE(override.group_id, user_highest_group.group_id) AS effective_group_id, + override.spend_limit_micros AS override_limit_micros + FROM candidate_users + LEFT JOIN user_ai_budget_overrides override ON override.user_id = candidate_users.user_id + LEFT JOIN user_highest_group ON user_highest_group.user_id = candidate_users.user_id +), +group_limits AS ( + -- Per attributed group, how many members take the group's own limit and the + -- combined limit of those carrying an override. + SELECT + effective.effective_group_id AS group_id, + count(*) FILTER (WHERE effective.override_limit_micros IS NULL) AS plain_member_count, + COALESCE(SUM(effective.override_limit_micros), 0)::BIGINT AS override_limit_sum + FROM effective + WHERE effective.effective_group_id IS NOT NULL + GROUP BY effective.effective_group_id +), +group_totals AS ( + -- Combined limit per budgeted group, counting members with no override at the + -- group's own limit and adding the overrides on top. Unbudgeted groups are + -- absent here, so the join below leaves their total null. + SELECT + queried_groups.id AS group_id, + (budget.spend_limit_micros * COALESCE(group_limits.plain_member_count, 0) + + COALESCE(group_limits.override_limit_sum, 0))::BIGINT AS total_spend_limit_micros + FROM queried_groups + JOIN group_ai_budgets budget ON budget.group_id = queried_groups.id + LEFT JOIN group_limits ON group_limits.group_id = queried_groups.id +), +group_spend AS ( + -- Spend per queried group over the period. + SELECT + spend.effective_group_id AS group_id, + COALESCE(SUM(spend.spend_micros), 0)::BIGINT AS current_spend_micros + FROM ai_user_daily_spend spend + WHERE spend.effective_group_id IN (SELECT id FROM queried_groups) + AND spend.day >= ((@period_start::timestamptz) AT TIME ZONE 'UTC')::date + GROUP BY spend.effective_group_id +) SELECT - groups.id AS group_id, - groups.organization_id AS organization_id, + queried_groups.id AS group_id, + queried_groups.organization_id AS organization_id, budget.spend_limit_micros AS spend_limit_micros, - COALESCE(SUM(spend.spend_micros), 0)::BIGINT AS current_spend_micros -FROM groups -LEFT JOIN group_ai_budgets budget ON budget.group_id = groups.id -LEFT JOIN ai_user_daily_spend spend - ON spend.effective_group_id = groups.id - AND spend.day >= ((@period_start::timestamptz) AT TIME ZONE 'UTC')::date -WHERE groups.organization_id = @organization_id - AND groups.id = ANY(@group_ids::uuid[]) -GROUP BY groups.id, budget.spend_limit_micros -ORDER BY groups.id; + group_totals.total_spend_limit_micros AS total_spend_limit_micros, + COALESCE(group_spend.current_spend_micros, 0)::BIGINT AS current_spend_micros +FROM queried_groups +LEFT JOIN group_ai_budgets budget ON budget.group_id = queried_groups.id +LEFT JOIN group_totals ON group_totals.group_id = queried_groups.id +LEFT JOIN group_spend ON group_spend.group_id = queried_groups.id +ORDER BY queried_groups.id; -- name: GetGroupMembersAISpend :many -- Returns each user's AI spend attributed to the queried group, on or after diff --git a/codersdk/aibridge.go b/codersdk/aibridge.go index 3215bed8b3..e5fb5aeadb 100644 --- a/codersdk/aibridge.go +++ b/codersdk/aibridge.go @@ -84,11 +84,15 @@ type OrganizationGroupsAISpend struct { // within the active budget period. type OrganizationGroupAISpend struct { GroupID uuid.UUID `json:"group_id" format:"uuid"` - // SpendLimitMicros is the group's configured AI spend limit. Null when - // the group has no configured budget. + // SpendLimitMicros is the group's configured AI spend budget per member. + // Null when the group has no configured budget. SpendLimitMicros *int64 `json:"spend_limit_micros"` - // CurrentSpendMicros is the group's spend over the current budget - // period. + // TotalSpendLimitMicros is the currently configured combined budget of the + // members attributed to this group, with each member's override replacing + // their share. Null when the group has no budget, and zero when no members + // are attributed to it. + TotalSpendLimitMicros *int64 `json:"total_spend_limit_micros"` + // CurrentSpendMicros is the group's spend over the current budget period. CurrentSpendMicros int64 `json:"current_spend_micros"` } diff --git a/docs/reference/api/enterprise.md b/docs/reference/api/enterprise.md index 622b36f39b..a4e8294cc4 100644 --- a/docs/reference/api/enterprise.md +++ b/docs/reference/api/enterprise.md @@ -1070,7 +1070,8 @@ Returns the AI spend limit and aggregate spend for the group. "group_id": "306db4e0-7449-4501-b76f-075576fe2d8f", "period_end": "2019-08-24T14:15:22Z", "period_start": "2019-08-24T14:15:22Z", - "spend_limit_micros": 0 + "spend_limit_micros": 0, + "total_spend_limit_micros": 0 } ``` @@ -2084,7 +2085,8 @@ Unknown or unreadable group IDs are silently omitted. { "current_spend_micros": 0, "group_id": "306db4e0-7449-4501-b76f-075576fe2d8f", - "spend_limit_micros": 0 + "spend_limit_micros": 0, + "total_spend_limit_micros": 0 } ], "period_end": "2019-08-24T14:15:22Z", diff --git a/docs/reference/api/schemas.md b/docs/reference/api/schemas.md index 8bfce5d3bd..b2d5a62c20 100644 --- a/docs/reference/api/schemas.md +++ b/docs/reference/api/schemas.md @@ -7756,19 +7756,21 @@ Only certain features set these fields: - FeatureManagedAgentLimit - FeatureAgen "group_id": "306db4e0-7449-4501-b76f-075576fe2d8f", "period_end": "2019-08-24T14:15:22Z", "period_start": "2019-08-24T14:15:22Z", - "spend_limit_micros": 0 + "spend_limit_micros": 0, + "total_spend_limit_micros": 0 } ``` ### Properties -| Name | Type | Required | Restrictions | Description | -|------------------------|---------|----------|--------------|------------------------------------------------------------------------------------------------------------| -| `current_spend_micros` | integer | false | | Current spend micros is the group's spend over the current budget period. | -| `group_id` | string | false | | | -| `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. | -| `spend_limit_micros` | integer | false | | Spend limit micros is the group's configured AI spend limit. Null when the group has no configured budget. | +| Name | Type | Required | Restrictions | Description | +|----------------------------|---------|----------|--------------|----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------| +| `current_spend_micros` | integer | false | | Current spend micros is the group's spend over the current budget period. | +| `group_id` | string | false | | | +| `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. | +| `spend_limit_micros` | integer | false | | Spend limit micros is the group's configured AI spend budget per member. Null when the group has no configured budget. | +| `total_spend_limit_micros` | integer | false | | Total spend limit micros is the currently configured combined budget of the members attributed to this group, with each member's override replacing their share. Null when the group has no budget, and zero when no members are attributed to it. | ## codersdk.GroupMemberAISpend @@ -9314,17 +9316,19 @@ Only certain features set these fields: - FeatureManagedAgentLimit - FeatureAgen { "current_spend_micros": 0, "group_id": "306db4e0-7449-4501-b76f-075576fe2d8f", - "spend_limit_micros": 0 + "spend_limit_micros": 0, + "total_spend_limit_micros": 0 } ``` ### Properties -| Name | Type | Required | Restrictions | Description | -|------------------------|---------|----------|--------------|------------------------------------------------------------------------------------------------------------| -| `current_spend_micros` | integer | false | | Current spend micros is the group's spend over the current budget period. | -| `group_id` | string | false | | | -| `spend_limit_micros` | integer | false | | Spend limit micros is the group's configured AI spend limit. Null when the group has no configured budget. | +| Name | Type | Required | Restrictions | Description | +|----------------------------|---------|----------|--------------|----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------| +| `current_spend_micros` | integer | false | | Current spend micros is the group's spend over the current budget period. | +| `group_id` | string | false | | | +| `spend_limit_micros` | integer | false | | Spend limit micros is the group's configured AI spend budget per member. Null when the group has no configured budget. | +| `total_spend_limit_micros` | integer | false | | Total spend limit micros is the currently configured combined budget of the members attributed to this group, with each member's override replacing their share. Null when the group has no budget, and zero when no members are attributed to it. | ## codersdk.OrganizationGroupsAISpend @@ -9334,7 +9338,8 @@ Only certain features set these fields: - FeatureManagedAgentLimit - FeatureAgen { "current_spend_micros": 0, "group_id": "306db4e0-7449-4501-b76f-075576fe2d8f", - "spend_limit_micros": 0 + "spend_limit_micros": 0, + "total_spend_limit_micros": 0 } ], "period_end": "2019-08-24T14:15:22Z", diff --git a/enterprise/coderd/aibridge_test.go b/enterprise/coderd/aibridge_test.go index 7410f9279c..9d9fc1776b 100644 --- a/enterprise/coderd/aibridge_test.go +++ b/enterprise/coderd/aibridge_test.go @@ -3544,46 +3544,53 @@ func TestOrganizationGroupsAISpend(t *testing.T) { require.Equal(t, group.ID, resp.Groups[0].GroupID) }) + // The group has a single member, so a budgeted group's total matches its + // per-member limit and an unbudgeted group reports null. tests := []struct { - name string - setBudget bool - spendLimit int64 - spent int64 - wantSpendLimit *int64 - wantCurrentSpend int64 + name string + setBudget bool + spendLimit int64 + spent int64 + wantSpendLimit *int64 + wantTotalSpendLimit *int64 + wantCurrentSpend int64 }{ { name: "NoBudgetNoSpend", }, { - name: "ZeroLimitBudget", - setBudget: true, - spendLimit: 0, - wantSpendLimit: ptr.Ref(int64(0)), - wantCurrentSpend: 0, + name: "ZeroLimitBudget", + setBudget: true, + spendLimit: 0, + wantSpendLimit: ptr.Ref(int64(0)), + wantTotalSpendLimit: ptr.Ref(int64(0)), + wantCurrentSpend: 0, }, { - name: "BudgetZeroSpend", - setBudget: true, - spendLimit: 1_000_000_000, - wantSpendLimit: ptr.Ref(int64(1_000_000_000)), - wantCurrentSpend: 0, + name: "BudgetZeroSpend", + setBudget: true, + spendLimit: 1_000_000_000, + wantSpendLimit: ptr.Ref(int64(1_000_000_000)), + wantTotalSpendLimit: ptr.Ref(int64(1_000_000_000)), + wantCurrentSpend: 0, }, { - name: "BudgetWithSpend", - setBudget: true, - spendLimit: 1_000_000_000, - spent: 250_000_000, - wantSpendLimit: ptr.Ref(int64(1_000_000_000)), - wantCurrentSpend: 250_000_000, + name: "BudgetWithSpend", + setBudget: true, + spendLimit: 1_000_000_000, + spent: 250_000_000, + wantSpendLimit: ptr.Ref(int64(1_000_000_000)), + wantTotalSpendLimit: ptr.Ref(int64(1_000_000_000)), + wantCurrentSpend: 250_000_000, }, { - name: "SpendExceedsLimit", - setBudget: true, - spendLimit: 1_000_000_000, - spent: 1_500_000_000, - wantSpendLimit: ptr.Ref(int64(1_000_000_000)), - wantCurrentSpend: 1_500_000_000, + name: "SpendExceedsLimit", + setBudget: true, + spendLimit: 1_000_000_000, + spent: 1_500_000_000, + wantSpendLimit: ptr.Ref(int64(1_000_000_000)), + wantTotalSpendLimit: ptr.Ref(int64(1_000_000_000)), + wantCurrentSpend: 1_500_000_000, }, } @@ -3631,9 +3638,38 @@ func TestOrganizationGroupsAISpend(t *testing.T) { require.Len(t, got.Groups, 1) require.Equal(t, group.ID, got.Groups[0].GroupID) require.Equal(t, tt.wantSpendLimit, got.Groups[0].SpendLimitMicros) + require.Equal(t, tt.wantTotalSpendLimit, got.Groups[0].TotalSpendLimitMicros) require.Equal(t, tt.wantCurrentSpend, got.Groups[0].CurrentSpendMicros) }) } + + t.Run("TotalCombinesAllMembers", func(t *testing.T) { + t.Parallel() + + adminClient, _, group := setupAICostControlTest(t, aiCostControlTestOptions{GroupName: "total-all-members-org-group"}) + ctx := testutil.Context(t, testutil.WaitLong) + + // Given: a second member in a group with a per-member budget. + _, second := coderdtest.CreateAnotherUser(t, adminClient, group.OrganizationID) + _, err := adminClient.PatchGroup(ctx, group.ID, codersdk.PatchGroupRequest{ + AddUsers: []string{second.ID.String()}, + }) + require.NoError(t, err) + _, err = adminClient.UpsertGroupAIBudget(ctx, group.ID, codersdk.UpsertGroupAIBudgetRequest{ + SpendLimitMicros: 1_000_000_000, + }) + require.NoError(t, err) + + // When: querying the group's spend. + got, err := adminClient.OrganizationGroupsAISpend(ctx, group.OrganizationID, []uuid.UUID{group.ID}) + require.NoError(t, err) + + // Then: the total covers both members while the per-member limit is unchanged. + require.Len(t, got.Groups, 1) + require.Equal(t, ptr.Ref(int64(1_000_000_000)), got.Groups[0].SpendLimitMicros) + require.Equal(t, ptr.Ref(int64(2_000_000_000)), got.Groups[0].TotalSpendLimitMicros) + require.Equal(t, int64(0), got.Groups[0].CurrentSpendMicros) + }) } func TestOrganizationGroupsAISpendRoleAccess(t *testing.T) { @@ -4660,46 +4696,53 @@ func TestGroupAISpend(t *testing.T) { require.Equal(t, http.StatusNotFound, sdkErr.StatusCode()) }) + // The group has a single member, so a budgeted group's total matches its + // per-member limit and an unbudgeted group reports null. tests := []struct { - name string - setBudget bool - spendLimit int64 - spent int64 - wantSpendLimit *int64 - wantCurrentSpend int64 + name string + setBudget bool + spendLimit int64 + spent int64 + wantSpendLimit *int64 + wantTotalSpendLimit *int64 + wantCurrentSpend int64 }{ { name: "NoBudgetNoSpend", }, { - name: "ZeroLimitBudget", - setBudget: true, - spendLimit: 0, - wantSpendLimit: ptr.Ref(int64(0)), - wantCurrentSpend: 0, + name: "ZeroLimitBudget", + setBudget: true, + spendLimit: 0, + wantSpendLimit: ptr.Ref(int64(0)), + wantTotalSpendLimit: ptr.Ref(int64(0)), + wantCurrentSpend: 0, }, { - name: "BudgetZeroSpend", - setBudget: true, - spendLimit: 1_000_000_000, - wantSpendLimit: ptr.Ref(int64(1_000_000_000)), - wantCurrentSpend: 0, + name: "BudgetZeroSpend", + setBudget: true, + spendLimit: 1_000_000_000, + wantSpendLimit: ptr.Ref(int64(1_000_000_000)), + wantTotalSpendLimit: ptr.Ref(int64(1_000_000_000)), + wantCurrentSpend: 0, }, { - name: "BudgetWithSpend", - setBudget: true, - spendLimit: 1_000_000_000, - spent: 250_000_000, - wantSpendLimit: ptr.Ref(int64(1_000_000_000)), - wantCurrentSpend: 250_000_000, + name: "BudgetWithSpend", + setBudget: true, + spendLimit: 1_000_000_000, + spent: 250_000_000, + wantSpendLimit: ptr.Ref(int64(1_000_000_000)), + wantTotalSpendLimit: ptr.Ref(int64(1_000_000_000)), + wantCurrentSpend: 250_000_000, }, { - name: "SpendExceedsLimit", - setBudget: true, - spendLimit: 1_000_000_000, - spent: 1_500_000_000, - wantSpendLimit: ptr.Ref(int64(1_000_000_000)), - wantCurrentSpend: 1_500_000_000, + name: "SpendExceedsLimit", + setBudget: true, + spendLimit: 1_000_000_000, + spent: 1_500_000_000, + wantSpendLimit: ptr.Ref(int64(1_000_000_000)), + wantTotalSpendLimit: ptr.Ref(int64(1_000_000_000)), + wantCurrentSpend: 1_500_000_000, }, } @@ -4746,9 +4789,37 @@ func TestGroupAISpend(t *testing.T) { require.Equal(t, wantPeriodEnd, got.PeriodEnd) require.Equal(t, group.ID, got.GroupID) require.Equal(t, tt.wantSpendLimit, got.SpendLimitMicros) + require.Equal(t, tt.wantTotalSpendLimit, got.TotalSpendLimitMicros) require.Equal(t, tt.wantCurrentSpend, got.CurrentSpendMicros) }) } + + t.Run("TotalCombinesAllMembers", func(t *testing.T) { + t.Parallel() + + adminClient, _, group := setupAICostControlTest(t, aiCostControlTestOptions{GroupName: "total-all-members-group"}) + ctx := testutil.Context(t, testutil.WaitLong) + + // Given: a second member in a group with a per-member budget. + _, second := coderdtest.CreateAnotherUser(t, adminClient, group.OrganizationID) + _, err := adminClient.PatchGroup(ctx, group.ID, codersdk.PatchGroupRequest{ + AddUsers: []string{second.ID.String()}, + }) + require.NoError(t, err) + _, err = adminClient.UpsertGroupAIBudget(ctx, group.ID, codersdk.UpsertGroupAIBudgetRequest{ + SpendLimitMicros: 1_000_000_000, + }) + require.NoError(t, err) + + // When: querying the group's spend. + got, err := adminClient.GroupAISpend(ctx, group.ID) + require.NoError(t, err) + + // Then: the total covers both members while the per-member limit is unchanged. + require.Equal(t, ptr.Ref(int64(1_000_000_000)), got.SpendLimitMicros) + require.Equal(t, ptr.Ref(int64(2_000_000_000)), got.TotalSpendLimitMicros) + require.Equal(t, int64(0), got.CurrentSpendMicros) + }) } func TestGroupAISpendRoleAccess(t *testing.T) { diff --git a/site/src/api/typesGenerated.ts b/site/src/api/typesGenerated.ts index b00be75244..f9fe468512 100644 --- a/site/src/api/typesGenerated.ts +++ b/site/src/api/typesGenerated.ts @@ -6921,13 +6921,19 @@ export interface Organization extends MinimalOrganization { export interface OrganizationGroupAISpend { readonly group_id: string; /** - * SpendLimitMicros is the group's configured AI spend limit. Null when - * the group has no configured budget. + * SpendLimitMicros is the group's configured AI spend budget per member. + * Null when the group has no configured budget. */ readonly spend_limit_micros: number | null; /** - * CurrentSpendMicros is the group's spend over the current budget - * period. + * TotalSpendLimitMicros is the currently configured combined budget of the + * members attributed to this group, with each member's override replacing + * their share. Null when the group has no budget, and zero when no members + * are attributed to it. + */ + readonly total_spend_limit_micros: number | null; + /** + * CurrentSpendMicros is the group's spend over the current budget period. */ readonly current_spend_micros: number; } diff --git a/site/src/pages/GroupsPage/GroupsPageView.stories.tsx b/site/src/pages/GroupsPage/GroupsPageView.stories.tsx index 598e1bfad1..454f184f62 100644 --- a/site/src/pages/GroupsPage/GroupsPageView.stories.tsx +++ b/site/src/pages/GroupsPage/GroupsPageView.stories.tsx @@ -51,6 +51,7 @@ export const WithAIBudgets: Story = { group_id: "ai-unlimited", current_spend_micros: 25_492_000_000, spend_limit_micros: null, + total_spend_limit_micros: null, }, }, { @@ -58,7 +59,8 @@ export const WithAIBudgets: Story = { spend: { group_id: "ai-under", current_spend_micros: 10_000_000, - spend_limit_micros: 50_000_000, + spend_limit_micros: 25_000_000, + total_spend_limit_micros: 50_000_000, }, }, { @@ -66,7 +68,8 @@ export const WithAIBudgets: Story = { spend: { group_id: "ai-warning", current_spend_micros: 46_000_000, - spend_limit_micros: 50_000_000, + spend_limit_micros: 25_000_000, + total_spend_limit_micros: 50_000_000, }, }, { @@ -74,7 +77,8 @@ export const WithAIBudgets: Story = { spend: { group_id: "ai-at-limit", current_spend_micros: 50_000_000, - spend_limit_micros: 50_000_000, + spend_limit_micros: 25_000_000, + total_spend_limit_micros: 50_000_000, }, }, { @@ -82,7 +86,8 @@ export const WithAIBudgets: Story = { spend: { group_id: "ai-over", current_spend_micros: 75_000_000, - spend_limit_micros: 50_000_000, + spend_limit_micros: 25_000_000, + total_spend_limit_micros: 50_000_000, }, }, { @@ -91,6 +96,7 @@ export const WithAIBudgets: Story = { group_id: "ai-zero-budget", current_spend_micros: 5_000_000, spend_limit_micros: 0, + total_spend_limit_micros: 0, }, }, { @@ -99,6 +105,7 @@ export const WithAIBudgets: Story = { group_id: "ai-zero-both", current_spend_micros: 0, spend_limit_micros: 0, + total_spend_limit_micros: 0, }, }, // No spend exercises the missing-spend em-dash fallback. diff --git a/site/src/pages/GroupsPage/GroupsPageView.tsx b/site/src/pages/GroupsPage/GroupsPageView.tsx index 2668f25417..f7ed9fffd1 100644 --- a/site/src/pages/GroupsPage/GroupsPageView.tsx +++ b/site/src/pages/GroupsPage/GroupsPageView.tsx @@ -218,7 +218,7 @@ const GroupRow: FC = ({ group, showAIBudget }) => { {group.spend ? ( ) : ( EM_DASH