feat(coderd): resolve effective user AI budget (#26142)

Closes https://linear.app/codercom/issue/AIGOV-287/add-effective-group-resolution Implements the effective AI budget resolution from the AI Governance cost-controls RFC: for a given user, a `user_ai_budget_overrides` row wins if present, otherwise the deployment budget policy (`highest`) picks the largest group budget across the user's groups, ties broken alphabetically. For now, I keep the logic under `coderd/aibridge/budget`, but that may change during the implementation of budget enforcement.
2026-09-24 15:04:27 +08:00 · 2026-06-10 14:01:20 +00:00
parent f9dfa18c46
commit 1cada0649c
9 changed files with 384 additions and 0 deletions
@@ -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)
+	}
+}
@@ -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)
+		})
+	}
+}
@@ -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)
 }
@@ -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{})
@@ -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)
@@ -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()
@@ -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
@@ -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
@@ -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;