From fcca7c451a5128a3a3e2cc877849400a8183fa27 Mon Sep 17 00:00:00 2001 From: Susana Cardoso Ferreira Date: Tue, 21 Jul 2026 09:22:26 +0000 Subject: [PATCH 1/4] feat: fall back to the Everyone group for AI spend attribution --- coderd/aibridge/budget/budget.go | 80 ++++-- coderd/aibridge/budget/budget_test.go | 169 +++++++++--- coderd/aibridgedserver/aibridgedserver.go | 19 +- .../aibridgedserver/aibridgedserver_test.go | 240 +++++++++++++----- coderd/aibridgedserver/cost.go | 25 +- coderd/apidoc/docs.go | 4 +- coderd/apidoc/swagger.json | 4 +- coderd/database/dbauthz/dbauthz.go | 7 + coderd/database/dbauthz/dbauthz_test.go | 8 + coderd/database/dbmetrics/querymetrics.go | 8 + coderd/database/dbmock/dbmock.go | 15 ++ coderd/database/querier.go | 21 +- coderd/database/querier_test.go | 128 ++++++++++ coderd/database/queries.sql.go | 82 ++++-- coderd/database/queries/aicostcontrol.sql | 74 ++++-- codersdk/aibridge.go | 16 +- docs/reference/api/schemas.md | 30 +-- enterprise/coderd/aibridge.go | 14 +- enterprise/coderd/aibridge_test.go | 114 +++++++-- site/src/api/typesGenerated.ts | 16 +- 20 files changed, 805 insertions(+), 269 deletions(-) diff --git a/coderd/aibridge/budget/budget.go b/coderd/aibridge/budget/budget.go index 5b86cdbdd6813..61a4c793173b3 100644 --- a/coderd/aibridge/budget/budget.go +++ b/coderd/aibridge/budget/budget.go @@ -20,37 +20,43 @@ import ( type Store interface { GetUserAIBudgetOverride(ctx context.Context, userID uuid.UUID) (database.UserAIBudgetOverride, error) GetHighestGroupAIBudgetByUser(ctx context.Context, userID uuid.UUID) (database.GetHighestGroupAIBudgetByUserRow, error) + GetUserEveryoneFallbackGroup(ctx context.Context, userID uuid.UUID) (uuid.UUID, error) } -// EffectiveBudget is the AI budget that applies to a user after override and -// policy resolution. -type EffectiveBudget struct { +// EffectiveGroup is a user's resolved effective group and, when a budget +// applies, its limit. Limit is nil for the Everyone fallback (unlimited). +type EffectiveGroup struct { // GroupID is the group the spend is attributed to. GroupID uuid.UUID - // SpendLimitMicros is the effective spend limit in micro-units - // (1 unit = 1,000,000). + Limit *Limit +} + +// Limit is an AI spend limit and the tier that produced it. +type Limit struct { + // SpendLimitMicros is the spend limit in micro-units (1 unit = 1,000,000). SpendLimitMicros int64 Source codersdk.AIBudgetLimitSource } -// ResolveUserAIBudget returns the effective AI budget for userID. The second -// return value is false when no budget is configured for the user. A per-user -// override wins unconditionally; otherwise the budget is selected from the -// user's groups according to policy. -// +// ResolveUserAIBudget returns the effective AI budget group for userID. The +// second return value is false when no budget is configured for the user. A +// per-user override wins unconditionally; otherwise the budget is selected from +// the user's groups according to policy. // TODO(AIGOV-527): unify effective group resolution in a single place. -func ResolveUserAIBudget(ctx context.Context, db Store, userID uuid.UUID, policy codersdk.AIBudgetPolicy) (EffectiveBudget, bool, error) { +func ResolveUserAIBudget(ctx context.Context, db Store, userID uuid.UUID, policy codersdk.AIBudgetPolicy) (EffectiveGroup, bool, error) { // A per-user override always wins. override, err := db.GetUserAIBudgetOverride(ctx, userID) if err == nil { - return EffectiveBudget{ - GroupID: override.GroupID, - SpendLimitMicros: override.SpendLimitMicros, - Source: codersdk.AIBudgetLimitSourceUserOverride, + return EffectiveGroup{ + GroupID: override.GroupID, + Limit: &Limit{ + SpendLimitMicros: override.SpendLimitMicros, + Source: codersdk.AIBudgetLimitSourceUserOverride, + }, }, true, nil } if !errors.Is(err, sql.ErrNoRows) { - return EffectiveBudget{}, false, xerrors.Errorf("get user AI budget override: %w", err) + return EffectiveGroup{}, false, xerrors.Errorf("get user AI budget override: %w", err) } // No override: select a group budget according to the deployment policy. @@ -58,17 +64,45 @@ func ResolveUserAIBudget(ctx context.Context, db Store, userID uuid.UUID, policy case codersdk.AIBudgetPolicyHighest: row, err := db.GetHighestGroupAIBudgetByUser(ctx, userID) if errors.Is(err, sql.ErrNoRows) { - return EffectiveBudget{}, false, nil + return EffectiveGroup{}, false, nil } if err != nil { - return EffectiveBudget{}, false, xerrors.Errorf("get highest group AI budget: %w", err) + return EffectiveGroup{}, false, xerrors.Errorf("get highest group AI budget: %w", err) } - return EffectiveBudget{ - GroupID: row.GroupID, - SpendLimitMicros: row.SpendLimitMicros, - Source: codersdk.AIBudgetLimitSourceGroup, + return EffectiveGroup{ + GroupID: row.GroupID, + Limit: &Limit{ + SpendLimitMicros: row.SpendLimitMicros, + Source: codersdk.AIBudgetLimitSourceGroup, + }, }, true, nil default: - return EffectiveBudget{}, false, xerrors.Errorf("unsupported AI budget policy: %q", policy) + return EffectiveGroup{}, false, xerrors.Errorf("unsupported AI budget policy: %q", policy) + } +} + +// ResolveUserEffectiveGroup resolves the user's effective group, falling back to +// the organization's Everyone group when no override or group budget applies. +// The second return value is false when no effective group was found for the +// user. +func ResolveUserEffectiveGroup(ctx context.Context, db Store, userID uuid.UUID, policy codersdk.AIBudgetPolicy) (EffectiveGroup, bool, error) { + group, ok, err := ResolveUserAIBudget(ctx, db, userID, policy) + if err != nil { + return EffectiveGroup{}, false, err + } + if ok { + return group, true, nil + } + + // No override or group budget: fall back to the Everyone group (unlimited). + groupID, err := db.GetUserEveryoneFallbackGroup(ctx, userID) + if errors.Is(err, sql.ErrNoRows) { + // This should not happen, as a user should always be a member of an + // organization and its associated Everyone group. + return EffectiveGroup{}, false, nil + } + if err != nil { + return EffectiveGroup{}, false, xerrors.Errorf("get everyone fallback group: %w", err) } + return EffectiveGroup{GroupID: groupID}, true, nil } diff --git a/coderd/aibridge/budget/budget_test.go b/coderd/aibridge/budget/budget_test.go index 4caaae02780a9..b931278bd54ba 100644 --- a/coderd/aibridge/budget/budget_test.go +++ b/coderd/aibridge/budget/budget_test.go @@ -1,7 +1,6 @@ package budget_test import ( - "bytes" "context" "testing" @@ -52,18 +51,18 @@ func TestResolveUserAIBudget(t *testing.T) { tests := []struct { name string policy codersdk.AIBudgetPolicy - setup func(t *testing.T, ctx context.Context, db database.Store) (userID uuid.UUID, want budget.EffectiveBudget, wantOK bool) + setup func(t *testing.T, ctx context.Context, db database.Store) (userID uuid.UUID, want budget.EffectiveGroup, wantOK bool) wantErr string }{ { name: "OverrideWins", policy: codersdk.AIBudgetPolicyHighest, - setup: func(t *testing.T, ctx context.Context, db database.Store) (uuid.UUID, budget.EffectiveBudget, bool) { + setup: func(t *testing.T, ctx context.Context, db database.Store) (uuid.UUID, budget.EffectiveGroup, bool) { org := dbgen.Organization(t, db, database.Organization{}) user := dbgen.User(t, db, database.User{}) // A higher group budget that the override must still beat. budgetedGroup(t, ctx, db, org.ID, user.ID, "rich-group", 9_000_000) - // The override names its own group; the user must be a member. + // The override names a group the user must be a member of. og := dbgen.Group(t, db, database.Group{OrganizationID: org.ID, Name: "override-group"}) dbgen.GroupMember(t, db, database.GroupMemberTable{UserID: user.ID, GroupID: og.ID}) _, err := db.UpsertUserAIBudgetOverride(ctx, database.UpsertUserAIBudgetOverrideParams{ @@ -72,92 +71,89 @@ func TestResolveUserAIBudget(t *testing.T) { SpendLimitMicros: 1_000_000, }) require.NoError(t, err) - return user.ID, budget.EffectiveBudget{GroupID: og.ID, SpendLimitMicros: 1_000_000, Source: codersdk.AIBudgetLimitSourceUserOverride}, true + return user.ID, budget.EffectiveGroup{GroupID: og.ID, Limit: &budget.Limit{SpendLimitMicros: 1_000_000, Source: codersdk.AIBudgetLimitSourceUserOverride}}, true }, }, { name: "SingleGroupBudget", policy: codersdk.AIBudgetPolicyHighest, - setup: func(t *testing.T, ctx context.Context, db database.Store) (uuid.UUID, budget.EffectiveBudget, bool) { + setup: func(t *testing.T, ctx context.Context, db database.Store) (uuid.UUID, budget.EffectiveGroup, bool) { org := dbgen.Organization(t, db, database.Organization{}) user := dbgen.User(t, db, database.User{}) gid := budgetedGroup(t, ctx, db, org.ID, user.ID, "only", 8_000_000) - return user.ID, budget.EffectiveBudget{GroupID: gid, SpendLimitMicros: 8_000_000, Source: codersdk.AIBudgetLimitSourceGroup}, true + return user.ID, budget.EffectiveGroup{GroupID: gid, Limit: &budget.Limit{SpendLimitMicros: 8_000_000, Source: codersdk.AIBudgetLimitSourceGroup}}, true }, }, { name: "HighestGroupWins", policy: codersdk.AIBudgetPolicyHighest, - setup: func(t *testing.T, ctx context.Context, db database.Store) (uuid.UUID, budget.EffectiveBudget, bool) { + setup: func(t *testing.T, ctx context.Context, db database.Store) (uuid.UUID, budget.EffectiveGroup, bool) { org := dbgen.Organization(t, db, database.Organization{}) user := dbgen.User(t, db, database.User{}) budgetedGroup(t, ctx, db, org.ID, user.ID, "low", 5_000_000) budgetedGroup(t, ctx, db, org.ID, user.ID, "mid", 20_000_000) high := budgetedGroup(t, ctx, db, org.ID, user.ID, "high", 50_000_000) - return user.ID, budget.EffectiveBudget{GroupID: high, SpendLimitMicros: 50_000_000, Source: codersdk.AIBudgetLimitSourceGroup}, true + return user.ID, budget.EffectiveGroup{GroupID: high, Limit: &budget.Limit{SpendLimitMicros: 50_000_000, Source: codersdk.AIBudgetLimitSourceGroup}}, true }, }, { - name: "TieBrokenByName", + name: "TieBrokenByGroupName", policy: codersdk.AIBudgetPolicyHighest, - setup: func(t *testing.T, ctx context.Context, db database.Store) (uuid.UUID, budget.EffectiveBudget, bool) { + setup: func(t *testing.T, ctx context.Context, db database.Store) (uuid.UUID, budget.EffectiveGroup, bool) { org := dbgen.Organization(t, db, database.Organization{}) user := dbgen.User(t, db, database.User{}) - // Equal limits; "alpha" must win over "beta" by name ascending. + // Equal limits in the same org, so "alpha" wins over "beta" by group + // name ascending. alpha := budgetedGroup(t, ctx, db, org.ID, user.ID, "alpha", 10_000_000) budgetedGroup(t, ctx, db, org.ID, user.ID, "beta", 10_000_000) - return user.ID, budget.EffectiveBudget{GroupID: alpha, SpendLimitMicros: 10_000_000, Source: codersdk.AIBudgetLimitSourceGroup}, true + return user.ID, budget.EffectiveGroup{GroupID: alpha, Limit: &budget.Limit{SpendLimitMicros: 10_000_000, Source: codersdk.AIBudgetLimitSourceGroup}}, true }, }, { - name: "TieBrokenByGroupID", + name: "TieBrokenByOrgName", policy: codersdk.AIBudgetPolicyHighest, - setup: func(t *testing.T, ctx context.Context, db database.Store) (uuid.UUID, budget.EffectiveBudget, bool) { + setup: func(t *testing.T, ctx context.Context, db database.Store) (uuid.UUID, budget.EffectiveGroup, bool) { user := dbgen.User(t, db, database.User{}) - // Two groups in different orgs share both name and limit. - // Group id breaks the tie, so resolution is deterministic. - org1 := dbgen.Organization(t, db, database.Organization{}) - org2 := dbgen.Organization(t, db, database.Organization{}) - g1 := budgetedGroup(t, ctx, db, org1.ID, user.ID, "dup", 10_000_000) - g2 := budgetedGroup(t, ctx, db, org2.ID, user.ID, "dup", 10_000_000) - winner := g1 - if bytes.Compare(g2[:], g1[:]) < 0 { - winner = g2 - } - return user.ID, budget.EffectiveBudget{GroupID: winner, SpendLimitMicros: 10_000_000, Source: codersdk.AIBudgetLimitSourceGroup}, true + // Two groups in different orgs share both name and limit. The + // organization name breaks the tie ascending, so "org-a" wins. + orgA := dbgen.Organization(t, db, database.Organization{Name: "org-a"}) + orgB := dbgen.Organization(t, db, database.Organization{Name: "org-b"}) + winner := budgetedGroup(t, ctx, db, orgA.ID, user.ID, "dup", 10_000_000) + budgetedGroup(t, ctx, db, orgB.ID, user.ID, "dup", 10_000_000) + return user.ID, budget.EffectiveGroup{GroupID: winner, Limit: &budget.Limit{SpendLimitMicros: 10_000_000, Source: codersdk.AIBudgetLimitSourceGroup}}, true }, }, { name: "GroupsButNoneBudgeted", policy: codersdk.AIBudgetPolicyHighest, - setup: func(t *testing.T, ctx context.Context, db database.Store) (uuid.UUID, budget.EffectiveBudget, bool) { + setup: func(t *testing.T, ctx context.Context, db database.Store) (uuid.UUID, budget.EffectiveGroup, bool) { org := dbgen.Organization(t, db, database.Organization{}) user := dbgen.User(t, db, database.User{}) g := dbgen.Group(t, db, database.Group{OrganizationID: org.ID, Name: "unbudgeted"}) dbgen.GroupMember(t, db, database.GroupMemberTable{UserID: user.ID, GroupID: g.ID}) - return user.ID, budget.EffectiveBudget{}, false + return user.ID, budget.EffectiveGroup{}, false }, }, { name: "EveryoneGroupBudget", policy: codersdk.AIBudgetPolicyHighest, - setup: func(t *testing.T, ctx context.Context, db database.Store) (uuid.UUID, budget.EffectiveBudget, bool) { + setup: func(t *testing.T, ctx context.Context, db database.Store) (uuid.UUID, budget.EffectiveGroup, bool) { org := dbgen.Organization(t, db, database.Organization{}) user := dbgen.User(t, db, database.User{}) // Membership is via organization_members only (no group_members row), // exercising the org-members half of group_members_expanded. everyoneID := budgetedEveryoneGroup(t, ctx, db, org.ID, user.ID, 7_000_000) - return user.ID, budget.EffectiveBudget{GroupID: everyoneID, SpendLimitMicros: 7_000_000, Source: codersdk.AIBudgetLimitSourceGroup}, true + return user.ID, budget.EffectiveGroup{GroupID: everyoneID, Limit: &budget.Limit{SpendLimitMicros: 7_000_000, Source: codersdk.AIBudgetLimitSourceGroup}}, true }, }, { name: "OverrideBeatsEveryoneBudget", policy: codersdk.AIBudgetPolicyHighest, - setup: func(t *testing.T, ctx context.Context, db database.Store) (uuid.UUID, budget.EffectiveBudget, bool) { + setup: func(t *testing.T, ctx context.Context, db database.Store) (uuid.UUID, budget.EffectiveGroup, bool) { org := dbgen.Organization(t, db, database.Organization{}) user := dbgen.User(t, db, database.User{}) everyoneID := budgetedEveryoneGroup(t, ctx, db, org.ID, user.ID, 7_000_000) - // Override attributed to the Everyone group; the user is a member + // Override attributed to the Everyone group. The user is a member // via organization_members, satisfying the membership trigger. _, err := db.UpsertUserAIBudgetOverride(ctx, database.UpsertUserAIBudgetOverrideParams{ UserID: user.ID, @@ -165,16 +161,16 @@ func TestResolveUserAIBudget(t *testing.T) { SpendLimitMicros: 2_000_000, }) require.NoError(t, err) - return user.ID, budget.EffectiveBudget{GroupID: everyoneID, SpendLimitMicros: 2_000_000, Source: codersdk.AIBudgetLimitSourceUserOverride}, true + return user.ID, budget.EffectiveGroup{GroupID: everyoneID, Limit: &budget.Limit{SpendLimitMicros: 2_000_000, Source: codersdk.AIBudgetLimitSourceUserOverride}}, true }, }, { name: "UnsupportedPolicy", policy: codersdk.AIBudgetPolicy("unsupported"), - setup: func(t *testing.T, ctx context.Context, db database.Store) (uuid.UUID, budget.EffectiveBudget, bool) { + setup: func(t *testing.T, ctx context.Context, db database.Store) (uuid.UUID, budget.EffectiveGroup, bool) { // No override, so resolution reaches the policy switch and errors. user := dbgen.User(t, db, database.User{}) - return user.ID, budget.EffectiveBudget{}, false + return user.ID, budget.EffectiveGroup{}, false }, wantErr: "unsupported AI budget policy", }, @@ -199,8 +195,107 @@ func TestResolveUserAIBudget(t *testing.T) { return } require.Equal(t, want.GroupID, got.GroupID) - require.Equal(t, want.SpendLimitMicros, got.SpendLimitMicros) - require.Equal(t, want.Source, got.Source) + require.Equal(t, want.Limit, got.Limit) + }) + } +} + +func TestResolveUserEffectiveGroup(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + policy codersdk.AIBudgetPolicy + setup func(t *testing.T, ctx context.Context, db database.Store) (userID uuid.UUID, want budget.EffectiveGroup, wantOK bool) + }{ + { + // The Everyone group has a budget, so it resolves via the budget + // path rather than the fallback. + name: "EveryoneGroupWithBudget", + policy: codersdk.AIBudgetPolicyHighest, + setup: func(t *testing.T, ctx context.Context, db database.Store) (uuid.UUID, budget.EffectiveGroup, bool) { + org := dbgen.Organization(t, db, database.Organization{}) + user := dbgen.User(t, db, database.User{}) + // The Everyone group's id equals the org id. + group := dbgen.Group(t, db, database.Group{ID: org.ID, OrganizationID: org.ID, Name: "Everyone"}) + dbgen.OrganizationMember(t, db, database.OrganizationMember{OrganizationID: org.ID, UserID: user.ID}) + _, err := db.UpsertGroupAIBudget(ctx, database.UpsertGroupAIBudgetParams{ + GroupID: group.ID, + SpendLimitMicros: 7_000_000, + }) + require.NoError(t, err) + return user.ID, budget.EffectiveGroup{GroupID: group.ID, Limit: &budget.Limit{SpendLimitMicros: 7_000_000, Source: codersdk.AIBudgetLimitSourceGroup}}, true + }, + }, + { + // With a single org and no budget, attribution falls back to that + // org's Everyone group with no limit. + name: "FallbackToEveryoneUnlimited", + policy: codersdk.AIBudgetPolicyHighest, + setup: func(t *testing.T, ctx context.Context, db database.Store) (uuid.UUID, budget.EffectiveGroup, bool) { + org := dbgen.Organization(t, db, database.Organization{}) + user := dbgen.User(t, db, database.User{}) + dbgen.OrganizationMember(t, db, database.OrganizationMember{OrganizationID: org.ID, UserID: user.ID}) + return user.ID, budget.EffectiveGroup{GroupID: org.ID}, true + }, + }, + { + // The fallback prefers the default org even over an org whose name + // sorts first. + name: "FallbackPrefersDefaultOrg", + policy: codersdk.AIBudgetPolicyHighest, + setup: func(t *testing.T, ctx context.Context, db database.Store) (uuid.UUID, budget.EffectiveGroup, bool) { + defaultOrg, err := db.GetDefaultOrganization(ctx) + require.NoError(t, err) + otherOrg := dbgen.Organization(t, db, database.Organization{Name: "0-sorts-first"}) + user := dbgen.User(t, db, database.User{}) + dbgen.OrganizationMember(t, db, database.OrganizationMember{OrganizationID: defaultOrg.ID, UserID: user.ID}) + dbgen.OrganizationMember(t, db, database.OrganizationMember{OrganizationID: otherOrg.ID, UserID: user.ID}) + return user.ID, budget.EffectiveGroup{GroupID: defaultOrg.ID}, true + }, + }, + { + // Among non-default orgs, the fallback breaks ties by org name + // ascending, so "org-a" wins. + name: "FallbackTieByOrgName", + policy: codersdk.AIBudgetPolicyHighest, + setup: func(t *testing.T, ctx context.Context, db database.Store) (uuid.UUID, budget.EffectiveGroup, bool) { + user := dbgen.User(t, db, database.User{}) + orgA := dbgen.Organization(t, db, database.Organization{Name: "org-a"}) + orgB := dbgen.Organization(t, db, database.Organization{Name: "org-b"}) + dbgen.OrganizationMember(t, db, database.OrganizationMember{OrganizationID: orgA.ID, UserID: user.ID}) + dbgen.OrganizationMember(t, db, database.OrganizationMember{OrganizationID: orgB.ID, UserID: user.ID}) + return user.ID, budget.EffectiveGroup{GroupID: orgA.ID}, true + }, + }, + { + // A user with no org membership has no effective group. + name: "NoOrgMembership", + policy: codersdk.AIBudgetPolicyHighest, + setup: func(t *testing.T, ctx context.Context, db database.Store) (uuid.UUID, budget.EffectiveGroup, bool) { + user := dbgen.User(t, db, database.User{}) + return user.ID, budget.EffectiveGroup{}, false + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + db, _ := dbtestutil.NewDB(t) + ctx := testutil.Context(t, testutil.WaitLong) + + userID, want, wantOK := tt.setup(t, ctx, db) + got, ok, err := budget.ResolveUserEffectiveGroup(ctx, db, userID, tt.policy) + require.NoError(t, err) + require.Equal(t, wantOK, ok) + if !wantOK { + return + } + + require.Equal(t, want.GroupID, got.GroupID) + require.Equal(t, want.Limit, got.Limit) }) } } diff --git a/coderd/aibridgedserver/aibridgedserver.go b/coderd/aibridgedserver/aibridgedserver.go index c17a087920076..6a0f68aab43fd 100644 --- a/coderd/aibridgedserver/aibridgedserver.go +++ b/coderd/aibridgedserver/aibridgedserver.go @@ -80,6 +80,7 @@ type store interface { GetAIModelPriceByProviderModel(ctx context.Context, arg database.GetAIModelPriceByProviderModelParams) (database.AIModelPrice, error) GetUserAIBudgetOverride(ctx context.Context, userID uuid.UUID) (database.UserAIBudgetOverride, error) GetHighestGroupAIBudgetByUser(ctx context.Context, userID uuid.UUID) (database.GetHighestGroupAIBudgetByUserRow, error) + GetUserEveryoneFallbackGroup(ctx context.Context, userID uuid.UUID) (uuid.UUID, error) GetUserAISpendSince(ctx context.Context, arg database.GetUserAISpendSinceParams) (database.GetUserAISpendSinceRow, error) // MCPConfigurator-related queries. @@ -812,32 +813,32 @@ type userAIBudget struct { // overages, not build an accounting system. // - Fail-open is acceptable for this case. func (s *Server) checkUserAIBudget(ctx context.Context, userID uuid.UUID, periodStart time.Time) (userAIBudget, error) { - effectiveBudget, ok, err := budget.ResolveUserAIBudget(ctx, s.store, userID, s.budgetPolicy) + effectiveGroup, ok, err := budget.ResolveUserAIBudget(ctx, s.store, userID, s.budgetPolicy) if err != nil { return userAIBudget{}, xerrors.Errorf("resolve effective AI budget for user %q with budget policy %q: %w", userID, s.budgetPolicy, err) } - if !ok { - // No budget configured for the user; return zero-valued status. + if !ok || effectiveGroup.Limit == nil { + // No enforceable spend limit for the user; return zero-valued status. return userAIBudget{}, nil } spend, err := s.store.GetUserAISpendSince(ctx, database.GetUserAISpendSinceParams{ UserID: userID, - EffectiveGroupID: effectiveBudget.GroupID, + EffectiveGroupID: effectiveGroup.GroupID, PeriodStart: periodStart, }) if err != nil { - return userAIBudget{}, xerrors.Errorf("get user AI spend for user %q in group %q: %w", userID, effectiveBudget.GroupID, err) + return userAIBudget{}, xerrors.Errorf("get user AI spend for user %q in group %q: %w", userID, effectiveGroup.GroupID, err) } - exceeded := spend.SpendMicros >= effectiveBudget.SpendLimitMicros + exceeded := spend.SpendMicros >= effectiveGroup.Limit.SpendLimitMicros logger := s.logger.With( slog.F("user_id", userID), - slog.F("effective_group_id", effectiveBudget.GroupID), + slog.F("effective_group_id", effectiveGroup.GroupID), slog.F("period_start", periodStart), slog.F("current_spend_micros", spend.SpendMicros), - slog.F("spend_limit_micros", effectiveBudget.SpendLimitMicros), + slog.F("spend_limit_micros", effectiveGroup.Limit.SpendLimitMicros), slog.F("exceeded", exceeded), ) logger.Debug(ctx, "user AI spend status") @@ -847,7 +848,7 @@ func (s *Server) checkUserAIBudget(ctx context.Context, userID uuid.UUID, period return userAIBudget{ Exceeded: exceeded, - SpendLimitMicros: ptr.Ref(effectiveBudget.SpendLimitMicros), + SpendLimitMicros: ptr.Ref(effectiveGroup.Limit.SpendLimitMicros), }, nil } diff --git a/coderd/aibridgedserver/aibridgedserver_test.go b/coderd/aibridgedserver/aibridgedserver_test.go index bf28dbdf9edac..fd983476400e9 100644 --- a/coderd/aibridgedserver/aibridgedserver_test.go +++ b/coderd/aibridgedserver/aibridgedserver_test.go @@ -732,6 +732,51 @@ func TestIsBudgetExceeded_Enforcement(t *testing.T) { require.False(t, afterResp.GetExceeded()) require.Equal(t, int64(overrideLimitMicros), afterResp.GetSpendLimitMicros()) }) + + t.Run("unbudgeted member is not blocked", func(t *testing.T) { + t.Parallel() + + clock := quartz.NewMock(t) + clock.Set(time.Date(2026, time.March, 15, 0, 0, 0, 0, time.UTC)) + + ctx := testutil.Context(t, testutil.WaitLong) + logger := testutil.Logger(t) + rawDB, _ := dbtestutil.NewDB(t) + authzDB := dbauthz.New(rawDB, rbac.NewStrictAuthorizer(prometheus.NewRegistry()), logger, coderdtest.AccessControlStorePointer()) + + // An org member with no group budget and no override. Their spend is + // unlimited, so enforcement never blocks them regardless of accrued + // spend, even though attribution would fall back to the Everyone group. + org := dbgen.Organization(t, rawDB, database.Organization{}) + user := dbgen.User(t, rawDB, database.User{}) + dbgen.OrganizationMember(t, rawDB, database.OrganizationMember{OrganizationID: org.ID, UserID: user.ID}) + + // Record spend attributed to the Everyone group. Without a configured + // limit it must not cause a block. + _, err := rawDB.IncrementUserAIDailySpend(ctx, database.IncrementUserAIDailySpendParams{ + UserID: user.ID, + EffectiveGroupID: org.ID, + Day: clock.Now(), + CostMicros: 1_000_000_000, // $1,000 USD + }) + require.NoError(t, err) + + srv, err := aibridgedserver.NewServer(ctx, aibridgedserver.Options{ + Store: authzDB, + AISeatTracker: agplaiseats.Noop{}, + AccessURL: "/", + GatewayCfg: codersdk.AIBridgeConfig{}, + Experiments: requiredExperiments, + Logger: logger, + Clock: clock, + }) + require.NoError(t, err) + + resp, err := srv.IsBudgetExceeded(ctx, &proto.IsBudgetExceededRequest{UserId: user.ID.String()}) + require.NoError(t, err) + require.False(t, resp.GetExceeded()) + require.Nil(t, resp.SpendLimitMicros) + }) } func TestGetMCPServerConfigs(t *testing.T) { @@ -1636,7 +1681,7 @@ func TestRecordTokenUsage(t *testing.T) { CacheWritePrice: sql.NullInt64{Int64: 4_000_000, Valid: true}, } // No override - expectTokenUsageCostLookups(db, intc, nil, group, price) + expectTokenUsageCostLookups(db, intc, nil, group, nil, price) // input 300 + output 1200 + cache read 15 + cache write 40. const wantCost int64 = 1555 @@ -1691,7 +1736,7 @@ func TestRecordTokenUsage(t *testing.T) { InputPrice: sql.NullInt64{Int64: 3_000_000, Valid: true}, } // No group - expectTokenUsageCostLookups(db, intc, override, nil, price) + expectTokenUsageCostLookups(db, intc, override, nil, nil, price) // input 300. const wantCost int64 = 300 @@ -1717,6 +1762,49 @@ func TestRecordTokenUsage(t *testing.T) { }).Return(database.AIUserDailySpend{}, nil) }, }, + { + // No override or group budget, so attribution falls back to the + // user's Everyone group. + name: "valid token usage falls back to the Everyone group", + request: &proto.RecordTokenUsageRequest{ + InterceptionId: uuid.NewString(), + MsgId: "msg_123", + InputTokens: 100, + CreatedAt: timestamppb.New(now), + }, + setupMocks: func(t *testing.T, db *dbmock.MockStore, req *proto.RecordTokenUsageRequest) { + interceptionID, err := uuid.Parse(req.GetInterceptionId()) + assert.NoError(t, err, "parse interception UUID") + + intc := newTestInterception(interceptionID) + everyoneID := uuid.New() + price := &database.AIModelPrice{ + Provider: intc.Provider, + Model: intc.Model, + InputPrice: sql.NullInt64{Int64: 3_000_000, Valid: true}, + } + expectTokenUsageCostLookups(db, intc, nil, nil, &everyoneID, price) + + // input 300. + const wantCost int64 = 300 + + db.EXPECT().InTx(gomock.Any(), nil).DoAndReturn( + func(fn func(database.Store) error, _ *database.TxOptions) error { return fn(db) }, + ) + + db.EXPECT().InsertAIBridgeTokenUsage(gomock.Any(), gomock.Cond(func(p database.InsertAIBridgeTokenUsageParams) bool { + return assert.Equal(t, uuid.NullUUID{UUID: everyoneID, Valid: true}, p.EffectiveGroupID, "effective group ID") && + assert.Equal(t, sql.NullInt64{Int64: wantCost, Valid: true}, p.CostMicros, "cost") + })).Return(database.AIBridgeTokenUsage{ID: uuid.New(), InterceptionID: interceptionID}, nil) + + db.EXPECT().IncrementUserAIDailySpend(gomock.Any(), database.IncrementUserAIDailySpendParams{ + UserID: intc.InitiatorID, + EffectiveGroupID: everyoneID, + Day: now.UTC().Truncate(24 * time.Hour), + CostMicros: wantCost, + }).Return(database.AIUserDailySpend{}, nil) + }, + }, { // Model has no price row, so cost is NULL. name: "valid token usage with effective group and no price", @@ -1739,7 +1827,7 @@ func TestRecordTokenUsage(t *testing.T) { // Budget resolves to a group, but the model has no price row. // The resolved group must survive the price lookup's early // return on sql.ErrNoRows, while prices and cost stay NULL. - expectTokenUsageCostLookups(db, intc, nil, group, nil) + expectTokenUsageCostLookups(db, intc, nil, group, nil, nil) db.EXPECT().InTx(gomock.Any(), nil).DoAndReturn( func(fn func(database.Store) error, _ *database.TxOptions) error { return fn(db) }, @@ -1793,7 +1881,7 @@ func TestRecordTokenUsage(t *testing.T) { CacheReadPrice: sql.NullInt64{Valid: false}, CacheWritePrice: sql.NullInt64{Valid: false}, } - expectTokenUsageCostLookups(db, intc, nil, group, price) + expectTokenUsageCostLookups(db, intc, nil, group, nil, price) db.EXPECT().InTx(gomock.Any(), nil).DoAndReturn( func(fn func(database.Store) error, _ *database.TxOptions) error { return fn(db) }, @@ -1846,7 +1934,7 @@ func TestRecordTokenUsage(t *testing.T) { CacheReadPrice: sql.NullInt64{Int64: 0, Valid: true}, CacheWritePrice: sql.NullInt64{Int64: 0, Valid: true}, } - expectTokenUsageCostLookups(db, intc, nil, group, price) + expectTokenUsageCostLookups(db, intc, nil, group, nil, price) db.EXPECT().InTx(gomock.Any(), nil).DoAndReturn( func(fn func(database.Store) error, _ *database.TxOptions) error { return fn(db) }, @@ -1871,60 +1959,7 @@ func TestRecordTokenUsage(t *testing.T) { }, }, { - // No budget configured, model is priced: group is NULL but cost is computed. - name: "valid token usage with no budget and cost", - request: &proto.RecordTokenUsageRequest{ - InterceptionId: uuid.NewString(), - MsgId: "msg_123", - InputTokens: 100, - OutputTokens: 200, - CacheReadInputTokens: 50, - CacheWriteInputTokens: 10, - CreatedAt: timestamppb.Now(), - }, - setupMocks: func(t *testing.T, db *dbmock.MockStore, req *proto.RecordTokenUsageRequest) { - interceptionID, err := uuid.Parse(req.GetInterceptionId()) - assert.NoError(t, err, "parse interception UUID") - - intc := newTestInterception(interceptionID) - price := &database.AIModelPrice{ - Provider: intc.Provider, - Model: intc.Model, - InputPrice: sql.NullInt64{Int64: 3_000_000, Valid: true}, - OutputPrice: sql.NullInt64{Int64: 6_000_000, Valid: true}, - CacheReadPrice: sql.NullInt64{Int64: 300_000, Valid: true}, - CacheWritePrice: sql.NullInt64{Int64: 4_000_000, Valid: true}, - } - // No budget configured, but the model is priced: cost is - // computed independently of budget resolution, and the group - // attribution stays NULL. - expectTokenUsageCostLookups(db, intc, nil, nil, price) - - db.EXPECT().InTx(gomock.Any(), nil).DoAndReturn( - func(fn func(database.Store) error, _ *database.TxOptions) error { return fn(db) }, - ) - - // input 300 + output 1200 + cache read 15 + cache write 40. - const wantCost int64 = 1555 - - db.EXPECT().InsertAIBridgeTokenUsage(gomock.Any(), gomock.Cond(func(p database.InsertAIBridgeTokenUsageParams) bool { - if !assert.False(t, p.EffectiveGroupID.Valid, "effective group ID null") || - !assert.Equal(t, price.InputPrice, p.InputPriceMicros, "input price") || - !assert.Equal(t, price.OutputPrice, p.OutputPriceMicros, "output price") || - !assert.Equal(t, price.CacheReadPrice, p.CacheReadPriceMicros, "cache read price") || - !assert.Equal(t, price.CacheWritePrice, p.CacheWritePriceMicros, "cache write price") || - !assert.Equal(t, sql.NullInt64{Int64: wantCost, Valid: true}, p.CostMicros, "cost") { - return false - } - return true - })).Return(database.AIBridgeTokenUsage{ID: uuid.New(), InterceptionID: interceptionID}, nil) - - // Spend update is skipped because the effective group is NULL. - db.EXPECT().IncrementUserAIDailySpend(gomock.Any(), gomock.Any()).Times(0) - }, - }, - { - // No budget and no price row: group and cost are NULL. + // No budget and no price row: attribution falls back to Everyone and cost is NULL. name: "valid token usage with no budget and no price", request: &proto.RecordTokenUsageRequest{ InterceptionId: uuid.NewString(), @@ -1940,10 +1975,12 @@ func TestRecordTokenUsage(t *testing.T) { interceptionID, err := uuid.Parse(req.GetInterceptionId()) assert.NoError(t, err, "parse interception UUID") - // No budget configured and no price row: tokens recorded - // with NULL cost, prices, and group attribution. + // No budget configured, so attribution falls back to the + // Everyone group. The model has no price row, so cost and + // prices stay NULL. intc := newTestInterception(interceptionID) - expectTokenUsageCostLookups(db, intc, nil, nil, nil) + everyoneID := uuid.New() + expectTokenUsageCostLookups(db, intc, nil, nil, &everyoneID, nil) db.EXPECT().InTx(gomock.Any(), nil).DoAndReturn( func(fn func(database.Store) error, _ *database.TxOptions) error { return fn(db) }, @@ -1959,7 +1996,7 @@ func TestRecordTokenUsage(t *testing.T) { !assert.Equal(t, req.GetCacheWriteInputTokens(), p.CacheWriteInputTokens, "cache write input tokens") || !assert.JSONEq(t, metadataJSON, string(p.Metadata), "metadata") || !assert.WithinDuration(t, req.GetCreatedAt().AsTime(), p.CreatedAt, time.Second, "created at") || - !assert.False(t, p.EffectiveGroupID.Valid, "effective group ID null") || + !assert.Equal(t, uuid.NullUUID{UUID: everyoneID, Valid: true}, p.EffectiveGroupID, "effective group ID") || !assert.False(t, p.InputPriceMicros.Valid, "input price null") || !assert.False(t, p.OutputPriceMicros.Valid, "output price null") || !assert.False(t, p.CacheReadPriceMicros.Valid, "cache read price null") || @@ -1983,7 +2020,57 @@ func TestRecordTokenUsage(t *testing.T) { CreatedAt: req.GetCreatedAt().AsTime(), }, nil) - // Spend update is skipped because the effective group and cost are NULL. + // Spend update is skipped because cost is NULL. + db.EXPECT().IncrementUserAIDailySpend(gomock.Any(), gomock.Any()).Times(0) + }, + }, + { + // A user with no organization has no effective group. This is + // not expected in practice. Spend is still recorded, but with a + // NULL group, and the daily spend update is skipped. + name: "valid token usage with no effective group", + request: &proto.RecordTokenUsageRequest{ + InterceptionId: uuid.NewString(), + MsgId: "msg_123", + InputTokens: 100, + OutputTokens: 200, + CacheReadInputTokens: 50, + CacheWriteInputTokens: 10, + CreatedAt: timestamppb.Now(), + }, + setupMocks: func(t *testing.T, db *dbmock.MockStore, req *proto.RecordTokenUsageRequest) { + interceptionID, err := uuid.Parse(req.GetInterceptionId()) + assert.NoError(t, err, "parse interception UUID") + + intc := newTestInterception(interceptionID) + price := &database.AIModelPrice{ + Provider: intc.Provider, + Model: intc.Model, + InputPrice: sql.NullInt64{Int64: 3_000_000, Valid: true}, + OutputPrice: sql.NullInt64{Int64: 6_000_000, Valid: true}, + CacheReadPrice: sql.NullInt64{Int64: 300_000, Valid: true}, + CacheWritePrice: sql.NullInt64{Int64: 4_000_000, Valid: true}, + } + // Every resolution lookup misses, including the Everyone + // fallback, so the group stays NULL while cost is computed. + expectTokenUsageCostLookups(db, intc, nil, nil, nil, price) + + db.EXPECT().InTx(gomock.Any(), nil).DoAndReturn( + func(fn func(database.Store) error, _ *database.TxOptions) error { return fn(db) }, + ) + + // input 300 + output 1200 + cache read 15 + cache write 40. + const wantCost int64 = 1555 + + db.EXPECT().InsertAIBridgeTokenUsage(gomock.Any(), gomock.Cond(func(p database.InsertAIBridgeTokenUsageParams) bool { + if !assert.False(t, p.EffectiveGroupID.Valid, "effective group ID null") || + !assert.Equal(t, sql.NullInt64{Int64: wantCost, Valid: true}, p.CostMicros, "cost") { + return false + } + return true + })).Return(database.AIBridgeTokenUsage{ID: uuid.New(), InterceptionID: interceptionID}, nil) + + // Spend update is skipped because the effective group is NULL. db.EXPECT().IncrementUserAIDailySpend(gomock.Any(), gomock.Any()).Times(0) }, }, @@ -2039,6 +2126,8 @@ func TestRecordTokenUsage(t *testing.T) { Return(database.UserAIBudgetOverride{}, sql.ErrNoRows) db.EXPECT().GetHighestGroupAIBudgetByUser(gomock.Any(), intc.InitiatorID). Return(database.GetHighestGroupAIBudgetByUserRow{}, sql.ErrNoRows) + db.EXPECT().GetUserEveryoneFallbackGroup(gomock.Any(), intc.InitiatorID). + Return(uuid.New(), nil) db.EXPECT().GetAIModelPriceByProviderModel(gomock.Any(), gomock.Any()). Return(database.AIModelPrice{}, sql.ErrConnDone) }, @@ -2057,7 +2146,8 @@ func TestRecordTokenUsage(t *testing.T) { interceptionID, err := uuid.Parse(req.GetInterceptionId()) assert.NoError(t, err, "parse interception UUID") - expectTokenUsageCostLookups(db, newTestInterception(interceptionID), nil, nil, nil) + everyoneID := uuid.New() + expectTokenUsageCostLookups(db, newTestInterception(interceptionID), nil, nil, &everyoneID, nil) db.EXPECT().InTx(gomock.Any(), nil).DoAndReturn( func(fn func(database.Store) error, _ *database.TxOptions) error { return fn(db) }, @@ -2085,7 +2175,7 @@ func TestRecordTokenUsage(t *testing.T) { Model: intc.Model, InputPrice: sql.NullInt64{Int64: 3_000_000, Valid: true}, } - expectTokenUsageCostLookups(db, intc, nil, group, price) + expectTokenUsageCostLookups(db, intc, nil, group, nil, price) db.EXPECT().InTx(gomock.Any(), nil).DoAndReturn( func(fn func(database.Store) error, _ *database.TxOptions) error { return fn(db) }, @@ -2213,14 +2303,16 @@ func newTestInterception(id uuid.UUID) database.AIBridgeInterception { } // expectTokenUsageCostLookups mocks the store lookups made by resolveTokenUsageCost -// (budget resolution and the price lookup). A nil override, group, or price makes that -// lookup return sql.ErrNoRows. Budget resolution mirrors production code: a non-nil override -// wins and skips the group lookup, so group is consulted only when override is nil. +// (budget resolution and the price lookup). A nil override, group, everyoneGroupID, or +// price makes that lookup return sql.ErrNoRows. Budget resolution mirrors production code: +// a non-nil override wins and skips the group lookup, and the Everyone fallback is consulted +// only when both override and group are nil. func expectTokenUsageCostLookups( db *dbmock.MockStore, intc database.AIBridgeInterception, override *database.UserAIBudgetOverride, group *database.GetHighestGroupAIBudgetByUserRow, + everyoneGroupID *uuid.UUID, price *database.AIModelPrice, ) { db.EXPECT().GetAIBridgeInterceptionByID(gomock.Any(), intc.ID).Return(intc, nil) @@ -2235,6 +2327,13 @@ func expectTokenUsageCostLookups( } else { db.EXPECT().GetHighestGroupAIBudgetByUser(gomock.Any(), intc.InitiatorID). Return(database.GetHighestGroupAIBudgetByUserRow{}, sql.ErrNoRows) + if everyoneGroupID != nil { + db.EXPECT().GetUserEveryoneFallbackGroup(gomock.Any(), intc.InitiatorID). + Return(*everyoneGroupID, nil) + } else { + db.EXPECT().GetUserEveryoneFallbackGroup(gomock.Any(), intc.InitiatorID). + Return(uuid.Nil, sql.ErrNoRows) + } } } @@ -2736,7 +2835,8 @@ func TestStructuredLogging(t *testing.T) { name: "RecordTokenUsage_logs_when_enabled", structuredLogging: true, setupMocks: func(db *dbmock.MockStore, intcID uuid.UUID) { - expectTokenUsageCostLookups(db, newTestInterception(intcID), nil, nil, nil) + everyoneID := uuid.New() + expectTokenUsageCostLookups(db, newTestInterception(intcID), nil, nil, &everyoneID, nil) db.EXPECT().InTx(gomock.Any(), nil).DoAndReturn( func(fn func(database.Store) error, _ *database.TxOptions) error { return fn(db) }, ) diff --git a/coderd/aibridgedserver/cost.go b/coderd/aibridgedserver/cost.go index 0a6b9ba5b63f2..416b03ea7c289 100644 --- a/coderd/aibridgedserver/cost.go +++ b/coderd/aibridgedserver/cost.go @@ -31,22 +31,25 @@ type tokenUsageCost struct { } // resolveTokenUsageCost resolves the effective group and per-token prices for an -// interception and computes its cost. Two outcomes are expected and yield NULL -// columns rather than an error: a user with no configured budget (yields a NULL -// group) and a model absent from the price table (yields NULL prices and cost). -// Any other error is returned. A NULL cost unambiguously means "model not priced". +// interception and computes its cost. A model absent from the price table yields +// NULL prices and cost rather than an error (a NULL cost unambiguously means +// "model not priced"). Any other error is returned. func (s *Server) resolveTokenUsageCost(ctx context.Context, intc database.AIBridgeInterception, in *proto.RecordTokenUsageRequest) (tokenUsageCost, error) { var result tokenUsageCost - // Resolve the effective group for attribution. This is independent of - // whether the model is priced. ok is false when no budget is configured, - // which leaves the group attribution NULL. - effectiveBudget, ok, err := budget.ResolveUserAIBudget(ctx, s.store, intc.InitiatorID, s.budgetPolicy) + // Resolve the effective group for attribution, independent of whether the + // model is priced. + effectiveGroup, ok, err := budget.ResolveUserEffectiveGroup(ctx, s.store, intc.InitiatorID, s.budgetPolicy) if err != nil { - return tokenUsageCost{}, xerrors.Errorf("resolve effective AI budget for user %q with policy %q: %w", intc.InitiatorID, s.budgetPolicy, err) + return tokenUsageCost{}, xerrors.Errorf("resolve effective AI group for user %q with policy %q: %w", intc.InitiatorID, s.budgetPolicy, err) } - if ok { - result.effectiveGroupID = uuid.NullUUID{UUID: effectiveBudget.GroupID, Valid: true} + if !ok { + // A user should always resolve to at least their Everyone group, so log + // this unexpected case. Spend is still recorded, with a NULL group. + s.logger.Warn(ctx, "no effective group for user, AI spend not attributed", + slog.F("user_id", intc.InitiatorID)) + } else { + result.effectiveGroupID = uuid.NullUUID{UUID: effectiveGroup.GroupID, Valid: true} } // Snapshot the price for this (provider, model) and compute cost. diff --git a/coderd/apidoc/docs.go b/coderd/apidoc/docs.go index 2c27365ac0797..8e01936e88737 100644 --- a/coderd/apidoc/docs.go +++ b/coderd/apidoc/docs.go @@ -20424,7 +20424,7 @@ const docTemplate = `{ "type": "object", "properties": { "effective_group_id": { - "description": "EffectiveGroupID is the user's effective budget group within the queried\ngroup's organization. Null when no effective budget group is visible in\nthis organization, including when the user's budget resolves to a group\nin another organization.", + "description": "EffectiveGroupID is the user's effective budget group within the queried\ngroup's organization, falling back to the Everyone group when no budget\napplies. Null when the effective group belongs to a different organization\nthan the queried group.", "type": "string", "format": "uuid" }, @@ -26027,7 +26027,7 @@ const docTemplate = `{ "type": "integer" }, "effective_group_id": { - "description": "EffectiveGroupID is the group the spend is attributed to. Null when\nno budget applies.", + "description": "EffectiveGroupID is the group the spend is attributed to, falling back to\nthe Everyone group when no budget applies. Null only when the user has no\norganization.", "type": "string", "format": "uuid" }, diff --git a/coderd/apidoc/swagger.json b/coderd/apidoc/swagger.json index 7f87f2bad29c0..0c47c56afe180 100644 --- a/coderd/apidoc/swagger.json +++ b/coderd/apidoc/swagger.json @@ -18590,7 +18590,7 @@ "type": "object", "properties": { "effective_group_id": { - "description": "EffectiveGroupID is the user's effective budget group within the queried\ngroup's organization. Null when no effective budget group is visible in\nthis organization, including when the user's budget resolves to a group\nin another organization.", + "description": "EffectiveGroupID is the user's effective budget group within the queried\ngroup's organization, falling back to the Everyone group when no budget\napplies. Null when the effective group belongs to a different organization\nthan the queried group.", "type": "string", "format": "uuid" }, @@ -23932,7 +23932,7 @@ "type": "integer" }, "effective_group_id": { - "description": "EffectiveGroupID is the group the spend is attributed to. Null when\nno budget applies.", + "description": "EffectiveGroupID is the group the spend is attributed to, falling back to\nthe Everyone group when no budget applies. Null only when the user has no\norganization.", "type": "string", "format": "uuid" }, diff --git a/coderd/database/dbauthz/dbauthz.go b/coderd/database/dbauthz/dbauthz.go index 3161f3c3e9db0..3cb39f78791b6 100644 --- a/coderd/database/dbauthz/dbauthz.go +++ b/coderd/database/dbauthz/dbauthz.go @@ -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) } diff --git a/coderd/database/dbauthz/dbauthz_test.go b/coderd/database/dbauthz/dbauthz_test.go index c029cd3761146..d064aeb9bf612 100644 --- a/coderd/database/dbauthz/dbauthz_test.go +++ b/coderd/database/dbauthz/dbauthz_test.go @@ -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{}) diff --git a/coderd/database/dbmetrics/querymetrics.go b/coderd/database/dbmetrics/querymetrics.go index 57c742457f0e4..16bfc7c80f10a 100644 --- a/coderd/database/dbmetrics/querymetrics.go +++ b/coderd/database/dbmetrics/querymetrics.go @@ -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) diff --git a/coderd/database/dbmock/dbmock.go b/coderd/database/dbmock/dbmock.go index 9103654b6d3e2..453a8798dbb1d 100644 --- a/coderd/database/dbmock/dbmock.go +++ b/coderd/database/dbmock/dbmock.go @@ -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() diff --git a/coderd/database/querier.go b/coderd/database/querier.go index 5c555cdd86f48..80b5c3e010882 100644 --- a/coderd/database/querier.go +++ b/coderd/database/querier.go @@ -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 organization name then 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". 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 organization name ascending. 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. diff --git a/coderd/database/querier_test.go b/coderd/database/querier_test.go index 53a477e9e6624..c94e8312b000f 100644 --- a/coderd/database/querier_test.go +++ b/coderd/database/querier_test.go @@ -13431,6 +13431,45 @@ 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("SpendWithDifferentEffectiveGroup", func(t *testing.T) { t.Parallel() db, _ := dbtestutil.NewDB(t) @@ -13686,6 +13725,95 @@ func TestGetGroupMembersAISpend(t *testing.T) { }) } +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 when another org's name sorts first. + 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{Name: "0-sorts-first"}) + user := dbgen.User(t, db, database.User{}) + dbgen.OrganizationMember(t, db, database.OrganizationMember{OrganizationID: defaultOrg.ID, UserID: user.ID}) + dbgen.OrganizationMember(t, db, database.OrganizationMember{OrganizationID: otherOrg.ID, UserID: user.ID}) + return user.ID, defaultOrg.ID + }, + }, + { + // Among non-default orgs, ties break by org name ascending. + name: "TieByOrgName", + setup: func(t *testing.T, ctx context.Context, db database.Store) (uuid.UUID, uuid.UUID) { + user := dbgen.User(t, db, database.User{}) + orgA := dbgen.Organization(t, db, database.Organization{Name: "org-a"}) + orgB := dbgen.Organization(t, db, database.Organization{Name: "org-b"}) + dbgen.OrganizationMember(t, db, database.OrganizationMember{OrganizationID: orgA.ID, UserID: user.ID}) + dbgen.OrganizationMember(t, db, database.OrganizationMember{OrganizationID: orgB.ID, UserID: user.ID}) + return user.ID, orgA.ID + }, + }, + { + // A soft-deleted org is excluded even when its name sorts first. + 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: "0-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}) + 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() { diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go index 876c952fc5975..a16487581912c 100644 --- a/coderd/database/queries.sql.go +++ b/coderd/database/queries.sql.go @@ -2594,15 +2594,29 @@ 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 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 + ORDER BY member.user_id, budget.spend_limit_micros DESC, organizations.name ASC, member.group_name 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 organization name ascending. + 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, organizations.name 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 +2625,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 +2678,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 +2718,16 @@ 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 +WHERE member.user_id = $1 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 + organizations.name ASC, -- organization name tiebreak + member.group_name ASC -- group name tiebreak LIMIT 1 ` @@ -2724,11 +2737,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 organization name then 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". func (q *sqlQuerier) GetHighestGroupAIBudgetByUser(ctx context.Context, userID uuid.UUID) (GetHighestGroupAIBudgetByUserRow, error) { row := q.db.QueryRowContext(ctx, getHighestGroupAIBudgetByUser, userID) var i GetHighestGroupAIBudgetByUserRow @@ -2856,6 +2869,29 @@ 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 + organizations.name ASC -- organization name 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 organization name ascending. 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) diff --git a/coderd/database/queries/aicostcontrol.sql b/coderd/database/queries/aicostcontrol.sql index c15a63d55fa32..4f48f558567b7 100644 --- a/coderd/database/queries/aicostcontrol.sql +++ b/coderd/database/queries/aicostcontrol.sql @@ -60,24 +60,37 @@ 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 organization name then 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". 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 +WHERE member.user_id = @user_id 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 + organizations.name ASC, -- organization name tiebreak + member.group_name ASC -- group name 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 organization name ascending. 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 + organizations.name ASC -- organization name tiebreak LIMIT 1; -- name: IncrementUserAIDailySpend :one @@ -126,9 +139,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 +167,29 @@ 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 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 + ORDER BY member.user_id, budget.spend_limit_micros DESC, organizations.name ASC, member.group_name 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 organization name ascending. + 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, organizations.name 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 +198,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 diff --git a/codersdk/aibridge.go b/codersdk/aibridge.go index 4598ac3d55535..ebbf7a2719962 100644 --- a/codersdk/aibridge.go +++ b/codersdk/aibridge.go @@ -34,12 +34,14 @@ type AIGroupBudget struct { LimitSource AIBudgetLimitSource `json:"limit_source"` } -// UserAIBudgetSummary is the effective AI budget for a user. When no -// budget applies, all fields except UserID are null. +// UserAIBudgetSummary is the effective AI budget for a user. When no budget +// applies, the effective group falls back to the Everyone group with a null +// limit and source. type UserAIBudgetSummary struct { UserID uuid.UUID `json:"user_id" format:"uuid"` - // EffectiveGroupID is the group the spend is attributed to. Null when - // no budget applies. + // EffectiveGroupID is the group the spend is attributed to, falling back to + // the Everyone group when no budget applies. Null only when the user has no + // organization. EffectiveGroupID *uuid.UUID `json:"effective_group_id" format:"uuid"` // SpendLimitMicros is the effective spend limit in micro-units. // Null when no budget applies to the user (unlimited). @@ -101,9 +103,9 @@ type GroupMembersAISpend struct { type GroupMemberAISpend struct { UserID uuid.UUID `json:"user_id" format:"uuid"` // EffectiveGroupID is the user's effective budget group within the queried - // group's organization. Null when no effective budget group is visible in - // this organization, including when the user's budget resolves to a group - // in another organization. + // group's organization, falling back to the Everyone group when no budget + // applies. Null when the effective group belongs to a different organization + // than the queried group. EffectiveGroupID *uuid.UUID `json:"effective_group_id" format:"uuid"` // GroupBudget is the budget when the queried group is this user's // effective budget source. Null when the user's budget resolves to another diff --git a/docs/reference/api/schemas.md b/docs/reference/api/schemas.md index ddba40ad34a87..0f2fbbfd0dca2 100644 --- a/docs/reference/api/schemas.md +++ b/docs/reference/api/schemas.md @@ -7742,12 +7742,12 @@ Only certain features set these fields: - FeatureManagedAgentLimit| ### Properties -| Name | Type | Required | Restrictions | Description | -|----------------------|--------------------------------------------------|----------|--------------|-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------| -| `effective_group_id` | string | false | | Effective group ID is the user's effective budget group within the queried group's organization. Null when no effective budget group is visible in this organization, including when the user's budget resolves to a group in another organization. | -| `group_budget` | [codersdk.AIGroupBudget](#codersdkaigroupbudget) | false | | Group budget is the budget when the queried group is this user's effective budget source. Null when the user's budget resolves to another group or no budget applies to the user. | -| `group_spend_micros` | integer | false | | Group spend micros is the user's spend attributed to the queried group over the current budget period. | -| `user_id` | string | false | | | +| Name | Type | Required | Restrictions | Description | +|----------------------|--------------------------------------------------|----------|--------------|-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------| +| `effective_group_id` | string | false | | Effective group ID is the user's effective budget group within the queried group's organization, falling back to the Everyone group when no budget applies. Null when the effective group belongs to a different organization than the queried group. | +| `group_budget` | [codersdk.AIGroupBudget](#codersdkaigroupbudget) | false | | Group budget is the budget when the queried group is this user's effective budget source. Null when the user's budget resolves to another group or no budget applies to the user. | +| `group_spend_micros` | integer | false | | Group spend micros is the user's spend attributed to the queried group over the current budget period. | +| `user_id` | string | false | | | ## codersdk.GroupMembersAISpend @@ -14110,15 +14110,15 @@ If the schedule is empty, the user will be updated to use the default schedule.| ### Properties -| Name | Type | Required | Restrictions | Description | -|------------------------|--------------------------------------------------------------|----------|--------------|----------------------------------------------------------------------------------------------------------------------| -| `current_spend_micros` | integer | false | | Current spend micros is the user's spend on their effective group over the current budget period. | -| `effective_group_id` | string | false | | Effective group ID is the group the spend is attributed to. Null when no budget applies. | -| `limit_source` | [codersdk.AIBudgetLimitSource](#codersdkaibudgetlimitsource) | false | | Limit source identifies which tier produced the limit. Null when no budget applies. | -| `period_end` | string | false | | Period end is the exclusive upper bound of the current budget period. | -| `period_start` | string | false | | Period start is the inclusive lower bound of the current budget period. | -| `spend_limit_micros` | integer | false | | Spend limit micros is the effective spend limit in micro-units. Null when no budget applies to the user (unlimited). | -| `user_id` | string | false | | | +| Name | Type | Required | Restrictions | Description | +|------------------------|--------------------------------------------------------------|----------|--------------|---------------------------------------------------------------------------------------------------------------------------------------------------------------------| +| `current_spend_micros` | integer | false | | Current spend micros is the user's spend on their effective group over the current budget period. | +| `effective_group_id` | string | false | | Effective group ID is the group the spend is attributed to, falling back to the Everyone group when no budget applies. Null only when the user has no organization. | +| `limit_source` | [codersdk.AIBudgetLimitSource](#codersdkaibudgetlimitsource) | false | | Limit source identifies which tier produced the limit. Null when no budget applies. | +| `period_end` | string | false | | Period end is the exclusive upper bound of the current budget period. | +| `period_start` | string | false | | Period start is the inclusive lower bound of the current budget period. | +| `spend_limit_micros` | integer | false | | Spend limit micros is the effective spend limit in micro-units. Null when no budget applies to the user (unlimited). | +| `user_id` | string | false | | | ## codersdk.UserActivity diff --git a/enterprise/coderd/aibridge.go b/enterprise/coderd/aibridge.go index 49ba9764b5d95..f07b6bb18d061 100644 --- a/enterprise/coderd/aibridge.go +++ b/enterprise/coderd/aibridge.go @@ -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 { diff --git a/enterprise/coderd/aibridge_test.go b/enterprise/coderd/aibridge_test.go index aa975b7da7e5f..f4fb5392bc6f6 100644 --- a/enterprise/coderd/aibridge_test.go +++ b/enterprise/coderd/aibridge_test.go @@ -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) @@ -3997,7 +4059,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) }) diff --git a/site/src/api/typesGenerated.ts b/site/src/api/typesGenerated.ts index 8a4c22d8dbb02..22945999560a1 100644 --- a/site/src/api/typesGenerated.ts +++ b/site/src/api/typesGenerated.ts @@ -5213,9 +5213,9 @@ export interface GroupMemberAISpend { readonly user_id: string; /** * EffectiveGroupID is the user's effective budget group within the queried - * group's organization. Null when no effective budget group is visible in - * this organization, including when the user's budget resolves to a group - * in another organization. + * group's organization, falling back to the Everyone group when no budget + * applies. Null when the effective group belongs to a different organization + * than the queried group. */ readonly effective_group_id: string | null; /** @@ -9907,14 +9907,16 @@ export interface UserAIBudgetOverride { // From codersdk/aibridge.go /** - * UserAIBudgetSummary is the effective AI budget for a user. When no - * budget applies, all fields except UserID are null. + * UserAIBudgetSummary is the effective AI budget for a user. When no budget + * applies, the effective group falls back to the Everyone group with a null + * limit and source. */ export interface UserAIBudgetSummary { readonly user_id: string; /** - * EffectiveGroupID is the group the spend is attributed to. Null when - * no budget applies. + * EffectiveGroupID is the group the spend is attributed to, falling back to + * the Everyone group when no budget applies. Null only when the user has no + * organization. */ readonly effective_group_id: string | null; /** From 5ec3a66f70a9a397344c3d5c20690528586eae0d Mon Sep 17 00:00:00 2001 From: Susana Cardoso Ferreira Date: Wed, 22 Jul 2026 10:11:30 +0000 Subject: [PATCH 2/4] chore: address review comments --- coderd/aibridge/budget/budget.go | 16 ++- coderd/aibridge/budget/budget_test.go | 21 ++- coderd/aibridgedserver/aibridgedserver.go | 3 + .../aibridgedserver/aibridgedserver_test.go | 11 +- coderd/aibridgedserver/cost.go | 8 +- coderd/apidoc/docs.go | 2 +- coderd/apidoc/swagger.json | 2 +- coderd/database/querier_test.go | 123 ++++++++++++++++++ coderd/database/queries.sql.go | 2 + coderd/database/queries/aicostcontrol.sql | 2 + codersdk/aibridge.go | 2 +- docs/reference/api/schemas.md | 18 +-- enterprise/coderd/aibridge_test.go | 59 +++++++++ site/src/api/typesGenerated.ts | 2 +- 14 files changed, 240 insertions(+), 31 deletions(-) diff --git a/coderd/aibridge/budget/budget.go b/coderd/aibridge/budget/budget.go index 61a4c793173b3..dc7d3acc3e0f6 100644 --- a/coderd/aibridge/budget/budget.go +++ b/coderd/aibridge/budget/budget.go @@ -28,20 +28,24 @@ type Store interface { type EffectiveGroup struct { // GroupID is the group the spend is attributed to. GroupID uuid.UUID - Limit *Limit + // Limit is the resolved spend limit, or nil for the unlimited Everyone + // fallback. + Limit *Limit } -// Limit is an AI spend limit and the tier that produced it. +// Limit is an AI spend limit and the source that produced it. type Limit struct { // SpendLimitMicros is the spend limit in micro-units (1 unit = 1,000,000). SpendLimitMicros int64 Source codersdk.AIBudgetLimitSource } -// ResolveUserAIBudget returns the effective AI budget group for userID. The -// second return value is false when no budget is configured for the user. A -// per-user override wins unconditionally; otherwise the budget is selected from -// the user's groups according to policy. +// ResolveUserAIBudget returns the effective AI budget group for userID, +// resolved in order: +// 1. A per-user override, if configured. +// 2. Otherwise, a group budget selected by the deployment policy. +// +// The second return value is false when no budget is configured for the user. // TODO(AIGOV-527): unify effective group resolution in a single place. func ResolveUserAIBudget(ctx context.Context, db Store, userID uuid.UUID, policy codersdk.AIBudgetPolicy) (EffectiveGroup, bool, error) { // A per-user override always wins. diff --git a/coderd/aibridge/budget/budget_test.go b/coderd/aibridge/budget/budget_test.go index b931278bd54ba..07c5731becc84 100644 --- a/coderd/aibridge/budget/budget_test.go +++ b/coderd/aibridge/budget/budget_test.go @@ -204,9 +204,10 @@ func TestResolveUserEffectiveGroup(t *testing.T) { t.Parallel() tests := []struct { - name string - policy codersdk.AIBudgetPolicy - setup func(t *testing.T, ctx context.Context, db database.Store) (userID uuid.UUID, want budget.EffectiveGroup, wantOK bool) + name string + policy codersdk.AIBudgetPolicy + setup func(t *testing.T, ctx context.Context, db database.Store) (userID uuid.UUID, want budget.EffectiveGroup, wantOK bool) + wantErr string }{ { // The Everyone group has a budget, so it resolves via the budget @@ -277,6 +278,16 @@ func TestResolveUserEffectiveGroup(t *testing.T) { return user.ID, budget.EffectiveGroup{}, false }, }, + { + // An unsupported policy surfaces the error from ResolveUserAIBudget. + name: "UnsupportedPolicy", + policy: codersdk.AIBudgetPolicy("unsupported"), + setup: func(t *testing.T, ctx context.Context, db database.Store) (uuid.UUID, budget.EffectiveGroup, bool) { + user := dbgen.User(t, db, database.User{}) + return user.ID, budget.EffectiveGroup{}, false + }, + wantErr: "unsupported AI budget policy", + }, } for _, tt := range tests { @@ -288,6 +299,10 @@ func TestResolveUserEffectiveGroup(t *testing.T) { userID, want, wantOK := tt.setup(t, ctx, db) got, ok, err := budget.ResolveUserEffectiveGroup(ctx, db, userID, tt.policy) + if tt.wantErr != "" { + require.ErrorContains(t, err, tt.wantErr) + return + } require.NoError(t, err) require.Equal(t, wantOK, ok) if !wantOK { diff --git a/coderd/aibridgedserver/aibridgedserver.go b/coderd/aibridgedserver/aibridgedserver.go index 6a0f68aab43fd..cf750944e93c8 100644 --- a/coderd/aibridgedserver/aibridgedserver.go +++ b/coderd/aibridgedserver/aibridgedserver.go @@ -817,6 +817,9 @@ func (s *Server) checkUserAIBudget(ctx context.Context, userID uuid.UUID, period if err != nil { return userAIBudget{}, xerrors.Errorf("resolve effective AI budget for user %q with budget policy %q: %w", userID, s.budgetPolicy, err) } + // ok is false when no budget is configured. The nil Limit check keeps + // enforcement failing open if a caller resolves via the unlimited + // Everyone fallback. if !ok || effectiveGroup.Limit == nil { // No enforceable spend limit for the user; return zero-valued status. return userAIBudget{}, nil diff --git a/coderd/aibridgedserver/aibridgedserver_test.go b/coderd/aibridgedserver/aibridgedserver_test.go index fd983476400e9..ab6a479b05289 100644 --- a/coderd/aibridgedserver/aibridgedserver_test.go +++ b/coderd/aibridgedserver/aibridgedserver_test.go @@ -744,9 +744,8 @@ func TestIsBudgetExceeded_Enforcement(t *testing.T) { rawDB, _ := dbtestutil.NewDB(t) authzDB := dbauthz.New(rawDB, rbac.NewStrictAuthorizer(prometheus.NewRegistry()), logger, coderdtest.AccessControlStorePointer()) - // An org member with no group budget and no override. Their spend is - // unlimited, so enforcement never blocks them regardless of accrued - // spend, even though attribution would fall back to the Everyone group. + // An org member with no group budget and no override: spend is + // unlimited, so enforcement never blocks them. org := dbgen.Organization(t, rawDB, database.Organization{}) user := dbgen.User(t, rawDB, database.User{}) dbgen.OrganizationMember(t, rawDB, database.OrganizationMember{OrganizationID: org.ID, UserID: user.ID}) @@ -2025,9 +2024,9 @@ func TestRecordTokenUsage(t *testing.T) { }, }, { - // A user with no organization has no effective group. This is - // not expected in practice. Spend is still recorded, but with a - // NULL group, and the daily spend update is skipped. + // A user with no organization has no effective group. Spend is + // still recorded, but with a NULL group, and the daily spend + // update is skipped. name: "valid token usage with no effective group", request: &proto.RecordTokenUsageRequest{ InterceptionId: uuid.NewString(), diff --git a/coderd/aibridgedserver/cost.go b/coderd/aibridgedserver/cost.go index 416b03ea7c289..87ba4ab64643a 100644 --- a/coderd/aibridgedserver/cost.go +++ b/coderd/aibridgedserver/cost.go @@ -31,9 +31,11 @@ type tokenUsageCost struct { } // resolveTokenUsageCost resolves the effective group and per-token prices for an -// interception and computes its cost. A model absent from the price table yields -// NULL prices and cost rather than an error (a NULL cost unambiguously means -// "model not priced"). Any other error is returned. +// interception and computes its cost. Two independent conditions yield a NULL +// column rather than an error: an unresolved effective group (the user has no +// org membership), and a model absent from the price table leaves prices and +// cost NULL (a NULL cost unambiguously means "model not priced"). +// Any other error is returned. func (s *Server) resolveTokenUsageCost(ctx context.Context, intc database.AIBridgeInterception, in *proto.RecordTokenUsageRequest) (tokenUsageCost, error) { var result tokenUsageCost diff --git a/coderd/apidoc/docs.go b/coderd/apidoc/docs.go index 8e01936e88737..69411ed4491f5 100644 --- a/coderd/apidoc/docs.go +++ b/coderd/apidoc/docs.go @@ -26027,7 +26027,7 @@ const docTemplate = `{ "type": "integer" }, "effective_group_id": { - "description": "EffectiveGroupID is the group the spend is attributed to, falling back to\nthe Everyone group when no budget applies. Null only when the user has no\norganization.", + "description": "EffectiveGroupID is the group the spend is attributed to, falling back to\nthe Everyone group when no budget applies. Null only when the user has no\norganization membership.", "type": "string", "format": "uuid" }, diff --git a/coderd/apidoc/swagger.json b/coderd/apidoc/swagger.json index 0c47c56afe180..57d188b92b23f 100644 --- a/coderd/apidoc/swagger.json +++ b/coderd/apidoc/swagger.json @@ -23932,7 +23932,7 @@ "type": "integer" }, "effective_group_id": { - "description": "EffectiveGroupID is the group the spend is attributed to, falling back to\nthe Everyone group when no budget applies. Null only when the user has no\norganization.", + "description": "EffectiveGroupID is the group the spend is attributed to, falling back to\nthe Everyone group when no budget applies. Null only when the user has no\norganization membership.", "type": "string", "format": "uuid" }, diff --git a/coderd/database/querier_test.go b/coderd/database/querier_test.go index c94e8312b000f..05a0020f77838 100644 --- a/coderd/database/querier_test.go +++ b/coderd/database/querier_test.go @@ -13725,6 +13725,129 @@ 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 org name ascending. + name: "TieByOrgName", + 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{Name: "org-a"}) + orgB := dbgen.Organization(t, db, database.Organization{Name: "org-b"}) + groupA := dbgen.Group(t, db, database.Group{OrganizationID: orgA.ID}) + groupB := 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: groupA.ID, UserID: user.ID}) + dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: groupB.ID, UserID: user.ID}) + _, err := db.UpsertGroupAIBudget(ctx, database.UpsertGroupAIBudgetParams{GroupID: groupA.ID, SpendLimitMicros: 1_000_000}) + require.NoError(t, err) + _, err = db.UpsertGroupAIBudget(ctx, database.UpsertGroupAIBudgetParams{GroupID: groupB.ID, SpendLimitMicros: 1_000_000}) + require.NoError(t, err) + return user.ID, database.GetHighestGroupAIBudgetByUserRow{GroupID: groupA.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() diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go index a16487581912c..1a0e00bde8021 100644 --- a/coderd/database/queries.sql.go +++ b/coderd/database/queries.sql.go @@ -2596,6 +2596,7 @@ user_highest_group AS ( JOIN group_members_expanded member ON member.group_id = budget.group_id JOIN organizations ON organizations.id = member.organization_id WHERE member.user_id IN (SELECT user_id FROM filtered_users) + AND organizations.deleted = false ORDER BY member.user_id, budget.spend_limit_micros DESC, organizations.name ASC, member.group_name ASC ), user_fallback_group AS ( @@ -2724,6 +2725,7 @@ 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 WHERE member.user_id = $1 + AND organizations.deleted = false ORDER BY budget.spend_limit_micros DESC, -- highest wins organizations.name ASC, -- organization name tiebreak diff --git a/coderd/database/queries/aicostcontrol.sql b/coderd/database/queries/aicostcontrol.sql index 4f48f558567b7..1f06d3b2ad77a 100644 --- a/coderd/database/queries/aicostcontrol.sql +++ b/coderd/database/queries/aicostcontrol.sql @@ -72,6 +72,7 @@ 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 WHERE member.user_id = @user_id + AND organizations.deleted = false ORDER BY budget.spend_limit_micros DESC, -- highest wins organizations.name ASC, -- organization name tiebreak @@ -169,6 +170,7 @@ user_highest_group AS ( JOIN group_members_expanded member ON member.group_id = budget.group_id JOIN organizations ON organizations.id = member.organization_id WHERE member.user_id IN (SELECT user_id FROM filtered_users) + AND organizations.deleted = false ORDER BY member.user_id, budget.spend_limit_micros DESC, organizations.name ASC, member.group_name ASC ), user_fallback_group AS ( diff --git a/codersdk/aibridge.go b/codersdk/aibridge.go index ebbf7a2719962..34572240ab195 100644 --- a/codersdk/aibridge.go +++ b/codersdk/aibridge.go @@ -41,7 +41,7 @@ type UserAIBudgetSummary struct { UserID uuid.UUID `json:"user_id" format:"uuid"` // EffectiveGroupID is the group the spend is attributed to, falling back to // the Everyone group when no budget applies. Null only when the user has no - // organization. + // organization membership. EffectiveGroupID *uuid.UUID `json:"effective_group_id" format:"uuid"` // SpendLimitMicros is the effective spend limit in micro-units. // Null when no budget applies to the user (unlimited). diff --git a/docs/reference/api/schemas.md b/docs/reference/api/schemas.md index 0f2fbbfd0dca2..ca6b0fad0d4fb 100644 --- a/docs/reference/api/schemas.md +++ b/docs/reference/api/schemas.md @@ -14110,15 +14110,15 @@ If the schedule is empty, the user will be updated to use the default schedule.| ### Properties -| Name | Type | Required | Restrictions | Description | -|------------------------|--------------------------------------------------------------|----------|--------------|---------------------------------------------------------------------------------------------------------------------------------------------------------------------| -| `current_spend_micros` | integer | false | | Current spend micros is the user's spend on their effective group over the current budget period. | -| `effective_group_id` | string | false | | Effective group ID is the group the spend is attributed to, falling back to the Everyone group when no budget applies. Null only when the user has no organization. | -| `limit_source` | [codersdk.AIBudgetLimitSource](#codersdkaibudgetlimitsource) | false | | Limit source identifies which tier produced the limit. Null when no budget applies. | -| `period_end` | string | false | | Period end is the exclusive upper bound of the current budget period. | -| `period_start` | string | false | | Period start is the inclusive lower bound of the current budget period. | -| `spend_limit_micros` | integer | false | | Spend limit micros is the effective spend limit in micro-units. Null when no budget applies to the user (unlimited). | -| `user_id` | string | false | | | +| Name | Type | Required | Restrictions | Description | +|------------------------|--------------------------------------------------------------|----------|--------------|--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------| +| `current_spend_micros` | integer | false | | Current spend micros is the user's spend on their effective group over the current budget period. | +| `effective_group_id` | string | false | | Effective group ID is the group the spend is attributed to, falling back to the Everyone group when no budget applies. Null only when the user has no organization membership. | +| `limit_source` | [codersdk.AIBudgetLimitSource](#codersdkaibudgetlimitsource) | false | | Limit source identifies which tier produced the limit. Null when no budget applies. | +| `period_end` | string | false | | Period end is the exclusive upper bound of the current budget period. | +| `period_start` | string | false | | Period start is the inclusive lower bound of the current budget period. | +| `spend_limit_micros` | integer | false | | Spend limit micros is the effective spend limit in micro-units. Null when no budget applies to the user (unlimited). | +| `user_id` | string | false | | | ## codersdk.UserActivity diff --git a/enterprise/coderd/aibridge_test.go b/enterprise/coderd/aibridge_test.go index f4fb5392bc6f6..ca270ece8d492 100644 --- a/enterprise/coderd/aibridge_test.go +++ b/enterprise/coderd/aibridge_test.go @@ -4034,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() diff --git a/site/src/api/typesGenerated.ts b/site/src/api/typesGenerated.ts index 22945999560a1..4158bd880bc73 100644 --- a/site/src/api/typesGenerated.ts +++ b/site/src/api/typesGenerated.ts @@ -9916,7 +9916,7 @@ export interface UserAIBudgetSummary { /** * EffectiveGroupID is the group the spend is attributed to, falling back to * the Everyone group when no budget applies. Null only when the user has no - * organization. + * organization membership. */ readonly effective_group_id: string | null; /** From db534d26445312898eac2d2a8b94d66282537757 Mon Sep 17 00:00:00 2001 From: Susana Cardoso Ferreira Date: Wed, 22 Jul 2026 10:34:45 +0000 Subject: [PATCH 3/4] tests: add CrossOrgFallbackMasked test --- coderd/database/querier_test.go | 45 +++++++++++++++++++++++++++++++++ 1 file changed, 45 insertions(+) diff --git a/coderd/database/querier_test.go b/coderd/database/querier_test.go index 05a0020f77838..5c7d80ce5c3b0 100644 --- a/coderd/database/querier_test.go +++ b/coderd/database/querier_test.go @@ -13470,6 +13470,51 @@ func TestGetGroupMembersAISpend(t *testing.T) { 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 also belongs to + // another org whose name sorts first, so the fallback prefers that org's + // Everyone group over the queried group's org. + user := dbgen.User(t, db, database.User{}) + queriedOrg := dbgen.Organization(t, db, database.Organization{Name: "z-queried-org"}) + otherOrg := dbgen.Organization(t, db, database.Organization{Name: "a-other-org"}) + queried := dbgen.Group(t, db, database.Group{OrganizationID: queriedOrg.ID}) + dbgen.OrganizationMember(t, db, database.OrganizationMember{UserID: user.ID, OrganizationID: queriedOrg.ID}) + dbgen.OrganizationMember(t, db, database.OrganizationMember{UserID: user.ID, OrganizationID: otherOrg.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) From d42be831c9dd5aeef6d35add039423031563de1e Mon Sep 17 00:00:00 2001 From: Susana Cardoso Ferreira Date: Wed, 22 Jul 2026 17:48:06 +0000 Subject: [PATCH 4/4] chore: update AI budget ties by earliest org membership --- coderd/aibridge/budget/budget_test.go | 63 ++++++++++-------- coderd/database/querier.go | 6 +- coderd/database/querier_test.go | 81 +++++++++++++---------- coderd/database/queries.sql.go | 29 +++++--- coderd/database/queries/aicostcontrol.sql | 29 +++++--- 5 files changed, 121 insertions(+), 87 deletions(-) diff --git a/coderd/aibridge/budget/budget_test.go b/coderd/aibridge/budget/budget_test.go index 07c5731becc84..82ca9c6ef9ea8 100644 --- a/coderd/aibridge/budget/budget_test.go +++ b/coderd/aibridge/budget/budget_test.go @@ -1,8 +1,10 @@ package budget_test import ( + "bytes" "context" "testing" + "time" "github.com/google/uuid" "github.com/stretchr/testify/require" @@ -80,6 +82,7 @@ func TestResolveUserAIBudget(t *testing.T) { setup: func(t *testing.T, ctx context.Context, db database.Store) (uuid.UUID, budget.EffectiveGroup, bool) { org := dbgen.Organization(t, db, database.Organization{}) user := dbgen.User(t, db, database.User{}) + dbgen.OrganizationMember(t, db, database.OrganizationMember{OrganizationID: org.ID, UserID: user.ID}) gid := budgetedGroup(t, ctx, db, org.ID, user.ID, "only", 8_000_000) return user.ID, budget.EffectiveGroup{GroupID: gid, Limit: &budget.Limit{SpendLimitMicros: 8_000_000, Source: codersdk.AIBudgetLimitSourceGroup}}, true }, @@ -90,6 +93,7 @@ func TestResolveUserAIBudget(t *testing.T) { setup: func(t *testing.T, ctx context.Context, db database.Store) (uuid.UUID, budget.EffectiveGroup, bool) { org := dbgen.Organization(t, db, database.Organization{}) user := dbgen.User(t, db, database.User{}) + dbgen.OrganizationMember(t, db, database.OrganizationMember{OrganizationID: org.ID, UserID: user.ID}) budgetedGroup(t, ctx, db, org.ID, user.ID, "low", 5_000_000) budgetedGroup(t, ctx, db, org.ID, user.ID, "mid", 20_000_000) high := budgetedGroup(t, ctx, db, org.ID, user.ID, "high", 50_000_000) @@ -97,29 +101,36 @@ func TestResolveUserAIBudget(t *testing.T) { }, }, { - name: "TieBrokenByGroupName", + name: "TieBrokenByEarliestOrgMembership", policy: codersdk.AIBudgetPolicyHighest, setup: func(t *testing.T, ctx context.Context, db database.Store) (uuid.UUID, budget.EffectiveGroup, bool) { - org := dbgen.Organization(t, db, database.Organization{}) user := dbgen.User(t, db, database.User{}) - // Equal limits in the same org, so "alpha" wins over "beta" by group - // name ascending. - alpha := budgetedGroup(t, ctx, db, org.ID, user.ID, "alpha", 10_000_000) - budgetedGroup(t, ctx, db, org.ID, user.ID, "beta", 10_000_000) - return user.ID, budget.EffectiveGroup{GroupID: alpha, Limit: &budget.Limit{SpendLimitMicros: 10_000_000, Source: codersdk.AIBudgetLimitSourceGroup}}, true + // Two groups in different orgs share the same limit. The earlier + // organization membership breaks the tie. + earlyOrg := dbgen.Organization(t, db, database.Organization{}) + lateOrg := dbgen.Organization(t, db, database.Organization{}) + dbgen.OrganizationMember(t, db, database.OrganizationMember{OrganizationID: earlyOrg.ID, UserID: user.ID, CreatedAt: time.Now().Add(-time.Hour)}) + dbgen.OrganizationMember(t, db, database.OrganizationMember{OrganizationID: lateOrg.ID, UserID: user.ID}) + winner := budgetedGroup(t, ctx, db, earlyOrg.ID, user.ID, "dup", 10_000_000) + budgetedGroup(t, ctx, db, lateOrg.ID, user.ID, "dup", 10_000_000) + return user.ID, budget.EffectiveGroup{GroupID: winner, Limit: &budget.Limit{SpendLimitMicros: 10_000_000, Source: codersdk.AIBudgetLimitSourceGroup}}, true }, }, { - name: "TieBrokenByOrgName", + name: "TieBrokenByGroupID", policy: codersdk.AIBudgetPolicyHighest, setup: func(t *testing.T, ctx context.Context, db database.Store) (uuid.UUID, budget.EffectiveGroup, bool) { + org := dbgen.Organization(t, db, database.Organization{}) user := dbgen.User(t, db, database.User{}) - // Two groups in different orgs share both name and limit. The - // organization name breaks the tie ascending, so "org-a" wins. - orgA := dbgen.Organization(t, db, database.Organization{Name: "org-a"}) - orgB := dbgen.Organization(t, db, database.Organization{Name: "org-b"}) - winner := budgetedGroup(t, ctx, db, orgA.ID, user.ID, "dup", 10_000_000) - budgetedGroup(t, ctx, db, orgB.ID, user.ID, "dup", 10_000_000) + dbgen.OrganizationMember(t, db, database.OrganizationMember{OrganizationID: org.ID, UserID: user.ID}) + // Both groups are in the same org, so both resolve to the same + // organization membership and the tie falls to the lowest group ID. + groupA := budgetedGroup(t, ctx, db, org.ID, user.ID, "alpha", 10_000_000) + groupB := budgetedGroup(t, ctx, db, org.ID, user.ID, "beta", 10_000_000) + winner := groupA + if bytes.Compare(groupB[:], groupA[:]) < 0 { + winner = groupB + } return user.ID, budget.EffectiveGroup{GroupID: winner, Limit: &budget.Limit{SpendLimitMicros: 10_000_000, Source: codersdk.AIBudgetLimitSourceGroup}}, true }, }, @@ -241,32 +252,32 @@ func TestResolveUserEffectiveGroup(t *testing.T) { }, }, { - // The fallback prefers the default org even over an org whose name - // sorts first. + // The fallback prefers the default org even over an org joined + // earlier. name: "FallbackPrefersDefaultOrg", policy: codersdk.AIBudgetPolicyHighest, setup: func(t *testing.T, ctx context.Context, db database.Store) (uuid.UUID, budget.EffectiveGroup, bool) { defaultOrg, err := db.GetDefaultOrganization(ctx) require.NoError(t, err) - otherOrg := dbgen.Organization(t, db, database.Organization{Name: "0-sorts-first"}) + otherOrg := dbgen.Organization(t, db, database.Organization{}) user := dbgen.User(t, db, database.User{}) + dbgen.OrganizationMember(t, db, database.OrganizationMember{OrganizationID: otherOrg.ID, UserID: user.ID, CreatedAt: time.Now().Add(-time.Hour)}) dbgen.OrganizationMember(t, db, database.OrganizationMember{OrganizationID: defaultOrg.ID, UserID: user.ID}) - dbgen.OrganizationMember(t, db, database.OrganizationMember{OrganizationID: otherOrg.ID, UserID: user.ID}) return user.ID, budget.EffectiveGroup{GroupID: defaultOrg.ID}, true }, }, { - // Among non-default orgs, the fallback breaks ties by org name - // ascending, so "org-a" wins. - name: "FallbackTieByOrgName", + // Among non-default orgs, the fallback breaks ties by the earliest + // organization membership. + name: "FallbackTieByEarliestOrgMembership", policy: codersdk.AIBudgetPolicyHighest, setup: func(t *testing.T, ctx context.Context, db database.Store) (uuid.UUID, budget.EffectiveGroup, bool) { user := dbgen.User(t, db, database.User{}) - orgA := dbgen.Organization(t, db, database.Organization{Name: "org-a"}) - orgB := dbgen.Organization(t, db, database.Organization{Name: "org-b"}) - dbgen.OrganizationMember(t, db, database.OrganizationMember{OrganizationID: orgA.ID, UserID: user.ID}) - dbgen.OrganizationMember(t, db, database.OrganizationMember{OrganizationID: orgB.ID, UserID: user.ID}) - return user.ID, budget.EffectiveGroup{GroupID: orgA.ID}, true + earlyOrg := dbgen.Organization(t, db, database.Organization{}) + lateOrg := dbgen.Organization(t, db, database.Organization{}) + dbgen.OrganizationMember(t, db, database.OrganizationMember{OrganizationID: earlyOrg.ID, UserID: user.ID, CreatedAt: time.Now().Add(-time.Hour)}) + dbgen.OrganizationMember(t, db, database.OrganizationMember{OrganizationID: lateOrg.ID, UserID: user.ID}) + return user.ID, budget.EffectiveGroup{GroupID: earlyOrg.ID}, true }, }, { diff --git a/coderd/database/querier.go b/coderd/database/querier.go index 80b5c3e010882..8e9b909de432a 100644 --- a/coderd/database/querier.go +++ b/coderd/database/querier.go @@ -604,7 +604,7 @@ 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 organization name then group name ascending. Implements the + // 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 @@ -877,8 +877,8 @@ type sqlcQuerier interface { 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 organization name ascending. Returns no rows when the user has no - // organization membership. + // 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. diff --git a/coderd/database/querier_test.go b/coderd/database/querier_test.go index 5c7d80ce5c3b0..d0467a86b3ee2 100644 --- a/coderd/database/querier_test.go +++ b/coderd/database/querier_test.go @@ -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) @@ -13475,15 +13483,14 @@ func TestGetGroupMembersAISpend(t *testing.T) { db, _ := dbtestutil.NewDB(t) ctx := testutil.Context(t, testutil.WaitShort) - // Given: an unbudgeted member of the queried group who also belongs to - // another org whose name sorts first, so the fallback prefers that org's - // Everyone group over the queried group's org. + // 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{Name: "z-queried-org"}) - otherOrg := dbgen.Organization(t, db, database.Organization{Name: "a-other-org"}) + 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.OrganizationMember(t, db, database.OrganizationMember{UserID: user.ID, OrganizationID: otherOrg.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. @@ -13843,23 +13850,24 @@ func TestGetHighestGroupAIBudgetByUser(t *testing.T) { }, }, { - // Equal limits across orgs break by org name ascending. - name: "TieByOrgName", + // 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{}) - orgA := dbgen.Organization(t, db, database.Organization{Name: "org-a"}) - orgB := dbgen.Organization(t, db, database.Organization{Name: "org-b"}) - groupA := dbgen.Group(t, db, database.Group{OrganizationID: orgA.ID}) - groupB := 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: groupA.ID, UserID: user.ID}) - dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: groupB.ID, UserID: user.ID}) - _, err := db.UpsertGroupAIBudget(ctx, database.UpsertGroupAIBudgetParams{GroupID: groupA.ID, SpendLimitMicros: 1_000_000}) + 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: groupB.ID, SpendLimitMicros: 1_000_000}) + _, err = db.UpsertGroupAIBudget(ctx, database.UpsertGroupAIBudgetParams{GroupID: lateGroup.ID, SpendLimitMicros: 1_000_000}) require.NoError(t, err) - return user.ID, database.GetHighestGroupAIBudgetByUserRow{GroupID: groupA.ID, SpendLimitMicros: 1_000_000} + return user.ID, database.GetHighestGroupAIBudgetByUserRow{GroupID: earlyGroup.ID, SpendLimitMicros: 1_000_000} }, }, { @@ -13912,39 +13920,40 @@ func TestGetUserEveryoneFallbackGroup(t *testing.T) { }, }, { - // The default org is preferred even when another org's name sorts first. + // 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{Name: "0-sorts-first"}) + 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}) - dbgen.OrganizationMember(t, db, database.OrganizationMember{OrganizationID: otherOrg.ID, UserID: user.ID}) return user.ID, defaultOrg.ID }, }, { - // Among non-default orgs, ties break by org name ascending. - name: "TieByOrgName", + // 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{}) - orgA := dbgen.Organization(t, db, database.Organization{Name: "org-a"}) - orgB := dbgen.Organization(t, db, database.Organization{Name: "org-b"}) - dbgen.OrganizationMember(t, db, database.OrganizationMember{OrganizationID: orgA.ID, UserID: user.ID}) - dbgen.OrganizationMember(t, db, database.OrganizationMember{OrganizationID: orgB.ID, UserID: user.ID}) - return user.ID, orgA.ID + 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 its name sorts first. + // 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: "0-deleted-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}) + 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(), diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go index 1a0e00bde8021..76168e903634a 100644 --- a/coderd/database/queries.sql.go +++ b/coderd/database/queries.sql.go @@ -2595,14 +2595,17 @@ user_highest_group AS ( 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) AND organizations.deleted = false - ORDER BY member.user_id, budget.spend_limit_micros DESC, organizations.name ASC, member.group_name ASC + 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 organization name ascending. + -- default org, then the earliest organization membership. SELECT DISTINCT ON (organization_members.user_id) organization_members.user_id, organizations.id AS group_id @@ -2610,7 +2613,7 @@ user_fallback_group AS ( 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, organizations.name ASC + 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 @@ -2724,12 +2727,15 @@ SELECT 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 - budget.spend_limit_micros DESC, -- highest wins - organizations.name ASC, -- organization name tiebreak - member.group_name ASC -- group name tiebreak + budget.spend_limit_micros DESC, -- highest wins + organization_members.created_at ASC, -- earliest organization membership + budget.group_id ASC -- deterministic tiebreak LIMIT 1 ` @@ -2739,7 +2745,7 @@ type GetHighestGroupAIBudgetByUserRow struct { } // Returns the highest group AI budget across the groups the user belongs to, -// breaking ties by organization name then group name ascending. Implements the +// 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 @@ -2878,15 +2884,16 @@ 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 - organizations.name ASC -- organization name tiebreak + 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 organization name ascending. Returns no rows when the user has no -// organization membership. +// 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 diff --git a/coderd/database/queries/aicostcontrol.sql b/coderd/database/queries/aicostcontrol.sql index 1f06d3b2ad77a..ece65a9eef613 100644 --- a/coderd/database/queries/aicostcontrol.sql +++ b/coderd/database/queries/aicostcontrol.sql @@ -60,7 +60,7 @@ 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 organization name then group name ascending. Implements the +-- 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 @@ -71,27 +71,31 @@ SELECT 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 - budget.spend_limit_micros DESC, -- highest wins - organizations.name ASC, -- organization name tiebreak - member.group_name ASC -- group name tiebreak + 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 organization name ascending. Returns no rows when the user has no --- organization membership. +-- 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 - organizations.name ASC -- organization name tiebreak + 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 @@ -169,14 +173,17 @@ user_highest_group AS ( 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) AND organizations.deleted = false - ORDER BY member.user_id, budget.spend_limit_micros DESC, organizations.name ASC, member.group_name ASC + 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 organization name ascending. + -- default org, then the earliest organization membership. SELECT DISTINCT ON (organization_members.user_id) organization_members.user_id, organizations.id AS group_id @@ -184,7 +191,7 @@ user_fallback_group AS ( 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, organizations.name ASC + 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