feat: fall back to the Everyone group for AI spend attribution (#27364)

## Description

Previously, a user with no per-user override and no membership in a budgeted group had no effective group, so their AI spend was attributed nowhere and was, therefore, untracked. This change falls back to the organization's Everyone group when no override or group budget applies.

Since every user in an organization is implicitly a member of that org's Everyone group, spend is now attributed and tracked for any user with organization membership. A user with no organization membership resolves to no group, so their daily spend is not incremented and a warning is logged.

The fallback is unlimited, so enforcement is unaffected: only override and group budgets can block requests. For users in multiple organizations, an existing budget on any Everyone group is still chosen by the "highest" policy; when none is budgeted, the fallback prefers the default org, then orders by organization name.

## Changes

- Add `ResolveUserEffectiveGroup` and the `GetUserEveryoneFallbackGroup` query: resolve override → group budget → Everyone group fallback.
- Attribute token-usage spend and the user AI spend endpoint via the fallback, so unbudgeted users resolve to their Everyone group instead of null.
- Update `GetGroupMembersAISpend` to surface the Everyone fallback as the effective group.
- Update `GetHighestGroupAIBudgetByUser` to break ties by organization name then group name, keeping multi-org resolution deterministic and consistent with the fallback.
- For multi-org users with no budget anywhere, the fallback picks the Everyone group deterministically: prefer the default org, then order by organization name.

Closes https://linear.app/codercom/issue/AIGOV-509/fall-back-to-the-everyone-group-for-spend-attribution

> [!NOTE]
> Initially generated by Claude Opus 4.7, modified and reviewed by @ssncferreira
This commit is contained in:
Susana Ferreira
2026-07-23 09:26:25 +01:00
committed by GitHub
parent 671173b498
commit c23f2c0223
20 changed files with 1096 additions and 272 deletions
+60 -22
View File
@@ -20,37 +20,47 @@ import (
type Store interface {
GetUserAIBudgetOverride(ctx context.Context, userID uuid.UUID) (database.UserAIBudgetOverride, error)
GetHighestGroupAIBudgetByUser(ctx context.Context, userID uuid.UUID) (database.GetHighestGroupAIBudgetByUserRow, error)
GetUserEveryoneFallbackGroup(ctx context.Context, userID uuid.UUID) (uuid.UUID, error)
}
// EffectiveBudget is the AI budget that applies to a user after override and
// policy resolution.
type EffectiveBudget struct {
// EffectiveGroup is a user's resolved effective group and, when a budget
// applies, its limit. Limit is nil for the Everyone fallback (unlimited).
type EffectiveGroup struct {
// GroupID is the group the spend is attributed to.
GroupID uuid.UUID
// SpendLimitMicros is the effective spend limit in micro-units
// (1 unit = 1,000,000).
// Limit is the resolved spend limit, or nil for the unlimited Everyone
// fallback.
Limit *Limit
}
// Limit is an AI spend limit and the source that produced it.
type Limit struct {
// SpendLimitMicros is the spend limit in micro-units (1 unit = 1,000,000).
SpendLimitMicros int64
Source codersdk.AIBudgetLimitSource
}
// ResolveUserAIBudget returns the effective AI budget for userID. The second
// return value is false when no budget is configured for the user. A per-user
// override wins unconditionally; otherwise the budget is selected from the
// user's groups according to policy.
// ResolveUserAIBudget returns the effective AI budget group for userID,
// resolved in order:
// 1. A per-user override, if configured.
// 2. Otherwise, a group budget selected by the deployment policy.
//
// The second return value is false when no budget is configured for the user.
// TODO(AIGOV-527): unify effective group resolution in a single place.
func ResolveUserAIBudget(ctx context.Context, db Store, userID uuid.UUID, policy codersdk.AIBudgetPolicy) (EffectiveBudget, bool, error) {
func ResolveUserAIBudget(ctx context.Context, db Store, userID uuid.UUID, policy codersdk.AIBudgetPolicy) (EffectiveGroup, bool, error) {
// A per-user override always wins.
override, err := db.GetUserAIBudgetOverride(ctx, userID)
if err == nil {
return EffectiveBudget{
GroupID: override.GroupID,
SpendLimitMicros: override.SpendLimitMicros,
Source: codersdk.AIBudgetLimitSourceUserOverride,
return EffectiveGroup{
GroupID: override.GroupID,
Limit: &Limit{
SpendLimitMicros: override.SpendLimitMicros,
Source: codersdk.AIBudgetLimitSourceUserOverride,
},
}, true, nil
}
if !errors.Is(err, sql.ErrNoRows) {
return EffectiveBudget{}, false, xerrors.Errorf("get user AI budget override: %w", err)
return EffectiveGroup{}, false, xerrors.Errorf("get user AI budget override: %w", err)
}
// No override: select a group budget according to the deployment policy.
@@ -58,17 +68,45 @@ func ResolveUserAIBudget(ctx context.Context, db Store, userID uuid.UUID, policy
case codersdk.AIBudgetPolicyHighest:
row, err := db.GetHighestGroupAIBudgetByUser(ctx, userID)
if errors.Is(err, sql.ErrNoRows) {
return EffectiveBudget{}, false, nil
return EffectiveGroup{}, false, nil
}
if err != nil {
return EffectiveBudget{}, false, xerrors.Errorf("get highest group AI budget: %w", err)
return EffectiveGroup{}, false, xerrors.Errorf("get highest group AI budget: %w", err)
}
return EffectiveBudget{
GroupID: row.GroupID,
SpendLimitMicros: row.SpendLimitMicros,
Source: codersdk.AIBudgetLimitSourceGroup,
return EffectiveGroup{
GroupID: row.GroupID,
Limit: &Limit{
SpendLimitMicros: row.SpendLimitMicros,
Source: codersdk.AIBudgetLimitSourceGroup,
},
}, true, nil
default:
return EffectiveBudget{}, false, xerrors.Errorf("unsupported AI budget policy: %q", policy)
return EffectiveGroup{}, false, xerrors.Errorf("unsupported AI budget policy: %q", policy)
}
}
// ResolveUserEffectiveGroup resolves the user's effective group, falling back to
// the organization's Everyone group when no override or group budget applies.
// The second return value is false when no effective group was found for the
// user.
func ResolveUserEffectiveGroup(ctx context.Context, db Store, userID uuid.UUID, policy codersdk.AIBudgetPolicy) (EffectiveGroup, bool, error) {
group, ok, err := ResolveUserAIBudget(ctx, db, userID, policy)
if err != nil {
return EffectiveGroup{}, false, err
}
if ok {
return group, true, nil
}
// No override or group budget: fall back to the Everyone group (unlimited).
groupID, err := db.GetUserEveryoneFallbackGroup(ctx, userID)
if errors.Is(err, sql.ErrNoRows) {
// This should not happen, as a user should always be a member of an
// organization and its associated Everyone group.
return EffectiveGroup{}, false, nil
}
if err != nil {
return EffectiveGroup{}, false, xerrors.Errorf("get everyone fallback group: %w", err)
}
return EffectiveGroup{GroupID: groupID}, true, nil
}
+158 -37
View File
@@ -4,6 +4,7 @@ import (
"bytes"
"context"
"testing"
"time"
"github.com/google/uuid"
"github.com/stretchr/testify/require"
@@ -52,18 +53,18 @@ func TestResolveUserAIBudget(t *testing.T) {
tests := []struct {
name string
policy codersdk.AIBudgetPolicy
setup func(t *testing.T, ctx context.Context, db database.Store) (userID uuid.UUID, want budget.EffectiveBudget, wantOK bool)
setup func(t *testing.T, ctx context.Context, db database.Store) (userID uuid.UUID, want budget.EffectiveGroup, wantOK bool)
wantErr string
}{
{
name: "OverrideWins",
policy: codersdk.AIBudgetPolicyHighest,
setup: func(t *testing.T, ctx context.Context, db database.Store) (uuid.UUID, budget.EffectiveBudget, bool) {
setup: func(t *testing.T, ctx context.Context, db database.Store) (uuid.UUID, budget.EffectiveGroup, bool) {
org := dbgen.Organization(t, db, database.Organization{})
user := dbgen.User(t, db, database.User{})
// A higher group budget that the override must still beat.
budgetedGroup(t, ctx, db, org.ID, user.ID, "rich-group", 9_000_000)
// The override names its own group; the user must be a member.
// The override names a group the user must be a member of.
og := dbgen.Group(t, db, database.Group{OrganizationID: org.ID, Name: "override-group"})
dbgen.GroupMember(t, db, database.GroupMemberTable{UserID: user.ID, GroupID: og.ID})
_, err := db.UpsertUserAIBudgetOverride(ctx, database.UpsertUserAIBudgetOverrideParams{
@@ -72,92 +73,98 @@ func TestResolveUserAIBudget(t *testing.T) {
SpendLimitMicros: 1_000_000,
})
require.NoError(t, err)
return user.ID, budget.EffectiveBudget{GroupID: og.ID, SpendLimitMicros: 1_000_000, Source: codersdk.AIBudgetLimitSourceUserOverride}, true
return user.ID, budget.EffectiveGroup{GroupID: og.ID, Limit: &budget.Limit{SpendLimitMicros: 1_000_000, Source: codersdk.AIBudgetLimitSourceUserOverride}}, true
},
},
{
name: "SingleGroupBudget",
policy: codersdk.AIBudgetPolicyHighest,
setup: func(t *testing.T, ctx context.Context, db database.Store) (uuid.UUID, budget.EffectiveBudget, bool) {
setup: func(t *testing.T, ctx context.Context, db database.Store) (uuid.UUID, budget.EffectiveGroup, bool) {
org := dbgen.Organization(t, db, database.Organization{})
user := dbgen.User(t, db, database.User{})
dbgen.OrganizationMember(t, db, database.OrganizationMember{OrganizationID: org.ID, UserID: user.ID})
gid := budgetedGroup(t, ctx, db, org.ID, user.ID, "only", 8_000_000)
return user.ID, budget.EffectiveBudget{GroupID: gid, SpendLimitMicros: 8_000_000, Source: codersdk.AIBudgetLimitSourceGroup}, true
return user.ID, budget.EffectiveGroup{GroupID: gid, Limit: &budget.Limit{SpendLimitMicros: 8_000_000, Source: codersdk.AIBudgetLimitSourceGroup}}, true
},
},
{
name: "HighestGroupWins",
policy: codersdk.AIBudgetPolicyHighest,
setup: func(t *testing.T, ctx context.Context, db database.Store) (uuid.UUID, budget.EffectiveBudget, bool) {
setup: func(t *testing.T, ctx context.Context, db database.Store) (uuid.UUID, budget.EffectiveGroup, bool) {
org := dbgen.Organization(t, db, database.Organization{})
user := dbgen.User(t, db, database.User{})
dbgen.OrganizationMember(t, db, database.OrganizationMember{OrganizationID: org.ID, UserID: user.ID})
budgetedGroup(t, ctx, db, org.ID, user.ID, "low", 5_000_000)
budgetedGroup(t, ctx, db, org.ID, user.ID, "mid", 20_000_000)
high := budgetedGroup(t, ctx, db, org.ID, user.ID, "high", 50_000_000)
return user.ID, budget.EffectiveBudget{GroupID: high, SpendLimitMicros: 50_000_000, Source: codersdk.AIBudgetLimitSourceGroup}, true
return user.ID, budget.EffectiveGroup{GroupID: high, Limit: &budget.Limit{SpendLimitMicros: 50_000_000, Source: codersdk.AIBudgetLimitSourceGroup}}, true
},
},
{
name: "TieBrokenByName",
name: "TieBrokenByEarliestOrgMembership",
policy: codersdk.AIBudgetPolicyHighest,
setup: func(t *testing.T, ctx context.Context, db database.Store) (uuid.UUID, budget.EffectiveBudget, bool) {
org := dbgen.Organization(t, db, database.Organization{})
setup: func(t *testing.T, ctx context.Context, db database.Store) (uuid.UUID, budget.EffectiveGroup, bool) {
user := dbgen.User(t, db, database.User{})
// Equal limits; "alpha" must win over "beta" by name ascending.
alpha := budgetedGroup(t, ctx, db, org.ID, user.ID, "alpha", 10_000_000)
budgetedGroup(t, ctx, db, org.ID, user.ID, "beta", 10_000_000)
return user.ID, budget.EffectiveBudget{GroupID: alpha, SpendLimitMicros: 10_000_000, Source: codersdk.AIBudgetLimitSourceGroup}, true
// Two groups in different orgs share the same limit. The earlier
// organization membership breaks the tie.
earlyOrg := dbgen.Organization(t, db, database.Organization{})
lateOrg := dbgen.Organization(t, db, database.Organization{})
dbgen.OrganizationMember(t, db, database.OrganizationMember{OrganizationID: earlyOrg.ID, UserID: user.ID, CreatedAt: time.Now().Add(-time.Hour)})
dbgen.OrganizationMember(t, db, database.OrganizationMember{OrganizationID: lateOrg.ID, UserID: user.ID})
winner := budgetedGroup(t, ctx, db, earlyOrg.ID, user.ID, "dup", 10_000_000)
budgetedGroup(t, ctx, db, lateOrg.ID, user.ID, "dup", 10_000_000)
return user.ID, budget.EffectiveGroup{GroupID: winner, Limit: &budget.Limit{SpendLimitMicros: 10_000_000, Source: codersdk.AIBudgetLimitSourceGroup}}, true
},
},
{
name: "TieBrokenByGroupID",
policy: codersdk.AIBudgetPolicyHighest,
setup: func(t *testing.T, ctx context.Context, db database.Store) (uuid.UUID, budget.EffectiveBudget, bool) {
setup: func(t *testing.T, ctx context.Context, db database.Store) (uuid.UUID, budget.EffectiveGroup, bool) {
org := dbgen.Organization(t, db, database.Organization{})
user := dbgen.User(t, db, database.User{})
// Two groups in different orgs share both name and limit.
// Group id breaks the tie, so resolution is deterministic.
org1 := dbgen.Organization(t, db, database.Organization{})
org2 := dbgen.Organization(t, db, database.Organization{})
g1 := budgetedGroup(t, ctx, db, org1.ID, user.ID, "dup", 10_000_000)
g2 := budgetedGroup(t, ctx, db, org2.ID, user.ID, "dup", 10_000_000)
winner := g1
if bytes.Compare(g2[:], g1[:]) < 0 {
winner = g2
dbgen.OrganizationMember(t, db, database.OrganizationMember{OrganizationID: org.ID, UserID: user.ID})
// Both groups are in the same org, so both resolve to the same
// organization membership and the tie falls to the lowest group ID.
groupA := budgetedGroup(t, ctx, db, org.ID, user.ID, "alpha", 10_000_000)
groupB := budgetedGroup(t, ctx, db, org.ID, user.ID, "beta", 10_000_000)
winner := groupA
if bytes.Compare(groupB[:], groupA[:]) < 0 {
winner = groupB
}
return user.ID, budget.EffectiveBudget{GroupID: winner, SpendLimitMicros: 10_000_000, Source: codersdk.AIBudgetLimitSourceGroup}, true
return user.ID, budget.EffectiveGroup{GroupID: winner, Limit: &budget.Limit{SpendLimitMicros: 10_000_000, Source: codersdk.AIBudgetLimitSourceGroup}}, true
},
},
{
name: "GroupsButNoneBudgeted",
policy: codersdk.AIBudgetPolicyHighest,
setup: func(t *testing.T, ctx context.Context, db database.Store) (uuid.UUID, budget.EffectiveBudget, bool) {
setup: func(t *testing.T, ctx context.Context, db database.Store) (uuid.UUID, budget.EffectiveGroup, bool) {
org := dbgen.Organization(t, db, database.Organization{})
user := dbgen.User(t, db, database.User{})
g := dbgen.Group(t, db, database.Group{OrganizationID: org.ID, Name: "unbudgeted"})
dbgen.GroupMember(t, db, database.GroupMemberTable{UserID: user.ID, GroupID: g.ID})
return user.ID, budget.EffectiveBudget{}, false
return user.ID, budget.EffectiveGroup{}, false
},
},
{
name: "EveryoneGroupBudget",
policy: codersdk.AIBudgetPolicyHighest,
setup: func(t *testing.T, ctx context.Context, db database.Store) (uuid.UUID, budget.EffectiveBudget, bool) {
setup: func(t *testing.T, ctx context.Context, db database.Store) (uuid.UUID, budget.EffectiveGroup, bool) {
org := dbgen.Organization(t, db, database.Organization{})
user := dbgen.User(t, db, database.User{})
// Membership is via organization_members only (no group_members row),
// exercising the org-members half of group_members_expanded.
everyoneID := budgetedEveryoneGroup(t, ctx, db, org.ID, user.ID, 7_000_000)
return user.ID, budget.EffectiveBudget{GroupID: everyoneID, SpendLimitMicros: 7_000_000, Source: codersdk.AIBudgetLimitSourceGroup}, true
return user.ID, budget.EffectiveGroup{GroupID: everyoneID, Limit: &budget.Limit{SpendLimitMicros: 7_000_000, Source: codersdk.AIBudgetLimitSourceGroup}}, true
},
},
{
name: "OverrideBeatsEveryoneBudget",
policy: codersdk.AIBudgetPolicyHighest,
setup: func(t *testing.T, ctx context.Context, db database.Store) (uuid.UUID, budget.EffectiveBudget, bool) {
setup: func(t *testing.T, ctx context.Context, db database.Store) (uuid.UUID, budget.EffectiveGroup, bool) {
org := dbgen.Organization(t, db, database.Organization{})
user := dbgen.User(t, db, database.User{})
everyoneID := budgetedEveryoneGroup(t, ctx, db, org.ID, user.ID, 7_000_000)
// Override attributed to the Everyone group; the user is a member
// Override attributed to the Everyone group. The user is a member
// via organization_members, satisfying the membership trigger.
_, err := db.UpsertUserAIBudgetOverride(ctx, database.UpsertUserAIBudgetOverrideParams{
UserID: user.ID,
@@ -165,16 +172,16 @@ func TestResolveUserAIBudget(t *testing.T) {
SpendLimitMicros: 2_000_000,
})
require.NoError(t, err)
return user.ID, budget.EffectiveBudget{GroupID: everyoneID, SpendLimitMicros: 2_000_000, Source: codersdk.AIBudgetLimitSourceUserOverride}, true
return user.ID, budget.EffectiveGroup{GroupID: everyoneID, Limit: &budget.Limit{SpendLimitMicros: 2_000_000, Source: codersdk.AIBudgetLimitSourceUserOverride}}, true
},
},
{
name: "UnsupportedPolicy",
policy: codersdk.AIBudgetPolicy("unsupported"),
setup: func(t *testing.T, ctx context.Context, db database.Store) (uuid.UUID, budget.EffectiveBudget, bool) {
setup: func(t *testing.T, ctx context.Context, db database.Store) (uuid.UUID, budget.EffectiveGroup, bool) {
// No override, so resolution reaches the policy switch and errors.
user := dbgen.User(t, db, database.User{})
return user.ID, budget.EffectiveBudget{}, false
return user.ID, budget.EffectiveGroup{}, false
},
wantErr: "unsupported AI budget policy",
},
@@ -199,8 +206,122 @@ func TestResolveUserAIBudget(t *testing.T) {
return
}
require.Equal(t, want.GroupID, got.GroupID)
require.Equal(t, want.SpendLimitMicros, got.SpendLimitMicros)
require.Equal(t, want.Source, got.Source)
require.Equal(t, want.Limit, got.Limit)
})
}
}
func TestResolveUserEffectiveGroup(t *testing.T) {
t.Parallel()
tests := []struct {
name string
policy codersdk.AIBudgetPolicy
setup func(t *testing.T, ctx context.Context, db database.Store) (userID uuid.UUID, want budget.EffectiveGroup, wantOK bool)
wantErr string
}{
{
// The Everyone group has a budget, so it resolves via the budget
// path rather than the fallback.
name: "EveryoneGroupWithBudget",
policy: codersdk.AIBudgetPolicyHighest,
setup: func(t *testing.T, ctx context.Context, db database.Store) (uuid.UUID, budget.EffectiveGroup, bool) {
org := dbgen.Organization(t, db, database.Organization{})
user := dbgen.User(t, db, database.User{})
// The Everyone group's id equals the org id.
group := dbgen.Group(t, db, database.Group{ID: org.ID, OrganizationID: org.ID, Name: "Everyone"})
dbgen.OrganizationMember(t, db, database.OrganizationMember{OrganizationID: org.ID, UserID: user.ID})
_, err := db.UpsertGroupAIBudget(ctx, database.UpsertGroupAIBudgetParams{
GroupID: group.ID,
SpendLimitMicros: 7_000_000,
})
require.NoError(t, err)
return user.ID, budget.EffectiveGroup{GroupID: group.ID, Limit: &budget.Limit{SpendLimitMicros: 7_000_000, Source: codersdk.AIBudgetLimitSourceGroup}}, true
},
},
{
// With a single org and no budget, attribution falls back to that
// org's Everyone group with no limit.
name: "FallbackToEveryoneUnlimited",
policy: codersdk.AIBudgetPolicyHighest,
setup: func(t *testing.T, ctx context.Context, db database.Store) (uuid.UUID, budget.EffectiveGroup, bool) {
org := dbgen.Organization(t, db, database.Organization{})
user := dbgen.User(t, db, database.User{})
dbgen.OrganizationMember(t, db, database.OrganizationMember{OrganizationID: org.ID, UserID: user.ID})
return user.ID, budget.EffectiveGroup{GroupID: org.ID}, true
},
},
{
// The fallback prefers the default org even over an org joined
// earlier.
name: "FallbackPrefersDefaultOrg",
policy: codersdk.AIBudgetPolicyHighest,
setup: func(t *testing.T, ctx context.Context, db database.Store) (uuid.UUID, budget.EffectiveGroup, bool) {
defaultOrg, err := db.GetDefaultOrganization(ctx)
require.NoError(t, err)
otherOrg := dbgen.Organization(t, db, database.Organization{})
user := dbgen.User(t, db, database.User{})
dbgen.OrganizationMember(t, db, database.OrganizationMember{OrganizationID: otherOrg.ID, UserID: user.ID, CreatedAt: time.Now().Add(-time.Hour)})
dbgen.OrganizationMember(t, db, database.OrganizationMember{OrganizationID: defaultOrg.ID, UserID: user.ID})
return user.ID, budget.EffectiveGroup{GroupID: defaultOrg.ID}, true
},
},
{
// Among non-default orgs, the fallback breaks ties by the earliest
// organization membership.
name: "FallbackTieByEarliestOrgMembership",
policy: codersdk.AIBudgetPolicyHighest,
setup: func(t *testing.T, ctx context.Context, db database.Store) (uuid.UUID, budget.EffectiveGroup, bool) {
user := dbgen.User(t, db, database.User{})
earlyOrg := dbgen.Organization(t, db, database.Organization{})
lateOrg := dbgen.Organization(t, db, database.Organization{})
dbgen.OrganizationMember(t, db, database.OrganizationMember{OrganizationID: earlyOrg.ID, UserID: user.ID, CreatedAt: time.Now().Add(-time.Hour)})
dbgen.OrganizationMember(t, db, database.OrganizationMember{OrganizationID: lateOrg.ID, UserID: user.ID})
return user.ID, budget.EffectiveGroup{GroupID: earlyOrg.ID}, true
},
},
{
// A user with no org membership has no effective group.
name: "NoOrgMembership",
policy: codersdk.AIBudgetPolicyHighest,
setup: func(t *testing.T, ctx context.Context, db database.Store) (uuid.UUID, budget.EffectiveGroup, bool) {
user := dbgen.User(t, db, database.User{})
return user.ID, budget.EffectiveGroup{}, false
},
},
{
// An unsupported policy surfaces the error from ResolveUserAIBudget.
name: "UnsupportedPolicy",
policy: codersdk.AIBudgetPolicy("unsupported"),
setup: func(t *testing.T, ctx context.Context, db database.Store) (uuid.UUID, budget.EffectiveGroup, bool) {
user := dbgen.User(t, db, database.User{})
return user.ID, budget.EffectiveGroup{}, false
},
wantErr: "unsupported AI budget policy",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
t.Parallel()
db, _ := dbtestutil.NewDB(t)
ctx := testutil.Context(t, testutil.WaitLong)
userID, want, wantOK := tt.setup(t, ctx, db)
got, ok, err := budget.ResolveUserEffectiveGroup(ctx, db, userID, tt.policy)
if tt.wantErr != "" {
require.ErrorContains(t, err, tt.wantErr)
return
}
require.NoError(t, err)
require.Equal(t, wantOK, ok)
if !wantOK {
return
}
require.Equal(t, want.GroupID, got.GroupID)
require.Equal(t, want.Limit, got.Limit)
})
}
}