mirror of
https://github.com/coder/coder.git
synced 2026-09-24 15:04:27 +08:00
fix: stop group spend limits from leaking across org boundaries (#24294)
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 > 🤖
This commit is contained in:
@@ -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 {
|
||||
|
||||
@@ -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{
|
||||
|
||||
@@ -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())
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
+6
-2
@@ -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.",
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
|
||||
@@ -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):
|
||||
|
||||
@@ -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")
|
||||
})
|
||||
}
|
||||
Reference in New Issue
Block a user