diff --git a/coderd/aibridge/budget/budget.go b/coderd/aibridge/budget/budget.go new file mode 100644 index 0000000000..b7aa3accba --- /dev/null +++ b/coderd/aibridge/budget/budget.go @@ -0,0 +1,76 @@ +// Package budget resolves the effective AI spend budget for a user. A +// per-user override always wins; otherwise the deployment budget policy selects +// a budget from the groups the user belongs to. +package budget + +import ( + "context" + "database/sql" + "errors" + + "github.com/google/uuid" + "golang.org/x/xerrors" + + "github.com/coder/coder/v2/coderd/database" + "github.com/coder/coder/v2/codersdk" +) + +// LimitSource identifies which tier produced an EffectiveBudget. +type LimitSource string + +const ( + // SourceUserOverride indicates the budget came from a per-user override. + SourceUserOverride LimitSource = "user_override" + // SourceGroup indicates the budget came from a group budget selected by the + // deployment policy. + SourceGroup LimitSource = "group" +) + +// EffectiveBudget is the AI budget that applies to a user after override and +// policy resolution. +type EffectiveBudget 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). + SpendLimitMicros int64 + Source LimitSource +} + +// 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. +func ResolveUserAIBudget(ctx context.Context, db database.Store, userID uuid.UUID, policy codersdk.AIBudgetPolicy) (EffectiveBudget, 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: SourceUserOverride, + }, true, nil + } + if !errors.Is(err, sql.ErrNoRows) { + return EffectiveBudget{}, false, xerrors.Errorf("get user AI budget override: %w", err) + } + + // No override: select a group budget according to the deployment policy. + switch policy { + case codersdk.AIBudgetPolicyHighest: + row, err := db.GetHighestGroupAIBudgetByUser(ctx, userID) + if errors.Is(err, sql.ErrNoRows) { + return EffectiveBudget{}, false, nil + } + if err != nil { + return EffectiveBudget{}, false, xerrors.Errorf("get highest group AI budget: %w", err) + } + return EffectiveBudget{ + GroupID: row.GroupID, + SpendLimitMicros: row.SpendLimitMicros, + Source: SourceGroup, + }, true, nil + default: + return EffectiveBudget{}, false, xerrors.Errorf("unsupported AI budget policy: %q", policy) + } +} diff --git a/coderd/aibridge/budget/budget_test.go b/coderd/aibridge/budget/budget_test.go new file mode 100644 index 0000000000..9171ac2e38 --- /dev/null +++ b/coderd/aibridge/budget/budget_test.go @@ -0,0 +1,206 @@ +package budget_test + +import ( + "bytes" + "context" + "testing" + + "github.com/google/uuid" + "github.com/stretchr/testify/require" + + "github.com/coder/coder/v2/coderd/aibridge/budget" + "github.com/coder/coder/v2/coderd/database" + "github.com/coder/coder/v2/coderd/database/dbgen" + "github.com/coder/coder/v2/coderd/database/dbtestutil" + "github.com/coder/coder/v2/codersdk" + "github.com/coder/coder/v2/testutil" +) + +func TestResolveUserAIBudget(t *testing.T) { + t.Parallel() + + // budgetedGroup creates a regular group in the org, adds the user to it, and + // sets a group AI budget. Returns the group ID. + budgetedGroup := func(t *testing.T, ctx context.Context, db database.Store, orgID, userID uuid.UUID, groupName string, spendLimit int64) uuid.UUID { + t.Helper() + g := dbgen.Group(t, db, database.Group{OrganizationID: orgID, Name: groupName}) + dbgen.GroupMember(t, db, database.GroupMemberTable{UserID: userID, GroupID: g.ID}) + _, err := db.UpsertGroupAIBudget(ctx, database.UpsertGroupAIBudgetParams{ + GroupID: g.ID, + SpendLimitMicros: spendLimit, + }) + require.NoError(t, err) + return g.ID + } + + // budgetedEveryoneGroup creates the org's "Everyone" group (id == org id), + // which is not auto-created for orgs built via dbgen, makes the user an org + // member so membership flows through organization_members, and sets a group + // AI budget. Returns the group ID. + budgetedEveryoneGroup := func(t *testing.T, ctx context.Context, db database.Store, orgID, userID uuid.UUID, spendLimit int64) uuid.UUID { + t.Helper() + g := dbgen.Group(t, db, database.Group{ID: orgID, OrganizationID: orgID, Name: "Everyone"}) + dbgen.OrganizationMember(t, db, database.OrganizationMember{OrganizationID: orgID, UserID: userID}) + _, err := db.UpsertGroupAIBudget(ctx, database.UpsertGroupAIBudgetParams{ + GroupID: g.ID, + SpendLimitMicros: spendLimit, + }) + require.NoError(t, err) + return g.ID + } + + 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) + wantErr string + }{ + { + name: "OverrideWins", + 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{}) + 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. + 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{ + UserID: user.ID, + GroupID: og.ID, + SpendLimitMicros: 1_000_000, + }) + require.NoError(t, err) + return user.ID, budget.EffectiveBudget{GroupID: og.ID, SpendLimitMicros: 1_000_000, Source: budget.SourceUserOverride}, true + }, + }, + { + name: "SingleGroupBudget", + 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{}) + user := dbgen.User(t, db, database.User{}) + gid := budgetedGroup(t, ctx, db, org.ID, user.ID, "only", 8_000_000) + return user.ID, budget.EffectiveBudget{GroupID: gid, SpendLimitMicros: 8_000_000, Source: budget.SourceGroup}, true + }, + }, + { + name: "HighestGroupWins", + 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{}) + user := dbgen.User(t, db, database.User{}) + budgetedGroup(t, ctx, db, org.ID, user.ID, "low", 5_000_000) + budgetedGroup(t, ctx, db, org.ID, user.ID, "mid", 20_000_000) + high := budgetedGroup(t, ctx, db, org.ID, user.ID, "high", 50_000_000) + return user.ID, budget.EffectiveBudget{GroupID: high, SpendLimitMicros: 50_000_000, Source: budget.SourceGroup}, true + }, + }, + { + name: "TieBrokenByName", + 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{}) + 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: budget.SourceGroup}, true + }, + }, + { + name: "TieBrokenByGroupID", + policy: codersdk.AIBudgetPolicyHighest, + setup: func(t *testing.T, ctx context.Context, db database.Store) (uuid.UUID, budget.EffectiveBudget, bool) { + user := dbgen.User(t, db, database.User{}) + // Two groups in different orgs share both name and limit. + // Group id breaks the tie, so resolution is deterministic. + org1 := dbgen.Organization(t, db, database.Organization{}) + org2 := dbgen.Organization(t, db, database.Organization{}) + g1 := budgetedGroup(t, ctx, db, org1.ID, user.ID, "dup", 10_000_000) + g2 := budgetedGroup(t, ctx, db, org2.ID, user.ID, "dup", 10_000_000) + winner := g1 + if bytes.Compare(g2[:], g1[:]) < 0 { + winner = g2 + } + return user.ID, budget.EffectiveBudget{GroupID: winner, SpendLimitMicros: 10_000_000, Source: budget.SourceGroup}, true + }, + }, + { + name: "GroupsButNoneBudgeted", + 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{}) + 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 + }, + }, + { + name: "EveryoneGroupBudget", + 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{}) + 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: budget.SourceGroup}, true + }, + }, + { + name: "OverrideBeatsEveryoneBudget", + 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{}) + 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 + // via organization_members, satisfying the membership trigger. + _, err := db.UpsertUserAIBudgetOverride(ctx, database.UpsertUserAIBudgetOverrideParams{ + UserID: user.ID, + GroupID: everyoneID, + SpendLimitMicros: 2_000_000, + }) + require.NoError(t, err) + return user.ID, budget.EffectiveBudget{GroupID: everyoneID, SpendLimitMicros: 2_000_000, Source: budget.SourceUserOverride}, true + }, + }, + { + name: "UnsupportedPolicy", + policy: codersdk.AIBudgetPolicy("unsupported"), + setup: func(t *testing.T, ctx context.Context, db database.Store) (uuid.UUID, budget.EffectiveBudget, bool) { + // No override, so resolution reaches the policy switch and errors. + user := dbgen.User(t, db, database.User{}) + return user.ID, budget.EffectiveBudget{}, 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.ResolveUserAIBudget(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.SpendLimitMicros, got.SpendLimitMicros) + require.Equal(t, want.Source, got.Source) + }) + } +} diff --git a/coderd/database/dbauthz/dbauthz.go b/coderd/database/dbauthz/dbauthz.go index 4b08644dec..65bd6174fd 100644 --- a/coderd/database/dbauthz/dbauthz.go +++ b/coderd/database/dbauthz/dbauthz.go @@ -3625,6 +3625,13 @@ func (q *querier) GetHealthSettings(ctx context.Context) (string, error) { return q.db.GetHealthSettings(ctx) } +func (q *querier) GetHighestGroupAIBudgetByUser(ctx context.Context, userID uuid.UUID) (database.GetHighestGroupAIBudgetByUserRow, error) { + if _, err := q.GetUserByID(ctx, userID); err != nil { // AuthZ check + return database.GetHighestGroupAIBudgetByUserRow{}, err + } + return q.db.GetHighestGroupAIBudgetByUser(ctx, userID) +} + func (q *querier) GetInboxNotificationByID(ctx context.Context, id uuid.UUID) (database.InboxNotification, error) { return fetchWithAction(q.log, q.auth, policy.ActionRead, q.db.GetInboxNotificationByID)(ctx, id) } diff --git a/coderd/database/dbauthz/dbauthz_test.go b/coderd/database/dbauthz/dbauthz_test.go index 916eca2319..1911ae7783 100644 --- a/coderd/database/dbauthz/dbauthz_test.go +++ b/coderd/database/dbauthz/dbauthz_test.go @@ -6502,6 +6502,14 @@ func (s *MethodTestSuite) TestAIBridge() { check.Args(user.ID).Asserts(user, policy.ActionRead).Returns(override) })) + s.Run("GetHighestGroupAIBudgetByUser", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { + user := testutil.Fake(s.T(), faker, database.User{}) + row := testutil.Fake(s.T(), faker, database.GetHighestGroupAIBudgetByUserRow{}) + dbm.EXPECT().GetUserByID(gomock.Any(), user.ID).Return(user, nil).AnyTimes() + dbm.EXPECT().GetHighestGroupAIBudgetByUser(gomock.Any(), user.ID).Return(row, nil).AnyTimes() + check.Args(user.ID).Asserts(user, policy.ActionRead).Returns(row) + })) + 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 cae6549e8d..e4acbec466 100644 --- a/coderd/database/dbmetrics/querymetrics.go +++ b/coderd/database/dbmetrics/querymetrics.go @@ -2001,6 +2001,14 @@ func (m queryMetricsStore) GetHealthSettings(ctx context.Context) (string, error return r0, r1 } +func (m queryMetricsStore) GetHighestGroupAIBudgetByUser(ctx context.Context, userID uuid.UUID) (database.GetHighestGroupAIBudgetByUserRow, error) { + start := time.Now() + r0, r1 := m.s.GetHighestGroupAIBudgetByUser(ctx, userID) + m.queryLatencies.WithLabelValues("GetHighestGroupAIBudgetByUser").Observe(time.Since(start).Seconds()) + m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "GetHighestGroupAIBudgetByUser").Inc() + return r0, r1 +} + func (m queryMetricsStore) GetInboxNotificationByID(ctx context.Context, id uuid.UUID) (database.InboxNotification, error) { start := time.Now() r0, r1 := m.s.GetInboxNotificationByID(ctx, id) diff --git a/coderd/database/dbmock/dbmock.go b/coderd/database/dbmock/dbmock.go index 80952fabee..ad8c0099d9 100644 --- a/coderd/database/dbmock/dbmock.go +++ b/coderd/database/dbmock/dbmock.go @@ -3720,6 +3720,21 @@ func (mr *MockStoreMockRecorder) GetHealthSettings(ctx any) *gomock.Call { return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetHealthSettings", reflect.TypeOf((*MockStore)(nil).GetHealthSettings), ctx) } +// GetHighestGroupAIBudgetByUser mocks base method. +func (m *MockStore) GetHighestGroupAIBudgetByUser(ctx context.Context, userID uuid.UUID) (database.GetHighestGroupAIBudgetByUserRow, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "GetHighestGroupAIBudgetByUser", ctx, userID) + ret0, _ := ret[0].(database.GetHighestGroupAIBudgetByUserRow) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// GetHighestGroupAIBudgetByUser indicates an expected call of GetHighestGroupAIBudgetByUser. +func (mr *MockStoreMockRecorder) GetHighestGroupAIBudgetByUser(ctx, userID any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetHighestGroupAIBudgetByUser", reflect.TypeOf((*MockStore)(nil).GetHighestGroupAIBudgetByUser), ctx, userID) +} + // GetInboxNotificationByID mocks base method. func (m *MockStore) GetInboxNotificationByID(ctx context.Context, id uuid.UUID) (database.InboxNotification, error) { m.ctrl.T.Helper() diff --git a/coderd/database/querier.go b/coderd/database/querier.go index 08a2b18155..39642dce7a 100644 --- a/coderd/database/querier.go +++ b/coderd/database/querier.go @@ -503,6 +503,13 @@ type sqlcQuerier interface { // A limit of 0 means "no limit". 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". + 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 // param user_id: The user ID diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go index 96258d521c..bddca21e0c 100644 --- a/coderd/database/queries.sql.go +++ b/coderd/database/queries.sql.go @@ -2617,6 +2617,41 @@ func (q *sqlQuerier) GetGroupAIBudget(ctx context.Context, groupID uuid.UUID) (G return i, err } +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 +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 +LIMIT 1 +` + +type GetHighestGroupAIBudgetByUserRow struct { + GroupID uuid.UUID `db:"group_id" json:"group_id"` + SpendLimitMicros int64 `db:"spend_limit_micros" json:"spend_limit_micros"` +} + +// 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". +func (q *sqlQuerier) GetHighestGroupAIBudgetByUser(ctx context.Context, userID uuid.UUID) (GetHighestGroupAIBudgetByUserRow, error) { + row := q.db.QueryRowContext(ctx, getHighestGroupAIBudgetByUser, userID) + var i GetHighestGroupAIBudgetByUserRow + err := row.Scan(&i.GroupID, &i.SpendLimitMicros) + return i, err +} + const getUserAIBudgetOverride = `-- name: GetUserAIBudgetOverride :one SELECT user_id, group_id, spend_limit_micros, created_at, updated_at FROM user_ai_budget_overrides diff --git a/coderd/database/queries/aicostcontrol.sql b/coderd/database/queries/aicostcontrol.sql index 188ec7357e..cad1ed6452 100644 --- a/coderd/database/queries/aicostcontrol.sql +++ b/coderd/database/queries/aicostcontrol.sql @@ -57,3 +57,25 @@ RETURNING *; -- name: DeleteUserAIBudgetOverride :one 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". +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 +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 +LIMIT 1;