opensource/coder
3
0
Fork 0
mirror of https://github.com/coder/coder.git synced 2026-09-24 15:04:27 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix: report combined member limit in group AI spend (#27589)

Browse Source 
## Problem

The organization groups page showed each group's AI budget as the
group's per-member limit, so the total it displayed was effectively
group members × group budget. That ignores per-user budget overrides
charged to the group, so a group where one member has an override
reported a limit that doesn't match what its members can actually spend.

## Changes

- Add `total_spend_limit_micros` to the organization groups AI spend
payload, the combined budget of the members attributed to the group,
with each member's override replacing their share.
- Return `null` for the total when the group has no budget, since its
members spend without a cap.
- Both the organization groups and single group spend endpoints report
the new field, as they share the same query.
- Use the total as the denominator on the groups page AI budget column.

Depends on #27568
This commit is contained in:
Susana Ferreira
2026-07-29 09:16:37 +01:00
committed by GitHub
parent e96d8646e2
commit 0b4095085e
14 changed files with 705 additions and 177 deletions
Download Patch File Download Diff File Expand all files Collapse all files
coderd/apidoc/docs.go
Generated
+12 -4
View File
@@ -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"
}
}
coderd/apidoc/swagger.json
Generated
+12 -4
View File
@@ -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"
}
}
coderd/database/db2sdk/db2sdk.go
+3
View File
@@ -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
}
coderd/database/querier.go
Generated
+6 -2
View File
@@ -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)
coderd/database/querier_test.go
+301 -47
View File
@@ -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) {
coderd/database/queries.sql.go
Generated
+99 -20
View File
@@ -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
coderd/database/queries/aicostcontrol.sql
+91 -14
View File
@@ -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