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:
@@ -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)
|
||||
}
|
||||
|
||||
@@ -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
@@ -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)
|
||||
|
||||
Generated
+15
@@ -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()
|
||||
|
||||
Generated
+13
-8
@@ -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.
|
||||
|
||||
@@ -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() {
|
||||
|
||||
Generated
+68
-23
@@ -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,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
|
||||
|
||||
Reference in New Issue
Block a user