From f9b438420030296a6d015d60633e9486828ada7e Mon Sep 17 00:00:00 2001 From: Susana Cardoso Ferreira Date: Tue, 28 Jul 2026 14:40:06 +0000 Subject: [PATCH 1/4] fix: report combined member limit in group AI spend --- coderd/apidoc/docs.go | 12 +- coderd/apidoc/swagger.json | 12 +- coderd/database/db2sdk/db2sdk.go | 3 + coderd/database/querier.go | 8 +- coderd/database/querier_test.go | 240 ++++++++++++++++-- coderd/database/queries.sql.go | 120 +++++++-- coderd/database/queries/aicostcontrol.sql | 106 +++++++- codersdk/aibridge.go | 9 +- docs/reference/api/enterprise.md | 6 +- docs/reference/api/schemas.md | 35 +-- enterprise/coderd/aibridge_test.go | 181 +++++++++---- site/src/api/typesGenerated.ts | 11 +- .../GroupsPage/GroupsPageView.stories.tsx | 7 + 13 files changed, 617 insertions(+), 133 deletions(-) diff --git a/coderd/apidoc/docs.go b/coderd/apidoc/docs.go index b95e78772e219..2402fc5299a8f 100644 --- a/coderd/apidoc/docs.go +++ b/coderd/apidoc/docs.go @@ -20700,7 +20700,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 per-member AI spend limit.\nNull when the group has no configured budget.", + "type": "integer" + }, + "total_spend_limit_micros": { + "description": "TotalSpendLimitMicros is the combined limit of the members whose spend is\nattributed to this group, applying each member's override where one\nexists. Null when any of those members has no limit, and zero when no\nmembers are attributed to the group.", "type": "integer" } } @@ -22188,7 +22192,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 per-member AI spend limit.\nNull when the group has no configured budget.", + "type": "integer" + }, + "total_spend_limit_micros": { + "description": "TotalSpendLimitMicros is the combined limit of the members whose spend is\nattributed to this group, applying each member's override where one\nexists. Null when any of those members has no limit, and zero when no\nmembers are attributed to the group.", "type": "integer" } } diff --git a/coderd/apidoc/swagger.json b/coderd/apidoc/swagger.json index 1a3f184e8d3a0..8f91477619b42 100644 --- a/coderd/apidoc/swagger.json +++ b/coderd/apidoc/swagger.json @@ -18844,7 +18844,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 per-member AI spend limit.\nNull when the group has no configured budget.", + "type": "integer" + }, + "total_spend_limit_micros": { + "description": "TotalSpendLimitMicros is the combined limit of the members whose spend is\nattributed to this group, applying each member's override where one\nexists. Null when any of those members has no limit, and zero when no\nmembers are attributed to the group.", "type": "integer" } } @@ -20254,7 +20258,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 per-member AI spend limit.\nNull when the group has no configured budget.", + "type": "integer" + }, + "total_spend_limit_micros": { + "description": "TotalSpendLimitMicros is the combined limit of the members whose spend is\nattributed to this group, applying each member's override where one\nexists. Null when any of those members has no limit, and zero when no\nmembers are attributed to the group.", "type": "integer" } } diff --git a/coderd/database/db2sdk/db2sdk.go b/coderd/database/db2sdk/db2sdk.go index c6e49c4dd933c..47dff8ac87268 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 f82ef6ea146b6..070052ff0c51d 100644 --- a/coderd/database/querier.go +++ b/coderd/database/querier.go @@ -684,9 +684,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 per-member limit applied to the members + // attributed to the group, replaced by their override where one exists. It is + // null when the group has no budget, because those members spend without a cap. // 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 5cc5c9bd1a77b..5ff2a4360a7b8 100644 --- a/coderd/database/querier_test.go +++ b/coderd/database/querier_test.go @@ -12888,28 +12888,50 @@ func TestGetOrganizationGroupsAISpend(t *testing.T) { spend int64 } + // overrideSpec charges a per-user budget override to one of the two groups. + type overrideSpec struct { + user int + onB bool + limit int64 + } + + // Every case builds an org with two users and two groups. The first group is + // the subject, and the second exists so a case can pull a member's effective + // budget away from the first. members and membersB hold indexes into the + // users slice, and seeded spend is always attributed to users[0] on the + // first group. tests := []struct { name string setBudget bool spendLimit int64 + setBudgetB bool + spendLimitB int64 + members []int + membersB []int + override *overrideSpec rows []seedRow wantCurrentSpend int64 + wantTotalLimit sql.NullInt64 + wantTotalLimitB sql.NullInt64 }{ { name: "NoBudgetNoSpend", wantCurrentSpend: 0, + wantTotalLimit: sql.NullInt64{}, }, { name: "ZeroLimitBudget", setBudget: true, spendLimit: 0, wantCurrentSpend: 0, + wantTotalLimit: sql.NullInt64{Int64: 0, Valid: true}, }, { name: "BudgetZeroSpend", setBudget: true, spendLimit: 1_000_000, wantCurrentSpend: 0, + wantTotalLimit: sql.NullInt64{Int64: 0, Valid: true}, }, { name: "BudgetWithSpend", @@ -12917,11 +12939,84 @@ func TestGetOrganizationGroupsAISpend(t *testing.T) { spendLimit: 1_000_000, rows: []seedRow{{now, 250}}, wantCurrentSpend: 250, + wantTotalLimit: sql.NullInt64{Int64: 0, Valid: true}, }, { name: "NoBudgetWithSpend", rows: []seedRow{{now, 100}}, wantCurrentSpend: 100, + wantTotalLimit: sql.NullInt64{}, + }, + { + name: "BudgetedTwoPlainMembers", + setBudget: true, + spendLimit: 100, + members: []int{0, 1}, + wantTotalLimit: sql.NullInt64{Int64: 200, Valid: true}, + }, + { + name: "BudgetedPlainMemberPlusOverride", + setBudget: true, + spendLimit: 100, + members: []int{0, 1}, + override: &overrideSpec{user: 0, limit: 1000}, + // The override replaces its holder's share of the group limit. + wantTotalLimit: sql.NullInt64{Int64: 1100, Valid: true}, + }, + { + name: "BudgetedOnlyOverrideMember", + setBudget: true, + spendLimit: 100, + members: []int{0}, + override: &overrideSpec{user: 0, limit: 1000}, + wantTotalLimit: sql.NullInt64{Int64: 1000, Valid: true}, + }, + { + name: "BudgetedZeroLimitWithMembers", + setBudget: true, + spendLimit: 0, + members: []int{0, 1}, + wantTotalLimit: sql.NullInt64{Int64: 0, Valid: true}, + }, + { + name: "UnbudgetedWithMembers", + members: []int{0, 1}, + wantTotalLimit: sql.NullInt64{}, + }, + { + name: "UnbudgetedWithOverride", + members: []int{0}, + override: &overrideSpec{user: 0, limit: 1000}, + // Members of a group with no budget spend without a cap, so the + // override does not make the group's total finite. + wantTotalLimit: sql.NullInt64{}, + }, + { + name: "MemberResolvesToHigherBudgetGroup", + setBudget: true, + spendLimit: 100, + setBudgetB: true, + spendLimitB: 500, + members: []int{0, 1}, + membersB: []int{0}, + // users[0] is attributed to the higher-limit second group, so only + // users[1] counts toward the first. + wantTotalLimit: sql.NullInt64{Int64: 100, Valid: true}, + wantTotalLimitB: sql.NullInt64{Int64: 500, Valid: true}, + }, + { + name: "OverrideChargedToOtherGroup", + setBudget: true, + spendLimit: 100, + setBudgetB: true, + spendLimitB: 50, + members: []int{0, 1}, + membersB: []int{0}, + override: &overrideSpec{user: 0, onB: true, limit: 1000}, + // The override attributes users[0] to the second group even though + // the first carries the higher group limit. + wantTotalLimit: sql.NullInt64{Int64: 100, Valid: true}, + wantTotalLimitB: sql.NullInt64{Int64: 1000, Valid: true}, }, } @@ -12931,20 +13026,55 @@ 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: an org with two groups and two members, wired up per the case. org := dbgen.Organization(t, db, database.Organization{}) + users := []database.User{ + dbgen.User(t, db, database.User{}), + dbgen.User(t, db, database.User{}), + } + for _, u := range users { + dbgen.OrganizationMember(t, db, database.OrganizationMember{ + UserID: u.ID, + OrganizationID: org.ID, + }) + } group := dbgen.Group(t, db, database.Group{OrganizationID: org.ID}) - if tt.setBudget { + groupB := dbgen.Group(t, db, database.Group{OrganizationID: org.ID}) + for _, b := range []struct { + groupID uuid.UUID + set bool + limit int64 + }{{group.ID, tt.setBudget, tt.spendLimit}, {groupB.ID, tt.setBudgetB, tt.spendLimitB}} { + if !b.set { + continue + } _, err := db.UpsertGroupAIBudget(ctx, database.UpsertGroupAIBudgetParams{ - GroupID: group.ID, - SpendLimitMicros: tt.spendLimit, + GroupID: b.groupID, + SpendLimitMicros: b.limit, + }) + require.NoError(t, err) + } + for _, i := range tt.members { + dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: group.ID, UserID: users[i].ID}) + } + for _, i := range tt.membersB { + dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: groupB.ID, UserID: users[i].ID}) + } + if tt.override != nil { + overrideGroup := group.ID + if tt.override.onB { + overrideGroup = groupB.ID + } + _, err := db.UpsertUserAIBudgetOverride(ctx, database.UpsertUserAIBudgetOverrideParams{ + UserID: users[tt.override.user].ID, + GroupID: overrideGroup, + SpendLimitMicros: tt.override.limit, }) require.NoError(t, err) } for _, r := range tt.rows { _, err := db.IncrementUserAIDailySpend(ctx, database.IncrementUserAIDailySpendParams{ - UserID: user.ID, + UserID: users[0].ID, EffectiveGroupID: group.ID, Day: r.day, CostMicros: r.spend, @@ -12952,25 +13082,33 @@ func TestGetOrganizationGroupsAISpend(t *testing.T) { require.NoError(t, err) } - // When: querying spend for the group since monthStart. + // When: querying spend for both groups since monthStart. got, err := db.GetOrganizationGroupsAISpend(ctx, database.GetOrganizationGroupsAISpendParams{ OrganizationID: org.ID, - GroupIds: []uuid.UUID{group.ID}, + GroupIds: []uuid.UUID{group.ID, groupB.ID}, PeriodStart: monthStart, }) require.NoError(t, err) - // Then: one row is returned with the group's limit and spend. - require.Len(t, got, 1) - require.Equal(t, group.ID, got[0].GroupID) - require.Equal(t, org.ID, got[0].OrganizationID) + // Then: each group reports its own limit, spend, and the combined + // limit of the members attributed to it. + require.Len(t, got, 2) + byID := make(map[uuid.UUID]database.GetOrganizationGroupsAISpendRow, len(got)) + for _, r := range got { + byID[r.GroupID] = r + } + row, ok := byID[group.ID] + require.True(t, ok, "queried group missing from response") + require.Equal(t, org.ID, row.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") + require.True(t, row.SpendLimitMicros.Valid, "expected configured budget") + require.Equal(t, tt.spendLimit, row.SpendLimitMicros.Int64, "spend_limit_micros") } else { - require.False(t, got[0].SpendLimitMicros.Valid, "expected no configured budget") + require.False(t, row.SpendLimitMicros.Valid, "expected no configured budget") } - require.Equal(t, tt.wantCurrentSpend, got[0].CurrentSpendMicros) + require.Equal(t, tt.wantCurrentSpend, row.CurrentSpendMicros) + require.Equal(t, tt.wantTotalLimit, row.TotalSpendLimitMicros, "total_spend_limit_micros") + require.Equal(t, tt.wantTotalLimitB, byID[groupB.ID].TotalSpendLimitMicros, "second group total_spend_limit_micros") }) } @@ -13179,6 +13317,76 @@ 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 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: 200, Valid: true}, got[0].TotalSpendLimitMicros) + }) + + t.Run("EveryoneGroupWithoutBudgetIsUnlimited", func(t *testing.T) { + t.Parallel() + db, _ := dbtestutil.NewDB(t) + ctx := testutil.Context(t, testutil.WaitShort) + + // Given: an org with members but 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.False(t, got[0].TotalSpendLimitMicros.Valid) + }) } func TestGetGroupMembersAISpend(t *testing.T) { diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go index ee88e6cdba8c8..aacd189bce3e6 100644 --- a/coderd/database/queries.sql.go +++ b/coderd/database/queries.sql.go @@ -2860,41 +2860,120 @@ 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. Narrowing to members loses nobody, because a + -- user's effective group is always a group they belong to. 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 spend without a cap, and the + -- unbudgeted group they fall back to reports unlimited regardless. + 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 per-member limit applied to the members +// attributed to the group, replaced by their override where one exists. It is +// null when the group has no budget, because those members spend without a cap. // 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 +2985,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 6d9753ff6b295..cbe4035eeb579 100644 --- a/coderd/database/queries/aicostcontrol.sql +++ b/coderd/database/queries/aicostcontrol.sql @@ -122,23 +122,101 @@ 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 per-member limit applied to the members +-- attributed to the group, replaced by their override where one exists. It is +-- null when the group has no budget, because those members spend without a cap. -- 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. Narrowing to members loses nobody, because a + -- user's effective group is always a group they belong to. 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 spend without a cap, and the + -- unbudgeted group they fall back to reports unlimited regardless. + 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 3215bed8b3c25..de712638f7f3e 100644 --- a/codersdk/aibridge.go +++ b/codersdk/aibridge.go @@ -84,9 +84,14 @@ 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 per-member AI spend limit. + // Null when the group has no configured budget. SpendLimitMicros *int64 `json:"spend_limit_micros"` + // TotalSpendLimitMicros is the combined limit of the members whose spend is + // attributed to this group, applying each member's override where one + // exists. Null when any of those members has no limit, and zero when no + // members are attributed to the group. + 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 88026e8599f6b..9a85a7a1718b2 100644 --- a/docs/reference/api/enterprise.md +++ b/docs/reference/api/enterprise.md @@ -1066,7 +1066,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 } ``` @@ -2003,7 +2004,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 cae7eb143c5dd..2105a8c7a53ad 100644 --- a/docs/reference/api/schemas.md +++ b/docs/reference/api/schemas.md @@ -7747,19 +7747,21 @@ Only certain features set these fields: - FeatureManagedAgentLimit| "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 per-member AI spend limit. Null when the group has no configured budget. | +| `total_spend_limit_micros` | integer | false | | Total spend limit micros is the combined limit of the members whose spend is attributed to this group, applying each member's override where one exists. Null when any of those members has no limit, and zero when no members are attributed to the group. | ## codersdk.GroupMemberAISpend @@ -9291,17 +9293,19 @@ Only certain features set these fields: - FeatureManagedAgentLimit| { "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 per-member AI spend limit. Null when the group has no configured budget. | +| `total_spend_limit_micros` | integer | false | | Total spend limit micros is the combined limit of the members whose spend is attributed to this group, applying each member's override where one exists. Null when any of those members has no limit, and zero when no members are attributed to the group. | ## codersdk.OrganizationGroupsAISpend @@ -9311,7 +9315,8 @@ Only certain features set these fields: - FeatureManagedAgentLimit| { "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 7410f9279c100..31874447a77f2 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,37 @@ 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) + }) } func TestOrganizationGroupsAISpendRoleAccess(t *testing.T) { @@ -4660,46 +4695,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 +4788,36 @@ 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) + }) } func TestGroupAISpendRoleAccess(t *testing.T) { diff --git a/site/src/api/typesGenerated.ts b/site/src/api/typesGenerated.ts index 2d3f42b2b37bd..224ee061cc002 100644 --- a/site/src/api/typesGenerated.ts +++ b/site/src/api/typesGenerated.ts @@ -6884,10 +6884,17 @@ 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 per-member AI spend limit. + * Null when the group has no configured budget. */ readonly spend_limit_micros: number | null; + /** + * TotalSpendLimitMicros is the combined limit of the members whose spend is + * attributed to this group, applying each member's override where one + * exists. Null when any of those members has no limit, and zero when no + * members are attributed to the group. + */ + readonly total_spend_limit_micros: number | null; /** * CurrentSpendMicros is the group's spend over the current budget * period. diff --git a/site/src/pages/GroupsPage/GroupsPageView.stories.tsx b/site/src/pages/GroupsPage/GroupsPageView.stories.tsx index 598e1bfad1d73..23dcf700ef188 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, }, }, { @@ -59,6 +60,7 @@ export const WithAIBudgets: Story = { group_id: "ai-under", current_spend_micros: 10_000_000, spend_limit_micros: 50_000_000, + total_spend_limit_micros: 50_000_000, }, }, { @@ -67,6 +69,7 @@ export const WithAIBudgets: Story = { group_id: "ai-warning", current_spend_micros: 46_000_000, spend_limit_micros: 50_000_000, + total_spend_limit_micros: 50_000_000, }, }, { @@ -75,6 +78,7 @@ export const WithAIBudgets: Story = { group_id: "ai-at-limit", current_spend_micros: 50_000_000, spend_limit_micros: 50_000_000, + total_spend_limit_micros: 50_000_000, }, }, { @@ -83,6 +87,7 @@ export const WithAIBudgets: Story = { group_id: "ai-over", current_spend_micros: 75_000_000, spend_limit_micros: 50_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. From 208b14b4739a2fa651d56fc57f60f363fa2694a1 Mon Sep 17 00:00:00 2001 From: Susana Cardoso Ferreira Date: Tue, 28 Jul 2026 14:47:45 +0000 Subject: [PATCH 2/4] chore: update groups page to show group's total spend --- site/src/pages/GroupsPage/GroupsPageView.stories.tsx | 8 ++++---- site/src/pages/GroupsPage/GroupsPageView.tsx | 2 +- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/site/src/pages/GroupsPage/GroupsPageView.stories.tsx b/site/src/pages/GroupsPage/GroupsPageView.stories.tsx index 23dcf700ef188..454f184f62006 100644 --- a/site/src/pages/GroupsPage/GroupsPageView.stories.tsx +++ b/site/src/pages/GroupsPage/GroupsPageView.stories.tsx @@ -59,7 +59,7 @@ 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, }, }, @@ -68,7 +68,7 @@ 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, }, }, @@ -77,7 +77,7 @@ 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, }, }, @@ -86,7 +86,7 @@ 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, }, }, diff --git a/site/src/pages/GroupsPage/GroupsPageView.tsx b/site/src/pages/GroupsPage/GroupsPageView.tsx index 2668f25417fd2..f7ed9fffd1dc2 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 From e768a87b1600a9706f8f6f2935d3cbdf86dc3d7c Mon Sep 17 00:00:00 2001 From: Susana Cardoso Ferreira Date: Tue, 28 Jul 2026 17:20:43 +0000 Subject: [PATCH 3/4] chore: address review comments --- coderd/apidoc/docs.go | 12 ++++++------ coderd/apidoc/swagger.json | 12 ++++++------ codersdk/aibridge.go | 13 ++++++------- docs/reference/api/schemas.md | 28 ++++++++++++++-------------- site/src/api/typesGenerated.ts | 13 ++++++------- 5 files changed, 38 insertions(+), 40 deletions(-) diff --git a/coderd/apidoc/docs.go b/coderd/apidoc/docs.go index 2402fc5299a8f..fb609c591c620 100644 --- a/coderd/apidoc/docs.go +++ b/coderd/apidoc/docs.go @@ -20682,7 +20682,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": { @@ -20700,11 +20700,11 @@ const docTemplate = `{ "format": "date-time" }, "spend_limit_micros": { - "description": "SpendLimitMicros is the group's configured per-member AI spend limit.\nNull when the 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 combined limit of the members whose spend is\nattributed to this group, applying each member's override where one\nexists. Null when any of those members has no limit, and zero when no\nmembers are attributed to the group.", + "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" } } @@ -22184,7 +22184,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": { @@ -22192,11 +22192,11 @@ const docTemplate = `{ "format": "uuid" }, "spend_limit_micros": { - "description": "SpendLimitMicros is the group's configured per-member AI spend limit.\nNull when the 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 combined limit of the members whose spend is\nattributed to this group, applying each member's override where one\nexists. Null when any of those members has no limit, and zero when no\nmembers are attributed to the group.", + "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 8f91477619b42..9a02023b8ce6f 100644 --- a/coderd/apidoc/swagger.json +++ b/coderd/apidoc/swagger.json @@ -18826,7 +18826,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": { @@ -18844,11 +18844,11 @@ "format": "date-time" }, "spend_limit_micros": { - "description": "SpendLimitMicros is the group's configured per-member AI spend limit.\nNull when the 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 combined limit of the members whose spend is\nattributed to this group, applying each member's override where one\nexists. Null when any of those members has no limit, and zero when no\nmembers are attributed to the group.", + "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" } } @@ -20250,7 +20250,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": { @@ -20258,11 +20258,11 @@ "format": "uuid" }, "spend_limit_micros": { - "description": "SpendLimitMicros is the group's configured per-member AI spend limit.\nNull when the 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 combined limit of the members whose spend is\nattributed to this group, applying each member's override where one\nexists. Null when any of those members has no limit, and zero when no\nmembers are attributed to the group.", + "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/codersdk/aibridge.go b/codersdk/aibridge.go index de712638f7f3e..e5fb5aeadb649 100644 --- a/codersdk/aibridge.go +++ b/codersdk/aibridge.go @@ -84,16 +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 per-member AI spend limit. + // 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"` - // TotalSpendLimitMicros is the combined limit of the members whose spend is - // attributed to this group, applying each member's override where one - // exists. Null when any of those members has no limit, and zero when no - // members are attributed to the group. + // 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 is the group's spend over the current budget period. CurrentSpendMicros int64 `json:"current_spend_micros"` } diff --git a/docs/reference/api/schemas.md b/docs/reference/api/schemas.md index 2105a8c7a53ad..01e16582d77d0 100644 --- a/docs/reference/api/schemas.md +++ b/docs/reference/api/schemas.md @@ -7754,14 +7754,14 @@ Only certain features set these fields: - FeatureManagedAgentLimit| ### 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 per-member AI spend limit. Null when the group has no configured budget. | -| `total_spend_limit_micros` | integer | false | | Total spend limit micros is the combined limit of the members whose spend is attributed to this group, applying each member's override where one exists. Null when any of those members has no limit, and zero when no members are attributed to the group. | +| 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 @@ -9300,12 +9300,12 @@ Only certain features set these fields: - FeatureManagedAgentLimit| ### 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 per-member AI spend limit. Null when the group has no configured budget. | -| `total_spend_limit_micros` | integer | false | | Total spend limit micros is the combined limit of the members whose spend is attributed to this group, applying each member's override where one exists. Null when any of those members has no limit, and zero when no members are attributed to the group. | +| 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 diff --git a/site/src/api/typesGenerated.ts b/site/src/api/typesGenerated.ts index 224ee061cc002..792ad79110688 100644 --- a/site/src/api/typesGenerated.ts +++ b/site/src/api/typesGenerated.ts @@ -6884,20 +6884,19 @@ export interface Organization extends MinimalOrganization { export interface OrganizationGroupAISpend { readonly group_id: string; /** - * SpendLimitMicros is the group's configured per-member AI spend limit. + * 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; /** - * TotalSpendLimitMicros is the combined limit of the members whose spend is - * attributed to this group, applying each member's override where one - * exists. Null when any of those members has no limit, and zero when no - * members are attributed to the group. + * 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. + * CurrentSpendMicros is the group's spend over the current budget period. */ readonly current_spend_micros: number; } From daf46508af14ea478e7ea5cbc32d8c3b0f8c7513 Mon Sep 17 00:00:00 2001 From: Susana Cardoso Ferreira Date: Tue, 28 Jul 2026 18:01:16 +0000 Subject: [PATCH 4/4] tests: improve tests --- coderd/database/querier.go | 6 +- coderd/database/querier_test.go | 350 ++++++++++++---------- coderd/database/queries.sql.go | 15 +- coderd/database/queries/aicostcontrol.sql | 15 +- enterprise/coderd/aibridge_test.go | 2 + 5 files changed, 217 insertions(+), 171 deletions(-) diff --git a/coderd/database/querier.go b/coderd/database/querier.go index 070052ff0c51d..e3030339624f0 100644 --- a/coderd/database/querier.go +++ b/coderd/database/querier.go @@ -686,9 +686,9 @@ type sqlcQuerier interface { // Returns AI spend limits and aggregate spend for groups in @group_ids that // 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 per-member limit applied to the members - // attributed to the group, replaced by their override where one exists. It is - // null when the group has no budget, because those members spend without a cap. + // 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) diff --git a/coderd/database/querier_test.go b/coderd/database/querier_test.go index 5ff2a4360a7b8..54c0bdc0f2f92 100644 --- a/coderd/database/querier_test.go +++ b/coderd/database/querier_test.go @@ -12883,140 +12883,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 { - day time.Time - spend int64 + // nullInt64 keeps the expectations below readable. + nullInt64 := func(v int64) sql.NullInt64 { + return sql.NullInt64{Int64: v, Valid: true} } - // overrideSpec charges a per-user budget override to one of the two groups. - type overrideSpec struct { - user int - onB bool + 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 + } - // Every case builds an org with two users and two groups. The first group is - // the subject, and the second exists so a case can pull a member's effective - // budget away from the first. members and membersB hold indexes into the - // users slice, and seeded spend is always attributed to users[0] on the - // first group. + // 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 - setBudgetB bool - spendLimitB int64 - members []int - membersB []int - override *overrideSpec - rows []seedRow - wantCurrentSpend int64 - wantTotalLimit sql.NullInt64 - wantTotalLimitB sql.NullInt64 + name string + budgets []groupBudget + members []membership + overrides []override + spend []spendRow + want map[string]wantGroup }{ { - name: "NoBudgetNoSpend", - wantCurrentSpend: 0, - wantTotalLimit: sql.NullInt64{}, + name: "NoBudgetNoSpend", + want: map[string]wantGroup{ + "eng": {spendLimit: sql.NullInt64{}, totalLimit: sql.NullInt64{}}, + }, }, { - name: "ZeroLimitBudget", - setBudget: true, - spendLimit: 0, - wantCurrentSpend: 0, - wantTotalLimit: sql.NullInt64{Int64: 0, Valid: true}, + 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, - wantTotalLimit: sql.NullInt64{Int64: 0, Valid: true}, + // 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, - wantTotalLimit: sql.NullInt64{Int64: 0, Valid: true}, + // 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, - wantTotalLimit: sql.NullInt64{}, + 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", - setBudget: true, - spendLimit: 100, - members: []int{0, 1}, - wantTotalLimit: sql.NullInt64{Int64: 200, Valid: true}, + 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", - setBudget: true, - spendLimit: 100, - members: []int{0, 1}, - override: &overrideSpec{user: 0, limit: 1000}, + 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. - wantTotalLimit: sql.NullInt64{Int64: 1100, Valid: true}, + want: map[string]wantGroup{ + "eng": {spendLimit: nullInt64(100), totalLimit: nullInt64(1100)}, + }, }, { - name: "BudgetedOnlyOverrideMember", - setBudget: true, - spendLimit: 100, - members: []int{0}, - override: &overrideSpec{user: 0, limit: 1000}, - wantTotalLimit: sql.NullInt64{Int64: 1000, Valid: true}, + 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", - setBudget: true, - spendLimit: 0, - members: []int{0, 1}, - wantTotalLimit: sql.NullInt64{Int64: 0, Valid: true}, + 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: []int{0, 1}, - wantTotalLimit: sql.NullInt64{}, + 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: []int{0}, - override: &overrideSpec{user: 0, limit: 1000}, + 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. - wantTotalLimit: sql.NullInt64{}, + want: map[string]wantGroup{ + "eng": {spendLimit: sql.NullInt64{}, totalLimit: sql.NullInt64{}}, + }, }, { - name: "MemberResolvesToHigherBudgetGroup", - setBudget: true, - spendLimit: 100, - setBudgetB: true, - spendLimitB: 500, - members: []int{0, 1}, - membersB: []int{0}, - // users[0] is attributed to the higher-limit second group, so only - // users[1] counts toward the first. - wantTotalLimit: sql.NullInt64{Int64: 100, Valid: true}, - wantTotalLimitB: sql.NullInt64{Int64: 500, Valid: true}, + 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", - setBudget: true, - spendLimit: 100, - setBudgetB: true, - spendLimitB: 50, - members: []int{0, 1}, - membersB: []int{0}, - override: &overrideSpec{user: 0, onB: true, limit: 1000}, - // The override attributes users[0] to the second group even though - // the first carries the higher group limit. - wantTotalLimit: sql.NullInt64{Int64: 100, Valid: true}, - wantTotalLimitB: sql.NullInt64{Int64: 1000, Valid: true}, + 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)}, + }, }, } @@ -13026,89 +13060,92 @@ func TestGetOrganizationGroupsAISpend(t *testing.T) { db, _ := dbtestutil.NewDB(t) ctx := testutil.Context(t, testutil.WaitShort) - // Given: an org with two groups and two members, wired up per the case. + // Given: the groups, members, overrides and spend the case declares. org := dbgen.Organization(t, db, database.Organization{}) - users := []database.User{ - dbgen.User(t, db, database.User{}), - dbgen.User(t, db, database.User{}), + 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 } - for _, u := range users { + 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: u.ID, + UserID: user.ID, OrganizationID: org.ID, }) + userIDs[name] = user.ID + return user.ID } - group := dbgen.Group(t, db, database.Group{OrganizationID: org.ID}) - groupB := dbgen.Group(t, db, database.Group{OrganizationID: org.ID}) - for _, b := range []struct { - groupID uuid.UUID - set bool - limit int64 - }{{group.ID, tt.setBudget, tt.spendLimit}, {groupB.ID, tt.setBudgetB, tt.spendLimitB}} { - if !b.set { - continue - } + + for name := range tt.want { + groupID(name) + } + for _, b := range tt.budgets { _, err := db.UpsertGroupAIBudget(ctx, database.UpsertGroupAIBudgetParams{ - GroupID: b.groupID, + GroupID: groupID(b.group), SpendLimitMicros: b.limit, }) require.NoError(t, err) } - for _, i := range tt.members { - dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: group.ID, UserID: users[i].ID}) - } - for _, i := range tt.membersB { - dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: groupB.ID, UserID: users[i].ID}) + for _, m := range tt.members { + dbgen.GroupMember(t, db, database.GroupMemberTable{ + GroupID: groupID(m.group), + UserID: userID(m.user), + }) } - if tt.override != nil { - overrideGroup := group.ID - if tt.override.onB { - overrideGroup = groupB.ID - } + for _, o := range tt.overrides { _, err := db.UpsertUserAIBudgetOverride(ctx, database.UpsertUserAIBudgetOverrideParams{ - UserID: users[tt.override.user].ID, - GroupID: overrideGroup, - SpendLimitMicros: tt.override.limit, + UserID: userID(o.user), + GroupID: groupID(o.group), + SpendLimitMicros: o.limit, }) require.NoError(t, err) } - for _, r := range tt.rows { + for _, s := range tt.spend { _, err := db.IncrementUserAIDailySpend(ctx, database.IncrementUserAIDailySpendParams{ - UserID: users[0].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 both groups 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, groupB.ID}, + GroupIds: queried, PeriodStart: monthStart, }) require.NoError(t, err) // Then: each group reports its own limit, spend, and the combined // limit of the members attributed to it. - require.Len(t, got, 2) + require.Len(t, got, len(groupIDs)) byID := make(map[uuid.UUID]database.GetOrganizationGroupsAISpendRow, len(got)) for _, r := range got { byID[r.GroupID] = r } - row, ok := byID[group.ID] - require.True(t, ok, "queried group missing from response") - require.Equal(t, org.ID, row.OrganizationID) - if tt.setBudget { - require.True(t, row.SpendLimitMicros.Valid, "expected configured budget") - require.Equal(t, tt.spendLimit, row.SpendLimitMicros.Int64, "spend_limit_micros") - } else { - require.False(t, row.SpendLimitMicros.Valid, "expected no configured budget") + 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, row.CurrentSpendMicros) - require.Equal(t, tt.wantTotalLimit, row.TotalSpendLimitMicros, "total_spend_limit_micros") - require.Equal(t, tt.wantTotalLimitB, byID[groupB.ID].TotalSpendLimitMicros, "second group total_spend_limit_micros") }) } @@ -13153,10 +13190,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) }) @@ -13323,7 +13364,8 @@ func TestGetOrganizationGroupsAISpend(t *testing.T) { db, _ := dbtestutil.NewDB(t) ctx := testutil.Context(t, testutil.WaitShort) - // Given: an org whose implicit Everyone group carries the only budget. + // 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{}) @@ -13353,7 +13395,9 @@ func TestGetOrganizationGroupsAISpend(t *testing.T) { // Then: every org member counts toward the total. require.Len(t, got, 1) - require.Equal(t, sql.NullInt64{Int64: 200, Valid: true}, got[0].TotalSpendLimitMicros) + 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) { @@ -13361,7 +13405,7 @@ func TestGetOrganizationGroupsAISpend(t *testing.T) { db, _ := dbtestutil.NewDB(t) ctx := testutil.Context(t, testutil.WaitShort) - // Given: an org with members but no budget anywhere, which is where + // 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 { @@ -13385,7 +13429,9 @@ func TestGetOrganizationGroupsAISpend(t *testing.T) { // Then: the total is null, so the group reads as unlimited. require.Len(t, got, 1) - require.False(t, got[0].TotalSpendLimitMicros.Valid) + 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") }) } diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go index aacd189bce3e6..277f2e8db9cbe 100644 --- a/coderd/database/queries.sql.go +++ b/coderd/database/queries.sql.go @@ -2868,9 +2868,8 @@ WITH queried_groups AS ( AND groups.id = ANY($2::uuid[]) ), candidate_users AS ( - -- Members of the queried groups. Narrowing to members loses nobody, because a - -- user's effective group is always a group they belong to. Uses - -- group_members_expanded so the implicit Everyone group counts. + -- 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) @@ -2894,8 +2893,8 @@ user_highest_group AS ( ), effective AS ( -- Effective budget group per user: an override wins over the highest-limit - -- group they belong to. Users with neither spend without a cap, and the - -- unbudgeted group they fall back to reports unlimited regardless. + -- 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, @@ -2967,9 +2966,9 @@ type GetOrganizationGroupsAISpendRow struct { // Returns AI spend limits and aggregate spend for groups in @group_ids that // 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 per-member limit applied to the members -// attributed to the group, replaced by their override where one exists. It is -// null when the group has no budget, because those members spend without a cap. +// 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) { diff --git a/coderd/database/queries/aicostcontrol.sql b/coderd/database/queries/aicostcontrol.sql index cbe4035eeb579..43c5970fe2dca 100644 --- a/coderd/database/queries/aicostcontrol.sql +++ b/coderd/database/queries/aicostcontrol.sql @@ -124,9 +124,9 @@ WHERE user_id = @user_id -- Returns AI spend limits and aggregate spend for groups in @group_ids that -- 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 per-member limit applied to the members --- attributed to the group, replaced by their override where one exists. It is --- null when the group has no budget, because those members spend without a cap. +-- 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 ( @@ -137,9 +137,8 @@ WITH queried_groups AS ( AND groups.id = ANY(@group_ids::uuid[]) ), candidate_users AS ( - -- Members of the queried groups. Narrowing to members loses nobody, because a - -- user's effective group is always a group they belong to. Uses - -- group_members_expanded so the implicit Everyone group counts. + -- 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) @@ -163,8 +162,8 @@ user_highest_group AS ( ), effective AS ( -- Effective budget group per user: an override wins over the highest-limit - -- group they belong to. Users with neither spend without a cap, and the - -- unbudgeted group they fall back to reports unlimited regardless. + -- 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, diff --git a/enterprise/coderd/aibridge_test.go b/enterprise/coderd/aibridge_test.go index 31874447a77f2..9d9fc1776b745 100644 --- a/enterprise/coderd/aibridge_test.go +++ b/enterprise/coderd/aibridge_test.go @@ -3668,6 +3668,7 @@ func TestOrganizationGroupsAISpend(t *testing.T) { 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) }) } @@ -4817,6 +4818,7 @@ func TestGroupAISpend(t *testing.T) { // 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) }) }