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
+7
View File
@@ -5111,6 +5111,13 @@ func (q *querier) GetUserCount(ctx context.Context, includeSystem bool) (int64,
return q.db.GetUserCount(ctx, includeSystem)
}
func (q *querier) GetUserEveryoneFallbackGroup(ctx context.Context, userID uuid.UUID) (uuid.UUID, error) {
if _, err := q.GetUserByID(ctx, userID); err != nil { // AuthZ check
return uuid.Nil, err
}
return q.db.GetUserEveryoneFallbackGroup(ctx, userID)
}
func (q *querier) GetUserForChatSyntheticAPIKeyByID(ctx context.Context, id uuid.UUID) (database.User, error) {
return fetchWithAction(q.log, q.auth, policy.ActionReadPersonal, q.db.GetUserForChatSyntheticAPIKeyByID)(ctx, id)
}
+8
View File
@@ -7011,6 +7011,14 @@ func (s *MethodTestSuite) TestAIBridge() {
check.Args(user.ID).Asserts(user, policy.ActionRead).Returns(row)
}))
s.Run("GetUserEveryoneFallbackGroup", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
user := testutil.Fake(s.T(), faker, database.User{})
group := testutil.Fake(s.T(), faker, database.Group{})
dbm.EXPECT().GetUserByID(gomock.Any(), user.ID).Return(user, nil).AnyTimes()
dbm.EXPECT().GetUserEveryoneFallbackGroup(gomock.Any(), user.ID).Return(group.ID, nil).AnyTimes()
check.Args(user.ID).Asserts(user, policy.ActionRead).Returns(group.ID)
}))
s.Run("UpsertUserAIBudgetOverride", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
user := testutil.Fake(s.T(), faker, database.User{})
group := testutil.Fake(s.T(), faker, database.Group{})
+8
View File
@@ -3297,6 +3297,14 @@ func (m queryMetricsStore) GetUserCount(ctx context.Context, includeSystem bool)
return r0, r1
}
func (m queryMetricsStore) GetUserEveryoneFallbackGroup(ctx context.Context, userID uuid.UUID) (uuid.UUID, error) {
start := time.Now()
r0, r1 := m.s.GetUserEveryoneFallbackGroup(ctx, userID)
m.queryLatencies.WithLabelValues("GetUserEveryoneFallbackGroup").Observe(time.Since(start).Seconds())
m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "GetUserEveryoneFallbackGroup").Inc()
return r0, r1
}
func (m queryMetricsStore) GetUserForChatSyntheticAPIKeyByID(ctx context.Context, id uuid.UUID) (database.User, error) {
start := time.Now()
r0, r1 := m.s.GetUserForChatSyntheticAPIKeyByID(ctx, id)
+15
View File
@@ -6163,6 +6163,21 @@ func (mr *MockStoreMockRecorder) GetUserCount(ctx, includeSystem any) *gomock.Ca
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetUserCount", reflect.TypeOf((*MockStore)(nil).GetUserCount), ctx, includeSystem)
}
// GetUserEveryoneFallbackGroup mocks base method.
func (m *MockStore) GetUserEveryoneFallbackGroup(ctx context.Context, userID uuid.UUID) (uuid.UUID, error) {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "GetUserEveryoneFallbackGroup", ctx, userID)
ret0, _ := ret[0].(uuid.UUID)
ret1, _ := ret[1].(error)
return ret0, ret1
}
// GetUserEveryoneFallbackGroup indicates an expected call of GetUserEveryoneFallbackGroup.
func (mr *MockStoreMockRecorder) GetUserEveryoneFallbackGroup(ctx, userID any) *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetUserEveryoneFallbackGroup", reflect.TypeOf((*MockStore)(nil).GetUserEveryoneFallbackGroup), ctx, userID)
}
// GetUserForChatSyntheticAPIKeyByID mocks base method.
func (m *MockStore) GetUserForChatSyntheticAPIKeyByID(ctx context.Context, id uuid.UUID) (database.User, error) {
m.ctrl.T.Helper()
+13 -8
View File
@@ -580,9 +580,9 @@ type sqlcQuerier interface {
// Returns each user's AI spend attributed to the queried group, on or after
// period_start until NOW. Only current members of the queried group are
// returned. spend_limit_micros and limit_source are populated only when the
// queried group is the user's effective budget source. The effective_group_id
// is null when the user has no configured budget or when the effective group
// belongs to a different organization than the queried group.
// queried group is the user's effective budget source. The effective group
// falls back to the Everyone group, and effective_group_id is null only when
// that group belongs to a different organization than the queried group.
// The period_start parameter is normalized to its UTC calendar day.
// TODO(AIGOV-527): unify effective group resolution in a single place.
// Spend is aggregated for the queried group, not the user's effective group.
@@ -604,11 +604,11 @@ type sqlcQuerier interface {
GetGroups(ctx context.Context, arg GetGroupsParams) ([]GetGroupsRow, error)
GetHealthSettings(ctx context.Context) (string, error)
// Returns the highest group AI budget across the groups the user belongs to,
// breaking ties by group name ascending. Implements the "highest" budget policy.
// group_members_expanded is a UNION of group_members and organization_members,
// so the implicit "Everyone" group (group_id == organization_id) is included.
// Returns no rows when the user has no budgeted groups; callers should treat
// sql.ErrNoRows as "no group budget".
// breaking ties by the earliest organization membership. Implements the
// "highest" budget policy. group_members_expanded is a UNION of group_members
// and organization_members, so the implicit "Everyone" group
// (group_id == organization_id) is included. Returns no rows when the user has
// no budgeted groups. Callers should treat sql.ErrNoRows as "no group budget".
GetHighestGroupAIBudgetByUser(ctx context.Context, userID uuid.UUID) (GetHighestGroupAIBudgetByUserRow, error)
GetInboxNotificationByID(ctx context.Context, id uuid.UUID) (InboxNotification, error)
// Fetches inbox notifications for a user filtered by templates and targets
@@ -875,6 +875,11 @@ type sqlcQuerier interface {
GetUserChatSpendInPeriod(ctx context.Context, arg GetUserChatSpendInPeriodParams) (int64, error)
GetUserCodeDiffDisplayMode(ctx context.Context, userID uuid.UUID) (string, error)
GetUserCount(ctx context.Context, includeSystem bool) (int64, error)
// Returns the "Everyone" group (id == organization_id) to attribute a user's
// spend to when no override or budgeted group applies. Prefers the default org,
// then the earliest organization membership. Returns no rows when the user has
// no organization membership.
GetUserEveryoneFallbackGroup(ctx context.Context, userID uuid.UUID) (uuid.UUID, error)
GetUserForChatSyntheticAPIKeyByID(ctx context.Context, id uuid.UUID) (User, error)
// Returns the minimum (most restrictive) group limit for a user.
// Returns -1 if no group limits match the specified scope.
+309 -4
View File
@@ -13359,8 +13359,8 @@ func TestGetGroupMembersAISpend(t *testing.T) {
user := dbgen.User(t, db, database.User{})
org := dbgen.Organization(t, db, database.Organization{})
queried := dbgen.Group(t, db, database.Group{OrganizationID: org.ID})
groupA := dbgen.Group(t, db, database.Group{OrganizationID: org.ID, Name: "aaa-tie-group"})
groupB := dbgen.Group(t, db, database.Group{OrganizationID: org.ID, Name: "bbb-tie-group"})
groupA := dbgen.Group(t, db, database.Group{OrganizationID: org.ID})
groupB := dbgen.Group(t, db, database.Group{OrganizationID: org.ID})
dbgen.OrganizationMember(t, db, database.OrganizationMember{UserID: user.ID, OrganizationID: org.ID})
dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: queried.ID, UserID: user.ID})
dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: groupA.ID, UserID: user.ID})
@@ -13376,6 +13376,14 @@ func TestGetGroupMembersAISpend(t *testing.T) {
})
require.NoError(t, err)
// Both groups are in the same org, so both resolve to the same
// organization membership and the tie falls to the lowest group ID.
winner := groupA.ID
// Postgres orders the uuid type by its bytes.
if bytes.Compare(groupB.ID[:], groupA.ID[:]) < 0 {
winner = groupB.ID
}
// When: querying spend for the user.
got, err := db.GetGroupMembersAISpend(ctx, database.GetGroupMembersAISpendParams{
GroupID: queried.ID,
@@ -13384,9 +13392,9 @@ func TestGetGroupMembersAISpend(t *testing.T) {
})
require.NoError(t, err)
// Then: the tie is broken by group name ascending, so groupA wins.
// Then: the tie falls to the lowest group ID.
require.Len(t, got, 1)
require.Equal(t, uuid.NullUUID{UUID: groupA.ID, Valid: true}, got[0].EffectiveGroupID)
require.Equal(t, uuid.NullUUID{UUID: winner, Valid: true}, got[0].EffectiveGroupID)
require.False(t, got[0].SpendLimitMicros.Valid)
require.False(t, got[0].LimitSource.Valid)
require.Equal(t, int64(0), got[0].GroupSpendMicros)
@@ -13431,6 +13439,89 @@ func TestGetGroupMembersAISpend(t *testing.T) {
require.Equal(t, int64(0), got[0].GroupSpendMicros)
})
t.Run("FallbackToEveryoneGroup", func(t *testing.T) {
t.Parallel()
db, _ := dbtestutil.NewDB(t)
ctx := testutil.Context(t, testutil.WaitShort)
// Given: an unbudgeted member of the queried group whose org has an
// Everyone group but no override or budgeted group.
user := dbgen.User(t, db, database.User{})
org := dbgen.Organization(t, db, database.Organization{})
queried := dbgen.Group(t, db, database.Group{OrganizationID: org.ID})
dbgen.OrganizationMember(t, db, database.OrganizationMember{UserID: user.ID, OrganizationID: org.ID})
dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: queried.ID, UserID: user.ID})
// The Everyone group (id == org id) must exist for the effective group
// join to resolve the fallback.
//nolint:gocritic // Requires system context.
_, err := db.InsertAllUsersGroup(dbauthz.AsSystemRestricted(ctx), org.ID)
require.NoError(t, err)
_, err = db.IncrementUserAIDailySpend(ctx, database.IncrementUserAIDailySpendParams{
UserID: user.ID, EffectiveGroupID: queried.ID, Day: now, CostMicros: 250,
})
require.NoError(t, err)
// When: querying spend for the user.
got, err := db.GetGroupMembersAISpend(ctx, database.GetGroupMembersAISpendParams{
GroupID: queried.ID,
UserIds: []uuid.UUID{user.ID},
PeriodStart: monthStart,
})
require.NoError(t, err)
// Then: with no budget, the effective group falls back to the Everyone
// group. The limit and source are null, and queried-group spend is returned.
require.Len(t, got, 1)
require.Equal(t, uuid.NullUUID{UUID: org.ID, Valid: true}, got[0].EffectiveGroupID)
require.False(t, got[0].SpendLimitMicros.Valid)
require.False(t, got[0].LimitSource.Valid)
require.Equal(t, int64(250), got[0].GroupSpendMicros)
})
t.Run("CrossOrgFallbackMasked", func(t *testing.T) {
t.Parallel()
db, _ := dbtestutil.NewDB(t)
ctx := testutil.Context(t, testutil.WaitShort)
// Given: an unbudgeted member of the queried group who joined another
// org earlier. The fallback picks the earlier org's Everyone group.
user := dbgen.User(t, db, database.User{})
queriedOrg := dbgen.Organization(t, db, database.Organization{})
otherOrg := dbgen.Organization(t, db, database.Organization{})
queried := dbgen.Group(t, db, database.Group{OrganizationID: queriedOrg.ID})
dbgen.OrganizationMember(t, db, database.OrganizationMember{UserID: user.ID, OrganizationID: otherOrg.ID, CreatedAt: dbtime.Now().Add(-time.Hour)})
dbgen.OrganizationMember(t, db, database.OrganizationMember{UserID: user.ID, OrganizationID: queriedOrg.ID})
dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: queried.ID, UserID: user.ID})
// Both orgs have an Everyone group (id == org id), as in production.
//nolint:gocritic // Requires system context.
_, err := db.InsertAllUsersGroup(dbauthz.AsSystemRestricted(ctx), queriedOrg.ID)
require.NoError(t, err)
//nolint:gocritic // Requires system context.
_, err = db.InsertAllUsersGroup(dbauthz.AsSystemRestricted(ctx), otherOrg.ID)
require.NoError(t, err)
_, err = db.IncrementUserAIDailySpend(ctx, database.IncrementUserAIDailySpendParams{
UserID: user.ID, EffectiveGroupID: queried.ID, Day: now, CostMicros: 250,
})
require.NoError(t, err)
// When: querying spend for the user.
got, err := db.GetGroupMembersAISpend(ctx, database.GetGroupMembersAISpendParams{
GroupID: queried.ID,
UserIds: []uuid.UUID{user.ID},
PeriodStart: monthStart,
})
require.NoError(t, err)
// Then: the fallback resolves to the other org's Everyone group, so
// effective_group_id is masked to null, while queried-group spend still
// returns.
require.Len(t, got, 1)
require.False(t, got[0].EffectiveGroupID.Valid, "cross-org effective group must be masked")
require.False(t, got[0].SpendLimitMicros.Valid)
require.False(t, got[0].LimitSource.Valid)
require.Equal(t, int64(250), got[0].GroupSpendMicros)
})
t.Run("SpendWithDifferentEffectiveGroup", func(t *testing.T) {
t.Parallel()
db, _ := dbtestutil.NewDB(t)
@@ -13686,6 +13777,220 @@ func TestGetGroupMembersAISpend(t *testing.T) {
})
}
func TestGetHighestGroupAIBudgetByUser(t *testing.T) {
t.Parallel()
tests := []struct {
name string
setup func(t *testing.T, ctx context.Context, db database.Store) (userID uuid.UUID, want database.GetHighestGroupAIBudgetByUserRow)
wantErr error
}{
{
// Among the user's budgeted groups, the highest limit wins.
name: "HighestWins",
setup: func(t *testing.T, ctx context.Context, db database.Store) (uuid.UUID, database.GetHighestGroupAIBudgetByUserRow) {
user := dbgen.User(t, db, database.User{})
org := dbgen.Organization(t, db, database.Organization{})
lower := dbgen.Group(t, db, database.Group{OrganizationID: org.ID})
higher := dbgen.Group(t, db, database.Group{OrganizationID: org.ID})
dbgen.OrganizationMember(t, db, database.OrganizationMember{OrganizationID: org.ID, UserID: user.ID})
dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: lower.ID, UserID: user.ID})
dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: higher.ID, UserID: user.ID})
_, err := db.UpsertGroupAIBudget(ctx, database.UpsertGroupAIBudgetParams{GroupID: lower.ID, SpendLimitMicros: 1_000_000})
require.NoError(t, err)
_, err = db.UpsertGroupAIBudget(ctx, database.UpsertGroupAIBudgetParams{GroupID: higher.ID, SpendLimitMicros: 2_000_000})
require.NoError(t, err)
return user.ID, database.GetHighestGroupAIBudgetByUserRow{GroupID: higher.ID, SpendLimitMicros: 2_000_000}
},
},
{
// The highest limit wins across the user's orgs, not just within one.
name: "HighestWinsAcrossOrgs",
setup: func(t *testing.T, ctx context.Context, db database.Store) (uuid.UUID, database.GetHighestGroupAIBudgetByUserRow) {
user := dbgen.User(t, db, database.User{})
orgA := dbgen.Organization(t, db, database.Organization{})
orgB := dbgen.Organization(t, db, database.Organization{})
lower := dbgen.Group(t, db, database.Group{OrganizationID: orgA.ID})
higher := dbgen.Group(t, db, database.Group{OrganizationID: orgB.ID})
dbgen.OrganizationMember(t, db, database.OrganizationMember{OrganizationID: orgA.ID, UserID: user.ID})
dbgen.OrganizationMember(t, db, database.OrganizationMember{OrganizationID: orgB.ID, UserID: user.ID})
dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: lower.ID, UserID: user.ID})
dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: higher.ID, UserID: user.ID})
_, err := db.UpsertGroupAIBudget(ctx, database.UpsertGroupAIBudgetParams{GroupID: lower.ID, SpendLimitMicros: 1_000_000})
require.NoError(t, err)
_, err = db.UpsertGroupAIBudget(ctx, database.UpsertGroupAIBudgetParams{GroupID: higher.ID, SpendLimitMicros: 2_000_000})
require.NoError(t, err)
return user.ID, database.GetHighestGroupAIBudgetByUserRow{GroupID: higher.ID, SpendLimitMicros: 2_000_000}
},
},
{
// A budgeted group in a soft-deleted org is excluded even when its
// limit is higher.
name: "ExcludesDeletedOrg",
setup: func(t *testing.T, ctx context.Context, db database.Store) (uuid.UUID, database.GetHighestGroupAIBudgetByUserRow) {
user := dbgen.User(t, db, database.User{})
liveOrg := dbgen.Organization(t, db, database.Organization{Name: "live-org"})
deletedOrg := dbgen.Organization(t, db, database.Organization{Name: "deleted-org"})
liveGroup := dbgen.Group(t, db, database.Group{OrganizationID: liveOrg.ID})
deletedGroup := dbgen.Group(t, db, database.Group{OrganizationID: deletedOrg.ID})
dbgen.OrganizationMember(t, db, database.OrganizationMember{OrganizationID: liveOrg.ID, UserID: user.ID})
dbgen.OrganizationMember(t, db, database.OrganizationMember{OrganizationID: deletedOrg.ID, UserID: user.ID})
dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: liveGroup.ID, UserID: user.ID})
dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: deletedGroup.ID, UserID: user.ID})
_, err := db.UpsertGroupAIBudget(ctx, database.UpsertGroupAIBudgetParams{GroupID: liveGroup.ID, SpendLimitMicros: 1_000_000})
require.NoError(t, err)
_, err = db.UpsertGroupAIBudget(ctx, database.UpsertGroupAIBudgetParams{GroupID: deletedGroup.ID, SpendLimitMicros: 5_000_000})
require.NoError(t, err)
err = db.UpdateOrganizationDeletedByID(ctx, database.UpdateOrganizationDeletedByIDParams{
ID: deletedOrg.ID,
UpdatedAt: dbtime.Now(),
})
require.NoError(t, err)
return user.ID, database.GetHighestGroupAIBudgetByUserRow{GroupID: liveGroup.ID, SpendLimitMicros: 1_000_000}
},
},
{
// Equal limits across orgs break by the earliest organization
// membership.
name: "TieByEarliestOrgMembership",
setup: func(t *testing.T, ctx context.Context, db database.Store) (uuid.UUID, database.GetHighestGroupAIBudgetByUserRow) {
user := dbgen.User(t, db, database.User{})
earlyOrg := dbgen.Organization(t, db, database.Organization{})
lateOrg := dbgen.Organization(t, db, database.Organization{})
earlyGroup := dbgen.Group(t, db, database.Group{OrganizationID: earlyOrg.ID})
lateGroup := dbgen.Group(t, db, database.Group{OrganizationID: lateOrg.ID})
dbgen.OrganizationMember(t, db, database.OrganizationMember{OrganizationID: earlyOrg.ID, UserID: user.ID, CreatedAt: dbtime.Now().Add(-time.Hour)})
dbgen.OrganizationMember(t, db, database.OrganizationMember{OrganizationID: lateOrg.ID, UserID: user.ID})
dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: earlyGroup.ID, UserID: user.ID})
dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: lateGroup.ID, UserID: user.ID})
_, err := db.UpsertGroupAIBudget(ctx, database.UpsertGroupAIBudgetParams{GroupID: earlyGroup.ID, SpendLimitMicros: 1_000_000})
require.NoError(t, err)
_, err = db.UpsertGroupAIBudget(ctx, database.UpsertGroupAIBudgetParams{GroupID: lateGroup.ID, SpendLimitMicros: 1_000_000})
require.NoError(t, err)
return user.ID, database.GetHighestGroupAIBudgetByUserRow{GroupID: earlyGroup.ID, SpendLimitMicros: 1_000_000}
},
},
{
// A user with no budgeted group has no highest budget.
name: "NoBudgetedGroup",
wantErr: sql.ErrNoRows,
setup: func(t *testing.T, ctx context.Context, db database.Store) (uuid.UUID, database.GetHighestGroupAIBudgetByUserRow) {
user := dbgen.User(t, db, database.User{})
org := dbgen.Organization(t, db, database.Organization{})
dbgen.OrganizationMember(t, db, database.OrganizationMember{OrganizationID: org.ID, UserID: user.ID})
return user.ID, database.GetHighestGroupAIBudgetByUserRow{}
},
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
t.Parallel()
db, _ := dbtestutil.NewDB(t)
ctx := testutil.Context(t, testutil.WaitShort)
userID, want := tt.setup(t, ctx, db)
got, err := db.GetHighestGroupAIBudgetByUser(ctx, userID)
if tt.wantErr != nil {
require.ErrorIs(t, err, tt.wantErr)
return
}
require.NoError(t, err)
require.Equal(t, want, got)
})
}
}
func TestGetUserEveryoneFallbackGroup(t *testing.T) {
t.Parallel()
tests := []struct {
name string
setup func(t *testing.T, ctx context.Context, db database.Store) (userID uuid.UUID, wantGroupID uuid.UUID)
wantErr error
}{
{
// A single-org member falls back to that org's Everyone group.
name: "SingleOrg",
setup: func(t *testing.T, ctx context.Context, db database.Store) (uuid.UUID, uuid.UUID) {
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, org.ID
},
},
{
// The default org is preferred even over an org joined earlier.
name: "PrefersDefaultOrg",
setup: func(t *testing.T, ctx context.Context, db database.Store) (uuid.UUID, uuid.UUID) {
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: dbtime.Now().Add(-time.Hour)})
dbgen.OrganizationMember(t, db, database.OrganizationMember{OrganizationID: defaultOrg.ID, UserID: user.ID})
return user.ID, defaultOrg.ID
},
},
{
// Among non-default orgs, ties break by the earliest organization
// membership.
name: "TieByEarliestOrgMembership",
setup: func(t *testing.T, ctx context.Context, db database.Store) (uuid.UUID, uuid.UUID) {
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: dbtime.Now().Add(-time.Hour)})
dbgen.OrganizationMember(t, db, database.OrganizationMember{OrganizationID: lateOrg.ID, UserID: user.ID})
return user.ID, earlyOrg.ID
},
},
{
// A soft-deleted org is excluded even when it was joined earlier.
name: "ExcludesDeletedOrg",
setup: func(t *testing.T, ctx context.Context, db database.Store) (uuid.UUID, uuid.UUID) {
user := dbgen.User(t, db, database.User{})
liveOrg := dbgen.Organization(t, db, database.Organization{Name: "live-org"})
deletedOrg := dbgen.Organization(t, db, database.Organization{Name: "deleted-org"})
dbgen.OrganizationMember(t, db, database.OrganizationMember{OrganizationID: liveOrg.ID, UserID: user.ID})
dbgen.OrganizationMember(t, db, database.OrganizationMember{OrganizationID: deletedOrg.ID, UserID: user.ID, CreatedAt: dbtime.Now().Add(-time.Hour)})
err := db.UpdateOrganizationDeletedByID(ctx, database.UpdateOrganizationDeletedByIDParams{
ID: deletedOrg.ID,
UpdatedAt: dbtime.Now(),
})
require.NoError(t, err)
return user.ID, liveOrg.ID
},
},
{
// A user with no org membership has no fallback group.
name: "NoOrgMembership",
wantErr: sql.ErrNoRows,
setup: func(t *testing.T, ctx context.Context, db database.Store) (uuid.UUID, uuid.UUID) {
user := dbgen.User(t, db, database.User{})
return user.ID, uuid.Nil
},
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
t.Parallel()
db, _ := dbtestutil.NewDB(t)
ctx := testutil.Context(t, testutil.WaitShort)
userID, wantGroupID := tt.setup(t, ctx, db)
got, err := db.GetUserEveryoneFallbackGroup(ctx, userID)
if tt.wantErr != nil {
require.ErrorIs(t, err, tt.wantErr)
return
}
require.NoError(t, err)
require.Equal(t, wantGroupID, got)
})
}
}
func TestChatPinOrderQueries(t *testing.T) {
t.Parallel()
if testing.Short() {
+68 -23
View File
@@ -2594,15 +2594,33 @@ user_highest_group AS (
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 filtered_users)
ORDER BY member.user_id, budget.spend_limit_micros DESC, member.group_name ASC, budget.group_id ASC
AND organizations.deleted = false
ORDER BY member.user_id, budget.spend_limit_micros DESC, organization_members.created_at ASC, budget.group_id ASC
),
user_fallback_group AS (
-- Per user, the Everyone group to fall back to when no override or budgeted
-- group applies. The Everyone group has id == organization_id. Prefers the
-- default org, then the earliest organization membership.
SELECT DISTINCT ON (organization_members.user_id)
organization_members.user_id,
organizations.id AS group_id
FROM organization_members
JOIN organizations ON organizations.id = organization_members.organization_id
WHERE organization_members.user_id IN (SELECT user_id FROM filtered_users)
AND organizations.deleted = false
ORDER BY organization_members.user_id, organizations.is_default DESC, organization_members.created_at ASC, organizations.id ASC
),
effective AS (
-- Effective budget per user: a per-user override wins over the
-- highest-limit group.
-- Effective budget per user: a per-user override wins over the highest-limit
-- group, which wins over the Everyone group fallback.
SELECT
filtered_users.user_id,
COALESCE(override.group_id, user_highest_group.group_id) AS raw_effective_group_id,
COALESCE(override.group_id, user_highest_group.group_id, user_fallback_group.group_id) AS raw_effective_group_id,
COALESCE(override.spend_limit_micros, user_highest_group.spend_limit_micros) AS spend_limit_micros,
(CASE
WHEN override.group_id IS NOT NULL THEN 'user_override'
@@ -2611,6 +2629,7 @@ effective AS (
FROM filtered_users
LEFT JOIN user_ai_budget_overrides override ON override.user_id = filtered_users.user_id
LEFT JOIN user_highest_group ON user_highest_group.user_id = filtered_users.user_id
LEFT JOIN user_fallback_group ON user_fallback_group.user_id = filtered_users.user_id
),
applied_budget AS (
-- The limit and source only for users whose effective budget source is the
@@ -2663,9 +2682,9 @@ type GetGroupMembersAISpendRow struct {
// Returns each user's AI spend attributed to the queried group, on or after
// period_start until NOW. Only current members of the queried group are
// returned. spend_limit_micros and limit_source are populated only when the
// queried group is the user's effective budget source. The effective_group_id
// is null when the user has no configured budget or when the effective group
// belongs to a different organization than the queried group.
// queried group is the user's effective budget source. The effective group
// falls back to the Everyone group, and effective_group_id is null only when
// that group belongs to a different organization than the queried group.
// The period_start parameter is normalized to its UTC calendar day.
// TODO(AIGOV-527): unify effective group resolution in a single place.
// Spend is aggregated for the queried group, not the user's effective group.
@@ -2703,18 +2722,20 @@ func (q *sqlQuerier) GetGroupMembersAISpend(ctx context.Context, arg GetGroupMem
const getHighestGroupAIBudgetByUser = `-- name: GetHighestGroupAIBudgetByUser :one
SELECT
gaib.group_id,
gaib.spend_limit_micros
FROM group_ai_budgets gaib
JOIN group_members_expanded gme ON gme.group_id = gaib.group_id
WHERE gme.user_id = $1
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 = $1
AND organizations.deleted = false
ORDER BY
gaib.spend_limit_micros DESC, -- highest wins
gme.group_name ASC, -- alphabetical tiebreak
-- Final tiebreak on the group id makes the result deterministic when two
-- groups share both name and limit, which is possible across organizations
-- (groups are unique on (organization_id, name), not name alone).
gaib.group_id ASC
budget.spend_limit_micros DESC, -- highest wins
organization_members.created_at ASC, -- earliest organization membership
budget.group_id ASC -- deterministic tiebreak
LIMIT 1
`
@@ -2724,11 +2745,11 @@ type GetHighestGroupAIBudgetByUserRow struct {
}
// Returns the highest group AI budget across the groups the user belongs to,
// breaking ties by group name ascending. Implements the "highest" budget policy.
// group_members_expanded is a UNION of group_members and organization_members,
// so the implicit "Everyone" group (group_id == organization_id) is included.
// Returns no rows when the user has no budgeted groups; callers should treat
// sql.ErrNoRows as "no group budget".
// breaking ties by the earliest organization membership. Implements the
// "highest" budget policy. group_members_expanded is a UNION of group_members
// and organization_members, so the implicit "Everyone" group
// (group_id == organization_id) is included. Returns no rows when the user has
// no budgeted groups. Callers should treat sql.ErrNoRows as "no group budget".
func (q *sqlQuerier) GetHighestGroupAIBudgetByUser(ctx context.Context, userID uuid.UUID) (GetHighestGroupAIBudgetByUserRow, error) {
row := q.db.QueryRowContext(ctx, getHighestGroupAIBudgetByUser, userID)
var i GetHighestGroupAIBudgetByUserRow
@@ -2856,6 +2877,30 @@ func (q *sqlQuerier) GetUserAISpendSince(ctx context.Context, arg GetUserAISpend
return i, err
}
const getUserEveryoneFallbackGroup = `-- name: GetUserEveryoneFallbackGroup :one
SELECT organizations.id AS group_id
FROM organization_members
JOIN organizations ON organizations.id = organization_members.organization_id
WHERE organization_members.user_id = $1
AND organizations.deleted = false
ORDER BY
organizations.is_default DESC, -- prefer the default org
organization_members.created_at ASC, -- earliest organization membership
organizations.id ASC -- deterministic tiebreak
LIMIT 1
`
// Returns the "Everyone" group (id == organization_id) to attribute a user's
// spend to when no override or budgeted group applies. Prefers the default org,
// then the earliest organization membership. Returns no rows when the user has
// no organization membership.
func (q *sqlQuerier) GetUserEveryoneFallbackGroup(ctx context.Context, userID uuid.UUID) (uuid.UUID, error) {
row := q.db.QueryRowContext(ctx, getUserEveryoneFallbackGroup, userID)
var group_id uuid.UUID
err := row.Scan(&group_id)
return group_id, err
}
const incrementUserAIDailySpend = `-- name: IncrementUserAIDailySpend :one
INSERT INTO ai_user_daily_spend (user_id, effective_group_id, day, spend_micros)
VALUES ($1, $2, (($3::timestamptz) AT TIME ZONE 'UTC')::date, $4)
+60 -23
View File
@@ -60,24 +60,42 @@ DELETE FROM user_ai_budget_overrides WHERE user_id = @user_id RETURNING *;
-- name: GetHighestGroupAIBudgetByUser :one
-- Returns the highest group AI budget across the groups the user belongs to,
-- breaking ties by group name ascending. Implements the "highest" budget policy.
-- group_members_expanded is a UNION of group_members and organization_members,
-- so the implicit "Everyone" group (group_id == organization_id) is included.
-- Returns no rows when the user has no budgeted groups; callers should treat
-- sql.ErrNoRows as "no group budget".
-- breaking ties by the earliest organization membership. Implements the
-- "highest" budget policy. group_members_expanded is a UNION of group_members
-- and organization_members, so the implicit "Everyone" group
-- (group_id == organization_id) is included. Returns no rows when the user has
-- no budgeted groups. Callers should treat sql.ErrNoRows as "no group budget".
SELECT
gaib.group_id,
gaib.spend_limit_micros
FROM group_ai_budgets gaib
JOIN group_members_expanded gme ON gme.group_id = gaib.group_id
WHERE gme.user_id = @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 = @user_id
AND organizations.deleted = false
ORDER BY
gaib.spend_limit_micros DESC, -- highest wins
gme.group_name ASC, -- alphabetical tiebreak
-- Final tiebreak on the group id makes the result deterministic when two
-- groups share both name and limit, which is possible across organizations
-- (groups are unique on (organization_id, name), not name alone).
gaib.group_id ASC
budget.spend_limit_micros DESC, -- highest wins
organization_members.created_at ASC, -- earliest organization membership
budget.group_id ASC -- deterministic tiebreak
LIMIT 1;
-- name: GetUserEveryoneFallbackGroup :one
-- Returns the "Everyone" group (id == organization_id) to attribute a user's
-- spend to when no override or budgeted group applies. Prefers the default org,
-- then the earliest organization membership. Returns no rows when the user has
-- no organization membership.
SELECT organizations.id AS group_id
FROM organization_members
JOIN organizations ON organizations.id = organization_members.organization_id
WHERE organization_members.user_id = @user_id
AND organizations.deleted = false
ORDER BY
organizations.is_default DESC, -- prefer the default org
organization_members.created_at ASC, -- earliest organization membership
organizations.id ASC -- deterministic tiebreak
LIMIT 1;
-- name: IncrementUserAIDailySpend :one
@@ -126,9 +144,9 @@ ORDER BY groups.id;
-- Returns each user's AI spend attributed to the queried group, on or after
-- period_start until NOW. Only current members of the queried group are
-- returned. spend_limit_micros and limit_source are populated only when the
-- queried group is the user's effective budget source. The effective_group_id
-- is null when the user has no configured budget or when the effective group
-- belongs to a different organization than the queried group.
-- queried group is the user's effective budget source. The effective group
-- falls back to the Everyone group, and effective_group_id is null only when
-- that group belongs to a different organization than the queried group.
-- The period_start parameter is normalized to its UTC calendar day.
-- TODO(AIGOV-527): unify effective group resolution in a single place.
WITH queried_group AS (
@@ -154,15 +172,33 @@ user_highest_group AS (
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 filtered_users)
ORDER BY member.user_id, budget.spend_limit_micros DESC, member.group_name ASC, budget.group_id ASC
AND organizations.deleted = false
ORDER BY member.user_id, budget.spend_limit_micros DESC, organization_members.created_at ASC, budget.group_id ASC
),
user_fallback_group AS (
-- Per user, the Everyone group to fall back to when no override or budgeted
-- group applies. The Everyone group has id == organization_id. Prefers the
-- default org, then the earliest organization membership.
SELECT DISTINCT ON (organization_members.user_id)
organization_members.user_id,
organizations.id AS group_id
FROM organization_members
JOIN organizations ON organizations.id = organization_members.organization_id
WHERE organization_members.user_id IN (SELECT user_id FROM filtered_users)
AND organizations.deleted = false
ORDER BY organization_members.user_id, organizations.is_default DESC, organization_members.created_at ASC, organizations.id ASC
),
effective AS (
-- Effective budget per user: a per-user override wins over the
-- highest-limit group.
-- Effective budget per user: a per-user override wins over the highest-limit
-- group, which wins over the Everyone group fallback.
SELECT
filtered_users.user_id,
COALESCE(override.group_id, user_highest_group.group_id) AS raw_effective_group_id,
COALESCE(override.group_id, user_highest_group.group_id, user_fallback_group.group_id) AS raw_effective_group_id,
COALESCE(override.spend_limit_micros, user_highest_group.spend_limit_micros) AS spend_limit_micros,
(CASE
WHEN override.group_id IS NOT NULL THEN 'user_override'
@@ -171,6 +207,7 @@ effective AS (
FROM filtered_users
LEFT JOIN user_ai_budget_overrides override ON override.user_id = filtered_users.user_id
LEFT JOIN user_highest_group ON user_highest_group.user_id = filtered_users.user_id
LEFT JOIN user_fallback_group ON user_fallback_group.user_id = filtered_users.user_id
),
applied_budget AS (
-- The limit and source only for users whose effective budget source is the