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:
Cian Johnston
2026-04-14 16:56:17 +01:00
committed by GitHub
parent 4d4266a4ad
commit c552f9f281
12 changed files with 542 additions and 79 deletions
+7 -7
View File
@@ -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 {
+18 -8
View File
@@ -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{
+2 -2
View File
@@ -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())
+9 -9
View File
@@ -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.
+18 -5
View File
@@ -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
+65 -22
View File
@@ -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
+29 -8
View File
@@ -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
View File
@@ -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.",
+6 -6
View File
@@ -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
}
+1 -1
View File
@@ -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
}
+33 -9
View File
@@ -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")
})
}