From c552f9f281270d18f7914623664a31ac33132f97 Mon Sep 17 00:00:00 2001 From: Cian Johnston Date: Tue, 14 Apr 2026 16:56:17 +0100 Subject: [PATCH] fix: stop group spend limits from leaking across org boundaries (#24294) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Three SQL queries (`GetUserGroupSpendLimit`, `ResolveUserChatSpendLimit`, `GetUserChatSpendInPeriod`) aggregated chat spend limits and usage globally across all organizations. A restrictive group limit in org A would bleed into org B. ## Changes - Add `organization_id` parameter to all three SQL queries in `coderd/database/queries/chats.sql` - When nil UUID is passed, queries fall back to global behavior (backward compat for HTTP dashboard endpoints) - When real org ID is passed, limits and spend are scoped to that organization - Thread `organizationID` through `ResolveUsageLimitStatus` → `checkUsageLimit` → all chatd call sites - Update dbauthz wrappers for new param structs - HTTP endpoints (`chatCostSummary`, `getMyChatUsageLimitStatus`) pass `uuid.Nil` with TODO for future org-scoped UI - Add `TestResolveUsageLimitStatus_OrgScoped` with 5 test cases covering org isolation, nil-UUID fallback, spend scoping, and user override priority Closes coder/internal#1466 > 🤖 --- coderd/database/dbauthz/dbauthz.go | 14 +- coderd/database/dbauthz/dbauthz_test.go | 26 +- coderd/database/dbmetrics/querymetrics.go | 4 +- coderd/database/dbmock/dbmock.go | 18 +- coderd/database/querier.go | 23 +- coderd/database/queries.sql.go | 87 +++-- coderd/database/queries/chats.sql | 37 +- coderd/exp_chats.go | 8 +- coderd/x/chatd/chatd.go | 12 +- coderd/x/chatd/subagent.go | 2 +- coderd/x/chatd/usagelimit.go | 42 ++- enterprise/coderd/x/chatd/usagelimit_test.go | 348 +++++++++++++++++++ 12 files changed, 542 insertions(+), 79 deletions(-) create mode 100644 enterprise/coderd/x/chatd/usagelimit_test.go diff --git a/coderd/database/dbauthz/dbauthz.go b/coderd/database/dbauthz/dbauthz.go index c336d696e7..628495e7c9 100644 --- a/coderd/database/dbauthz/dbauthz.go +++ b/coderd/database/dbauthz/dbauthz.go @@ -4222,11 +4222,11 @@ func (q *querier) GetUserCount(ctx context.Context, includeSystem bool) (int64, return q.db.GetUserCount(ctx, includeSystem) } -func (q *querier) GetUserGroupSpendLimit(ctx context.Context, userID uuid.UUID) (int64, error) { - if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceChat.WithOwner(userID.String())); err != nil { +func (q *querier) GetUserGroupSpendLimit(ctx context.Context, arg database.GetUserGroupSpendLimitParams) (int64, error) { + if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceChat.WithOwner(arg.UserID.String())); err != nil { return 0, err } - return q.db.GetUserGroupSpendLimit(ctx, userID) + return q.db.GetUserGroupSpendLimit(ctx, arg) } func (q *querier) GetUserLatencyInsights(ctx context.Context, arg database.GetUserLatencyInsightsParams) ([]database.GetUserLatencyInsightsRow, error) { @@ -5816,11 +5816,11 @@ func (q *querier) RemoveUserFromGroups(ctx context.Context, arg database.RemoveU return q.db.RemoveUserFromGroups(ctx, arg) } -func (q *querier) ResolveUserChatSpendLimit(ctx context.Context, userID uuid.UUID) (int64, error) { - if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceChat.WithOwner(userID.String())); err != nil { - return 0, err +func (q *querier) ResolveUserChatSpendLimit(ctx context.Context, arg database.ResolveUserChatSpendLimitParams) (database.ResolveUserChatSpendLimitRow, error) { + if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceChat.WithOwner(arg.UserID.String())); err != nil { + return database.ResolveUserChatSpendLimitRow{}, err } - return q.db.ResolveUserChatSpendLimit(ctx, userID) + return q.db.ResolveUserChatSpendLimit(ctx, arg) } func (q *querier) RevokeDBCryptKey(ctx context.Context, activeKeyDigest string) error { diff --git a/coderd/database/dbauthz/dbauthz_test.go b/coderd/database/dbauthz/dbauthz_test.go index 7a6a5ec158..c9b74458b6 100644 --- a/coderd/database/dbauthz/dbauthz_test.go +++ b/coderd/database/dbauthz/dbauthz_test.go @@ -1125,7 +1125,9 @@ func (s *MethodTestSuite) TestChats() { })) s.Run("GetUserChatSpendInPeriod", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) { arg := database.GetUserChatSpendInPeriodParams{ - UserID: uuid.New(), + UserID: uuid.New(), + OrganizationID: uuid.NullUUID{UUID: uuid.New(), Valid: true}, + StartTime: time.Date(2025, 1, 1, 0, 0, 0, 0, time.UTC), EndTime: time.Date(2025, 2, 1, 0, 0, 0, 0, time.UTC), } @@ -1134,17 +1136,25 @@ func (s *MethodTestSuite) TestChats() { check.Args(arg).Asserts(rbac.ResourceChat.WithOwner(arg.UserID.String()), policy.ActionRead).Returns(spend) })) s.Run("GetUserGroupSpendLimit", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) { - userID := uuid.New() + arg := database.GetUserGroupSpendLimitParams{ + UserID: uuid.New(), + OrganizationID: uuid.NullUUID{UUID: uuid.New(), Valid: true}, + } limit := int64(456) - dbm.EXPECT().GetUserGroupSpendLimit(gomock.Any(), userID).Return(limit, nil).AnyTimes() - check.Args(userID).Asserts(rbac.ResourceChat.WithOwner(userID.String()), policy.ActionRead).Returns(limit) + dbm.EXPECT().GetUserGroupSpendLimit(gomock.Any(), arg).Return(limit, nil).AnyTimes() + check.Args(arg).Asserts(rbac.ResourceChat.WithOwner(arg.UserID.String()), policy.ActionRead).Returns(limit) })) + s.Run("ResolveUserChatSpendLimit", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) { - userID := uuid.New() - limit := int64(789) - dbm.EXPECT().ResolveUserChatSpendLimit(gomock.Any(), userID).Return(limit, nil).AnyTimes() - check.Args(userID).Asserts(rbac.ResourceChat.WithOwner(userID.String()), policy.ActionRead).Returns(limit) + arg := database.ResolveUserChatSpendLimitParams{ + UserID: uuid.New(), + OrganizationID: uuid.NullUUID{UUID: uuid.New(), Valid: true}, + } + row := database.ResolveUserChatSpendLimitRow{EffectiveLimitMicros: 789, LimitSource: "group"} + dbm.EXPECT().ResolveUserChatSpendLimit(gomock.Any(), arg).Return(row, nil).AnyTimes() + check.Args(arg).Asserts(rbac.ResourceChat.WithOwner(arg.UserID.String()), policy.ActionRead).Returns(row) })) + s.Run("GetChatUsageLimitConfig", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) { now := dbtime.Now() config := database.ChatUsageLimitConfig{ diff --git a/coderd/database/dbmetrics/querymetrics.go b/coderd/database/dbmetrics/querymetrics.go index a07cd8fa96..4150bdb129 100644 --- a/coderd/database/dbmetrics/querymetrics.go +++ b/coderd/database/dbmetrics/querymetrics.go @@ -2704,7 +2704,7 @@ func (m queryMetricsStore) GetUserCount(ctx context.Context, includeSystem bool) return r0, r1 } -func (m queryMetricsStore) GetUserGroupSpendLimit(ctx context.Context, userID uuid.UUID) (int64, error) { +func (m queryMetricsStore) GetUserGroupSpendLimit(ctx context.Context, userID database.GetUserGroupSpendLimitParams) (int64, error) { start := time.Now() r0, r1 := m.s.GetUserGroupSpendLimit(ctx, userID) m.queryLatencies.WithLabelValues("GetUserGroupSpendLimit").Observe(time.Since(start).Seconds()) @@ -4160,7 +4160,7 @@ func (m queryMetricsStore) RemoveUserFromGroups(ctx context.Context, arg databas return r0, r1 } -func (m queryMetricsStore) ResolveUserChatSpendLimit(ctx context.Context, userID uuid.UUID) (int64, error) { +func (m queryMetricsStore) ResolveUserChatSpendLimit(ctx context.Context, userID database.ResolveUserChatSpendLimitParams) (database.ResolveUserChatSpendLimitRow, error) { start := time.Now() r0, r1 := m.s.ResolveUserChatSpendLimit(ctx, userID) m.queryLatencies.WithLabelValues("ResolveUserChatSpendLimit").Observe(time.Since(start).Seconds()) diff --git a/coderd/database/dbmock/dbmock.go b/coderd/database/dbmock/dbmock.go index 141197ee32..4e6bd174bf 100644 --- a/coderd/database/dbmock/dbmock.go +++ b/coderd/database/dbmock/dbmock.go @@ -5058,18 +5058,18 @@ func (mr *MockStoreMockRecorder) GetUserCount(ctx, includeSystem any) *gomock.Ca } // GetUserGroupSpendLimit mocks base method. -func (m *MockStore) GetUserGroupSpendLimit(ctx context.Context, userID uuid.UUID) (int64, error) { +func (m *MockStore) GetUserGroupSpendLimit(ctx context.Context, arg database.GetUserGroupSpendLimitParams) (int64, error) { m.ctrl.T.Helper() - ret := m.ctrl.Call(m, "GetUserGroupSpendLimit", ctx, userID) + ret := m.ctrl.Call(m, "GetUserGroupSpendLimit", ctx, arg) ret0, _ := ret[0].(int64) ret1, _ := ret[1].(error) return ret0, ret1 } // GetUserGroupSpendLimit indicates an expected call of GetUserGroupSpendLimit. -func (mr *MockStoreMockRecorder) GetUserGroupSpendLimit(ctx, userID any) *gomock.Call { +func (mr *MockStoreMockRecorder) GetUserGroupSpendLimit(ctx, arg any) *gomock.Call { mr.mock.ctrl.T.Helper() - return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetUserGroupSpendLimit", reflect.TypeOf((*MockStore)(nil).GetUserGroupSpendLimit), ctx, userID) + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetUserGroupSpendLimit", reflect.TypeOf((*MockStore)(nil).GetUserGroupSpendLimit), ctx, arg) } // GetUserLatencyInsights mocks base method. @@ -7889,18 +7889,18 @@ func (mr *MockStoreMockRecorder) RemoveUserFromGroups(ctx, arg any) *gomock.Call } // ResolveUserChatSpendLimit mocks base method. -func (m *MockStore) ResolveUserChatSpendLimit(ctx context.Context, userID uuid.UUID) (int64, error) { +func (m *MockStore) ResolveUserChatSpendLimit(ctx context.Context, arg database.ResolveUserChatSpendLimitParams) (database.ResolveUserChatSpendLimitRow, error) { m.ctrl.T.Helper() - ret := m.ctrl.Call(m, "ResolveUserChatSpendLimit", ctx, userID) - ret0, _ := ret[0].(int64) + ret := m.ctrl.Call(m, "ResolveUserChatSpendLimit", ctx, arg) + ret0, _ := ret[0].(database.ResolveUserChatSpendLimitRow) ret1, _ := ret[1].(error) return ret0, ret1 } // ResolveUserChatSpendLimit indicates an expected call of ResolveUserChatSpendLimit. -func (mr *MockStoreMockRecorder) ResolveUserChatSpendLimit(ctx, userID any) *gomock.Call { +func (mr *MockStoreMockRecorder) ResolveUserChatSpendLimit(ctx, arg any) *gomock.Call { mr.mock.ctrl.T.Helper() - return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ResolveUserChatSpendLimit", reflect.TypeOf((*MockStore)(nil).ResolveUserChatSpendLimit), ctx, userID) + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ResolveUserChatSpendLimit", reflect.TypeOf((*MockStore)(nil).ResolveUserChatSpendLimit), ctx, arg) } // RevokeDBCryptKey mocks base method. diff --git a/coderd/database/querier.go b/coderd/database/querier.go index 54fd2b0132..3d00618488 100644 --- a/coderd/database/querier.go +++ b/coderd/database/querier.go @@ -642,11 +642,18 @@ type sqlcQuerier interface { GetUserChatCustomPrompt(ctx context.Context, userID uuid.UUID) (string, error) GetUserChatDebugLoggingEnabled(ctx context.Context, userID uuid.UUID) (bool, error) GetUserChatProviderKeys(ctx context.Context, userID uuid.UUID) ([]UserChatProviderKey, error) + // Returns the total spend for a user in the given period. + // When organization_id is NULL, spend across all organizations is + // returned (global behavior). Otherwise only spend within the + // specified organization is included. GetUserChatSpendInPeriod(ctx context.Context, arg GetUserChatSpendInPeriodParams) (int64, error) GetUserCount(ctx context.Context, includeSystem bool) (int64, error) // Returns the minimum (most restrictive) group limit for a user. - // Returns -1 if the user has no group limits applied. - GetUserGroupSpendLimit(ctx context.Context, userID uuid.UUID) (int64, error) + // Returns -1 if no group limits match the specified scope. + // When organization_id is NULL, groups across all organizations are + // considered (global behavior). Otherwise only groups within the + // specified organization are considered. + GetUserGroupSpendLimit(ctx context.Context, arg GetUserGroupSpendLimitParams) (int64, error) // GetUserLatencyInsights returns the median and 95th percentile connection // latency that users have experienced. The result can be filtered on // template_ids, meaning only user data from workspaces based on those templates @@ -907,11 +914,17 @@ type sqlcQuerier interface { RegisterWorkspaceProxy(ctx context.Context, arg RegisterWorkspaceProxyParams) (WorkspaceProxy, error) RemoveUserFromGroups(ctx context.Context, arg RemoveUserFromGroupsParams) ([]uuid.UUID, error) // Resolves the effective spend limit for a user using the hierarchy: - // 1. Individual user override (highest priority) - // 2. Minimum group limit across all user's groups + // 1. Individual user override (highest priority, applies globally across + // all organizations since it lives on the users table) + // 2. Minimum group limit across the user's groups // 3. Global default from config // Returns -1 if limits are not enabled. - ResolveUserChatSpendLimit(ctx context.Context, userID uuid.UUID) (int64, error) + // When organization_id is NULL, groups across all organizations are + // considered (global behavior). Otherwise only groups within the + // specified organization are considered. + // limit_source indicates which tier won: 'user', 'group', 'default', + // or 'disabled'. + ResolveUserChatSpendLimit(ctx context.Context, arg ResolveUserChatSpendLimitParams) (ResolveUserChatSpendLimitRow, error) RevokeDBCryptKey(ctx context.Context, activeKeyDigest string) error // Note that this selects from the CTE, not the original table. The CTE is named // the same as the original table to trick sqlc into reusing the existing struct diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go index bb542ef974..01b80b65ea 100644 --- a/coderd/database/queries.sql.go +++ b/coderd/database/queries.sql.go @@ -6747,19 +6747,31 @@ SELECT COALESCE(SUM(cm.total_cost_micros), 0)::bigint AS total_spend_micros FROM chat_messages cm JOIN chats c ON c.id = cm.chat_id WHERE c.owner_id = $1::uuid - AND cm.created_at >= $2::timestamptz - AND cm.created_at < $3::timestamptz + AND ($2::uuid IS NULL + OR c.organization_id = $2::uuid) + AND cm.created_at >= $3::timestamptz + AND cm.created_at < $4::timestamptz AND cm.total_cost_micros IS NOT NULL ` type GetUserChatSpendInPeriodParams struct { - UserID uuid.UUID `db:"user_id" json:"user_id"` - StartTime time.Time `db:"start_time" json:"start_time"` - EndTime time.Time `db:"end_time" json:"end_time"` + UserID uuid.UUID `db:"user_id" json:"user_id"` + OrganizationID uuid.NullUUID `db:"organization_id" json:"organization_id"` + StartTime time.Time `db:"start_time" json:"start_time"` + EndTime time.Time `db:"end_time" json:"end_time"` } +// Returns the total spend for a user in the given period. +// When organization_id is NULL, spend across all organizations is +// returned (global behavior). Otherwise only spend within the +// specified organization is included. func (q *sqlQuerier) GetUserChatSpendInPeriod(ctx context.Context, arg GetUserChatSpendInPeriodParams) (int64, error) { - row := q.db.QueryRowContext(ctx, getUserChatSpendInPeriod, arg.UserID, arg.StartTime, arg.EndTime) + row := q.db.QueryRowContext(ctx, getUserChatSpendInPeriod, + arg.UserID, + arg.OrganizationID, + arg.StartTime, + arg.EndTime, + ) var total_spend_micros int64 err := row.Scan(&total_spend_micros) return total_spend_micros, err @@ -6770,13 +6782,23 @@ SELECT COALESCE(MIN(g.chat_spend_limit_micros), -1)::bigint AS limit_micros FROM groups g JOIN group_members_expanded gme ON gme.group_id = g.id WHERE gme.user_id = $1::uuid + AND ($2::uuid IS NULL + OR g.organization_id = $2::uuid) AND g.chat_spend_limit_micros IS NOT NULL ` +type GetUserGroupSpendLimitParams struct { + UserID uuid.UUID `db:"user_id" json:"user_id"` + OrganizationID uuid.NullUUID `db:"organization_id" json:"organization_id"` +} + // Returns the minimum (most restrictive) group limit for a user. -// Returns -1 if the user has no group limits applied. -func (q *sqlQuerier) GetUserGroupSpendLimit(ctx context.Context, userID uuid.UUID) (int64, error) { - row := q.db.QueryRowContext(ctx, getUserGroupSpendLimit, userID) +// Returns -1 if no group limits match the specified scope. +// When organization_id is NULL, groups across all organizations are +// considered (global behavior). Otherwise only groups within the +// specified organization are considered. +func (q *sqlQuerier) GetUserGroupSpendLimit(ctx context.Context, arg GetUserGroupSpendLimitParams) (int64, error) { + row := q.db.QueryRowContext(ctx, getUserGroupSpendLimit, arg.UserID, arg.OrganizationID) var limit_micros int64 err := row.Scan(&limit_micros) return limit_micros, err @@ -7300,15 +7322,17 @@ func (q *sqlQuerier) PopNextQueuedMessage(ctx context.Context, chatID uuid.UUID) const resolveUserChatSpendLimit = `-- name: ResolveUserChatSpendLimit :one SELECT CASE - -- If limits are disabled, return -1. WHEN NOT cfg.enabled THEN -1 - -- Individual override takes priority. WHEN u.chat_spend_limit_micros IS NOT NULL THEN u.chat_spend_limit_micros - -- Group limit (minimum across all user's groups) is next. WHEN gl.limit_micros IS NOT NULL THEN gl.limit_micros - -- Fall back to global default. ELSE cfg.default_limit_micros -END::bigint AS effective_limit_micros +END::bigint AS effective_limit_micros, +CASE + WHEN NOT cfg.enabled THEN 'disabled' + WHEN u.chat_spend_limit_micros IS NOT NULL THEN 'user' + WHEN gl.limit_micros IS NOT NULL THEN 'group' + ELSE 'default' +END AS limit_source FROM chat_usage_limit_config cfg CROSS JOIN users u LEFT JOIN LATERAL ( @@ -7316,22 +7340,41 @@ LEFT JOIN LATERAL ( FROM groups g JOIN group_members_expanded gme ON gme.group_id = g.id WHERE gme.user_id = $1::uuid + AND ($2::uuid IS NULL + OR g.organization_id = $2::uuid) AND g.chat_spend_limit_micros IS NOT NULL ) gl ON TRUE WHERE u.id = $1::uuid LIMIT 1 ` +type ResolveUserChatSpendLimitParams struct { + UserID uuid.UUID `db:"user_id" json:"user_id"` + OrganizationID uuid.NullUUID `db:"organization_id" json:"organization_id"` +} + +type ResolveUserChatSpendLimitRow struct { + EffectiveLimitMicros int64 `db:"effective_limit_micros" json:"effective_limit_micros"` + LimitSource string `db:"limit_source" json:"limit_source"` +} + // Resolves the effective spend limit for a user using the hierarchy: -// 1. Individual user override (highest priority) -// 2. Minimum group limit across all user's groups -// 3. Global default from config +// 1. Individual user override (highest priority, applies globally across +// all organizations since it lives on the users table) +// 2. Minimum group limit across the user's groups +// 3. Global default from config +// // Returns -1 if limits are not enabled. -func (q *sqlQuerier) ResolveUserChatSpendLimit(ctx context.Context, userID uuid.UUID) (int64, error) { - row := q.db.QueryRowContext(ctx, resolveUserChatSpendLimit, userID) - var effective_limit_micros int64 - err := row.Scan(&effective_limit_micros) - return effective_limit_micros, err +// When organization_id is NULL, groups across all organizations are +// considered (global behavior). Otherwise only groups within the +// specified organization are considered. +// limit_source indicates which tier won: 'user', 'group', 'default', +// or 'disabled'. +func (q *sqlQuerier) ResolveUserChatSpendLimit(ctx context.Context, arg ResolveUserChatSpendLimitParams) (ResolveUserChatSpendLimitRow, error) { + row := q.db.QueryRowContext(ctx, resolveUserChatSpendLimit, arg.UserID, arg.OrganizationID) + var i ResolveUserChatSpendLimitRow + err := row.Scan(&i.EffectiveLimitMicros, &i.LimitSource) + return i, err } const softDeleteChatMessageByID = `-- name: SoftDeleteChatMessageByID :exec diff --git a/coderd/database/queries/chats.sql b/coderd/database/queries/chats.sql index caa3552b3f..04887c91e0 100644 --- a/coderd/database/queries/chats.sql +++ b/coderd/database/queries/chats.sql @@ -1134,10 +1134,16 @@ FROM users WHERE id = @user_id::uuid AND chat_spend_limit_micros IS NOT NULL; -- name: GetUserChatSpendInPeriod :one +-- Returns the total spend for a user in the given period. +-- When organization_id is NULL, spend across all organizations is +-- returned (global behavior). Otherwise only spend within the +-- specified organization is included. SELECT COALESCE(SUM(cm.total_cost_micros), 0)::bigint AS total_spend_micros FROM chat_messages cm JOIN chats c ON c.id = cm.chat_id WHERE c.owner_id = @user_id::uuid + AND (sqlc.narg('organization_id')::uuid IS NULL + OR c.organization_id = sqlc.narg('organization_id')::uuid) AND cm.created_at >= @start_time::timestamptz AND cm.created_at < @end_time::timestamptz AND cm.total_cost_micros IS NOT NULL; @@ -1189,11 +1195,16 @@ WHERE id = @group_id::uuid AND chat_spend_limit_micros IS NOT NULL; -- name: GetUserGroupSpendLimit :one -- Returns the minimum (most restrictive) group limit for a user. --- Returns -1 if the user has no group limits applied. +-- Returns -1 if no group limits match the specified scope. +-- When organization_id is NULL, groups across all organizations are +-- considered (global behavior). Otherwise only groups within the +-- specified organization are considered. SELECT COALESCE(MIN(g.chat_spend_limit_micros), -1)::bigint AS limit_micros FROM groups g JOIN group_members_expanded gme ON gme.group_id = g.id WHERE gme.user_id = @user_id::uuid + AND (sqlc.narg('organization_id')::uuid IS NULL + OR g.organization_id = sqlc.narg('organization_id')::uuid) AND g.chat_spend_limit_micros IS NOT NULL; -- name: GetChatsByWorkspaceIDs :many @@ -1205,20 +1216,28 @@ ORDER BY workspace_id, updated_at DESC; -- name: ResolveUserChatSpendLimit :one -- Resolves the effective spend limit for a user using the hierarchy: --- 1. Individual user override (highest priority) --- 2. Minimum group limit across all user's groups +-- 1. Individual user override (highest priority, applies globally across +-- all organizations since it lives on the users table) +-- 2. Minimum group limit across the user's groups -- 3. Global default from config -- Returns -1 if limits are not enabled. +-- When organization_id is NULL, groups across all organizations are +-- considered (global behavior). Otherwise only groups within the +-- specified organization are considered. +-- limit_source indicates which tier won: 'user', 'group', 'default', +-- or 'disabled'. SELECT CASE - -- If limits are disabled, return -1. WHEN NOT cfg.enabled THEN -1 - -- Individual override takes priority. WHEN u.chat_spend_limit_micros IS NOT NULL THEN u.chat_spend_limit_micros - -- Group limit (minimum across all user's groups) is next. WHEN gl.limit_micros IS NOT NULL THEN gl.limit_micros - -- Fall back to global default. ELSE cfg.default_limit_micros -END::bigint AS effective_limit_micros +END::bigint AS effective_limit_micros, +CASE + WHEN NOT cfg.enabled THEN 'disabled' + WHEN u.chat_spend_limit_micros IS NOT NULL THEN 'user' + WHEN gl.limit_micros IS NOT NULL THEN 'group' + ELSE 'default' +END AS limit_source FROM chat_usage_limit_config cfg CROSS JOIN users u LEFT JOIN LATERAL ( @@ -1226,6 +1245,8 @@ LEFT JOIN LATERAL ( FROM groups g JOIN group_members_expanded gme ON gme.group_id = g.id WHERE gme.user_id = @user_id::uuid + AND (sqlc.narg('organization_id')::uuid IS NULL + OR g.organization_id = sqlc.narg('organization_id')::uuid) AND g.chat_spend_limit_micros IS NOT NULL ) gl ON TRUE WHERE u.id = @user_id::uuid diff --git a/coderd/exp_chats.go b/coderd/exp_chats.go index 23b86c7856..c626120371 100644 --- a/coderd/exp_chats.go +++ b/coderd/exp_chats.go @@ -804,7 +804,9 @@ func (api *API) chatCostSummary(rw http.ResponseWriter, r *http.Request) { chatBreakdowns = append(chatBreakdowns, convertChatCostChatBreakdown(chat)) } - usageStatus, err := chatd.ResolveUsageLimitStatus(ctx, api.Database, targetUser.ID, time.Now()) + // TODO(CODAGT-161): pass real organization ID + // when the HTTP endpoint supports org-scoped queries. + usageStatus, err := chatd.ResolveUsageLimitStatus(ctx, api.Database, targetUser.ID, uuid.NullUUID{}, time.Now()) if err != nil { api.Logger.Warn(ctx, "failed to resolve usage limit status", slog.Error(err)) } @@ -1105,7 +1107,9 @@ func (api *API) updateChatUsageLimitConfig(rw http.ResponseWriter, r *http.Reque //nolint:revive // HTTP handler writes to ResponseWriter. func (api *API) getMyChatUsageLimitStatus(rw http.ResponseWriter, r *http.Request) { ctx := r.Context() - status, err := chatd.ResolveUsageLimitStatus(ctx, api.Database, httpmw.APIKey(r).UserID, time.Now()) + // TODO(CODAGT-161): pass real organization ID + // when the HTTP endpoint supports org-scoped queries. + status, err := chatd.ResolveUsageLimitStatus(ctx, api.Database, httpmw.APIKey(r).UserID, uuid.NullUUID{}, time.Now()) if err != nil { httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{ Message: "Failed to get chat usage limit status.", diff --git a/coderd/x/chatd/chatd.go b/coderd/x/chatd/chatd.go index f17d7e8c9f..1a2bdd425a 100644 --- a/coderd/x/chatd/chatd.go +++ b/coderd/x/chatd/chatd.go @@ -878,7 +878,7 @@ func (p *Server) CreateChat(ctx context.Context, opts CreateOptions) (database.C var chat database.Chat txErr := p.db.InTx(func(tx database.Store) error { - if limitErr := p.checkUsageLimit(ctx, tx, opts.OwnerID); limitErr != nil { + if limitErr := p.checkUsageLimit(ctx, tx, opts.OwnerID, uuid.NullUUID{UUID: opts.OrganizationID, Valid: true}); limitErr != nil { return limitErr } @@ -1047,7 +1047,7 @@ func (p *Server) SendMessage( } // Enforce usage limits before queueing or inserting. - if limitErr := p.checkUsageLimit(ctx, tx, lockedChat.OwnerID); limitErr != nil { + if limitErr := p.checkUsageLimit(ctx, tx, lockedChat.OwnerID, uuid.NullUUID{UUID: lockedChat.OrganizationID, Valid: true}); limitErr != nil { return limitErr } @@ -1169,8 +1169,8 @@ func (p *Server) SendMessage( return result, nil } -func (p *Server) checkUsageLimit(ctx context.Context, store database.Store, ownerID uuid.UUID) error { - status, err := ResolveUsageLimitStatus(ctx, store, ownerID, time.Now()) +func (p *Server) checkUsageLimit(ctx context.Context, store database.Store, ownerID uuid.UUID, organizationID uuid.NullUUID) error { + status, err := ResolveUsageLimitStatus(ctx, store, ownerID, organizationID, time.Now()) if err != nil { // Fail open: never block chat due to a limit-resolution failure. p.logger.Warn(ctx, "usage limit check failed, allowing message", @@ -1223,7 +1223,7 @@ func (p *Server) EditMessage( return xerrors.Errorf("lock chat: %w", err) } - if limitErr := p.checkUsageLimit(ctx, tx, lockedChat.OwnerID); limitErr != nil { + if limitErr := p.checkUsageLimit(ctx, tx, lockedChat.OwnerID, uuid.NullUUID{UUID: lockedChat.OrganizationID, Valid: true}); limitErr != nil { return limitErr } @@ -2028,7 +2028,7 @@ func (p *Server) regenerateChatTitleWithStore( chat database.Chat, keys chatprovider.ProviderAPIKeys, ) (database.Chat, error) { - if limitErr := p.checkUsageLimit(ctx, store, chat.OwnerID); limitErr != nil { + if limitErr := p.checkUsageLimit(ctx, store, chat.OwnerID, uuid.NullUUID{UUID: chat.OrganizationID, Valid: true}); limitErr != nil { return database.Chat{}, limitErr } diff --git a/coderd/x/chatd/subagent.go b/coderd/x/chatd/subagent.go index 300ce08aeb..4792df2b66 100644 --- a/coderd/x/chatd/subagent.go +++ b/coderd/x/chatd/subagent.go @@ -455,7 +455,7 @@ func (p *Server) createChildSubagentChatWithOptions( var child database.Chat txErr := p.db.InTx(func(tx database.Store) error { - if limitErr := p.checkUsageLimit(ctx, tx, parent.OwnerID); limitErr != nil { + if limitErr := p.checkUsageLimit(ctx, tx, parent.OwnerID, uuid.NullUUID{UUID: parent.OrganizationID, Valid: true}); limitErr != nil { return limitErr } diff --git a/coderd/x/chatd/usagelimit.go b/coderd/x/chatd/usagelimit.go index 12535421d4..cbe67f50e1 100644 --- a/coderd/x/chatd/usagelimit.go +++ b/coderd/x/chatd/usagelimit.go @@ -43,7 +43,10 @@ func ComputeUsagePeriodBounds(now time.Time, period codersdk.ChatUsageLimitPerio return start, end } -// ResolveUsageLimitStatus resolves the current usage-limit status for userID. +// ResolveUsageLimitStatus resolves the current usage-limit status for +// userID within organizationID. When organizationID is invalid (Valid +// == false), limits and spend are computed globally across all +// organizations (legacy behavior). // // Note: There is a potential race condition where two concurrent messages // from the same user can both pass the limit check if processed in @@ -60,7 +63,7 @@ func ComputeUsagePeriodBounds(now time.Time, period codersdk.ChatUsageLimitPerio // Then scan spend once over the widest active window with conditional SUMs // for each period and compare each spend/limit pair Go-side, blocking on // whichever period is tightest. -func ResolveUsageLimitStatus(ctx context.Context, db database.Store, userID uuid.UUID, now time.Time) (*codersdk.ChatUsageLimitStatus, error) { +func ResolveUsageLimitStatus(ctx context.Context, db database.Store, userID uuid.UUID, organizationID uuid.NullUUID, now time.Time) (*codersdk.ChatUsageLimitStatus, error) { //nolint:gocritic // AsChatd provides narrowly-scoped daemon access for // deployment config reads and cross-user chat spend aggregation. authCtx := dbauthz.AsChatd(ctx) @@ -83,27 +86,41 @@ func ResolveUsageLimitStatus(ctx context.Context, db database.Store, userID uuid // Resolve effective limit in a single query: // individual override > group limit > global default. - effectiveLimit, err := db.ResolveUserChatSpendLimit(authCtx, userID) + limitResult, err := db.ResolveUserChatSpendLimit(authCtx, database.ResolveUserChatSpendLimitParams{ + UserID: userID, + OrganizationID: organizationID, + }) if err != nil { return nil, err } - // -1 means limits are disabled (shouldn't happen since we checked above, - // but handle gracefully). - if effectiveLimit < 0 { + // -1 means limits are disabled (shouldn't happen since we checked + // above, but handle gracefully). + if limitResult.EffectiveLimitMicros < 0 { return nil, nil //nolint:nilnil // Nil status cleanly signals disabled limits. } start, end := ComputeUsagePeriodBounds(now, period) + // When the winning limit tier is org-scoped (group), scope spend + // to the same org. When the limit is global (user override or + // deployment default), check spend globally to prevent a user + // from exceeding their limit by spreading spend across orgs. + spendOrgID := organizationID + if limitResult.LimitSource != limitSourceGroup { + spendOrgID = uuid.NullUUID{} + } + spendTotal, err := db.GetUserChatSpendInPeriod(authCtx, database.GetUserChatSpendInPeriodParams{ - UserID: userID, - StartTime: start, - EndTime: end, + UserID: userID, + OrganizationID: spendOrgID, + StartTime: start, + EndTime: end, }) if err != nil { return nil, err } + effectiveLimit := limitResult.EffectiveLimitMicros return &codersdk.ChatUsageLimitStatus{ IsLimited: true, Period: period, @@ -114,6 +131,13 @@ func ResolveUsageLimitStatus(ctx context.Context, db database.Store, userID uuid }, nil } +// Limit source constants returned by ResolveUserChatSpendLimit. +const ( + limitSourceUser = "user" + limitSourceGroup = "group" + limitSourceDefault = "default" +) + func mapDBPeriodToSDK(dbPeriod string) (codersdk.ChatUsageLimitPeriod, bool) { switch dbPeriod { case string(codersdk.ChatUsageLimitPeriodDay): diff --git a/enterprise/coderd/x/chatd/usagelimit_test.go b/enterprise/coderd/x/chatd/usagelimit_test.go new file mode 100644 index 0000000000..7dbfa0cbb0 --- /dev/null +++ b/enterprise/coderd/x/chatd/usagelimit_test.go @@ -0,0 +1,348 @@ +package chatd_test + +import ( + "encoding/json" + "testing" + "time" + + "github.com/google/uuid" + "github.com/stretchr/testify/require" + + "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/coderd/x/chatd" + "github.com/coder/coder/v2/codersdk" + "github.com/coder/coder/v2/testutil" +) + +func TestResolveUsageLimitStatus_OrgScoped(t *testing.T) { + t.Parallel() + + db, _ := dbtestutil.NewDB(t) + ctx := testutil.Context(t, testutil.WaitLong) + + // Create two orgs and a user in both. + orgA := dbgen.Organization(t, db, database.Organization{}) + orgB := dbgen.Organization(t, db, database.Organization{}) + user := dbgen.User(t, db, database.User{}) + dbgen.OrganizationMember(t, db, database.OrganizationMember{ + UserID: user.ID, + OrganizationID: orgA.ID, + }) + dbgen.OrganizationMember(t, db, database.OrganizationMember{ + UserID: user.ID, + OrganizationID: orgB.ID, + }) + + // Create groups with different spend limits. + // groupA ($5) and groupA2 ($20) are both in orgA to exercise + // MIN aggregation within a single org. + groupA := dbgen.Group(t, db, database.Group{ + OrganizationID: orgA.ID, + }) + groupA2 := dbgen.Group(t, db, database.Group{ + OrganizationID: orgA.ID, + }) + groupB := dbgen.Group(t, db, database.Group{ + OrganizationID: orgB.ID, + }) + dbgen.GroupMember(t, db, database.GroupMemberTable{ + UserID: user.ID, + GroupID: groupA.ID, + }) + dbgen.GroupMember(t, db, database.GroupMemberTable{ + UserID: user.ID, + GroupID: groupA2.ID, + }) + dbgen.GroupMember(t, db, database.GroupMemberTable{ + UserID: user.ID, + GroupID: groupB.ID, + }) + + // Set group spend limits: groupA=$5, groupA2=$20, groupB=$50. + _, err := db.UpsertChatUsageLimitGroupOverride(ctx, database.UpsertChatUsageLimitGroupOverrideParams{ + GroupID: groupA.ID, + SpendLimitMicros: 5_000_000, + }) + require.NoError(t, err) + _, err = db.UpsertChatUsageLimitGroupOverride(ctx, database.UpsertChatUsageLimitGroupOverrideParams{ + GroupID: groupA2.ID, + SpendLimitMicros: 20_000_000, + }) + require.NoError(t, err) + _, err = db.UpsertChatUsageLimitGroupOverride(ctx, database.UpsertChatUsageLimitGroupOverrideParams{ + GroupID: groupB.ID, + SpendLimitMicros: 50_000_000, + }) + require.NoError(t, err) + + // Enable usage limits with a high default so group limits win. + _, err = db.UpsertChatUsageLimitConfig(ctx, database.UpsertChatUsageLimitConfigParams{ + Enabled: true, + DefaultLimitMicros: 100_000_000, + Period: string(codersdk.ChatUsageLimitPeriodMonth), + }) + require.NoError(t, err) + + // We need a chat provider + model config for inserting chats. + _, err = db.InsertChatProvider(ctx, database.InsertChatProviderParams{ + Provider: "openai", + DisplayName: "openai", + APIKey: "test-key", + CreatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, + Enabled: true, + CentralApiKeyEnabled: true, + }) + require.NoError(t, err) + modelConfig, err := db.InsertChatModelConfig(ctx, database.InsertChatModelConfigParams{ + Provider: "openai", + Model: "gpt-4o-mini", + DisplayName: "Test Model", + CreatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, + UpdatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, + Enabled: true, + IsDefault: true, + ContextLimit: 128000, + CompressionThreshold: 70, + Options: json.RawMessage(`{}`), + }) + require.NoError(t, err) + + now := time.Now().UTC() + + // insertChatWithSpend is a test helper that creates a chat in the + // given org and inserts a single message with the specified cost. + insertChatWithSpend := func(t *testing.T, ownerID, orgID, modelCfgID uuid.UUID, costMicros int64) { + t.Helper() + tctx := testutil.Context(t, testutil.WaitLong) + c, err := db.InsertChat(tctx, database.InsertChatParams{ + OrganizationID: orgID, + OwnerID: ownerID, + LastModelConfigID: modelCfgID, + Title: "test chat", + Status: database.ChatStatusWaiting, + MCPServerIDs: []uuid.UUID{}, + }) + require.NoError(t, err) + _, err = db.InsertChatMessages(tctx, database.InsertChatMessagesParams{ + ChatID: c.ID, + CreatedBy: []uuid.UUID{uuid.Nil}, + ModelConfigID: []uuid.UUID{modelCfgID}, + Role: []database.ChatMessageRole{database.ChatMessageRoleAssistant}, + Content: []string{`[{"type":"text","text":"hello"}]`}, + ContentVersion: []int16{1}, + Visibility: []database.ChatMessageVisibility{database.ChatMessageVisibilityBoth}, + InputTokens: []int64{100}, + OutputTokens: []int64{50}, + TotalTokens: []int64{150}, + ReasoningTokens: []int64{0}, + CacheCreationTokens: []int64{0}, + CacheReadTokens: []int64{0}, + ContextLimit: []int64{128000}, + Compressed: []bool{false}, + TotalCostMicros: []int64{costMicros}, + RuntimeMs: []int64{500}, + ProviderResponseID: []string{uuid.NewString()}, + }) + require.NoError(t, err) + } + + t.Run("OrgA_gets_orgA_limit", func(t *testing.T) { + t.Parallel() + ctx := testutil.Context(t, testutil.WaitLong) + // orgA has groupA ($5) and groupA2 ($20). MIN($5, $20) = $5. + status, err := chatd.ResolveUsageLimitStatus(ctx, db, user.ID, uuid.NullUUID{UUID: orgA.ID, Valid: true}, now) + require.NoError(t, err) + require.NotNil(t, status) + require.NotNil(t, status.SpendLimitMicros) + require.Equal(t, int64(5_000_000), *status.SpendLimitMicros, + "orgA should resolve to MIN of both groups ($5, $20) = $5") + }) + + t.Run("OrgB_gets_orgB_limit", func(t *testing.T) { + t.Parallel() + ctx := testutil.Context(t, testutil.WaitLong) + status, err := chatd.ResolveUsageLimitStatus(ctx, db, user.ID, uuid.NullUUID{UUID: orgB.ID, Valid: true}, now) + require.NoError(t, err) + require.NotNil(t, status) + require.NotNil(t, status.SpendLimitMicros) + require.Equal(t, int64(50_000_000), *status.SpendLimitMicros, + "orgB should resolve to groupB's $50 limit, not global MIN") + }) + + t.Run("UnknownOrg_gets_global_default", func(t *testing.T) { + t.Parallel() + ctx := testutil.Context(t, testutil.WaitLong) + // When the org ID does not match any group the user belongs + // to, MIN() over an empty set returns NULL, the CASE sees + // gl.limit_micros IS NOT NULL as false, and falls through + // to the global default. This subtest guards that contract: + // if someone changes the NULL-handling in + // ResolveUserChatSpendLimit, this will catch it. + randomOrg := uuid.NullUUID{UUID: uuid.New(), Valid: true} + status, err := chatd.ResolveUsageLimitStatus(ctx, db, user.ID, randomOrg, now) + require.NoError(t, err) + require.NotNil(t, status) + require.NotNil(t, status.SpendLimitMicros) + require.Equal(t, int64(100_000_000), *status.SpendLimitMicros, + "org with no matching groups should fall through to global default ($100)") + }) + + t.Run("NilOrg_gets_global_min", func(t *testing.T) { + t.Parallel() + ctx := testutil.Context(t, testutil.WaitLong) + // NULL org = global behavior: MIN across all groups. + status, err := chatd.ResolveUsageLimitStatus(ctx, db, user.ID, uuid.NullUUID{}, now) + require.NoError(t, err) + require.NotNil(t, status) + require.NotNil(t, status.SpendLimitMicros) + require.Equal(t, int64(5_000_000), *status.SpendLimitMicros, + "nil org should fall back to global MIN($5, $20, $50) = $5") + }) + + t.Run("Spend_scoped_to_org", func(t *testing.T) { + t.Parallel() + ctx := testutil.Context(t, testutil.WaitLong) + + // Dedicated user so spend insertion doesn't affect sibling subtests. + spendUser := dbgen.User(t, db, database.User{}) + dbgen.OrganizationMember(t, db, database.OrganizationMember{ + UserID: spendUser.ID, + OrganizationID: orgA.ID, + }) + dbgen.OrganizationMember(t, db, database.OrganizationMember{ + UserID: spendUser.ID, + OrganizationID: orgB.ID, + }) + dbgen.GroupMember(t, db, database.GroupMemberTable{ + UserID: spendUser.ID, + GroupID: groupA.ID, + }) + dbgen.GroupMember(t, db, database.GroupMemberTable{ + UserID: spendUser.ID, + GroupID: groupB.ID, + }) + + insertChatWithSpend(t, spendUser.ID, orgA.ID, modelConfig.ID, 3_000_000) + + // Resolve for orgB: should see zero spend (orgA's $3 not counted). + statusB, err := chatd.ResolveUsageLimitStatus(ctx, db, spendUser.ID, uuid.NullUUID{UUID: orgB.ID, Valid: true}, now) + require.NoError(t, err) + require.NotNil(t, statusB) + require.Equal(t, int64(0), statusB.CurrentSpend, + "orgB should not include orgA's spend") + + // Resolve for orgA: should see $3 spend. + statusA, err := chatd.ResolveUsageLimitStatus(ctx, db, spendUser.ID, uuid.NullUUID{UUID: orgA.ID, Valid: true}, now) + require.NoError(t, err) + require.NotNil(t, statusA) + require.Equal(t, int64(3_000_000), statusA.CurrentSpend, + "orgA should include its own spend") + + // Nil org: should see $3 (global). + statusNil, err := chatd.ResolveUsageLimitStatus(ctx, db, spendUser.ID, uuid.NullUUID{}, now) + require.NoError(t, err) + require.NotNil(t, statusNil) + require.Equal(t, int64(3_000_000), statusNil.CurrentSpend, + "nil org should include all spend globally") + }) + + t.Run("User_override_beats_group", func(t *testing.T) { + t.Parallel() + ctx := testutil.Context(t, testutil.WaitLong) + // Create a separate user with a personal override. + user2 := dbgen.User(t, db, database.User{}) + dbgen.OrganizationMember(t, db, database.OrganizationMember{ + UserID: user2.ID, + OrganizationID: orgA.ID, + }) + dbgen.GroupMember(t, db, database.GroupMemberTable{ + UserID: user2.ID, + GroupID: groupA.ID, + }) + + // Set $10 user override (beats groupA's $5 limit). + _, err := db.UpsertChatUsageLimitUserOverride(ctx, database.UpsertChatUsageLimitUserOverrideParams{ + UserID: user2.ID, + SpendLimitMicros: 10_000_000, + }) + require.NoError(t, err) + + status, err := chatd.ResolveUsageLimitStatus(ctx, db, user2.ID, uuid.NullUUID{UUID: orgA.ID, Valid: true}, now) + require.NoError(t, err) + require.NotNil(t, status) + require.NotNil(t, status.SpendLimitMicros) + require.Equal(t, int64(10_000_000), *status.SpendLimitMicros, + "user override should take priority over group limit") + }) + + t.Run("UserOverride_spend_is_global", func(t *testing.T) { + t.Parallel() + // When user override wins, spend should be checked globally, + // not per-org. Otherwise a user in N orgs can spend limit*N. + user3 := dbgen.User(t, db, database.User{}) + dbgen.OrganizationMember(t, db, database.OrganizationMember{ + UserID: user3.ID, + OrganizationID: orgA.ID, + }) + dbgen.OrganizationMember(t, db, database.OrganizationMember{ + UserID: user3.ID, + OrganizationID: orgB.ID, + }) + + // Set $10 user override. + _, err := db.UpsertChatUsageLimitUserOverride(testutil.Context(t, testutil.WaitLong), database.UpsertChatUsageLimitUserOverrideParams{ + UserID: user3.ID, + SpendLimitMicros: 10_000_000, + }) + require.NoError(t, err) + + // $6 in orgA + $6 in orgB = $12 total. + insertChatWithSpend(t, user3.ID, orgA.ID, modelConfig.ID, 6_000_000) + insertChatWithSpend(t, user3.ID, orgB.ID, modelConfig.ID, 6_000_000) + + ctx := testutil.Context(t, testutil.WaitLong) + status, err := chatd.ResolveUsageLimitStatus(ctx, db, user3.ID, uuid.NullUUID{UUID: orgA.ID, Valid: true}, now) + require.NoError(t, err) + require.NotNil(t, status) + require.NotNil(t, status.SpendLimitMicros) + require.Equal(t, int64(10_000_000), *status.SpendLimitMicros) + // Spend should be global ($12), not org-scoped ($6). + require.Equal(t, int64(12_000_000), status.CurrentSpend, + "user override should check global spend to prevent cross-org evasion") + }) + + t.Run("GlobalDefault_spend_is_global", func(t *testing.T) { + t.Parallel() + // When global default wins (no groups in the target org, + // no user override), spend should also be checked globally. + user4 := dbgen.User(t, db, database.User{}) + orgC := dbgen.Organization(t, db, database.Organization{}) + dbgen.OrganizationMember(t, db, database.OrganizationMember{ + UserID: user4.ID, + OrganizationID: orgA.ID, + }) + dbgen.OrganizationMember(t, db, database.OrganizationMember{ + UserID: user4.ID, + OrganizationID: orgC.ID, + }) + + // $30 in orgA + $40 in orgC = $70 total. + insertChatWithSpend(t, user4.ID, orgA.ID, modelConfig.ID, 30_000_000) + insertChatWithSpend(t, user4.ID, orgC.ID, modelConfig.ID, 40_000_000) + + ctx := testutil.Context(t, testutil.WaitLong) + // user4 has no groups in orgC, no override: falls through + // to global default ($100). + status, err := chatd.ResolveUsageLimitStatus(ctx, db, user4.ID, uuid.NullUUID{UUID: orgC.ID, Valid: true}, now) + require.NoError(t, err) + require.NotNil(t, status) + require.NotNil(t, status.SpendLimitMicros) + require.Equal(t, int64(100_000_000), *status.SpendLimitMicros, + "should fall through to global default ($100)") + // Spend should be global ($70), not org-scoped ($40). + require.Equal(t, int64(70_000_000), status.CurrentSpend, + "global default should check global spend") + }) +}