diff --git a/coderd/aibridge/budget/budget.go b/coderd/aibridge/budget/budget.go index 5b86cdbdd68..dc7d3acc3e0 100644 --- a/coderd/aibridge/budget/budget.go +++ b/coderd/aibridge/budget/budget.go @@ -20,37 +20,47 @@ import ( type Store interface { GetUserAIBudgetOverride(ctx context.Context, userID uuid.UUID) (database.UserAIBudgetOverride, error) GetHighestGroupAIBudgetByUser(ctx context.Context, userID uuid.UUID) (database.GetHighestGroupAIBudgetByUserRow, error) + GetUserEveryoneFallbackGroup(ctx context.Context, userID uuid.UUID) (uuid.UUID, error) } -// EffectiveBudget is the AI budget that applies to a user after override and -// policy resolution. -type EffectiveBudget struct { +// EffectiveGroup is a user's resolved effective group and, when a budget +// applies, its limit. Limit is nil for the Everyone fallback (unlimited). +type EffectiveGroup struct { // GroupID is the group the spend is attributed to. GroupID uuid.UUID - // SpendLimitMicros is the effective spend limit in micro-units - // (1 unit = 1,000,000). + // Limit is the resolved spend limit, or nil for the unlimited Everyone + // fallback. + Limit *Limit +} + +// Limit is an AI spend limit and the source that produced it. +type Limit struct { + // SpendLimitMicros is the spend limit in micro-units (1 unit = 1,000,000). SpendLimitMicros int64 Source codersdk.AIBudgetLimitSource } -// ResolveUserAIBudget returns the effective AI budget for userID. The second -// return value is false when no budget is configured for the user. A per-user -// override wins unconditionally; otherwise the budget is selected from the -// user's groups according to policy. +// ResolveUserAIBudget returns the effective AI budget group for userID, +// resolved in order: +// 1. A per-user override, if configured. +// 2. Otherwise, a group budget selected by the deployment policy. // +// The second return value is false when no budget is configured for the user. // TODO(AIGOV-527): unify effective group resolution in a single place. -func ResolveUserAIBudget(ctx context.Context, db Store, userID uuid.UUID, policy codersdk.AIBudgetPolicy) (EffectiveBudget, bool, error) { +func ResolveUserAIBudget(ctx context.Context, db Store, userID uuid.UUID, policy codersdk.AIBudgetPolicy) (EffectiveGroup, bool, error) { // A per-user override always wins. override, err := db.GetUserAIBudgetOverride(ctx, userID) if err == nil { - return EffectiveBudget{ - GroupID: override.GroupID, - SpendLimitMicros: override.SpendLimitMicros, - Source: codersdk.AIBudgetLimitSourceUserOverride, + return EffectiveGroup{ + GroupID: override.GroupID, + Limit: &Limit{ + SpendLimitMicros: override.SpendLimitMicros, + Source: codersdk.AIBudgetLimitSourceUserOverride, + }, }, true, nil } if !errors.Is(err, sql.ErrNoRows) { - return EffectiveBudget{}, false, xerrors.Errorf("get user AI budget override: %w", err) + return EffectiveGroup{}, false, xerrors.Errorf("get user AI budget override: %w", err) } // No override: select a group budget according to the deployment policy. @@ -58,17 +68,45 @@ func ResolveUserAIBudget(ctx context.Context, db Store, userID uuid.UUID, policy case codersdk.AIBudgetPolicyHighest: row, err := db.GetHighestGroupAIBudgetByUser(ctx, userID) if errors.Is(err, sql.ErrNoRows) { - return EffectiveBudget{}, false, nil + return EffectiveGroup{}, false, nil } if err != nil { - return EffectiveBudget{}, false, xerrors.Errorf("get highest group AI budget: %w", err) + return EffectiveGroup{}, false, xerrors.Errorf("get highest group AI budget: %w", err) } - return EffectiveBudget{ - GroupID: row.GroupID, - SpendLimitMicros: row.SpendLimitMicros, - Source: codersdk.AIBudgetLimitSourceGroup, + return EffectiveGroup{ + GroupID: row.GroupID, + Limit: &Limit{ + SpendLimitMicros: row.SpendLimitMicros, + Source: codersdk.AIBudgetLimitSourceGroup, + }, }, true, nil default: - return EffectiveBudget{}, false, xerrors.Errorf("unsupported AI budget policy: %q", policy) + return EffectiveGroup{}, false, xerrors.Errorf("unsupported AI budget policy: %q", policy) + } +} + +// ResolveUserEffectiveGroup resolves the user's effective group, falling back to +// the organization's Everyone group when no override or group budget applies. +// The second return value is false when no effective group was found for the +// user. +func ResolveUserEffectiveGroup(ctx context.Context, db Store, userID uuid.UUID, policy codersdk.AIBudgetPolicy) (EffectiveGroup, bool, error) { + group, ok, err := ResolveUserAIBudget(ctx, db, userID, policy) + if err != nil { + return EffectiveGroup{}, false, err + } + if ok { + return group, true, nil + } + + // No override or group budget: fall back to the Everyone group (unlimited). + groupID, err := db.GetUserEveryoneFallbackGroup(ctx, userID) + if errors.Is(err, sql.ErrNoRows) { + // This should not happen, as a user should always be a member of an + // organization and its associated Everyone group. + return EffectiveGroup{}, false, nil + } + if err != nil { + return EffectiveGroup{}, false, xerrors.Errorf("get everyone fallback group: %w", err) } + return EffectiveGroup{GroupID: groupID}, true, nil } diff --git a/coderd/aibridge/budget/budget_test.go b/coderd/aibridge/budget/budget_test.go index 4caaae02780..82ca9c6ef9e 100644 --- a/coderd/aibridge/budget/budget_test.go +++ b/coderd/aibridge/budget/budget_test.go @@ -4,6 +4,7 @@ import ( "bytes" "context" "testing" + "time" "github.com/google/uuid" "github.com/stretchr/testify/require" @@ -52,18 +53,18 @@ func TestResolveUserAIBudget(t *testing.T) { tests := []struct { name string policy codersdk.AIBudgetPolicy - setup func(t *testing.T, ctx context.Context, db database.Store) (userID uuid.UUID, want budget.EffectiveBudget, wantOK bool) + setup func(t *testing.T, ctx context.Context, db database.Store) (userID uuid.UUID, want budget.EffectiveGroup, wantOK bool) wantErr string }{ { name: "OverrideWins", policy: codersdk.AIBudgetPolicyHighest, - setup: func(t *testing.T, ctx context.Context, db database.Store) (uuid.UUID, budget.EffectiveBudget, bool) { + setup: func(t *testing.T, ctx context.Context, db database.Store) (uuid.UUID, budget.EffectiveGroup, bool) { org := dbgen.Organization(t, db, database.Organization{}) user := dbgen.User(t, db, database.User{}) // A higher group budget that the override must still beat. budgetedGroup(t, ctx, db, org.ID, user.ID, "rich-group", 9_000_000) - // The override names its own group; the user must be a member. + // The override names a group the user must be a member of. og := dbgen.Group(t, db, database.Group{OrganizationID: org.ID, Name: "override-group"}) dbgen.GroupMember(t, db, database.GroupMemberTable{UserID: user.ID, GroupID: og.ID}) _, err := db.UpsertUserAIBudgetOverride(ctx, database.UpsertUserAIBudgetOverrideParams{ @@ -72,92 +73,98 @@ func TestResolveUserAIBudget(t *testing.T) { SpendLimitMicros: 1_000_000, }) require.NoError(t, err) - return user.ID, budget.EffectiveBudget{GroupID: og.ID, SpendLimitMicros: 1_000_000, Source: codersdk.AIBudgetLimitSourceUserOverride}, true + return user.ID, budget.EffectiveGroup{GroupID: og.ID, Limit: &budget.Limit{SpendLimitMicros: 1_000_000, Source: codersdk.AIBudgetLimitSourceUserOverride}}, true }, }, { name: "SingleGroupBudget", policy: codersdk.AIBudgetPolicyHighest, - setup: func(t *testing.T, ctx context.Context, db database.Store) (uuid.UUID, budget.EffectiveBudget, bool) { + setup: func(t *testing.T, ctx context.Context, db database.Store) (uuid.UUID, budget.EffectiveGroup, bool) { org := dbgen.Organization(t, db, database.Organization{}) user := dbgen.User(t, db, database.User{}) + dbgen.OrganizationMember(t, db, database.OrganizationMember{OrganizationID: org.ID, UserID: user.ID}) gid := budgetedGroup(t, ctx, db, org.ID, user.ID, "only", 8_000_000) - return user.ID, budget.EffectiveBudget{GroupID: gid, SpendLimitMicros: 8_000_000, Source: codersdk.AIBudgetLimitSourceGroup}, true + return user.ID, budget.EffectiveGroup{GroupID: gid, Limit: &budget.Limit{SpendLimitMicros: 8_000_000, Source: codersdk.AIBudgetLimitSourceGroup}}, true }, }, { name: "HighestGroupWins", policy: codersdk.AIBudgetPolicyHighest, - setup: func(t *testing.T, ctx context.Context, db database.Store) (uuid.UUID, budget.EffectiveBudget, bool) { + setup: func(t *testing.T, ctx context.Context, db database.Store) (uuid.UUID, budget.EffectiveGroup, bool) { org := dbgen.Organization(t, db, database.Organization{}) user := dbgen.User(t, db, database.User{}) + dbgen.OrganizationMember(t, db, database.OrganizationMember{OrganizationID: org.ID, UserID: user.ID}) budgetedGroup(t, ctx, db, org.ID, user.ID, "low", 5_000_000) budgetedGroup(t, ctx, db, org.ID, user.ID, "mid", 20_000_000) high := budgetedGroup(t, ctx, db, org.ID, user.ID, "high", 50_000_000) - return user.ID, budget.EffectiveBudget{GroupID: high, SpendLimitMicros: 50_000_000, Source: codersdk.AIBudgetLimitSourceGroup}, true + return user.ID, budget.EffectiveGroup{GroupID: high, Limit: &budget.Limit{SpendLimitMicros: 50_000_000, Source: codersdk.AIBudgetLimitSourceGroup}}, true }, }, { - name: "TieBrokenByName", + name: "TieBrokenByEarliestOrgMembership", policy: codersdk.AIBudgetPolicyHighest, - setup: func(t *testing.T, ctx context.Context, db database.Store) (uuid.UUID, budget.EffectiveBudget, bool) { - org := dbgen.Organization(t, db, database.Organization{}) + setup: func(t *testing.T, ctx context.Context, db database.Store) (uuid.UUID, budget.EffectiveGroup, bool) { user := dbgen.User(t, db, database.User{}) - // Equal limits; "alpha" must win over "beta" by name ascending. - alpha := budgetedGroup(t, ctx, db, org.ID, user.ID, "alpha", 10_000_000) - budgetedGroup(t, ctx, db, org.ID, user.ID, "beta", 10_000_000) - return user.ID, budget.EffectiveBudget{GroupID: alpha, SpendLimitMicros: 10_000_000, Source: codersdk.AIBudgetLimitSourceGroup}, true + // Two groups in different orgs share the same limit. The earlier + // organization membership breaks the tie. + earlyOrg := dbgen.Organization(t, db, database.Organization{}) + lateOrg := dbgen.Organization(t, db, database.Organization{}) + dbgen.OrganizationMember(t, db, database.OrganizationMember{OrganizationID: earlyOrg.ID, UserID: user.ID, CreatedAt: time.Now().Add(-time.Hour)}) + dbgen.OrganizationMember(t, db, database.OrganizationMember{OrganizationID: lateOrg.ID, UserID: user.ID}) + winner := budgetedGroup(t, ctx, db, earlyOrg.ID, user.ID, "dup", 10_000_000) + budgetedGroup(t, ctx, db, lateOrg.ID, user.ID, "dup", 10_000_000) + return user.ID, budget.EffectiveGroup{GroupID: winner, Limit: &budget.Limit{SpendLimitMicros: 10_000_000, Source: codersdk.AIBudgetLimitSourceGroup}}, true }, }, { name: "TieBrokenByGroupID", policy: codersdk.AIBudgetPolicyHighest, - setup: func(t *testing.T, ctx context.Context, db database.Store) (uuid.UUID, budget.EffectiveBudget, bool) { + setup: func(t *testing.T, ctx context.Context, db database.Store) (uuid.UUID, budget.EffectiveGroup, bool) { + org := dbgen.Organization(t, db, database.Organization{}) user := dbgen.User(t, db, database.User{}) - // Two groups in different orgs share both name and limit. - // Group id breaks the tie, so resolution is deterministic. - org1 := dbgen.Organization(t, db, database.Organization{}) - org2 := dbgen.Organization(t, db, database.Organization{}) - g1 := budgetedGroup(t, ctx, db, org1.ID, user.ID, "dup", 10_000_000) - g2 := budgetedGroup(t, ctx, db, org2.ID, user.ID, "dup", 10_000_000) - winner := g1 - if bytes.Compare(g2[:], g1[:]) < 0 { - winner = g2 + dbgen.OrganizationMember(t, db, database.OrganizationMember{OrganizationID: org.ID, UserID: user.ID}) + // Both groups are in the same org, so both resolve to the same + // organization membership and the tie falls to the lowest group ID. + groupA := budgetedGroup(t, ctx, db, org.ID, user.ID, "alpha", 10_000_000) + groupB := budgetedGroup(t, ctx, db, org.ID, user.ID, "beta", 10_000_000) + winner := groupA + if bytes.Compare(groupB[:], groupA[:]) < 0 { + winner = groupB } - return user.ID, budget.EffectiveBudget{GroupID: winner, SpendLimitMicros: 10_000_000, Source: codersdk.AIBudgetLimitSourceGroup}, true + return user.ID, budget.EffectiveGroup{GroupID: winner, Limit: &budget.Limit{SpendLimitMicros: 10_000_000, Source: codersdk.AIBudgetLimitSourceGroup}}, true }, }, { name: "GroupsButNoneBudgeted", policy: codersdk.AIBudgetPolicyHighest, - setup: func(t *testing.T, ctx context.Context, db database.Store) (uuid.UUID, budget.EffectiveBudget, bool) { + setup: func(t *testing.T, ctx context.Context, db database.Store) (uuid.UUID, budget.EffectiveGroup, bool) { org := dbgen.Organization(t, db, database.Organization{}) user := dbgen.User(t, db, database.User{}) g := dbgen.Group(t, db, database.Group{OrganizationID: org.ID, Name: "unbudgeted"}) dbgen.GroupMember(t, db, database.GroupMemberTable{UserID: user.ID, GroupID: g.ID}) - return user.ID, budget.EffectiveBudget{}, false + return user.ID, budget.EffectiveGroup{}, false }, }, { name: "EveryoneGroupBudget", policy: codersdk.AIBudgetPolicyHighest, - setup: func(t *testing.T, ctx context.Context, db database.Store) (uuid.UUID, budget.EffectiveBudget, bool) { + setup: func(t *testing.T, ctx context.Context, db database.Store) (uuid.UUID, budget.EffectiveGroup, bool) { org := dbgen.Organization(t, db, database.Organization{}) user := dbgen.User(t, db, database.User{}) // Membership is via organization_members only (no group_members row), // exercising the org-members half of group_members_expanded. everyoneID := budgetedEveryoneGroup(t, ctx, db, org.ID, user.ID, 7_000_000) - return user.ID, budget.EffectiveBudget{GroupID: everyoneID, SpendLimitMicros: 7_000_000, Source: codersdk.AIBudgetLimitSourceGroup}, true + return user.ID, budget.EffectiveGroup{GroupID: everyoneID, Limit: &budget.Limit{SpendLimitMicros: 7_000_000, Source: codersdk.AIBudgetLimitSourceGroup}}, true }, }, { name: "OverrideBeatsEveryoneBudget", policy: codersdk.AIBudgetPolicyHighest, - setup: func(t *testing.T, ctx context.Context, db database.Store) (uuid.UUID, budget.EffectiveBudget, bool) { + setup: func(t *testing.T, ctx context.Context, db database.Store) (uuid.UUID, budget.EffectiveGroup, bool) { org := dbgen.Organization(t, db, database.Organization{}) user := dbgen.User(t, db, database.User{}) everyoneID := budgetedEveryoneGroup(t, ctx, db, org.ID, user.ID, 7_000_000) - // Override attributed to the Everyone group; the user is a member + // Override attributed to the Everyone group. The user is a member // via organization_members, satisfying the membership trigger. _, err := db.UpsertUserAIBudgetOverride(ctx, database.UpsertUserAIBudgetOverrideParams{ UserID: user.ID, @@ -165,16 +172,16 @@ func TestResolveUserAIBudget(t *testing.T) { SpendLimitMicros: 2_000_000, }) require.NoError(t, err) - return user.ID, budget.EffectiveBudget{GroupID: everyoneID, SpendLimitMicros: 2_000_000, Source: codersdk.AIBudgetLimitSourceUserOverride}, true + return user.ID, budget.EffectiveGroup{GroupID: everyoneID, Limit: &budget.Limit{SpendLimitMicros: 2_000_000, Source: codersdk.AIBudgetLimitSourceUserOverride}}, true }, }, { name: "UnsupportedPolicy", policy: codersdk.AIBudgetPolicy("unsupported"), - setup: func(t *testing.T, ctx context.Context, db database.Store) (uuid.UUID, budget.EffectiveBudget, bool) { + setup: func(t *testing.T, ctx context.Context, db database.Store) (uuid.UUID, budget.EffectiveGroup, bool) { // No override, so resolution reaches the policy switch and errors. user := dbgen.User(t, db, database.User{}) - return user.ID, budget.EffectiveBudget{}, false + return user.ID, budget.EffectiveGroup{}, false }, wantErr: "unsupported AI budget policy", }, @@ -199,8 +206,122 @@ func TestResolveUserAIBudget(t *testing.T) { return } require.Equal(t, want.GroupID, got.GroupID) - require.Equal(t, want.SpendLimitMicros, got.SpendLimitMicros) - require.Equal(t, want.Source, got.Source) + require.Equal(t, want.Limit, got.Limit) + }) + } +} + +func TestResolveUserEffectiveGroup(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + policy codersdk.AIBudgetPolicy + setup func(t *testing.T, ctx context.Context, db database.Store) (userID uuid.UUID, want budget.EffectiveGroup, wantOK bool) + wantErr string + }{ + { + // The Everyone group has a budget, so it resolves via the budget + // path rather than the fallback. + name: "EveryoneGroupWithBudget", + policy: codersdk.AIBudgetPolicyHighest, + setup: func(t *testing.T, ctx context.Context, db database.Store) (uuid.UUID, budget.EffectiveGroup, bool) { + org := dbgen.Organization(t, db, database.Organization{}) + user := dbgen.User(t, db, database.User{}) + // The Everyone group's id equals the org id. + group := dbgen.Group(t, db, database.Group{ID: org.ID, OrganizationID: org.ID, Name: "Everyone"}) + dbgen.OrganizationMember(t, db, database.OrganizationMember{OrganizationID: org.ID, UserID: user.ID}) + _, err := db.UpsertGroupAIBudget(ctx, database.UpsertGroupAIBudgetParams{ + GroupID: group.ID, + SpendLimitMicros: 7_000_000, + }) + require.NoError(t, err) + return user.ID, budget.EffectiveGroup{GroupID: group.ID, Limit: &budget.Limit{SpendLimitMicros: 7_000_000, Source: codersdk.AIBudgetLimitSourceGroup}}, true + }, + }, + { + // With a single org and no budget, attribution falls back to that + // org's Everyone group with no limit. + name: "FallbackToEveryoneUnlimited", + policy: codersdk.AIBudgetPolicyHighest, + setup: func(t *testing.T, ctx context.Context, db database.Store) (uuid.UUID, budget.EffectiveGroup, bool) { + org := dbgen.Organization(t, db, database.Organization{}) + user := dbgen.User(t, db, database.User{}) + dbgen.OrganizationMember(t, db, database.OrganizationMember{OrganizationID: org.ID, UserID: user.ID}) + return user.ID, budget.EffectiveGroup{GroupID: org.ID}, true + }, + }, + { + // The fallback prefers the default org even over an org joined + // earlier. + name: "FallbackPrefersDefaultOrg", + policy: codersdk.AIBudgetPolicyHighest, + setup: func(t *testing.T, ctx context.Context, db database.Store) (uuid.UUID, budget.EffectiveGroup, bool) { + defaultOrg, err := db.GetDefaultOrganization(ctx) + require.NoError(t, err) + otherOrg := dbgen.Organization(t, db, database.Organization{}) + user := dbgen.User(t, db, database.User{}) + dbgen.OrganizationMember(t, db, database.OrganizationMember{OrganizationID: otherOrg.ID, UserID: user.ID, CreatedAt: time.Now().Add(-time.Hour)}) + dbgen.OrganizationMember(t, db, database.OrganizationMember{OrganizationID: defaultOrg.ID, UserID: user.ID}) + return user.ID, budget.EffectiveGroup{GroupID: defaultOrg.ID}, true + }, + }, + { + // Among non-default orgs, the fallback breaks ties by the earliest + // organization membership. + name: "FallbackTieByEarliestOrgMembership", + policy: codersdk.AIBudgetPolicyHighest, + setup: func(t *testing.T, ctx context.Context, db database.Store) (uuid.UUID, budget.EffectiveGroup, bool) { + user := dbgen.User(t, db, database.User{}) + earlyOrg := dbgen.Organization(t, db, database.Organization{}) + lateOrg := dbgen.Organization(t, db, database.Organization{}) + dbgen.OrganizationMember(t, db, database.OrganizationMember{OrganizationID: earlyOrg.ID, UserID: user.ID, CreatedAt: time.Now().Add(-time.Hour)}) + dbgen.OrganizationMember(t, db, database.OrganizationMember{OrganizationID: lateOrg.ID, UserID: user.ID}) + return user.ID, budget.EffectiveGroup{GroupID: earlyOrg.ID}, true + }, + }, + { + // A user with no org membership has no effective group. + name: "NoOrgMembership", + policy: codersdk.AIBudgetPolicyHighest, + setup: func(t *testing.T, ctx context.Context, db database.Store) (uuid.UUID, budget.EffectiveGroup, bool) { + user := dbgen.User(t, db, database.User{}) + return user.ID, budget.EffectiveGroup{}, false + }, + }, + { + // An unsupported policy surfaces the error from ResolveUserAIBudget. + name: "UnsupportedPolicy", + policy: codersdk.AIBudgetPolicy("unsupported"), + setup: func(t *testing.T, ctx context.Context, db database.Store) (uuid.UUID, budget.EffectiveGroup, bool) { + user := dbgen.User(t, db, database.User{}) + return user.ID, budget.EffectiveGroup{}, false + }, + wantErr: "unsupported AI budget policy", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + db, _ := dbtestutil.NewDB(t) + ctx := testutil.Context(t, testutil.WaitLong) + + userID, want, wantOK := tt.setup(t, ctx, db) + got, ok, err := budget.ResolveUserEffectiveGroup(ctx, db, userID, tt.policy) + if tt.wantErr != "" { + require.ErrorContains(t, err, tt.wantErr) + return + } + require.NoError(t, err) + require.Equal(t, wantOK, ok) + if !wantOK { + return + } + + require.Equal(t, want.GroupID, got.GroupID) + require.Equal(t, want.Limit, got.Limit) }) } } diff --git a/coderd/aibridgedserver/aibridgedserver.go b/coderd/aibridgedserver/aibridgedserver.go index c17a0879200..cf750944e93 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,35 @@ 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. + // 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 } 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 +851,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 bf28dbdf9ed..ab6a479b052 100644 --- a/coderd/aibridgedserver/aibridgedserver_test.go +++ b/coderd/aibridgedserver/aibridgedserver_test.go @@ -732,6 +732,50 @@ 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: 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}) + + // 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 +1680,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 +1735,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 +1761,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 +1826,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 +1880,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 +1933,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 +1958,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 +1974,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 +1995,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 +2019,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. 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 +2125,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 +2145,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 +2174,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 +2302,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 +2326,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 +2834,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 0a6b9ba5b63..87ba4ab6464 100644 --- a/coderd/aibridgedserver/cost.go +++ b/coderd/aibridgedserver/cost.go @@ -31,22 +31,27 @@ 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. 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 - // 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 2c27365ac07..69411ed4491 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 membership.", "type": "string", "format": "uuid" }, diff --git a/coderd/apidoc/swagger.json b/coderd/apidoc/swagger.json index 7f87f2bad29..57d188b92b2 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 membership.", "type": "string", "format": "uuid" }, diff --git a/coderd/database/dbauthz/dbauthz.go b/coderd/database/dbauthz/dbauthz.go index 3161f3c3e9d..3cb39f78791 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 c029cd37611..d064aeb9bf6 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 57c742457f0..16bfc7c80f1 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 9103654b6d3..453a8798dbb 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 5c555cdd86f..8e9b909de43 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 the earliest organization membership. Implements the + // "highest" budget policy. group_members_expanded is a UNION of group_members + // and organization_members, so the implicit "Everyone" group + // (group_id == organization_id) is included. Returns no rows when the user has + // no budgeted groups. Callers should treat sql.ErrNoRows as "no group budget". GetHighestGroupAIBudgetByUser(ctx context.Context, userID uuid.UUID) (GetHighestGroupAIBudgetByUserRow, error) GetInboxNotificationByID(ctx context.Context, id uuid.UUID) (InboxNotification, error) // Fetches inbox notifications for a user filtered by templates and targets @@ -875,6 +875,11 @@ type sqlcQuerier interface { GetUserChatSpendInPeriod(ctx context.Context, arg GetUserChatSpendInPeriodParams) (int64, error) GetUserCodeDiffDisplayMode(ctx context.Context, userID uuid.UUID) (string, error) GetUserCount(ctx context.Context, includeSystem bool) (int64, error) + // Returns the "Everyone" group (id == organization_id) to attribute a user's + // spend to when no override or budgeted group applies. Prefers the default org, + // then the earliest organization membership. Returns no rows when the user has + // no organization membership. + GetUserEveryoneFallbackGroup(ctx context.Context, userID uuid.UUID) (uuid.UUID, error) GetUserForChatSyntheticAPIKeyByID(ctx context.Context, id uuid.UUID) (User, error) // Returns the minimum (most restrictive) group limit for a user. // Returns -1 if no group limits match the specified scope. diff --git a/coderd/database/querier_test.go b/coderd/database/querier_test.go index 53a477e9e66..d0467a86b3e 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) @@ -13431,6 +13439,89 @@ func TestGetGroupMembersAISpend(t *testing.T) { require.Equal(t, int64(0), got[0].GroupSpendMicros) }) + t.Run("FallbackToEveryoneGroup", func(t *testing.T) { + t.Parallel() + db, _ := dbtestutil.NewDB(t) + ctx := testutil.Context(t, testutil.WaitShort) + + // Given: an unbudgeted member of the queried group whose org has an + // Everyone group but no override or budgeted group. + user := dbgen.User(t, db, database.User{}) + org := dbgen.Organization(t, db, database.Organization{}) + queried := dbgen.Group(t, db, database.Group{OrganizationID: org.ID}) + dbgen.OrganizationMember(t, db, database.OrganizationMember{UserID: user.ID, OrganizationID: org.ID}) + dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: queried.ID, UserID: user.ID}) + // The Everyone group (id == org id) must exist for the effective group + // join to resolve the fallback. + //nolint:gocritic // Requires system context. + _, err := db.InsertAllUsersGroup(dbauthz.AsSystemRestricted(ctx), org.ID) + require.NoError(t, err) + _, err = db.IncrementUserAIDailySpend(ctx, database.IncrementUserAIDailySpendParams{ + UserID: user.ID, EffectiveGroupID: queried.ID, Day: now, CostMicros: 250, + }) + require.NoError(t, err) + + // When: querying spend for the user. + got, err := db.GetGroupMembersAISpend(ctx, database.GetGroupMembersAISpendParams{ + GroupID: queried.ID, + UserIds: []uuid.UUID{user.ID}, + PeriodStart: monthStart, + }) + require.NoError(t, err) + + // Then: with no budget, the effective group falls back to the Everyone + // group. The limit and source are null, and queried-group spend is returned. + require.Len(t, got, 1) + require.Equal(t, uuid.NullUUID{UUID: org.ID, Valid: true}, got[0].EffectiveGroupID) + require.False(t, got[0].SpendLimitMicros.Valid) + require.False(t, got[0].LimitSource.Valid) + require.Equal(t, int64(250), got[0].GroupSpendMicros) + }) + + t.Run("CrossOrgFallbackMasked", func(t *testing.T) { + t.Parallel() + db, _ := dbtestutil.NewDB(t) + ctx := testutil.Context(t, testutil.WaitShort) + + // Given: an unbudgeted member of the queried group who joined another + // org earlier. The fallback picks the earlier org's Everyone group. + user := dbgen.User(t, db, database.User{}) + queriedOrg := dbgen.Organization(t, db, database.Organization{}) + otherOrg := dbgen.Organization(t, db, database.Organization{}) + queried := dbgen.Group(t, db, database.Group{OrganizationID: queriedOrg.ID}) + dbgen.OrganizationMember(t, db, database.OrganizationMember{UserID: user.ID, OrganizationID: otherOrg.ID, CreatedAt: dbtime.Now().Add(-time.Hour)}) + dbgen.OrganizationMember(t, db, database.OrganizationMember{UserID: user.ID, OrganizationID: queriedOrg.ID}) + dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: queried.ID, UserID: user.ID}) + // Both orgs have an Everyone group (id == org id), as in production. + //nolint:gocritic // Requires system context. + _, err := db.InsertAllUsersGroup(dbauthz.AsSystemRestricted(ctx), queriedOrg.ID) + require.NoError(t, err) + //nolint:gocritic // Requires system context. + _, err = db.InsertAllUsersGroup(dbauthz.AsSystemRestricted(ctx), otherOrg.ID) + require.NoError(t, err) + _, err = db.IncrementUserAIDailySpend(ctx, database.IncrementUserAIDailySpendParams{ + UserID: user.ID, EffectiveGroupID: queried.ID, Day: now, CostMicros: 250, + }) + require.NoError(t, err) + + // When: querying spend for the user. + got, err := db.GetGroupMembersAISpend(ctx, database.GetGroupMembersAISpendParams{ + GroupID: queried.ID, + UserIds: []uuid.UUID{user.ID}, + PeriodStart: monthStart, + }) + require.NoError(t, err) + + // Then: the fallback resolves to the other org's Everyone group, so + // effective_group_id is masked to null, while queried-group spend still + // returns. + require.Len(t, got, 1) + require.False(t, got[0].EffectiveGroupID.Valid, "cross-org effective group must be masked") + require.False(t, got[0].SpendLimitMicros.Valid) + require.False(t, got[0].LimitSource.Valid) + require.Equal(t, int64(250), got[0].GroupSpendMicros) + }) + t.Run("SpendWithDifferentEffectiveGroup", func(t *testing.T) { t.Parallel() db, _ := dbtestutil.NewDB(t) @@ -13686,6 +13777,220 @@ func TestGetGroupMembersAISpend(t *testing.T) { }) } +func TestGetHighestGroupAIBudgetByUser(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + setup func(t *testing.T, ctx context.Context, db database.Store) (userID uuid.UUID, want database.GetHighestGroupAIBudgetByUserRow) + wantErr error + }{ + { + // Among the user's budgeted groups, the highest limit wins. + name: "HighestWins", + setup: func(t *testing.T, ctx context.Context, db database.Store) (uuid.UUID, database.GetHighestGroupAIBudgetByUserRow) { + user := dbgen.User(t, db, database.User{}) + org := dbgen.Organization(t, db, database.Organization{}) + lower := dbgen.Group(t, db, database.Group{OrganizationID: org.ID}) + higher := dbgen.Group(t, db, database.Group{OrganizationID: org.ID}) + dbgen.OrganizationMember(t, db, database.OrganizationMember{OrganizationID: org.ID, UserID: user.ID}) + dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: lower.ID, UserID: user.ID}) + dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: higher.ID, UserID: user.ID}) + _, err := db.UpsertGroupAIBudget(ctx, database.UpsertGroupAIBudgetParams{GroupID: lower.ID, SpendLimitMicros: 1_000_000}) + require.NoError(t, err) + _, err = db.UpsertGroupAIBudget(ctx, database.UpsertGroupAIBudgetParams{GroupID: higher.ID, SpendLimitMicros: 2_000_000}) + require.NoError(t, err) + return user.ID, database.GetHighestGroupAIBudgetByUserRow{GroupID: higher.ID, SpendLimitMicros: 2_000_000} + }, + }, + { + // The highest limit wins across the user's orgs, not just within one. + name: "HighestWinsAcrossOrgs", + setup: func(t *testing.T, ctx context.Context, db database.Store) (uuid.UUID, database.GetHighestGroupAIBudgetByUserRow) { + user := dbgen.User(t, db, database.User{}) + orgA := dbgen.Organization(t, db, database.Organization{}) + orgB := dbgen.Organization(t, db, database.Organization{}) + lower := dbgen.Group(t, db, database.Group{OrganizationID: orgA.ID}) + higher := dbgen.Group(t, db, database.Group{OrganizationID: orgB.ID}) + dbgen.OrganizationMember(t, db, database.OrganizationMember{OrganizationID: orgA.ID, UserID: user.ID}) + dbgen.OrganizationMember(t, db, database.OrganizationMember{OrganizationID: orgB.ID, UserID: user.ID}) + dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: lower.ID, UserID: user.ID}) + dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: higher.ID, UserID: user.ID}) + _, err := db.UpsertGroupAIBudget(ctx, database.UpsertGroupAIBudgetParams{GroupID: lower.ID, SpendLimitMicros: 1_000_000}) + require.NoError(t, err) + _, err = db.UpsertGroupAIBudget(ctx, database.UpsertGroupAIBudgetParams{GroupID: higher.ID, SpendLimitMicros: 2_000_000}) + require.NoError(t, err) + return user.ID, database.GetHighestGroupAIBudgetByUserRow{GroupID: higher.ID, SpendLimitMicros: 2_000_000} + }, + }, + { + // A budgeted group in a soft-deleted org is excluded even when its + // limit is higher. + name: "ExcludesDeletedOrg", + setup: func(t *testing.T, ctx context.Context, db database.Store) (uuid.UUID, database.GetHighestGroupAIBudgetByUserRow) { + user := dbgen.User(t, db, database.User{}) + liveOrg := dbgen.Organization(t, db, database.Organization{Name: "live-org"}) + deletedOrg := dbgen.Organization(t, db, database.Organization{Name: "deleted-org"}) + liveGroup := dbgen.Group(t, db, database.Group{OrganizationID: liveOrg.ID}) + deletedGroup := dbgen.Group(t, db, database.Group{OrganizationID: deletedOrg.ID}) + dbgen.OrganizationMember(t, db, database.OrganizationMember{OrganizationID: liveOrg.ID, UserID: user.ID}) + dbgen.OrganizationMember(t, db, database.OrganizationMember{OrganizationID: deletedOrg.ID, UserID: user.ID}) + dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: liveGroup.ID, UserID: user.ID}) + dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: deletedGroup.ID, UserID: user.ID}) + _, err := db.UpsertGroupAIBudget(ctx, database.UpsertGroupAIBudgetParams{GroupID: liveGroup.ID, SpendLimitMicros: 1_000_000}) + require.NoError(t, err) + _, err = db.UpsertGroupAIBudget(ctx, database.UpsertGroupAIBudgetParams{GroupID: deletedGroup.ID, SpendLimitMicros: 5_000_000}) + require.NoError(t, err) + err = db.UpdateOrganizationDeletedByID(ctx, database.UpdateOrganizationDeletedByIDParams{ + ID: deletedOrg.ID, + UpdatedAt: dbtime.Now(), + }) + require.NoError(t, err) + return user.ID, database.GetHighestGroupAIBudgetByUserRow{GroupID: liveGroup.ID, SpendLimitMicros: 1_000_000} + }, + }, + { + // Equal limits across orgs break by the earliest organization + // membership. + name: "TieByEarliestOrgMembership", + setup: func(t *testing.T, ctx context.Context, db database.Store) (uuid.UUID, database.GetHighestGroupAIBudgetByUserRow) { + user := dbgen.User(t, db, database.User{}) + earlyOrg := dbgen.Organization(t, db, database.Organization{}) + lateOrg := dbgen.Organization(t, db, database.Organization{}) + earlyGroup := dbgen.Group(t, db, database.Group{OrganizationID: earlyOrg.ID}) + lateGroup := dbgen.Group(t, db, database.Group{OrganizationID: lateOrg.ID}) + dbgen.OrganizationMember(t, db, database.OrganizationMember{OrganizationID: earlyOrg.ID, UserID: user.ID, CreatedAt: dbtime.Now().Add(-time.Hour)}) + dbgen.OrganizationMember(t, db, database.OrganizationMember{OrganizationID: lateOrg.ID, UserID: user.ID}) + dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: earlyGroup.ID, UserID: user.ID}) + dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: lateGroup.ID, UserID: user.ID}) + _, err := db.UpsertGroupAIBudget(ctx, database.UpsertGroupAIBudgetParams{GroupID: earlyGroup.ID, SpendLimitMicros: 1_000_000}) + require.NoError(t, err) + _, err = db.UpsertGroupAIBudget(ctx, database.UpsertGroupAIBudgetParams{GroupID: lateGroup.ID, SpendLimitMicros: 1_000_000}) + require.NoError(t, err) + return user.ID, database.GetHighestGroupAIBudgetByUserRow{GroupID: earlyGroup.ID, SpendLimitMicros: 1_000_000} + }, + }, + { + // A user with no budgeted group has no highest budget. + name: "NoBudgetedGroup", + wantErr: sql.ErrNoRows, + setup: func(t *testing.T, ctx context.Context, db database.Store) (uuid.UUID, database.GetHighestGroupAIBudgetByUserRow) { + user := dbgen.User(t, db, database.User{}) + org := dbgen.Organization(t, db, database.Organization{}) + dbgen.OrganizationMember(t, db, database.OrganizationMember{OrganizationID: org.ID, UserID: user.ID}) + return user.ID, database.GetHighestGroupAIBudgetByUserRow{} + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + db, _ := dbtestutil.NewDB(t) + ctx := testutil.Context(t, testutil.WaitShort) + + userID, want := tt.setup(t, ctx, db) + got, err := db.GetHighestGroupAIBudgetByUser(ctx, userID) + if tt.wantErr != nil { + require.ErrorIs(t, err, tt.wantErr) + return + } + require.NoError(t, err) + require.Equal(t, want, got) + }) + } +} + +func TestGetUserEveryoneFallbackGroup(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + setup func(t *testing.T, ctx context.Context, db database.Store) (userID uuid.UUID, wantGroupID uuid.UUID) + wantErr error + }{ + { + // A single-org member falls back to that org's Everyone group. + name: "SingleOrg", + setup: func(t *testing.T, ctx context.Context, db database.Store) (uuid.UUID, uuid.UUID) { + org := dbgen.Organization(t, db, database.Organization{}) + user := dbgen.User(t, db, database.User{}) + dbgen.OrganizationMember(t, db, database.OrganizationMember{OrganizationID: org.ID, UserID: user.ID}) + return user.ID, org.ID + }, + }, + { + // The default org is preferred even over an org joined earlier. + name: "PrefersDefaultOrg", + setup: func(t *testing.T, ctx context.Context, db database.Store) (uuid.UUID, uuid.UUID) { + defaultOrg, err := db.GetDefaultOrganization(ctx) + require.NoError(t, err) + otherOrg := dbgen.Organization(t, db, database.Organization{}) + user := dbgen.User(t, db, database.User{}) + dbgen.OrganizationMember(t, db, database.OrganizationMember{OrganizationID: otherOrg.ID, UserID: user.ID, CreatedAt: dbtime.Now().Add(-time.Hour)}) + dbgen.OrganizationMember(t, db, database.OrganizationMember{OrganizationID: defaultOrg.ID, UserID: user.ID}) + return user.ID, defaultOrg.ID + }, + }, + { + // Among non-default orgs, ties break by the earliest organization + // membership. + name: "TieByEarliestOrgMembership", + setup: func(t *testing.T, ctx context.Context, db database.Store) (uuid.UUID, uuid.UUID) { + user := dbgen.User(t, db, database.User{}) + earlyOrg := dbgen.Organization(t, db, database.Organization{}) + lateOrg := dbgen.Organization(t, db, database.Organization{}) + dbgen.OrganizationMember(t, db, database.OrganizationMember{OrganizationID: earlyOrg.ID, UserID: user.ID, CreatedAt: dbtime.Now().Add(-time.Hour)}) + dbgen.OrganizationMember(t, db, database.OrganizationMember{OrganizationID: lateOrg.ID, UserID: user.ID}) + return user.ID, earlyOrg.ID + }, + }, + { + // A soft-deleted org is excluded even when it was joined earlier. + name: "ExcludesDeletedOrg", + setup: func(t *testing.T, ctx context.Context, db database.Store) (uuid.UUID, uuid.UUID) { + user := dbgen.User(t, db, database.User{}) + liveOrg := dbgen.Organization(t, db, database.Organization{Name: "live-org"}) + deletedOrg := dbgen.Organization(t, db, database.Organization{Name: "deleted-org"}) + dbgen.OrganizationMember(t, db, database.OrganizationMember{OrganizationID: liveOrg.ID, UserID: user.ID}) + dbgen.OrganizationMember(t, db, database.OrganizationMember{OrganizationID: deletedOrg.ID, UserID: user.ID, CreatedAt: dbtime.Now().Add(-time.Hour)}) + err := db.UpdateOrganizationDeletedByID(ctx, database.UpdateOrganizationDeletedByIDParams{ + ID: deletedOrg.ID, + UpdatedAt: dbtime.Now(), + }) + require.NoError(t, err) + return user.ID, liveOrg.ID + }, + }, + { + // A user with no org membership has no fallback group. + name: "NoOrgMembership", + wantErr: sql.ErrNoRows, + setup: func(t *testing.T, ctx context.Context, db database.Store) (uuid.UUID, uuid.UUID) { + user := dbgen.User(t, db, database.User{}) + return user.ID, uuid.Nil + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + db, _ := dbtestutil.NewDB(t) + ctx := testutil.Context(t, testutil.WaitShort) + + userID, wantGroupID := tt.setup(t, ctx, db) + got, err := db.GetUserEveryoneFallbackGroup(ctx, userID) + if tt.wantErr != nil { + require.ErrorIs(t, err, tt.wantErr) + return + } + require.NoError(t, err) + require.Equal(t, wantGroupID, got) + }) + } +} + func TestChatPinOrderQueries(t *testing.T) { t.Parallel() if testing.Short() { diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go index 876c952fc59..76168e90363 100644 --- a/coderd/database/queries.sql.go +++ b/coderd/database/queries.sql.go @@ -2594,15 +2594,33 @@ user_highest_group AS ( budget.spend_limit_micros FROM group_ai_budgets budget JOIN group_members_expanded member ON member.group_id = budget.group_id + JOIN organizations ON organizations.id = member.organization_id + JOIN organization_members + ON organization_members.user_id = member.user_id + AND organization_members.organization_id = member.organization_id WHERE member.user_id IN (SELECT user_id FROM filtered_users) - ORDER BY member.user_id, budget.spend_limit_micros DESC, member.group_name ASC, budget.group_id ASC + AND organizations.deleted = false + ORDER BY member.user_id, budget.spend_limit_micros DESC, organization_members.created_at ASC, budget.group_id ASC +), +user_fallback_group AS ( + -- Per user, the Everyone group to fall back to when no override or budgeted + -- group applies. The Everyone group has id == organization_id. Prefers the + -- default org, then the earliest organization membership. + SELECT DISTINCT ON (organization_members.user_id) + organization_members.user_id, + organizations.id AS group_id + FROM organization_members + JOIN organizations ON organizations.id = organization_members.organization_id + WHERE organization_members.user_id IN (SELECT user_id FROM filtered_users) + AND organizations.deleted = false + ORDER BY organization_members.user_id, organizations.is_default DESC, organization_members.created_at ASC, organizations.id ASC ), effective AS ( - -- Effective budget per user: a per-user override wins over the - -- highest-limit group. + -- Effective budget per user: a per-user override wins over the highest-limit + -- group, which wins over the Everyone group fallback. SELECT filtered_users.user_id, - COALESCE(override.group_id, user_highest_group.group_id) AS raw_effective_group_id, + COALESCE(override.group_id, user_highest_group.group_id, user_fallback_group.group_id) AS raw_effective_group_id, COALESCE(override.spend_limit_micros, user_highest_group.spend_limit_micros) AS spend_limit_micros, (CASE WHEN override.group_id IS NOT NULL THEN 'user_override' @@ -2611,6 +2629,7 @@ effective AS ( FROM filtered_users LEFT JOIN user_ai_budget_overrides override ON override.user_id = filtered_users.user_id LEFT JOIN user_highest_group ON user_highest_group.user_id = filtered_users.user_id + LEFT JOIN user_fallback_group ON user_fallback_group.user_id = filtered_users.user_id ), applied_budget AS ( -- The limit and source only for users whose effective budget source is the @@ -2663,9 +2682,9 @@ type GetGroupMembersAISpendRow struct { // Returns each user's AI spend attributed to the queried group, on or after // period_start until NOW. Only current members of the queried group are // returned. spend_limit_micros and limit_source are populated only when the -// queried group is the user's effective budget source. The effective_group_id -// is null when the user has no configured budget or when the effective group -// belongs to a different organization than the queried group. +// queried group is the user's effective budget source. The effective group +// falls back to the Everyone group, and effective_group_id is null only when +// that group belongs to a different organization than the queried group. // The period_start parameter is normalized to its UTC calendar day. // TODO(AIGOV-527): unify effective group resolution in a single place. // Spend is aggregated for the queried group, not the user's effective group. @@ -2703,18 +2722,20 @@ func (q *sqlQuerier) GetGroupMembersAISpend(ctx context.Context, arg GetGroupMem const getHighestGroupAIBudgetByUser = `-- name: GetHighestGroupAIBudgetByUser :one SELECT - gaib.group_id, - gaib.spend_limit_micros -FROM group_ai_budgets gaib -JOIN group_members_expanded gme ON gme.group_id = gaib.group_id -WHERE gme.user_id = $1 + budget.group_id, + budget.spend_limit_micros +FROM group_ai_budgets budget +JOIN group_members_expanded member ON member.group_id = budget.group_id +JOIN organizations ON organizations.id = member.organization_id +JOIN organization_members + ON organization_members.user_id = member.user_id + AND organization_members.organization_id = member.organization_id +WHERE member.user_id = $1 + AND organizations.deleted = false ORDER BY - gaib.spend_limit_micros DESC, -- highest wins - gme.group_name ASC, -- alphabetical tiebreak - -- Final tiebreak on the group id makes the result deterministic when two - -- groups share both name and limit, which is possible across organizations - -- (groups are unique on (organization_id, name), not name alone). - gaib.group_id ASC + budget.spend_limit_micros DESC, -- highest wins + organization_members.created_at ASC, -- earliest organization membership + budget.group_id ASC -- deterministic tiebreak LIMIT 1 ` @@ -2724,11 +2745,11 @@ type GetHighestGroupAIBudgetByUserRow struct { } // Returns the highest group AI budget across the groups the user belongs to, -// breaking ties by group name ascending. Implements the "highest" budget policy. -// group_members_expanded is a UNION of group_members and organization_members, -// so the implicit "Everyone" group (group_id == organization_id) is included. -// Returns no rows when the user has no budgeted groups; callers should treat -// sql.ErrNoRows as "no group budget". +// breaking ties by the earliest organization membership. Implements the +// "highest" budget policy. group_members_expanded is a UNION of group_members +// and organization_members, so the implicit "Everyone" group +// (group_id == organization_id) is included. Returns no rows when the user has +// no budgeted groups. Callers should treat sql.ErrNoRows as "no group budget". func (q *sqlQuerier) GetHighestGroupAIBudgetByUser(ctx context.Context, userID uuid.UUID) (GetHighestGroupAIBudgetByUserRow, error) { row := q.db.QueryRowContext(ctx, getHighestGroupAIBudgetByUser, userID) var i GetHighestGroupAIBudgetByUserRow @@ -2856,6 +2877,30 @@ func (q *sqlQuerier) GetUserAISpendSince(ctx context.Context, arg GetUserAISpend return i, err } +const getUserEveryoneFallbackGroup = `-- name: GetUserEveryoneFallbackGroup :one +SELECT organizations.id AS group_id +FROM organization_members +JOIN organizations ON organizations.id = organization_members.organization_id +WHERE organization_members.user_id = $1 + AND organizations.deleted = false +ORDER BY + organizations.is_default DESC, -- prefer the default org + organization_members.created_at ASC, -- earliest organization membership + organizations.id ASC -- deterministic tiebreak +LIMIT 1 +` + +// Returns the "Everyone" group (id == organization_id) to attribute a user's +// spend to when no override or budgeted group applies. Prefers the default org, +// then the earliest organization membership. Returns no rows when the user has +// no organization membership. +func (q *sqlQuerier) GetUserEveryoneFallbackGroup(ctx context.Context, userID uuid.UUID) (uuid.UUID, error) { + row := q.db.QueryRowContext(ctx, getUserEveryoneFallbackGroup, userID) + var group_id uuid.UUID + err := row.Scan(&group_id) + return group_id, err +} + const incrementUserAIDailySpend = `-- name: IncrementUserAIDailySpend :one INSERT INTO ai_user_daily_spend (user_id, effective_group_id, day, spend_micros) VALUES ($1, $2, (($3::timestamptz) AT TIME ZONE 'UTC')::date, $4) diff --git a/coderd/database/queries/aicostcontrol.sql b/coderd/database/queries/aicostcontrol.sql index c15a63d55fa..ece65a9eef6 100644 --- a/coderd/database/queries/aicostcontrol.sql +++ b/coderd/database/queries/aicostcontrol.sql @@ -60,24 +60,42 @@ DELETE FROM user_ai_budget_overrides WHERE user_id = @user_id RETURNING *; -- name: GetHighestGroupAIBudgetByUser :one -- Returns the highest group AI budget across the groups the user belongs to, --- breaking ties by group name ascending. Implements the "highest" budget policy. --- group_members_expanded is a UNION of group_members and organization_members, --- so the implicit "Everyone" group (group_id == organization_id) is included. --- Returns no rows when the user has no budgeted groups; callers should treat --- sql.ErrNoRows as "no group budget". +-- breaking ties by the earliest organization membership. Implements the +-- "highest" budget policy. group_members_expanded is a UNION of group_members +-- and organization_members, so the implicit "Everyone" group +-- (group_id == organization_id) is included. Returns no rows when the user has +-- no budgeted groups. Callers should treat sql.ErrNoRows as "no group budget". SELECT - gaib.group_id, - gaib.spend_limit_micros -FROM group_ai_budgets gaib -JOIN group_members_expanded gme ON gme.group_id = gaib.group_id -WHERE gme.user_id = @user_id + budget.group_id, + budget.spend_limit_micros +FROM group_ai_budgets budget +JOIN group_members_expanded member ON member.group_id = budget.group_id +JOIN organizations ON organizations.id = member.organization_id +JOIN organization_members + ON organization_members.user_id = member.user_id + AND organization_members.organization_id = member.organization_id +WHERE member.user_id = @user_id + AND organizations.deleted = false ORDER BY - gaib.spend_limit_micros DESC, -- highest wins - gme.group_name ASC, -- alphabetical tiebreak - -- Final tiebreak on the group id makes the result deterministic when two - -- groups share both name and limit, which is possible across organizations - -- (groups are unique on (organization_id, name), not name alone). - gaib.group_id ASC + budget.spend_limit_micros DESC, -- highest wins + organization_members.created_at ASC, -- earliest organization membership + budget.group_id ASC -- deterministic tiebreak +LIMIT 1; + +-- name: GetUserEveryoneFallbackGroup :one +-- Returns the "Everyone" group (id == organization_id) to attribute a user's +-- spend to when no override or budgeted group applies. Prefers the default org, +-- then the earliest organization membership. Returns no rows when the user has +-- no organization membership. +SELECT organizations.id AS group_id +FROM organization_members +JOIN organizations ON organizations.id = organization_members.organization_id +WHERE organization_members.user_id = @user_id + AND organizations.deleted = false +ORDER BY + organizations.is_default DESC, -- prefer the default org + organization_members.created_at ASC, -- earliest organization membership + organizations.id ASC -- deterministic tiebreak LIMIT 1; -- name: IncrementUserAIDailySpend :one @@ -126,9 +144,9 @@ ORDER BY groups.id; -- Returns each user's AI spend attributed to the queried group, on or after -- period_start until NOW. Only current members of the queried group are -- returned. spend_limit_micros and limit_source are populated only when the --- queried group is the user's effective budget source. The effective_group_id --- is null when the user has no configured budget or when the effective group --- belongs to a different organization than the queried group. +-- queried group is the user's effective budget source. The effective group +-- falls back to the Everyone group, and effective_group_id is null only when +-- that group belongs to a different organization than the queried group. -- The period_start parameter is normalized to its UTC calendar day. -- TODO(AIGOV-527): unify effective group resolution in a single place. WITH queried_group AS ( @@ -154,15 +172,33 @@ user_highest_group AS ( budget.spend_limit_micros FROM group_ai_budgets budget JOIN group_members_expanded member ON member.group_id = budget.group_id + JOIN organizations ON organizations.id = member.organization_id + JOIN organization_members + ON organization_members.user_id = member.user_id + AND organization_members.organization_id = member.organization_id WHERE member.user_id IN (SELECT user_id FROM filtered_users) - ORDER BY member.user_id, budget.spend_limit_micros DESC, member.group_name ASC, budget.group_id ASC + AND organizations.deleted = false + ORDER BY member.user_id, budget.spend_limit_micros DESC, organization_members.created_at ASC, budget.group_id ASC +), +user_fallback_group AS ( + -- Per user, the Everyone group to fall back to when no override or budgeted + -- group applies. The Everyone group has id == organization_id. Prefers the + -- default org, then the earliest organization membership. + SELECT DISTINCT ON (organization_members.user_id) + organization_members.user_id, + organizations.id AS group_id + FROM organization_members + JOIN organizations ON organizations.id = organization_members.organization_id + WHERE organization_members.user_id IN (SELECT user_id FROM filtered_users) + AND organizations.deleted = false + ORDER BY organization_members.user_id, organizations.is_default DESC, organization_members.created_at ASC, organizations.id ASC ), effective AS ( - -- Effective budget per user: a per-user override wins over the - -- highest-limit group. + -- Effective budget per user: a per-user override wins over the highest-limit + -- group, which wins over the Everyone group fallback. SELECT filtered_users.user_id, - COALESCE(override.group_id, user_highest_group.group_id) AS raw_effective_group_id, + COALESCE(override.group_id, user_highest_group.group_id, user_fallback_group.group_id) AS raw_effective_group_id, COALESCE(override.spend_limit_micros, user_highest_group.spend_limit_micros) AS spend_limit_micros, (CASE WHEN override.group_id IS NOT NULL THEN 'user_override' @@ -171,6 +207,7 @@ effective AS ( FROM filtered_users LEFT JOIN user_ai_budget_overrides override ON override.user_id = filtered_users.user_id LEFT JOIN user_highest_group ON user_highest_group.user_id = filtered_users.user_id + LEFT JOIN user_fallback_group ON user_fallback_group.user_id = filtered_users.user_id ), applied_budget AS ( -- The limit and source only for users whose effective budget source is the diff --git a/codersdk/aibridge.go b/codersdk/aibridge.go index 4598ac3d555..34572240ab1 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 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). @@ -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 ddba40ad34a..ca6b0fad0d4 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 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.go b/enterprise/coderd/aibridge.go index 49ba9764b5d..f07b6bb18d0 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 aa975b7da7e..ca270ece8d4 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) @@ -3972,6 +4034,65 @@ func TestGroupMembersAISpend(t *testing.T) { require.Equal(t, int64(0), resp.Members[0].GroupSpendMicros) }) + t.Run("CrossOrgFallbackEveryoneMasked", func(t *testing.T) { + t.Parallel() + + dv := coderdtest.DeploymentValues(t) + dv.AI.BridgeConfig.Enabled = serpent.Bool(true) + dv.Experiments = []string{string(codersdk.ExperimentAIGatewayCostControl)} + db, ps := dbtestutil.NewDB(t) + ownerClient, owner := coderdenttest.New(t, &coderdenttest.Options{ + Options: &coderdtest.Options{DeploymentValues: dv, Database: db, Pubsub: ps}, + LicenseOptions: &coderdenttest.LicenseOptions{ + Features: license.Features{ + codersdk.FeatureTemplateRBAC: 1, + codersdk.FeatureAIBridge: 1, + codersdk.FeatureMultipleOrganizations: 1, + }, + }, + }) + userAdminClient, _ := coderdtest.CreateAnotherUser(t, ownerClient, owner.OrganizationID, rbac.RoleUserAdmin()) + ctx := testutil.Context(t, testutil.WaitLong) + + // Given: a member of the default org whose queried group lives in a + // non-default org, with no budget or override. The fallback prefers the + // default org's Everyone group, which lives in a different org than the + // queried group. + queriedOrg := coderdenttest.CreateOrganization(t, ownerClient, coderdenttest.CreateOrganizationOptions{}) + _, targetUser := coderdtest.CreateAnotherUser(t, ownerClient, owner.OrganizationID) + dbgen.OrganizationMember(t, db, database.OrganizationMember{UserID: targetUser.ID, OrganizationID: queriedOrg.ID}) + queried, err := userAdminClient.CreateGroup(ctx, queriedOrg.ID, codersdk.CreateGroupRequest{ + Name: "queried-cross-org-fallback-group", + }) + require.NoError(t, err) + _, err = userAdminClient.PatchGroup(ctx, queried.ID, codersdk.PatchGroupRequest{ + AddUsers: []string{targetUser.ID.String()}, + }) + require.NoError(t, err) + + // Spend is attributed to the queried group. + _, err = db.IncrementUserAIDailySpend(ctx, database.IncrementUserAIDailySpendParams{ + UserID: targetUser.ID, + EffectiveGroupID: queried.ID, + Day: dbtime.Now(), + CostMicros: 100_000_000, + }) + require.NoError(t, err) + + // When: the owner, who can read both orgs, queries the group. + //nolint:gocritic // The test asserts that even an owner sees the mask. + resp, err := ownerClient.GroupMembersAISpend(ctx, queried.ID, []uuid.UUID{targetUser.ID}) + require.NoError(t, err) + + // Then: effective_group_id is masked because the fallback Everyone group + // lives in the default org, while the queried group's spend still returns. + require.Len(t, resp.Members, 1) + require.Equal(t, targetUser.ID, resp.Members[0].UserID) + require.Nil(t, resp.Members[0].EffectiveGroupID, "cross-org fallback effective group must be masked") + require.Nil(t, resp.Members[0].GroupBudget) + require.Equal(t, int64(100_000_000), resp.Members[0].GroupSpendMicros) + }) + t.Run("OrgScopedRoute", func(t *testing.T) { t.Parallel() @@ -3997,7 +4118,7 @@ func TestGroupMembersAISpend(t *testing.T) { require.NoError(t, json.NewDecoder(res.Body).Decode(&got)) require.Len(t, got.Members, 1) require.Equal(t, targetUser.ID, got.Members[0].UserID) - require.Nil(t, got.Members[0].EffectiveGroupID) + require.Equal(t, &group.OrganizationID, got.Members[0].EffectiveGroupID) require.Nil(t, got.Members[0].GroupBudget) require.Equal(t, int64(0), got.Members[0].GroupSpendMicros) }) diff --git a/site/src/api/typesGenerated.ts b/site/src/api/typesGenerated.ts index 8a4c22d8dbb..4158bd880bc 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 membership. */ readonly effective_group_id: string | null; /**