mirror of
https://github.com/coder/coder.git
synced 2026-09-24 15:04:27 +08:00
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:
@@ -911,7 +911,7 @@ func (api *API) userAISpendStatus(rw http.ResponseWriter, r *http.Request) {
|
||||
)
|
||||
|
||||
policy := codersdk.NewAIBudgetPolicyFromString(api.DeploymentValues.AI.BridgeConfig.BudgetPolicy)
|
||||
effectiveBudget, ok, err := budget.ResolveUserAIBudget(ctx, api.Database, user.ID, policy)
|
||||
effectiveGroup, ok, err := budget.ResolveUserEffectiveGroup(ctx, api.Database, user.ID, policy)
|
||||
if err != nil {
|
||||
logger.Error(ctx, "failed to resolve user AI budget", slog.Error(err))
|
||||
httpapi.InternalServerError(rw, err)
|
||||
@@ -929,14 +929,16 @@ func (api *API) userAISpendStatus(rw http.ResponseWriter, r *http.Request) {
|
||||
}
|
||||
|
||||
if ok {
|
||||
resp.EffectiveGroupID = &effectiveBudget.GroupID
|
||||
resp.SpendLimitMicros = &effectiveBudget.SpendLimitMicros
|
||||
resp.LimitSource = &effectiveBudget.Source
|
||||
logger = logger.With(slog.F("effective_group_id", effectiveBudget.GroupID))
|
||||
resp.EffectiveGroupID = &effectiveGroup.GroupID
|
||||
if effectiveGroup.Limit != nil {
|
||||
resp.SpendLimitMicros = &effectiveGroup.Limit.SpendLimitMicros
|
||||
resp.LimitSource = &effectiveGroup.Limit.Source
|
||||
}
|
||||
logger = logger.With(slog.F("effective_group_id", effectiveGroup.GroupID))
|
||||
|
||||
spend, err := api.Database.GetUserAISpendSince(ctx, database.GetUserAISpendSinceParams{
|
||||
UserID: user.ID,
|
||||
EffectiveGroupID: effectiveBudget.GroupID,
|
||||
EffectiveGroupID: effectiveGroup.GroupID,
|
||||
PeriodStart: periodWindow.Start,
|
||||
})
|
||||
if err != nil {
|
||||
|
||||
@@ -3159,13 +3159,6 @@ func TestUserAISpendStatus(t *testing.T) {
|
||||
wantLimitSource *codersdk.AIBudgetLimitSource
|
||||
wantCurrentSpendMicros int64
|
||||
}{
|
||||
{
|
||||
name: "NoEffectiveGroup",
|
||||
wantHasEffectiveGroup: false,
|
||||
wantSpendLimitMicros: nil,
|
||||
wantLimitSource: nil,
|
||||
wantCurrentSpendMicros: 0,
|
||||
},
|
||||
{
|
||||
name: "GroupBudget/ZeroSpend",
|
||||
groupBudget: ptr.Ref(int64(1_000_000_000)),
|
||||
@@ -3279,6 +3272,65 @@ func TestUserAISpendStatus(t *testing.T) {
|
||||
require.Equal(t, tt.wantLimitSource, got.LimitSource)
|
||||
})
|
||||
}
|
||||
|
||||
t.Run("UnbudgetedFallsBackToEveryone", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
clock := quartz.NewMock(t)
|
||||
db, ps := dbtestutil.NewDB(t)
|
||||
adminClient, targetUser, group := setupAICostControlTest(t, aiCostControlTestOptions{
|
||||
GroupName: "spend-test-group",
|
||||
Clock: clock,
|
||||
Database: db,
|
||||
Pubsub: ps,
|
||||
})
|
||||
ctx := testutil.Context(t, testutil.WaitLong)
|
||||
clock.Set(time.Date(2026, time.March, 15, 12, 0, 0, 0, time.UTC))
|
||||
|
||||
// With no override or group budget, the effective group is the org's
|
||||
// Everyone group (id == org id) with no limit. The reported current
|
||||
// spend is the amount attributed to that Everyone group.
|
||||
everyoneGroupID := group.OrganizationID
|
||||
_, err := db.IncrementUserAIDailySpend(ctx, database.IncrementUserAIDailySpendParams{
|
||||
UserID: targetUser.ID,
|
||||
EffectiveGroupID: everyoneGroupID,
|
||||
Day: clock.Now(),
|
||||
CostMicros: 100_000_000,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
got, err := adminClient.UserAISpendStatus(ctx, targetUser.ID)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, &everyoneGroupID, got.EffectiveGroupID)
|
||||
require.Nil(t, got.SpendLimitMicros)
|
||||
require.Nil(t, got.LimitSource)
|
||||
require.Equal(t, int64(100_000_000), got.CurrentSpendMicros)
|
||||
})
|
||||
|
||||
t.Run("NoOrgReturnsNull", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
clock := quartz.NewMock(t)
|
||||
db, ps := dbtestutil.NewDB(t)
|
||||
adminClient, _, _ := setupAICostControlTest(t, aiCostControlTestOptions{
|
||||
GroupName: "spend-test-group",
|
||||
Clock: clock,
|
||||
Database: db,
|
||||
Pubsub: ps,
|
||||
})
|
||||
ctx := testutil.Context(t, testutil.WaitLong)
|
||||
clock.Set(time.Date(2026, time.March, 15, 12, 0, 0, 0, time.UTC))
|
||||
|
||||
// A user with no organization membership resolves to no effective group.
|
||||
orglessUser := dbgen.User(t, db, database.User{})
|
||||
|
||||
got, err := adminClient.UserAISpendStatus(ctx, orglessUser.ID)
|
||||
require.NoError(t, err)
|
||||
require.Nil(t, got.EffectiveGroupID)
|
||||
require.Nil(t, got.SpendLimitMicros)
|
||||
require.Nil(t, got.LimitSource)
|
||||
require.Equal(t, int64(0), got.CurrentSpendMicros)
|
||||
})
|
||||
}
|
||||
|
||||
func TestUserAISpendStatusRoleAccess(t *testing.T) {
|
||||
@@ -3798,23 +3850,21 @@ func TestGroupMembersAISpend(t *testing.T) {
|
||||
// Then: only the primary-org user is returned.
|
||||
require.Len(t, resp.Members, 1)
|
||||
require.Equal(t, targetUser.ID, resp.Members[0].UserID)
|
||||
require.Nil(t, resp.Members[0].EffectiveGroupID)
|
||||
require.Equal(t, &group.OrganizationID, resp.Members[0].EffectiveGroupID)
|
||||
require.Nil(t, resp.Members[0].GroupBudget)
|
||||
require.Equal(t, int64(0), resp.Members[0].GroupSpendMicros)
|
||||
})
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
groupLimit int64
|
||||
overrideLimit int64
|
||||
spent int64
|
||||
wantEffectiveGroup bool
|
||||
wantGroupBudget *codersdk.AIGroupBudget
|
||||
wantSpendMicros int64
|
||||
name string
|
||||
groupLimit int64
|
||||
overrideLimit int64
|
||||
spent int64
|
||||
wantEffectiveGroup bool
|
||||
wantEffectiveEveryone bool
|
||||
wantGroupBudget *codersdk.AIGroupBudget
|
||||
wantSpendMicros int64
|
||||
}{
|
||||
{
|
||||
name: "NoBudgetNoSpend",
|
||||
},
|
||||
{
|
||||
name: "BudgetZeroSpend",
|
||||
groupLimit: 1_000_000_000,
|
||||
@@ -3835,11 +3885,6 @@ func TestGroupMembersAISpend(t *testing.T) {
|
||||
},
|
||||
wantSpendMicros: 250_000_000,
|
||||
},
|
||||
{
|
||||
name: "NoBudgetWithSpend",
|
||||
spent: 100_000_000,
|
||||
wantSpendMicros: 100_000_000,
|
||||
},
|
||||
{
|
||||
name: "OverrideBudget",
|
||||
overrideLimit: 500_000_000,
|
||||
@@ -3849,6 +3894,19 @@ func TestGroupMembersAISpend(t *testing.T) {
|
||||
LimitSource: codersdk.AIBudgetLimitSourceUserOverride,
|
||||
},
|
||||
},
|
||||
{
|
||||
// With no budget, an in-org member falls back to the Everyone group.
|
||||
name: "FallbackToEveryoneNoSpend",
|
||||
wantEffectiveEveryone: true,
|
||||
},
|
||||
{
|
||||
// The fallback effective group is the Everyone group, while spend is
|
||||
// still attributed to the queried group.
|
||||
name: "FallbackToEveryoneWithSpend",
|
||||
spent: 100_000_000,
|
||||
wantEffectiveEveryone: true,
|
||||
wantSpendMicros: 100_000_000,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
@@ -3902,10 +3960,14 @@ func TestGroupMembersAISpend(t *testing.T) {
|
||||
require.Equal(t, wantPeriodEnd, got.PeriodEnd)
|
||||
require.Len(t, got.Members, 1)
|
||||
require.Equal(t, targetUser.ID, got.Members[0].UserID)
|
||||
if tt.wantEffectiveGroup {
|
||||
switch {
|
||||
case tt.wantEffectiveGroup:
|
||||
require.NotNil(t, got.Members[0].EffectiveGroupID)
|
||||
require.Equal(t, group.ID, *got.Members[0].EffectiveGroupID)
|
||||
} else {
|
||||
case tt.wantEffectiveEveryone:
|
||||
require.NotNil(t, got.Members[0].EffectiveGroupID)
|
||||
require.Equal(t, group.OrganizationID, *got.Members[0].EffectiveGroupID)
|
||||
default:
|
||||
require.Nil(t, got.Members[0].EffectiveGroupID)
|
||||
}
|
||||
require.Equal(t, tt.wantGroupBudget, got.Members[0].GroupBudget)
|
||||
@@ -3972,6 +4034,65 @@ func TestGroupMembersAISpend(t *testing.T) {
|
||||
require.Equal(t, int64(0), resp.Members[0].GroupSpendMicros)
|
||||
})
|
||||
|
||||
t.Run("CrossOrgFallbackEveryoneMasked", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
dv := coderdtest.DeploymentValues(t)
|
||||
dv.AI.BridgeConfig.Enabled = serpent.Bool(true)
|
||||
dv.Experiments = []string{string(codersdk.ExperimentAIGatewayCostControl)}
|
||||
db, ps := dbtestutil.NewDB(t)
|
||||
ownerClient, owner := coderdenttest.New(t, &coderdenttest.Options{
|
||||
Options: &coderdtest.Options{DeploymentValues: dv, Database: db, Pubsub: ps},
|
||||
LicenseOptions: &coderdenttest.LicenseOptions{
|
||||
Features: license.Features{
|
||||
codersdk.FeatureTemplateRBAC: 1,
|
||||
codersdk.FeatureAIBridge: 1,
|
||||
codersdk.FeatureMultipleOrganizations: 1,
|
||||
},
|
||||
},
|
||||
})
|
||||
userAdminClient, _ := coderdtest.CreateAnotherUser(t, ownerClient, owner.OrganizationID, rbac.RoleUserAdmin())
|
||||
ctx := testutil.Context(t, testutil.WaitLong)
|
||||
|
||||
// Given: a member of the default org whose queried group lives in a
|
||||
// non-default org, with no budget or override. The fallback prefers the
|
||||
// default org's Everyone group, which lives in a different org than the
|
||||
// queried group.
|
||||
queriedOrg := coderdenttest.CreateOrganization(t, ownerClient, coderdenttest.CreateOrganizationOptions{})
|
||||
_, targetUser := coderdtest.CreateAnotherUser(t, ownerClient, owner.OrganizationID)
|
||||
dbgen.OrganizationMember(t, db, database.OrganizationMember{UserID: targetUser.ID, OrganizationID: queriedOrg.ID})
|
||||
queried, err := userAdminClient.CreateGroup(ctx, queriedOrg.ID, codersdk.CreateGroupRequest{
|
||||
Name: "queried-cross-org-fallback-group",
|
||||
})
|
||||
require.NoError(t, err)
|
||||
_, err = userAdminClient.PatchGroup(ctx, queried.ID, codersdk.PatchGroupRequest{
|
||||
AddUsers: []string{targetUser.ID.String()},
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
// Spend is attributed to the queried group.
|
||||
_, err = db.IncrementUserAIDailySpend(ctx, database.IncrementUserAIDailySpendParams{
|
||||
UserID: targetUser.ID,
|
||||
EffectiveGroupID: queried.ID,
|
||||
Day: dbtime.Now(),
|
||||
CostMicros: 100_000_000,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
// When: the owner, who can read both orgs, queries the group.
|
||||
//nolint:gocritic // The test asserts that even an owner sees the mask.
|
||||
resp, err := ownerClient.GroupMembersAISpend(ctx, queried.ID, []uuid.UUID{targetUser.ID})
|
||||
require.NoError(t, err)
|
||||
|
||||
// Then: effective_group_id is masked because the fallback Everyone group
|
||||
// lives in the default org, while the queried group's spend still returns.
|
||||
require.Len(t, resp.Members, 1)
|
||||
require.Equal(t, targetUser.ID, resp.Members[0].UserID)
|
||||
require.Nil(t, resp.Members[0].EffectiveGroupID, "cross-org fallback effective group must be masked")
|
||||
require.Nil(t, resp.Members[0].GroupBudget)
|
||||
require.Equal(t, int64(100_000_000), resp.Members[0].GroupSpendMicros)
|
||||
})
|
||||
|
||||
t.Run("OrgScopedRoute", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
@@ -3997,7 +4118,7 @@ func TestGroupMembersAISpend(t *testing.T) {
|
||||
require.NoError(t, json.NewDecoder(res.Body).Decode(&got))
|
||||
require.Len(t, got.Members, 1)
|
||||
require.Equal(t, targetUser.ID, got.Members[0].UserID)
|
||||
require.Nil(t, got.Members[0].EffectiveGroupID)
|
||||
require.Equal(t, &group.OrganizationID, got.Members[0].EffectiveGroupID)
|
||||
require.Nil(t, got.Members[0].GroupBudget)
|
||||
require.Equal(t, int64(0), got.Members[0].GroupSpendMicros)
|
||||
})
|
||||
|
||||
Reference in New Issue
Block a user