opensource/coder
3
0
Fork 0
mirror of https://github.com/coder/coder.git synced 2026-09-24 15:04:27 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat: remove native chat usage limits in favor of AI Gateway budgets (#27329)

Browse Source 
## Stack Context

This stack makes AI Gateway data and budgets the source of truth for AI
spend controls.

1. Re-back the per-chat cost endpoint with AI Gateway data (#27328,
merged).
2. **This PR:** remove native chat usage limits.
3. Remove native chat cost tracking and its dedicated admin UI (#27330).

## Summary

Removes the native usage-limit API, SDK types, SQL, and chat enforcement
for deployment, user, and group chat limits. Compact AI Gateway budget
indicators remain in the Agents sidebar, user menu, and group settings.
Gateway budget rejections and provider quota failures continue to
classify as usage-limit errors, including a 409 response for synchronous
title generation.

Budget-period labels now use the API's UTC boundaries, so users see the
same dates in every browser timezone. The documentation explains the AI
Gateway replacement, its licensing requirements, and the differences
from native limits.

No schema is dropped in this release. The usage-limit table, index, user
and group columns, constraints, audit mappings, and generated scan
fields remain for mixed-version rolling upgrades. #27600 tracks their
removal after the compatibility window.

## Breaking change

Native day, week, and month chat spend limits are removed and are not
migrated. AI Gateway budgets are month-based, group-scoped with per-user
overrides, and require the AI Gateway entitlement. Deployments without
that entitlement no longer have chat spend enforcement.

> Mux prepared this PR on Mike's behalf.
This commit is contained in:
Michael Suchacz
2026-08-04 11:36:49 +02:00
committed by GitHub
parent a2287d6739
commit f0e6ac64b3
78 changed files with 611 additions and 6850 deletions
Download Patch File Download Diff File Expand all files Collapse all files
coderd/coderd.go
-13
View File
@@ -1420,19 +1420,6 @@ func New(options *Options) *API {
r.Delete("/", api.deleteChatModelConfig)
})
})
r.Route("/usage-limits", func(r chi.Router) {
r.Get("/", api.getChatUsageLimitConfig)
r.Put("/", api.updateChatUsageLimitConfig)
r.Get("/status", api.getMyChatUsageLimitStatus)
r.Route("/overrides/{user}", func(r chi.Router) {
r.Put("/", api.upsertChatUsageLimitOverride)
r.Delete("/", api.deleteChatUsageLimitOverride)
})
r.Route("/group-overrides/{group}", func(r chi.Router) {
r.Put("/", api.upsertChatUsageLimitGroupOverride)
r.Delete("/", api.deleteChatUsageLimitGroupOverride)
})
})
r.Route("/user-provider-configs", func(r chi.Router) {
r.Get("/", api.listUserChatProviderConfigs)
r.Route("/{providerConfig}", func(r chi.Router) {
coderd/database/dbauthz/dbauthz.go
-91
View File
@@ -2213,20 +2213,6 @@ func (q *querier) DeleteChatQueuedMessageReturningCount(ctx context.Context, arg
return q.db.DeleteChatQueuedMessageReturningCount(ctx, arg)
}
func (q *querier) DeleteChatUsageLimitGroupOverride(ctx context.Context, groupID uuid.UUID) error {
if err := q.authorizeContext(ctx, policy.ActionUpdate, rbac.ResourceDeploymentConfig); err != nil {
return err
}
return q.db.DeleteChatUsageLimitGroupOverride(ctx, groupID)
}
func (q *querier) DeleteChatUsageLimitUserOverride(ctx context.Context, userID uuid.UUID) error {
if err := q.authorizeContext(ctx, policy.ActionUpdate, rbac.ResourceDeploymentConfig); err != nil {
return err
}
return q.db.DeleteChatUsageLimitUserOverride(ctx, userID)
}
func (q *querier) DeleteCryptoKey(ctx context.Context, arg database.DeleteCryptoKeyParams) (database.CryptoKey, error) {
if err := q.authorizeContext(ctx, policy.ActionDelete, rbac.ResourceCryptoKey); err != nil {
return database.CryptoKey{}, err
@@ -3636,27 +3622,6 @@ func (q *querier) GetChatTitleGenerationModelOverride(ctx context.Context) (stri
return q.db.GetChatTitleGenerationModelOverride(ctx)
}
func (q *querier) GetChatUsageLimitConfig(ctx context.Context) (database.ChatUsageLimitConfig, error) {
if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceDeploymentConfig); err != nil {
return database.ChatUsageLimitConfig{}, err
}
return q.db.GetChatUsageLimitConfig(ctx)
}
func (q *querier) GetChatUsageLimitGroupOverride(ctx context.Context, groupID uuid.UUID) (database.GetChatUsageLimitGroupOverrideRow, error) {
if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceDeploymentConfig); err != nil {
return database.GetChatUsageLimitGroupOverrideRow{}, err
}
return q.db.GetChatUsageLimitGroupOverride(ctx, groupID)
}
func (q *querier) GetChatUsageLimitUserOverride(ctx context.Context, userID uuid.UUID) (database.GetChatUsageLimitUserOverrideRow, error) {
if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceDeploymentConfig); err != nil {
return database.GetChatUsageLimitUserOverrideRow{}, err
}
return q.db.GetChatUsageLimitUserOverride(ctx, userID)
}
func (q *querier) GetChatUserPromptsByChatID(ctx context.Context, arg database.GetChatUserPromptsByChatIDParams) ([]database.GetChatUserPromptsByChatIDRow, error) {
// Authorize read on the parent chat.
_, err := q.GetChatByID(ctx, arg.ChatID)
@@ -5179,13 +5144,6 @@ func (q *querier) GetUserChatPersonalModelOverride(ctx context.Context, arg data
return q.db.GetUserChatPersonalModelOverride(ctx, arg)
}
func (q *querier) GetUserChatSpendInPeriod(ctx context.Context, arg database.GetUserChatSpendInPeriodParams) (int64, error) {
if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceChat.WithOwner(arg.UserID.String())); err != nil {
return 0, err
}
return q.db.GetUserChatSpendInPeriod(ctx, arg)
}
func (q *querier) GetUserCodeDiffDisplayMode(ctx context.Context, userID uuid.UUID) (string, error) {
user, err := q.db.GetUserByID(ctx, userID)
if err != nil {
@@ -5216,13 +5174,6 @@ func (q *querier) GetUserForChatSyntheticAPIKeyByID(ctx context.Context, id uuid
return fetchWithAction(q.log, q.auth, policy.ActionReadPersonal, q.db.GetUserForChatSyntheticAPIKeyByID)(ctx, id)
}
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, arg)
}
func (q *querier) GetUserLatencyInsights(ctx context.Context, arg database.GetUserLatencyInsightsParams) ([]database.GetUserLatencyInsightsRow, error) {
// Used by insights endpoints. Need to check both for auditors and for regular users with template acl perms.
if err := q.authorizeContext(ctx, policy.ActionViewInsights, rbac.ResourceTemplate); err != nil {
@@ -6908,20 +6859,6 @@ func (q *querier) ListChatContextResourcesByChatID(ctx context.Context, chatID u
return q.db.ListChatContextResourcesByChatID(ctx, chatID)
}
func (q *querier) ListChatUsageLimitGroupOverrides(ctx context.Context) ([]database.ListChatUsageLimitGroupOverridesRow, error) {
if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceDeploymentConfig); err != nil {
return nil, err
}
return q.db.ListChatUsageLimitGroupOverrides(ctx)
}
func (q *querier) ListChatUsageLimitOverrides(ctx context.Context) ([]database.ListChatUsageLimitOverridesRow, error) {
if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceDeploymentConfig); err != nil {
return nil, err
}
return q.db.ListChatUsageLimitOverrides(ctx)
}
func (q *querier) ListProvisionerKeysByOrganization(ctx context.Context, organizationID uuid.UUID) ([]database.ProvisionerKey, error) {
return fetchWithPostFilter(q.auth, policy.ActionRead, q.db.ListProvisionerKeysByOrganization)(ctx, organizationID)
}
@@ -7157,13 +7094,6 @@ func (q *querier) ReorderChatQueuedMessageToHead(ctx context.Context, arg databa
return q.db.ReorderChatQueuedMessageToHead(ctx, arg)
}
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, arg)
}
func (q *querier) RevokeDBCryptKey(ctx context.Context, activeKeyDigest string) error {
if err := q.authorizeContext(ctx, policy.ActionUpdate, rbac.ResourceSystem); err != nil {
return err
@@ -9047,27 +8977,6 @@ func (q *querier) UpsertChatTitleGenerationModelOverride(ctx context.Context, va
return q.db.UpsertChatTitleGenerationModelOverride(ctx, value)
}
func (q *querier) UpsertChatUsageLimitConfig(ctx context.Context, arg database.UpsertChatUsageLimitConfigParams) (database.ChatUsageLimitConfig, error) {
if err := q.authorizeContext(ctx, policy.ActionUpdate, rbac.ResourceDeploymentConfig); err != nil {
return database.ChatUsageLimitConfig{}, err
}
return q.db.UpsertChatUsageLimitConfig(ctx, arg)
}
func (q *querier) UpsertChatUsageLimitGroupOverride(ctx context.Context, arg database.UpsertChatUsageLimitGroupOverrideParams) (database.UpsertChatUsageLimitGroupOverrideRow, error) {
if err := q.authorizeContext(ctx, policy.ActionUpdate, rbac.ResourceDeploymentConfig); err != nil {
return database.UpsertChatUsageLimitGroupOverrideRow{}, err
}
return q.db.UpsertChatUsageLimitGroupOverride(ctx, arg)
}
func (q *querier) UpsertChatUsageLimitUserOverride(ctx context.Context, arg database.UpsertChatUsageLimitUserOverrideParams) (database.UpsertChatUsageLimitUserOverrideRow, error) {
if err := q.authorizeContext(ctx, policy.ActionUpdate, rbac.ResourceDeploymentConfig); err != nil {
return database.UpsertChatUsageLimitUserOverrideRow{}, err
}
return q.db.UpsertChatUsageLimitUserOverride(ctx, arg)
}
//nolint:revive // Parameter name matches the generated querier interface.
func (q *querier) UpsertChatWorkspaceTTL(ctx context.Context, workspaceTtl string) error {
if err := q.authorizeContext(ctx, policy.ActionUpdate, rbac.ResourceDeploymentConfig); err != nil {
coderd/database/dbauthz/dbauthz_test.go
-146
View File
@@ -1688,152 +1688,6 @@ func (s *MethodTestSuite) TestChats() {
dbm.EXPECT().UpsertChatWorkspaceTTL(gomock.Any(), "1h").Return(nil).AnyTimes()
check.Args("1h").Asserts(rbac.ResourceDeploymentConfig, policy.ActionUpdate)
}))
s.Run("GetUserChatSpendInPeriod", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) {
arg := database.GetUserChatSpendInPeriodParams{
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),
}
spend := int64(123)
dbm.EXPECT().GetUserChatSpendInPeriod(gomock.Any(), arg).Return(spend, nil).AnyTimes()
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) {
arg := database.GetUserGroupSpendLimitParams{
UserID: uuid.New(),
OrganizationID: uuid.NullUUID{UUID: uuid.New(), Valid: true},
}
limit := int64(456)
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) {
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{
ID: 1,
Singleton: true,
Enabled: true,
DefaultLimitMicros: 1_000_000,
Period: "monthly",
CreatedAt: now,
UpdatedAt: now,
}
dbm.EXPECT().GetChatUsageLimitConfig(gomock.Any()).Return(config, nil).AnyTimes()
check.Args().Asserts(rbac.ResourceDeploymentConfig, policy.ActionRead).Returns(config)
}))
s.Run("GetChatUsageLimitGroupOverride", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) {
groupID := uuid.New()
override := database.GetChatUsageLimitGroupOverrideRow{
GroupID: groupID,
SpendLimitMicros: sql.NullInt64{Int64: 2_000_000, Valid: true},
}
dbm.EXPECT().GetChatUsageLimitGroupOverride(gomock.Any(), groupID).Return(override, nil).AnyTimes()
check.Args(groupID).Asserts(rbac.ResourceDeploymentConfig, policy.ActionRead).Returns(override)
}))
s.Run("GetChatUsageLimitUserOverride", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) {
userID := uuid.New()
override := database.GetChatUsageLimitUserOverrideRow{
UserID: userID,
SpendLimitMicros: sql.NullInt64{Int64: 3_000_000, Valid: true},
}
dbm.EXPECT().GetChatUsageLimitUserOverride(gomock.Any(), userID).Return(override, nil).AnyTimes()
check.Args(userID).Asserts(rbac.ResourceDeploymentConfig, policy.ActionRead).Returns(override)
}))
s.Run("ListChatUsageLimitGroupOverrides", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) {
overrides := []database.ListChatUsageLimitGroupOverridesRow{{
GroupID: uuid.New(),
GroupName: "group-name",
GroupDisplayName: "Group Name",
GroupAvatarUrl: "https://example.com/group.png",
SpendLimitMicros: sql.NullInt64{Int64: 4_000_000, Valid: true},
MemberCount: 5,
}}
dbm.EXPECT().ListChatUsageLimitGroupOverrides(gomock.Any()).Return(overrides, nil).AnyTimes()
check.Args().Asserts(rbac.ResourceDeploymentConfig, policy.ActionRead).Returns(overrides)
}))
s.Run("ListChatUsageLimitOverrides", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) {
overrides := []database.ListChatUsageLimitOverridesRow{{
UserID: uuid.New(),
Username: "usage-limit-user",
Name: "Usage Limit User",
AvatarURL: "https://example.com/avatar.png",
SpendLimitMicros: sql.NullInt64{Int64: 5_000_000, Valid: true},
}}
dbm.EXPECT().ListChatUsageLimitOverrides(gomock.Any()).Return(overrides, nil).AnyTimes()
check.Args().Asserts(rbac.ResourceDeploymentConfig, policy.ActionRead).Returns(overrides)
}))
s.Run("UpsertChatUsageLimitConfig", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) {
now := dbtime.Now()
arg := database.UpsertChatUsageLimitConfigParams{
Enabled: true,
DefaultLimitMicros: 6_000_000,
Period: "monthly",
}
config := database.ChatUsageLimitConfig{
ID: 1,
Singleton: true,
Enabled: arg.Enabled,
DefaultLimitMicros: arg.DefaultLimitMicros,
Period: arg.Period,
CreatedAt: now,
UpdatedAt: now,
}
dbm.EXPECT().UpsertChatUsageLimitConfig(gomock.Any(), arg).Return(config, nil).AnyTimes()
check.Args(arg).Asserts(rbac.ResourceDeploymentConfig, policy.ActionUpdate).Returns(config)
}))
s.Run("UpsertChatUsageLimitGroupOverride", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) {
arg := database.UpsertChatUsageLimitGroupOverrideParams{
SpendLimitMicros: 7_000_000,
GroupID: uuid.New(),
}
override := database.UpsertChatUsageLimitGroupOverrideRow{
GroupID: arg.GroupID,
Name: "group",
DisplayName: "Group",
AvatarURL: "",
SpendLimitMicros: sql.NullInt64{Int64: arg.SpendLimitMicros, Valid: true},
}
dbm.EXPECT().UpsertChatUsageLimitGroupOverride(gomock.Any(), arg).Return(override, nil).AnyTimes()
check.Args(arg).Asserts(rbac.ResourceDeploymentConfig, policy.ActionUpdate).Returns(override)
}))
s.Run("UpsertChatUsageLimitUserOverride", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) {
arg := database.UpsertChatUsageLimitUserOverrideParams{
SpendLimitMicros: 8_000_000,
UserID: uuid.New(),
}
override := database.UpsertChatUsageLimitUserOverrideRow{
UserID: arg.UserID,
Username: "user",
Name: "User",
AvatarURL: "",
SpendLimitMicros: sql.NullInt64{Int64: arg.SpendLimitMicros, Valid: true},
}
dbm.EXPECT().UpsertChatUsageLimitUserOverride(gomock.Any(), arg).Return(override, nil).AnyTimes()
check.Args(arg).Asserts(rbac.ResourceDeploymentConfig, policy.ActionUpdate).Returns(override)
}))
s.Run("DeleteChatUsageLimitGroupOverride", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) {
groupID := uuid.New()
dbm.EXPECT().DeleteChatUsageLimitGroupOverride(gomock.Any(), groupID).Return(nil).AnyTimes()
check.Args(groupID).Asserts(rbac.ResourceDeploymentConfig, policy.ActionUpdate)
}))
s.Run("DeleteChatUsageLimitUserOverride", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) {
userID := uuid.New()
dbm.EXPECT().DeleteChatUsageLimitUserOverride(gomock.Any(), userID).Return(nil).AnyTimes()
check.Args(userID).Asserts(rbac.ResourceDeploymentConfig, policy.ActionUpdate)
}))
s.Run("CleanupDeletedMCPServerIDsFromChats", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) {
dbm.EXPECT().CleanupDeletedMCPServerIDsFromChats(gomock.Any()).Return(nil).AnyTimes()
check.Args().Asserts(rbac.ResourceChat, policy.ActionUpdate)
coderd/database/dbmetrics/querymetrics.go
Generated
-104
View File
@@ -545,22 +545,6 @@ func (m queryMetricsStore) DeleteChatQueuedMessageReturningCount(ctx context.Con
return r0, r1
}
func (m queryMetricsStore) DeleteChatUsageLimitGroupOverride(ctx context.Context, groupID uuid.UUID) error {
start := time.Now()
r0 := m.s.DeleteChatUsageLimitGroupOverride(ctx, groupID)
m.queryLatencies.WithLabelValues("DeleteChatUsageLimitGroupOverride").Observe(time.Since(start).Seconds())
m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "DeleteChatUsageLimitGroupOverride").Inc()
return r0
}
func (m queryMetricsStore) DeleteChatUsageLimitUserOverride(ctx context.Context, userID uuid.UUID) error {
start := time.Now()
r0 := m.s.DeleteChatUsageLimitUserOverride(ctx, userID)
m.queryLatencies.WithLabelValues("DeleteChatUsageLimitUserOverride").Observe(time.Since(start).Seconds())
m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "DeleteChatUsageLimitUserOverride").Inc()
return r0
}
func (m queryMetricsStore) DeleteCryptoKey(ctx context.Context, arg database.DeleteCryptoKeyParams) (database.CryptoKey, error) {
start := time.Now()
r0, r1 := m.s.DeleteCryptoKey(ctx, arg)
@@ -1849,30 +1833,6 @@ func (m queryMetricsStore) GetChatTitleGenerationModelOverride(ctx context.Conte
return r0, r1
}
func (m queryMetricsStore) GetChatUsageLimitConfig(ctx context.Context) (database.ChatUsageLimitConfig, error) {
start := time.Now()
r0, r1 := m.s.GetChatUsageLimitConfig(ctx)
m.queryLatencies.WithLabelValues("GetChatUsageLimitConfig").Observe(time.Since(start).Seconds())
m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "GetChatUsageLimitConfig").Inc()
return r0, r1
}
func (m queryMetricsStore) GetChatUsageLimitGroupOverride(ctx context.Context, groupID uuid.UUID) (database.GetChatUsageLimitGroupOverrideRow, error) {
start := time.Now()
r0, r1 := m.s.GetChatUsageLimitGroupOverride(ctx, groupID)
m.queryLatencies.WithLabelValues("GetChatUsageLimitGroupOverride").Observe(time.Since(start).Seconds())
m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "GetChatUsageLimitGroupOverride").Inc()
return r0, r1
}
func (m queryMetricsStore) GetChatUsageLimitUserOverride(ctx context.Context, userID uuid.UUID) (database.GetChatUsageLimitUserOverrideRow, error) {
start := time.Now()
r0, r1 := m.s.GetChatUsageLimitUserOverride(ctx, userID)
m.queryLatencies.WithLabelValues("GetChatUsageLimitUserOverride").Observe(time.Since(start).Seconds())
m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "GetChatUsageLimitUserOverride").Inc()
return r0, r1
}
func (m queryMetricsStore) GetChatUserPromptsByChatID(ctx context.Context, arg database.GetChatUserPromptsByChatIDParams) ([]database.GetChatUserPromptsByChatIDRow, error) {
start := time.Now()
r0, r1 := m.s.GetChatUserPromptsByChatID(ctx, arg)
@@ -3337,14 +3297,6 @@ func (m queryMetricsStore) GetUserChatPersonalModelOverride(ctx context.Context,
return r0, r1
}
func (m queryMetricsStore) GetUserChatSpendInPeriod(ctx context.Context, arg database.GetUserChatSpendInPeriodParams) (int64, error) {
start := time.Now()
r0, r1 := m.s.GetUserChatSpendInPeriod(ctx, arg)
m.queryLatencies.WithLabelValues("GetUserChatSpendInPeriod").Observe(time.Since(start).Seconds())
m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "GetUserChatSpendInPeriod").Inc()
return r0, r1
}
func (m queryMetricsStore) GetUserCodeDiffDisplayMode(ctx context.Context, userID uuid.UUID) (string, error) {
start := time.Now()
r0, r1 := m.s.GetUserCodeDiffDisplayMode(ctx, userID)
@@ -3377,14 +3329,6 @@ func (m queryMetricsStore) GetUserForChatSyntheticAPIKeyByID(ctx context.Context
return r0, r1
}
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())
m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "GetUserGroupSpendLimit").Inc()
return r0, r1
}
func (m queryMetricsStore) GetUserLatencyInsights(ctx context.Context, arg database.GetUserLatencyInsightsParams) ([]database.GetUserLatencyInsightsRow, error) {
start := time.Now()
r0, r1 := m.s.GetUserLatencyInsights(ctx, arg)
@@ -4849,22 +4793,6 @@ func (m queryMetricsStore) ListChatContextResourcesByChatID(ctx context.Context,
return r0, r1
}
func (m queryMetricsStore) ListChatUsageLimitGroupOverrides(ctx context.Context) ([]database.ListChatUsageLimitGroupOverridesRow, error) {
start := time.Now()
r0, r1 := m.s.ListChatUsageLimitGroupOverrides(ctx)
m.queryLatencies.WithLabelValues("ListChatUsageLimitGroupOverrides").Observe(time.Since(start).Seconds())
m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "ListChatUsageLimitGroupOverrides").Inc()
return r0, r1
}
func (m queryMetricsStore) ListChatUsageLimitOverrides(ctx context.Context) ([]database.ListChatUsageLimitOverridesRow, error) {
start := time.Now()
r0, r1 := m.s.ListChatUsageLimitOverrides(ctx)
m.queryLatencies.WithLabelValues("ListChatUsageLimitOverrides").Observe(time.Since(start).Seconds())
m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "ListChatUsageLimitOverrides").Inc()
return r0, r1
}
func (m queryMetricsStore) ListProvisionerKeysByOrganization(ctx context.Context, organizationID uuid.UUID) ([]database.ProvisionerKey, error) {
start := time.Now()
r0, r1 := m.s.ListProvisionerKeysByOrganization(ctx, organizationID)
@@ -5073,14 +5001,6 @@ func (m queryMetricsStore) ReorderChatQueuedMessageToHead(ctx context.Context, a
return r0, r1
}
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())
m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "ResolveUserChatSpendLimit").Inc()
return r0, r1
}
func (m queryMetricsStore) RevokeDBCryptKey(ctx context.Context, activeKeyDigest string) error {
start := time.Now()
r0 := m.s.RevokeDBCryptKey(ctx, activeKeyDigest)
@@ -6449,30 +6369,6 @@ func (m queryMetricsStore) UpsertChatTitleGenerationModelOverride(ctx context.Co
return r0
}
func (m queryMetricsStore) UpsertChatUsageLimitConfig(ctx context.Context, arg database.UpsertChatUsageLimitConfigParams) (database.ChatUsageLimitConfig, error) {
start := time.Now()
r0, r1 := m.s.UpsertChatUsageLimitConfig(ctx, arg)
m.queryLatencies.WithLabelValues("UpsertChatUsageLimitConfig").Observe(time.Since(start).Seconds())
m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "UpsertChatUsageLimitConfig").Inc()
return r0, r1
}
func (m queryMetricsStore) UpsertChatUsageLimitGroupOverride(ctx context.Context, arg database.UpsertChatUsageLimitGroupOverrideParams) (database.UpsertChatUsageLimitGroupOverrideRow, error) {
start := time.Now()
r0, r1 := m.s.UpsertChatUsageLimitGroupOverride(ctx, arg)
m.queryLatencies.WithLabelValues("UpsertChatUsageLimitGroupOverride").Observe(time.Since(start).Seconds())
m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "UpsertChatUsageLimitGroupOverride").Inc()
return r0, r1
}
func (m queryMetricsStore) UpsertChatUsageLimitUserOverride(ctx context.Context, arg database.UpsertChatUsageLimitUserOverrideParams) (database.UpsertChatUsageLimitUserOverrideRow, error) {
start := time.Now()
r0, r1 := m.s.UpsertChatUsageLimitUserOverride(ctx, arg)
m.queryLatencies.WithLabelValues("UpsertChatUsageLimitUserOverride").Observe(time.Since(start).Seconds())
m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "UpsertChatUsageLimitUserOverride").Inc()
return r0, r1
}
func (m queryMetricsStore) UpsertChatWorkspaceTTL(ctx context.Context, workspaceTtl string) error {
start := time.Now()
r0 := m.s.UpsertChatWorkspaceTTL(ctx, workspaceTtl)
coderd/database/dbmock/dbmock.go
Generated
-193
View File
@@ -891,34 +891,6 @@ func (mr *MockStoreMockRecorder) DeleteChatQueuedMessageReturningCount(ctx, arg
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DeleteChatQueuedMessageReturningCount", reflect.TypeOf((*MockStore)(nil).DeleteChatQueuedMessageReturningCount), ctx, arg)
}
// DeleteChatUsageLimitGroupOverride mocks base method.
func (m *MockStore) DeleteChatUsageLimitGroupOverride(ctx context.Context, groupID uuid.UUID) error {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "DeleteChatUsageLimitGroupOverride", ctx, groupID)
ret0, _ := ret[0].(error)
return ret0
}
// DeleteChatUsageLimitGroupOverride indicates an expected call of DeleteChatUsageLimitGroupOverride.
func (mr *MockStoreMockRecorder) DeleteChatUsageLimitGroupOverride(ctx, groupID any) *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DeleteChatUsageLimitGroupOverride", reflect.TypeOf((*MockStore)(nil).DeleteChatUsageLimitGroupOverride), ctx, groupID)
}
// DeleteChatUsageLimitUserOverride mocks base method.
func (m *MockStore) DeleteChatUsageLimitUserOverride(ctx context.Context, userID uuid.UUID) error {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "DeleteChatUsageLimitUserOverride", ctx, userID)
ret0, _ := ret[0].(error)
return ret0
}
// DeleteChatUsageLimitUserOverride indicates an expected call of DeleteChatUsageLimitUserOverride.
func (mr *MockStoreMockRecorder) DeleteChatUsageLimitUserOverride(ctx, userID any) *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DeleteChatUsageLimitUserOverride", reflect.TypeOf((*MockStore)(nil).DeleteChatUsageLimitUserOverride), ctx, userID)
}
// DeleteCryptoKey mocks base method.
func (m *MockStore) DeleteCryptoKey(ctx context.Context, arg database.DeleteCryptoKeyParams) (database.CryptoKey, error) {
m.ctrl.T.Helper()
@@ -3418,51 +3390,6 @@ func (mr *MockStoreMockRecorder) GetChatTitleGenerationModelOverride(ctx any) *g
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetChatTitleGenerationModelOverride", reflect.TypeOf((*MockStore)(nil).GetChatTitleGenerationModelOverride), ctx)
}
// GetChatUsageLimitConfig mocks base method.
func (m *MockStore) GetChatUsageLimitConfig(ctx context.Context) (database.ChatUsageLimitConfig, error) {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "GetChatUsageLimitConfig", ctx)
ret0, _ := ret[0].(database.ChatUsageLimitConfig)
ret1, _ := ret[1].(error)
return ret0, ret1
}
// GetChatUsageLimitConfig indicates an expected call of GetChatUsageLimitConfig.
func (mr *MockStoreMockRecorder) GetChatUsageLimitConfig(ctx any) *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetChatUsageLimitConfig", reflect.TypeOf((*MockStore)(nil).GetChatUsageLimitConfig), ctx)
}
// GetChatUsageLimitGroupOverride mocks base method.
func (m *MockStore) GetChatUsageLimitGroupOverride(ctx context.Context, groupID uuid.UUID) (database.GetChatUsageLimitGroupOverrideRow, error) {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "GetChatUsageLimitGroupOverride", ctx, groupID)
ret0, _ := ret[0].(database.GetChatUsageLimitGroupOverrideRow)
ret1, _ := ret[1].(error)
return ret0, ret1
}
// GetChatUsageLimitGroupOverride indicates an expected call of GetChatUsageLimitGroupOverride.
func (mr *MockStoreMockRecorder) GetChatUsageLimitGroupOverride(ctx, groupID any) *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetChatUsageLimitGroupOverride", reflect.TypeOf((*MockStore)(nil).GetChatUsageLimitGroupOverride), ctx, groupID)
}
// GetChatUsageLimitUserOverride mocks base method.
func (m *MockStore) GetChatUsageLimitUserOverride(ctx context.Context, userID uuid.UUID) (database.GetChatUsageLimitUserOverrideRow, error) {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "GetChatUsageLimitUserOverride", ctx, userID)
ret0, _ := ret[0].(database.GetChatUsageLimitUserOverrideRow)
ret1, _ := ret[1].(error)
return ret0, ret1
}
// GetChatUsageLimitUserOverride indicates an expected call of GetChatUsageLimitUserOverride.
func (mr *MockStoreMockRecorder) GetChatUsageLimitUserOverride(ctx, userID any) *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetChatUsageLimitUserOverride", reflect.TypeOf((*MockStore)(nil).GetChatUsageLimitUserOverride), ctx, userID)
}
// GetChatUserPromptsByChatID mocks base method.
func (m *MockStore) GetChatUserPromptsByChatID(ctx context.Context, arg database.GetChatUserPromptsByChatIDParams) ([]database.GetChatUserPromptsByChatIDRow, error) {
m.ctrl.T.Helper()
@@ -6238,21 +6165,6 @@ func (mr *MockStoreMockRecorder) GetUserChatPersonalModelOverride(ctx, arg any)
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetUserChatPersonalModelOverride", reflect.TypeOf((*MockStore)(nil).GetUserChatPersonalModelOverride), ctx, arg)
}
// GetUserChatSpendInPeriod mocks base method.
func (m *MockStore) GetUserChatSpendInPeriod(ctx context.Context, arg database.GetUserChatSpendInPeriodParams) (int64, error) {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "GetUserChatSpendInPeriod", ctx, arg)
ret0, _ := ret[0].(int64)
ret1, _ := ret[1].(error)
return ret0, ret1
}
// GetUserChatSpendInPeriod indicates an expected call of GetUserChatSpendInPeriod.
func (mr *MockStoreMockRecorder) GetUserChatSpendInPeriod(ctx, arg any) *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetUserChatSpendInPeriod", reflect.TypeOf((*MockStore)(nil).GetUserChatSpendInPeriod), ctx, arg)
}
// GetUserCodeDiffDisplayMode mocks base method.
func (m *MockStore) GetUserCodeDiffDisplayMode(ctx context.Context, userID uuid.UUID) (string, error) {
m.ctrl.T.Helper()
@@ -6313,21 +6225,6 @@ func (mr *MockStoreMockRecorder) GetUserForChatSyntheticAPIKeyByID(ctx, id any)
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetUserForChatSyntheticAPIKeyByID", reflect.TypeOf((*MockStore)(nil).GetUserForChatSyntheticAPIKeyByID), ctx, id)
}
// GetUserGroupSpendLimit mocks base method.
func (m *MockStore) GetUserGroupSpendLimit(ctx context.Context, arg database.GetUserGroupSpendLimitParams) (int64, error) {
m.ctrl.T.Helper()
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, arg any) *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetUserGroupSpendLimit", reflect.TypeOf((*MockStore)(nil).GetUserGroupSpendLimit), ctx, arg)
}
// GetUserLatencyInsights mocks base method.
func (m *MockStore) GetUserLatencyInsights(ctx context.Context, arg database.GetUserLatencyInsightsParams) ([]database.GetUserLatencyInsightsRow, error) {
m.ctrl.T.Helper()
@@ -9131,36 +9028,6 @@ func (mr *MockStoreMockRecorder) ListChatContextResourcesByChatID(ctx, chatID an
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ListChatContextResourcesByChatID", reflect.TypeOf((*MockStore)(nil).ListChatContextResourcesByChatID), ctx, chatID)
}
// ListChatUsageLimitGroupOverrides mocks base method.
func (m *MockStore) ListChatUsageLimitGroupOverrides(ctx context.Context) ([]database.ListChatUsageLimitGroupOverridesRow, error) {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "ListChatUsageLimitGroupOverrides", ctx)
ret0, _ := ret[0].([]database.ListChatUsageLimitGroupOverridesRow)
ret1, _ := ret[1].(error)
return ret0, ret1
}
// ListChatUsageLimitGroupOverrides indicates an expected call of ListChatUsageLimitGroupOverrides.
func (mr *MockStoreMockRecorder) ListChatUsageLimitGroupOverrides(ctx any) *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ListChatUsageLimitGroupOverrides", reflect.TypeOf((*MockStore)(nil).ListChatUsageLimitGroupOverrides), ctx)
}
// ListChatUsageLimitOverrides mocks base method.
func (m *MockStore) ListChatUsageLimitOverrides(ctx context.Context) ([]database.ListChatUsageLimitOverridesRow, error) {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "ListChatUsageLimitOverrides", ctx)
ret0, _ := ret[0].([]database.ListChatUsageLimitOverridesRow)
ret1, _ := ret[1].(error)
return ret0, ret1
}
// ListChatUsageLimitOverrides indicates an expected call of ListChatUsageLimitOverrides.
func (mr *MockStoreMockRecorder) ListChatUsageLimitOverrides(ctx any) *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ListChatUsageLimitOverrides", reflect.TypeOf((*MockStore)(nil).ListChatUsageLimitOverrides), ctx)
}
// ListProvisionerKeysByOrganization mocks base method.
func (m *MockStore) ListProvisionerKeysByOrganization(ctx context.Context, organizationID uuid.UUID) ([]database.ProvisionerKey, error) {
m.ctrl.T.Helper()
@@ -9578,21 +9445,6 @@ func (mr *MockStoreMockRecorder) ReorderChatQueuedMessageToHead(ctx, arg any) *g
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ReorderChatQueuedMessageToHead", reflect.TypeOf((*MockStore)(nil).ReorderChatQueuedMessageToHead), ctx, arg)
}
// ResolveUserChatSpendLimit mocks base method.
func (m *MockStore) ResolveUserChatSpendLimit(ctx context.Context, arg database.ResolveUserChatSpendLimitParams) (database.ResolveUserChatSpendLimitRow, error) {
m.ctrl.T.Helper()
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, arg any) *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ResolveUserChatSpendLimit", reflect.TypeOf((*MockStore)(nil).ResolveUserChatSpendLimit), ctx, arg)
}
// RevokeDBCryptKey mocks base method.
func (m *MockStore) RevokeDBCryptKey(ctx context.Context, activeKeyDigest string) error {
m.ctrl.T.Helper()
@@ -12070,51 +11922,6 @@ func (mr *MockStoreMockRecorder) UpsertChatTitleGenerationModelOverride(ctx, val
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpsertChatTitleGenerationModelOverride", reflect.TypeOf((*MockStore)(nil).UpsertChatTitleGenerationModelOverride), ctx, value)
}
// UpsertChatUsageLimitConfig mocks base method.
func (m *MockStore) UpsertChatUsageLimitConfig(ctx context.Context, arg database.UpsertChatUsageLimitConfigParams) (database.ChatUsageLimitConfig, error) {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "UpsertChatUsageLimitConfig", ctx, arg)
ret0, _ := ret[0].(database.ChatUsageLimitConfig)
ret1, _ := ret[1].(error)
return ret0, ret1
}
// UpsertChatUsageLimitConfig indicates an expected call of UpsertChatUsageLimitConfig.
func (mr *MockStoreMockRecorder) UpsertChatUsageLimitConfig(ctx, arg any) *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpsertChatUsageLimitConfig", reflect.TypeOf((*MockStore)(nil).UpsertChatUsageLimitConfig), ctx, arg)
}
// UpsertChatUsageLimitGroupOverride mocks base method.
func (m *MockStore) UpsertChatUsageLimitGroupOverride(ctx context.Context, arg database.UpsertChatUsageLimitGroupOverrideParams) (database.UpsertChatUsageLimitGroupOverrideRow, error) {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "UpsertChatUsageLimitGroupOverride", ctx, arg)
ret0, _ := ret[0].(database.UpsertChatUsageLimitGroupOverrideRow)
ret1, _ := ret[1].(error)
return ret0, ret1
}
// UpsertChatUsageLimitGroupOverride indicates an expected call of UpsertChatUsageLimitGroupOverride.
func (mr *MockStoreMockRecorder) UpsertChatUsageLimitGroupOverride(ctx, arg any) *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpsertChatUsageLimitGroupOverride", reflect.TypeOf((*MockStore)(nil).UpsertChatUsageLimitGroupOverride), ctx, arg)
}
// UpsertChatUsageLimitUserOverride mocks base method.
func (m *MockStore) UpsertChatUsageLimitUserOverride(ctx context.Context, arg database.UpsertChatUsageLimitUserOverrideParams) (database.UpsertChatUsageLimitUserOverrideRow, error) {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "UpsertChatUsageLimitUserOverride", ctx, arg)
ret0, _ := ret[0].(database.UpsertChatUsageLimitUserOverrideRow)
ret1, _ := ret[1].(error)
return ret0, ret1
}
// UpsertChatUsageLimitUserOverride indicates an expected call of UpsertChatUsageLimitUserOverride.
func (mr *MockStoreMockRecorder) UpsertChatUsageLimitUserOverride(ctx, arg any) *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpsertChatUsageLimitUserOverride", reflect.TypeOf((*MockStore)(nil).UpsertChatUsageLimitUserOverride), ctx, arg)
}
// UpsertChatWorkspaceTTL mocks base method.
func (m *MockStore) UpsertChatWorkspaceTTL(ctx context.Context, workspaceTtl string) error {
m.ctrl.T.Helper()
coderd/database/querier.go
Generated
-33
View File
@@ -154,8 +154,6 @@ type sqlcQuerier interface {
// number of affected rows so callers can detect missing rows without
// a follow-up read.
DeleteChatQueuedMessageReturningCount(ctx context.Context, arg DeleteChatQueuedMessageReturningCountParams) (int64, error)
DeleteChatUsageLimitGroupOverride(ctx context.Context, groupID uuid.UUID) error
DeleteChatUsageLimitUserOverride(ctx context.Context, userID uuid.UUID) error
DeleteCryptoKey(ctx context.Context, arg DeleteCryptoKeyParams) (CryptoKey, error)
DeleteCustomRole(ctx context.Context, arg DeleteCustomRoleParams) error
DeleteExpiredAPIKeys(ctx context.Context, arg DeleteExpiredAPIKeysParams) (int64, error)
@@ -531,9 +529,6 @@ type sqlcQuerier interface {
// Returns an empty string when no allowlist has been configured (all templates allowed).
GetChatTemplateAllowlist(ctx context.Context) (string, error)
GetChatTitleGenerationModelOverride(ctx context.Context) (string, error)
GetChatUsageLimitConfig(ctx context.Context) (ChatUsageLimitConfig, error)
GetChatUsageLimitGroupOverride(ctx context.Context, groupID uuid.UUID) (GetChatUsageLimitGroupOverrideRow, error)
GetChatUsageLimitUserOverride(ctx context.Context, userID uuid.UUID) (GetChatUsageLimitUserOverrideRow, error)
// Returns the concatenated text of each user-visible user prompt in a
// chat, newest first. Used by the composer to populate the up/down
// arrow prompt-history cycle. Non-text parts (tool calls, files,
@@ -931,11 +926,6 @@ type sqlcQuerier interface {
GetUserChatCustomPrompt(ctx context.Context, userID uuid.UUID) (string, error)
GetUserChatDebugLoggingEnabled(ctx context.Context, userID uuid.UUID) (bool, error)
GetUserChatPersonalModelOverride(ctx context.Context, arg GetUserChatPersonalModelOverrideParams) (string, 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)
GetUserCodeDiffDisplayMode(ctx context.Context, userID uuid.UUID) (string, error)
GetUserCount(ctx context.Context, includeSystem bool) (int64, error)
// Returns the "Everyone" group (id == organization_id) to attribute a user's
@@ -944,12 +934,6 @@ type sqlcQuerier interface {
// no organization membership.
GetUserEveryoneFallbackGroup(ctx context.Context, userID uuid.UUID) (uuid.UUID, error)
GetUserForChatSyntheticAPIKeyByID(ctx context.Context, id uuid.UUID) (User, error)
// Returns the minimum (most restrictive) group limit for a user.
// 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
@@ -1303,8 +1287,6 @@ type sqlcQuerier interface {
// Lists a chat's pinned context resources, ordered deterministically by
// source.
ListChatContextResourcesByChatID(ctx context.Context, chatID uuid.UUID) ([]ChatContextResource, error)
ListChatUsageLimitGroupOverrides(ctx context.Context) ([]ListChatUsageLimitGroupOverridesRow, error)
ListChatUsageLimitOverrides(ctx context.Context) ([]ListChatUsageLimitOverridesRow, error)
ListProvisionerKeysByOrganization(ctx context.Context, organizationID uuid.UUID) ([]ProvisionerKey, error)
ListProvisionerKeysByOrganizationExcludeReserved(ctx context.Context, organizationID uuid.UUID) ([]ProvisionerKey, error)
ListTasks(ctx context.Context, arg ListTasksParams) ([]Task, error)
@@ -1366,18 +1348,6 @@ type sqlcQuerier interface {
// Sets the target queued message's position to one less than the
// current minimum position for that chat, moving it to the head.
ReorderChatQueuedMessageToHead(ctx context.Context, arg ReorderChatQueuedMessageToHeadParams) (int64, error)
// Resolves the effective spend limit for a user using the hierarchy:
// 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'.
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
@@ -1697,9 +1667,6 @@ type sqlcQuerier interface {
UpsertChatSystemPrompt(ctx context.Context, value string) error
UpsertChatTemplateAllowlist(ctx context.Context, templateAllowlist string) error
UpsertChatTitleGenerationModelOverride(ctx context.Context, value string) error
UpsertChatUsageLimitConfig(ctx context.Context, arg UpsertChatUsageLimitConfigParams) (ChatUsageLimitConfig, error)
UpsertChatUsageLimitGroupOverride(ctx context.Context, arg UpsertChatUsageLimitGroupOverrideParams) (UpsertChatUsageLimitGroupOverrideRow, error)
UpsertChatUsageLimitUserOverride(ctx context.Context, arg UpsertChatUsageLimitUserOverrideParams) (UpsertChatUsageLimitUserOverrideRow, error)
UpsertChatWorkspaceTTL(ctx context.Context, workspaceTtl string) error
// The default proxy is implied and not actually stored in the database.
// So we need to store it's configuration here for display purposes.
coderd/database/queries.sql.go
Generated
-390
View File
@@ -7097,24 +7097,6 @@ func (q *sqlQuerier) DeleteChatQueuedMessageReturningCount(ctx context.Context,
return result.RowsAffected()
}
const deleteChatUsageLimitGroupOverride = `-- name: DeleteChatUsageLimitGroupOverride :exec
UPDATE groups SET chat_spend_limit_micros = NULL WHERE id = $1::uuid
`
func (q *sqlQuerier) DeleteChatUsageLimitGroupOverride(ctx context.Context, groupID uuid.UUID) error {
_, err := q.db.ExecContext(ctx, deleteChatUsageLimitGroupOverride, groupID)
return err
}
const deleteChatUsageLimitUserOverride = `-- name: DeleteChatUsageLimitUserOverride :exec
UPDATE users SET chat_spend_limit_micros = NULL WHERE id = $1::uuid
`
func (q *sqlQuerier) DeleteChatUsageLimitUserOverride(ctx context.Context, userID uuid.UUID) error {
_, err := q.db.ExecContext(ctx, deleteChatUsageLimitUserOverride, userID)
return err
}
const deleteOldChats = `-- name: DeleteOldChats :execrows
WITH deletable AS (
SELECT id
@@ -9119,61 +9101,6 @@ func (q *sqlQuerier) GetChatStreamSyncRows(ctx context.Context, ids []uuid.UUID)
return items, nil
}
const getChatUsageLimitConfig = `-- name: GetChatUsageLimitConfig :one
SELECT id, singleton, enabled, default_limit_micros, period, created_at, updated_at FROM chat_usage_limit_config WHERE singleton = TRUE LIMIT 1
`
func (q *sqlQuerier) GetChatUsageLimitConfig(ctx context.Context) (ChatUsageLimitConfig, error) {
row := q.db.QueryRowContext(ctx, getChatUsageLimitConfig)
var i ChatUsageLimitConfig
err := row.Scan(
&i.ID,
&i.Singleton,
&i.Enabled,
&i.DefaultLimitMicros,
&i.Period,
&i.CreatedAt,
&i.UpdatedAt,
)
return i, err
}
const getChatUsageLimitGroupOverride = `-- name: GetChatUsageLimitGroupOverride :one
SELECT id AS group_id, chat_spend_limit_micros AS spend_limit_micros
FROM groups
WHERE id = $1::uuid AND chat_spend_limit_micros IS NOT NULL
`
type GetChatUsageLimitGroupOverrideRow struct {
GroupID uuid.UUID `db:"group_id" json:"group_id"`
SpendLimitMicros sql.NullInt64 `db:"spend_limit_micros" json:"spend_limit_micros"`
}
func (q *sqlQuerier) GetChatUsageLimitGroupOverride(ctx context.Context, groupID uuid.UUID) (GetChatUsageLimitGroupOverrideRow, error) {
row := q.db.QueryRowContext(ctx, getChatUsageLimitGroupOverride, groupID)
var i GetChatUsageLimitGroupOverrideRow
err := row.Scan(&i.GroupID, &i.SpendLimitMicros)
return i, err
}
const getChatUsageLimitUserOverride = `-- name: GetChatUsageLimitUserOverride :one
SELECT id AS user_id, chat_spend_limit_micros AS spend_limit_micros
FROM users
WHERE id = $1::uuid AND chat_spend_limit_micros IS NOT NULL
`
type GetChatUsageLimitUserOverrideRow struct {
UserID uuid.UUID `db:"user_id" json:"user_id"`
SpendLimitMicros sql.NullInt64 `db:"spend_limit_micros" json:"spend_limit_micros"`
}
func (q *sqlQuerier) GetChatUsageLimitUserOverride(ctx context.Context, userID uuid.UUID) (GetChatUsageLimitUserOverrideRow, error) {
row := q.db.QueryRowContext(ctx, getChatUsageLimitUserOverride, userID)
var i GetChatUsageLimitUserOverrideRow
err := row.Scan(&i.UserID, &i.SpendLimitMicros)
return i, err
}
const getChatUserPromptsByChatID = `-- name: GetChatUserPromptsByChatID :many
SELECT
cm.id,
@@ -10346,68 +10273,6 @@ func (q *sqlQuerier) GetTotalChatMessageRuntimeMsInRange(ctx context.Context, ar
return total_runtime_ms, err
}
const getUserChatSpendInPeriod = `-- name: GetUserChatSpendInPeriod :one
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 ($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"`
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.OrganizationID,
arg.StartTime,
arg.EndTime,
)
var total_spend_micros int64
err := row.Scan(&total_spend_micros)
return total_spend_micros, err
}
const getUserGroupSpendLimit = `-- name: GetUserGroupSpendLimit :one
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 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
}
const hydrateAgentChatsContext = `-- name: HydrateAgentChatsContext :many
WITH hydrated AS (
UPDATE chats
@@ -11138,106 +11003,6 @@ func (q *sqlQuerier) ListChatContextResourcesByChatID(ctx context.Context, chatI
return items, nil
}
const listChatUsageLimitGroupOverrides = `-- name: ListChatUsageLimitGroupOverrides :many
SELECT
g.id AS group_id,
g.name AS group_name,
g.display_name AS group_display_name,
g.avatar_url AS group_avatar_url,
g.chat_spend_limit_micros AS spend_limit_micros,
(SELECT COUNT(*)
FROM group_members_expanded gme
WHERE gme.group_id = g.id
AND gme.user_is_system = FALSE) AS member_count
FROM groups g
WHERE g.chat_spend_limit_micros IS NOT NULL
ORDER BY g.name ASC
`
type ListChatUsageLimitGroupOverridesRow struct {
GroupID uuid.UUID `db:"group_id" json:"group_id"`
GroupName string `db:"group_name" json:"group_name"`
GroupDisplayName string `db:"group_display_name" json:"group_display_name"`
GroupAvatarUrl string `db:"group_avatar_url" json:"group_avatar_url"`
SpendLimitMicros sql.NullInt64 `db:"spend_limit_micros" json:"spend_limit_micros"`
MemberCount int64 `db:"member_count" json:"member_count"`
}
func (q *sqlQuerier) ListChatUsageLimitGroupOverrides(ctx context.Context) ([]ListChatUsageLimitGroupOverridesRow, error) {
rows, err := q.db.QueryContext(ctx, listChatUsageLimitGroupOverrides)
if err != nil {
return nil, err
}
defer rows.Close()
var items []ListChatUsageLimitGroupOverridesRow
for rows.Next() {
var i ListChatUsageLimitGroupOverridesRow
if err := rows.Scan(
&i.GroupID,
&i.GroupName,
&i.GroupDisplayName,
&i.GroupAvatarUrl,
&i.SpendLimitMicros,
&i.MemberCount,
); err != nil {
return nil, err
}
items = append(items, i)
}
if err := rows.Close(); err != nil {
return nil, err
}
if err := rows.Err(); err != nil {
return nil, err
}
return items, nil
}
const listChatUsageLimitOverrides = `-- name: ListChatUsageLimitOverrides :many
SELECT u.id AS user_id, u.username, u.name, u.avatar_url,
u.chat_spend_limit_micros AS spend_limit_micros
FROM users u
WHERE u.chat_spend_limit_micros IS NOT NULL
ORDER BY u.username ASC
`
type ListChatUsageLimitOverridesRow struct {
UserID uuid.UUID `db:"user_id" json:"user_id"`
Username string `db:"username" json:"username"`
Name string `db:"name" json:"name"`
AvatarURL string `db:"avatar_url" json:"avatar_url"`
SpendLimitMicros sql.NullInt64 `db:"spend_limit_micros" json:"spend_limit_micros"`
}
func (q *sqlQuerier) ListChatUsageLimitOverrides(ctx context.Context) ([]ListChatUsageLimitOverridesRow, error) {
rows, err := q.db.QueryContext(ctx, listChatUsageLimitOverrides)
if err != nil {
return nil, err
}
defer rows.Close()
var items []ListChatUsageLimitOverridesRow
for rows.Next() {
var i ListChatUsageLimitOverridesRow
if err := rows.Scan(
&i.UserID,
&i.Username,
&i.Name,
&i.AvatarURL,
&i.SpendLimitMicros,
); err != nil {
return nil, err
}
items = append(items, i)
}
if err := rows.Close(); err != nil {
return nil, err
}
if err := rows.Err(); err != nil {
return nil, err
}
return items, nil
}
const lockChatAndBumpSnapshotVersion = `-- name: LockChatAndBumpSnapshotVersion :one
WITH bumped_chat AS (
UPDATE chats
@@ -11558,63 +11323,6 @@ func (q *sqlQuerier) ReorderChatQueuedMessageToHead(ctx context.Context, arg Reo
return result.RowsAffected()
}
const resolveUserChatSpendLimit = `-- name: ResolveUserChatSpendLimit :one
SELECT CASE
WHEN NOT cfg.enabled THEN -1
WHEN u.chat_spend_limit_micros IS NOT NULL THEN u.chat_spend_limit_micros
WHEN gl.limit_micros IS NOT NULL THEN gl.limit_micros
ELSE cfg.default_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 (
SELECT MIN(g.chat_spend_limit_micros) 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
) 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, 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'.
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 setChatContextSnapshot = `-- name: SetChatContextSnapshot :exec
UPDATE chats
SET
@@ -13823,104 +13531,6 @@ func (q *sqlQuerier) UpsertChatHeartbeat(ctx context.Context, arg UpsertChatHear
return err
}
const upsertChatUsageLimitConfig = `-- name: UpsertChatUsageLimitConfig :one
INSERT INTO chat_usage_limit_config (singleton, enabled, default_limit_micros, period, updated_at)
VALUES (TRUE, $1::boolean, $2::bigint, $3::text, NOW())
ON CONFLICT (singleton) DO UPDATE SET
enabled = EXCLUDED.enabled,
default_limit_micros = EXCLUDED.default_limit_micros,
period = EXCLUDED.period,
updated_at = NOW()
RETURNING id, singleton, enabled, default_limit_micros, period, created_at, updated_at
`
type UpsertChatUsageLimitConfigParams struct {
Enabled bool `db:"enabled" json:"enabled"`
DefaultLimitMicros int64 `db:"default_limit_micros" json:"default_limit_micros"`
Period string `db:"period" json:"period"`
}
func (q *sqlQuerier) UpsertChatUsageLimitConfig(ctx context.Context, arg UpsertChatUsageLimitConfigParams) (ChatUsageLimitConfig, error) {
row := q.db.QueryRowContext(ctx, upsertChatUsageLimitConfig, arg.Enabled, arg.DefaultLimitMicros, arg.Period)
var i ChatUsageLimitConfig
err := row.Scan(
&i.ID,
&i.Singleton,
&i.Enabled,
&i.DefaultLimitMicros,
&i.Period,
&i.CreatedAt,
&i.UpdatedAt,
)
return i, err
}
const upsertChatUsageLimitGroupOverride = `-- name: UpsertChatUsageLimitGroupOverride :one
UPDATE groups
SET chat_spend_limit_micros = $1::bigint
WHERE id = $2::uuid
RETURNING id AS group_id, name, display_name, avatar_url, chat_spend_limit_micros AS spend_limit_micros
`
type UpsertChatUsageLimitGroupOverrideParams struct {
SpendLimitMicros int64 `db:"spend_limit_micros" json:"spend_limit_micros"`
GroupID uuid.UUID `db:"group_id" json:"group_id"`
}
type UpsertChatUsageLimitGroupOverrideRow struct {
GroupID uuid.UUID `db:"group_id" json:"group_id"`
Name string `db:"name" json:"name"`
DisplayName string `db:"display_name" json:"display_name"`
AvatarURL string `db:"avatar_url" json:"avatar_url"`
SpendLimitMicros sql.NullInt64 `db:"spend_limit_micros" json:"spend_limit_micros"`
}
func (q *sqlQuerier) UpsertChatUsageLimitGroupOverride(ctx context.Context, arg UpsertChatUsageLimitGroupOverrideParams) (UpsertChatUsageLimitGroupOverrideRow, error) {
row := q.db.QueryRowContext(ctx, upsertChatUsageLimitGroupOverride, arg.SpendLimitMicros, arg.GroupID)
var i UpsertChatUsageLimitGroupOverrideRow
err := row.Scan(
&i.GroupID,
&i.Name,
&i.DisplayName,
&i.AvatarURL,
&i.SpendLimitMicros,
)
return i, err
}
const upsertChatUsageLimitUserOverride = `-- name: UpsertChatUsageLimitUserOverride :one
UPDATE users
SET chat_spend_limit_micros = $1::bigint
WHERE id = $2::uuid
RETURNING id AS user_id, username, name, avatar_url, chat_spend_limit_micros AS spend_limit_micros
`
type UpsertChatUsageLimitUserOverrideParams struct {
SpendLimitMicros int64 `db:"spend_limit_micros" json:"spend_limit_micros"`
UserID uuid.UUID `db:"user_id" json:"user_id"`
}
type UpsertChatUsageLimitUserOverrideRow struct {
UserID uuid.UUID `db:"user_id" json:"user_id"`
Username string `db:"username" json:"username"`
Name string `db:"name" json:"name"`
AvatarURL string `db:"avatar_url" json:"avatar_url"`
SpendLimitMicros sql.NullInt64 `db:"spend_limit_micros" json:"spend_limit_micros"`
}
func (q *sqlQuerier) UpsertChatUsageLimitUserOverride(ctx context.Context, arg UpsertChatUsageLimitUserOverrideParams) (UpsertChatUsageLimitUserOverrideRow, error) {
row := q.db.QueryRowContext(ctx, upsertChatUsageLimitUserOverride, arg.SpendLimitMicros, arg.UserID)
var i UpsertChatUsageLimitUserOverrideRow
err := row.Scan(
&i.UserID,
&i.Username,
&i.Name,
&i.AvatarURL,
&i.SpendLimitMicros,
)
return i, err
}
const batchUpsertConnectionLogs = `-- name: BatchUpsertConnectionLogs :exec
INSERT INTO connection_logs (
id, connect_time, organization_id, workspace_owner_id, workspace_id,
coderd/database/queries/chats.sql
-130
View File
@@ -2434,55 +2434,6 @@ LIMIT
OFFSET
sqlc.arg('page_offset')::int;
-- name: GetChatUsageLimitConfig :one
SELECT * FROM chat_usage_limit_config WHERE singleton = TRUE LIMIT 1;
-- name: UpsertChatUsageLimitConfig :one
INSERT INTO chat_usage_limit_config (singleton, enabled, default_limit_micros, period, updated_at)
VALUES (TRUE, @enabled::boolean, @default_limit_micros::bigint, @period::text, NOW())
ON CONFLICT (singleton) DO UPDATE SET
enabled = EXCLUDED.enabled,
default_limit_micros = EXCLUDED.default_limit_micros,
period = EXCLUDED.period,
updated_at = NOW()
RETURNING *;
-- name: ListChatUsageLimitOverrides :many
SELECT u.id AS user_id, u.username, u.name, u.avatar_url,
u.chat_spend_limit_micros AS spend_limit_micros
FROM users u
WHERE u.chat_spend_limit_micros IS NOT NULL
ORDER BY u.username ASC;
-- name: UpsertChatUsageLimitUserOverride :one
UPDATE users
SET chat_spend_limit_micros = @spend_limit_micros::bigint
WHERE id = @user_id::uuid
RETURNING id AS user_id, username, name, avatar_url, chat_spend_limit_micros AS spend_limit_micros;
-- name: DeleteChatUsageLimitUserOverride :exec
UPDATE users SET chat_spend_limit_micros = NULL WHERE id = @user_id::uuid;
-- name: GetChatUsageLimitUserOverride :one
SELECT id AS user_id, chat_spend_limit_micros AS spend_limit_micros
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;
-- name: GetTotalChatMessageRuntimeMsInRange :one
-- Computes hb_agent_runtime_v1 usage event payloads. Deliberately includes
-- soft-deleted messages and messages from all chats.
@@ -2508,49 +2459,6 @@ WHERE enabled = TRUE
)
);
-- name: ListChatUsageLimitGroupOverrides :many
SELECT
g.id AS group_id,
g.name AS group_name,
g.display_name AS group_display_name,
g.avatar_url AS group_avatar_url,
g.chat_spend_limit_micros AS spend_limit_micros,
(SELECT COUNT(*)
FROM group_members_expanded gme
WHERE gme.group_id = g.id
AND gme.user_is_system = FALSE) AS member_count
FROM groups g
WHERE g.chat_spend_limit_micros IS NOT NULL
ORDER BY g.name ASC;
-- name: UpsertChatUsageLimitGroupOverride :one
UPDATE groups
SET chat_spend_limit_micros = @spend_limit_micros::bigint
WHERE id = @group_id::uuid
RETURNING id AS group_id, name, display_name, avatar_url, chat_spend_limit_micros AS spend_limit_micros;
-- name: DeleteChatUsageLimitGroupOverride :exec
UPDATE groups SET chat_spend_limit_micros = NULL WHERE id = @group_id::uuid;
-- name: GetChatUsageLimitGroupOverride :one
SELECT id AS group_id, chat_spend_limit_micros AS spend_limit_micros
FROM groups
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 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
SELECT *
FROM chats_expanded
@@ -2558,44 +2466,6 @@ WHERE archived = false
AND workspace_id = ANY(@ids::uuid[])
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, 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
WHEN NOT cfg.enabled THEN -1
WHEN u.chat_spend_limit_micros IS NOT NULL THEN u.chat_spend_limit_micros
WHEN gl.limit_micros IS NOT NULL THEN gl.limit_micros
ELSE cfg.default_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 (
SELECT MIN(g.chat_spend_limit_micros) 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
) gl ON TRUE
WHERE u.id = @user_id::uuid
LIMIT 1;
-- name: UpdateChatLastReadMessageID :exec
-- Updates the last read message ID for a chat. This is used to track
-- which messages the owner has seen, enabling unread indicators.
coderd/exp_chats.go
+16 -537
View File
@@ -97,21 +97,6 @@ type chatDiffReference struct {
RepositoryRef *chatRepositoryRef
}
func writeChatUsageLimitExceeded(
ctx context.Context,
rw http.ResponseWriter,
limitErr *chatd.UsageLimitExceededError,
) {
httpapi.Write(ctx, rw, http.StatusConflict, codersdk.ChatUsageLimitExceededResponse{
Response: codersdk.Response{
Message: "Chat usage limit exceeded.",
},
SpentMicros: limitErr.ConsumedMicros,
LimitMicros: limitErr.LimitMicros,
ResetsAt: limitErr.PeriodEnd,
})
}
// Avoid returning raw dispatch errors, which may expose deployment internals.
func writeChatHookDispatchFailed(ctx context.Context, rw http.ResponseWriter, hookErr *dispatch.Error) {
httpapi.Write(ctx, rw, http.StatusBadGateway, codersdk.ChatHookDispatchFailedResponse{
@@ -145,13 +130,17 @@ func writeChatHookErr(ctx context.Context, rw http.ResponseWriter, err error, de
return false
}
func maybeWriteLimitErr(ctx context.Context, rw http.ResponseWriter, err error) bool {
var limitErr *chatd.UsageLimitExceededError
if errors.As(err, &limitErr) {
writeChatUsageLimitExceeded(ctx, rw, limitErr)
return true
// AI Gateway budget rejections and provider quota failures classify as usage
// limits; synchronous generation reports them as conflicts instead of 500s.
func maybeWriteChatUsageLimitError(ctx context.Context, rw http.ResponseWriter, err error) bool {
classified := chaterror.Classify(err)
if classified.Kind != codersdk.ChatErrorKindUsageLimit {
return false
}
return false
httpapi.Write(ctx, rw, http.StatusConflict, codersdk.Response{
Message: classified.Message,
})
return true
}
// requireChatDaemon reports whether the chat daemon exists, writing a 503
@@ -1479,9 +1468,6 @@ func (api *API) postChats(rw http.ResponseWriter, r *http.Request) {
if writeChatHookErr(ctx, rw, err, "Chat creation denied by lifecycle hook.") {
return
}
if maybeWriteLimitErr(ctx, rw, err) {
return
}
if xerrors.Is(err, chatd.ErrInvalidModelConfigID) {
httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{
Message: "Invalid model config ID.",
@@ -1691,13 +1677,6 @@ func (api *API) chatCostSummary(rw http.ResponseWriter, r *http.Request) {
chatBreakdowns = append(chatBreakdowns, convertChatCostChatBreakdown(chat))
}
// 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))
}
response := codersdk.ChatCostSummary{
StartDate: startDate,
EndDate: endDate,
@@ -1712,10 +1691,6 @@ func (api *API) chatCostSummary(rw http.ResponseWriter, r *http.Request) {
ByModel: modelBreakdowns,
ByChat: chatBreakdowns,
}
if usageStatus != nil {
response.UsageLimit = usageStatus
}
httpapi.Write(ctx, rw, http.StatusOK, response)
}
@@ -1820,447 +1795,6 @@ func (api *API) chatCostUsers(rw http.ResponseWriter, r *http.Request) {
})
}
// @Summary Get chat usage limit config
// @x-apidocgen {"skip": true}
// EXPERIMENTAL: this endpoint is experimental and is subject to change.
//
//nolint:revive // HTTP handler writes to ResponseWriter.
func (api *API) getChatUsageLimitConfig(rw http.ResponseWriter, r *http.Request) {
ctx := r.Context()
if !api.Authorize(r, policy.ActionRead, rbac.ResourceDeploymentConfig) {
httpapi.Forbidden(rw)
return
}
config, configErr := api.Database.GetChatUsageLimitConfig(ctx)
if configErr != nil && !errors.Is(configErr, sql.ErrNoRows) {
httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
Message: "Failed to get chat usage limit config.",
Detail: configErr.Error(),
})
return
}
overrideRows, err := api.Database.ListChatUsageLimitOverrides(ctx)
if err != nil {
httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
Message: "Failed to list chat usage limit overrides.",
Detail: err.Error(),
})
return
}
groupOverrides, err := api.Database.ListChatUsageLimitGroupOverrides(ctx)
if err != nil {
httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
Message: "Failed to list group usage limit overrides.",
Detail: err.Error(),
})
return
}
unpricedModelCount, err := api.Database.CountEnabledModelsWithoutPricing(ctx)
if err != nil {
httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
Message: "Failed to count unpriced chat models.",
Detail: err.Error(),
})
return
}
response := codersdk.ChatUsageLimitConfigResponse{
ChatUsageLimitConfig: codersdk.ChatUsageLimitConfig{},
UnpricedModelCount: unpricedModelCount,
Overrides: make([]codersdk.ChatUsageLimitOverride, 0, len(overrideRows)),
GroupOverrides: make([]codersdk.ChatUsageLimitGroupOverride, 0, len(groupOverrides)),
}
if configErr == nil {
response.Period = codersdk.ChatUsageLimitPeriod(config.Period)
response.UpdatedAt = config.UpdatedAt
if config.Enabled {
response.SpendLimitMicros = ptr.Ref(config.DefaultLimitMicros)
}
}
for _, row := range overrideRows {
response.Overrides = append(response.Overrides, codersdk.ChatUsageLimitOverride{
UserID: row.UserID,
Username: row.Username,
Name: row.Name,
AvatarURL: row.AvatarURL,
SpendLimitMicros: nullInt64Ptr(row.SpendLimitMicros),
})
}
for _, glo := range groupOverrides {
response.GroupOverrides = append(response.GroupOverrides, codersdk.ChatUsageLimitGroupOverride{
GroupID: glo.GroupID,
GroupName: glo.GroupName,
GroupDisplayName: glo.GroupDisplayName,
GroupAvatarURL: glo.GroupAvatarUrl,
MemberCount: glo.MemberCount,
SpendLimitMicros: nullInt64Ptr(glo.SpendLimitMicros),
})
}
httpapi.Write(ctx, rw, http.StatusOK, response)
}
// @Summary Update chat usage limit config
// @x-apidocgen {"skip": true}
// EXPERIMENTAL: this endpoint is experimental and is subject to change.
func (api *API) updateChatUsageLimitConfig(rw http.ResponseWriter, r *http.Request) {
ctx := r.Context()
if !api.Authorize(r, policy.ActionUpdate, rbac.ResourceDeploymentConfig) {
httpapi.Forbidden(rw)
return
}
var req codersdk.ChatUsageLimitConfig
if !httpapi.Read(ctx, rw, r, &req) {
return
}
params := database.UpsertChatUsageLimitConfigParams{
Enabled: false,
DefaultLimitMicros: 0,
Period: "",
}
if req.SpendLimitMicros == nil {
if req.Period != "" && !req.Period.Valid() {
httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{
Message: "Invalid chat usage limit period.",
Detail: "Period must be one of: day, week, month.",
})
return
}
params.Enabled = false
params.DefaultLimitMicros = 0
params.Period = string(req.Period)
if params.Period == "" {
params.Period = string(codersdk.ChatUsageLimitPeriodMonth)
}
} else {
if *req.SpendLimitMicros <= 0 {
httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{
Message: "Invalid chat usage limit spend limit.",
Detail: "Spend limit must be greater than 0.",
})
return
}
if !req.Period.Valid() {
httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{
Message: "Invalid chat usage limit period.",
Detail: "Period must be one of: day, week, month.",
})
return
}
params.Enabled = true
params.DefaultLimitMicros = *req.SpendLimitMicros
params.Period = string(req.Period)
}
config, err := api.Database.UpsertChatUsageLimitConfig(ctx, params)
if err != nil {
httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
Message: "Failed to update chat usage limit config.",
Detail: err.Error(),
})
return
}
response := codersdk.ChatUsageLimitConfig{
Period: codersdk.ChatUsageLimitPeriod(config.Period),
UpdatedAt: config.UpdatedAt,
}
if config.Enabled {
response.SpendLimitMicros = ptr.Ref(config.DefaultLimitMicros)
}
httpapi.Write(ctx, rw, http.StatusOK, response)
}
// @Summary Get my chat usage limit status
// @x-apidocgen {"skip": true}
// EXPERIMENTAL: this endpoint is experimental and is subject to change.
//
// getMyChatUsageLimitStatus returns the current usage-limit status for the
// authenticated user. No additional RBAC check is required because the
// endpoint always operates on the requesting user's own data via
// httpmw.APIKey(r).UserID.
//
//nolint:revive // HTTP handler writes to ResponseWriter.
func (api *API) getMyChatUsageLimitStatus(rw http.ResponseWriter, r *http.Request) {
ctx := r.Context()
// 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.",
Detail: err.Error(),
})
return
}
if status == nil {
httpapi.Write(ctx, rw, http.StatusOK, codersdk.ChatUsageLimitStatus{IsLimited: false})
return
}
httpapi.Write(ctx, rw, http.StatusOK, status)
}
// @Summary Upsert chat usage limit override
// @x-apidocgen {"skip": true}
// EXPERIMENTAL: this endpoint is experimental and is subject to change.
func (api *API) upsertChatUsageLimitOverride(rw http.ResponseWriter, r *http.Request) {
ctx := r.Context()
if !api.Authorize(r, policy.ActionUpdate, rbac.ResourceDeploymentConfig) {
httpapi.Forbidden(rw)
return
}
userID, ok := parseChatUsageLimitUserID(rw, r)
if !ok {
return
}
var req codersdk.UpsertChatUsageLimitOverrideRequest
if !httpapi.Read(ctx, rw, r, &req) {
return
}
if req.SpendLimitMicros <= 0 {
httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{
Message: "Invalid chat usage limit override.",
Detail: "Spend limit must be greater than 0.",
})
return
}
user, err := api.Database.GetUserByID(ctx, userID)
if err != nil {
if errors.Is(err, sql.ErrNoRows) {
httpapi.Write(ctx, rw, http.StatusNotFound, codersdk.Response{
Message: "User not found.",
})
return
}
httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
Message: "Failed to look up chat usage limit user.",
Detail: err.Error(),
})
return
}
_, err = api.Database.UpsertChatUsageLimitUserOverride(ctx, database.UpsertChatUsageLimitUserOverrideParams{
UserID: userID,
SpendLimitMicros: req.SpendLimitMicros,
})
if err != nil {
httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
Message: "Failed to upsert chat usage limit override.",
Detail: err.Error(),
})
return
}
httpapi.Write(ctx, rw, http.StatusOK, codersdk.ChatUsageLimitOverride{
UserID: user.ID,
Username: user.Username,
Name: user.Name,
AvatarURL: user.AvatarURL,
SpendLimitMicros: nullInt64Ptr(sql.NullInt64{Int64: req.SpendLimitMicros, Valid: true}),
})
}
// @Summary Delete chat usage limit override
// @x-apidocgen {"skip": true}
// EXPERIMENTAL: this endpoint is experimental and is subject to change.
func (api *API) deleteChatUsageLimitOverride(rw http.ResponseWriter, r *http.Request) {
ctx := r.Context()
if !api.Authorize(r, policy.ActionUpdate, rbac.ResourceDeploymentConfig) {
httpapi.Forbidden(rw)
return
}
userID, ok := parseChatUsageLimitUserID(rw, r)
if !ok {
return
}
if _, err := api.Database.GetUserByID(ctx, userID); err != nil {
if errors.Is(err, sql.ErrNoRows) {
writeChatUsageLimitUserNotFound(ctx, rw)
return
}
httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
Message: "Failed to look up chat usage limit user.",
Detail: err.Error(),
})
return
}
if _, err := api.Database.GetChatUsageLimitUserOverride(ctx, userID); err != nil {
if errors.Is(err, sql.ErrNoRows) {
writeChatUsageLimitOverrideNotFound(ctx, rw)
return
}
httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
Message: "Failed to look up chat usage limit override.",
Detail: err.Error(),
})
return
}
if err := api.Database.DeleteChatUsageLimitUserOverride(ctx, userID); err != nil {
httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
Message: "Failed to delete chat usage limit override.",
Detail: err.Error(),
})
return
}
rw.WriteHeader(http.StatusNoContent)
}
// @Summary Upsert chat usage limit group override
// @x-apidocgen {"skip": true}
// EXPERIMENTAL: this endpoint is experimental and is subject to change.
func (api *API) upsertChatUsageLimitGroupOverride(rw http.ResponseWriter, r *http.Request) {
ctx := r.Context()
if !api.Authorize(r, policy.ActionUpdate, rbac.ResourceDeploymentConfig) {
httpapi.Forbidden(rw)
return
}
groupIDStr := chi.URLParam(r, "group")
groupID, err := uuid.Parse(groupIDStr)
if err != nil {
httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{
Message: "Invalid group ID.",
Detail: err.Error(),
})
return
}
var req codersdk.UpdateChatUsageLimitGroupOverrideRequest
if !httpapi.Read(ctx, rw, r, &req) {
return
}
if req.SpendLimitMicros <= 0 {
httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{
Message: "Invalid chat usage limit group override.",
Detail: "Spend limit (in microdollars) must be greater than 0.",
})
return
}
group, err := api.Database.GetGroupByID(ctx, groupID)
if err != nil {
if errors.Is(err, sql.ErrNoRows) {
httpapi.Write(ctx, rw, http.StatusNotFound, codersdk.Response{
Message: "Group not found.",
})
return
}
httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
Message: "Failed to look up group details.",
Detail: err.Error(),
})
return
}
_, err = api.Database.UpsertChatUsageLimitGroupOverride(ctx, database.UpsertChatUsageLimitGroupOverrideParams{
GroupID: groupID,
SpendLimitMicros: req.SpendLimitMicros,
})
if err != nil {
httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
Message: "Failed to upsert group usage limit override.",
Detail: err.Error(),
})
return
}
memberCount, err := api.Database.GetGroupMembersCountByGroupID(ctx, database.GetGroupMembersCountByGroupIDParams{
GroupID: groupID,
IncludeSystem: false,
})
if err != nil {
if errors.Is(err, sql.ErrNoRows) {
writeChatUsageLimitGroupNotFound(ctx, rw)
return
}
httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
Message: "Failed to fetch group member count.",
Detail: err.Error(),
})
return
}
httpapi.Write(ctx, rw, http.StatusOK, codersdk.ChatUsageLimitGroupOverride{
GroupID: group.ID,
GroupName: group.Name,
GroupDisplayName: group.DisplayName,
GroupAvatarURL: group.AvatarURL,
MemberCount: memberCount,
SpendLimitMicros: nullInt64Ptr(sql.NullInt64{Int64: req.SpendLimitMicros, Valid: true}),
})
}
// @Summary Delete chat usage limit group override
// @x-apidocgen {"skip": true}
// EXPERIMENTAL: this endpoint is experimental and is subject to change.
func (api *API) deleteChatUsageLimitGroupOverride(rw http.ResponseWriter, r *http.Request) {
ctx := r.Context()
if !api.Authorize(r, policy.ActionUpdate, rbac.ResourceDeploymentConfig) {
httpapi.Forbidden(rw)
return
}
groupIDStr := chi.URLParam(r, "group")
groupID, err := uuid.Parse(groupIDStr)
if err != nil {
httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{
Message: "Invalid group ID.",
Detail: err.Error(),
})
return
}
if _, err := api.Database.GetGroupByID(ctx, groupID); err != nil {
if errors.Is(err, sql.ErrNoRows) {
writeChatUsageLimitGroupNotFound(ctx, rw)
return
}
httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
Message: "Failed to look up group details.",
Detail: err.Error(),
})
return
}
if _, err := api.Database.GetChatUsageLimitGroupOverride(ctx, groupID); err != nil {
if errors.Is(err, sql.ErrNoRows) {
writeChatUsageLimitGroupOverrideNotFound(ctx, rw)
return
}
httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
Message: "Failed to look up group usage limit override.",
Detail: err.Error(),
})
return
}
if err := api.Database.DeleteChatUsageLimitGroupOverride(ctx, groupID); err != nil {
httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
Message: "Failed to delete group usage limit override.",
Detail: err.Error(),
})
return
}
rw.WriteHeader(http.StatusNoContent)
}
// EXPERIMENTAL: this endpoint is experimental and is subject to change.
//
// @Summary Get chat by ID
@@ -3497,9 +3031,6 @@ func (api *API) postChatMessages(rw http.ResponseWriter, r *http.Request) {
if writeChatHookErr(ctx, rw, sendErr, "Chat message denied by lifecycle hook.") {
return
}
if maybeWriteLimitErr(ctx, rw, sendErr) {
return
}
if xerrors.Is(sendErr, chatd.ErrChatArchived) {
httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{
Message: "Cannot send messages to an archived chat.",
@@ -3687,9 +3218,6 @@ func (api *API) patchChatMessage(rw http.ResponseWriter, r *http.Request) {
if writeChatHookErr(ctx, rw, editErr, "Chat message denied by lifecycle hook.") {
return
}
if maybeWriteLimitErr(ctx, rw, editErr) {
return
}
switch {
case xerrors.Is(editErr, chatd.ErrChatArchived):
@@ -3858,9 +3386,6 @@ func (api *API) promoteChatQueuedMessage(rw http.ResponseWriter, r *http.Request
})
if txErr != nil {
if maybeWriteLimitErr(ctx, rw, txErr) {
return
}
switch {
case xerrors.Is(txErr, chatd.ErrChatArchived):
httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{
@@ -4163,9 +3688,6 @@ func (api *API) compactChat(rw http.ResponseWriter, r *http.Request) {
updated, err := api.chatDaemon.CompactChat(ctx, chat)
if err != nil {
if maybeWriteLimitErr(ctx, rw, err) {
return
}
if writeCommonChatMutationError(ctx, rw, err, "Cannot compact an archived chat.") {
return
}
@@ -4292,13 +3814,13 @@ func (api *API) regenerateChatTitle(rw http.ResponseWriter, r *http.Request) {
})
return
}
if maybeWriteLimitErr(ctx, rw, err) {
return
}
if httpapi.Is404Error(err) {
httpapi.ResourceNotFound(rw)
return
}
if maybeWriteChatUsageLimitError(ctx, rw, err) {
return
}
httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
Message: "Failed to regenerate chat title.",
Detail: err.Error(),
@@ -4341,13 +3863,13 @@ func (api *API) proposeChatTitle(rw http.ResponseWriter, r *http.Request) {
})
return
}
if maybeWriteLimitErr(ctx, rw, err) {
return
}
if httpapi.Is404Error(err) {
httpapi.ResourceNotFound(rw)
return
}
if maybeWriteChatUsageLimitError(ctx, rw, err) {
return
}
httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
Message: "Failed to generate chat title.",
Detail: err.Error(),
@@ -7992,49 +7514,6 @@ func chatModelConfigToUpdateParams(
}
}
func nullInt64Ptr(n sql.NullInt64) *int64 {
if !n.Valid {
return nil
}
return &n.Int64
}
func writeChatUsageLimitUserNotFound(ctx context.Context, rw http.ResponseWriter) {
httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{
Message: "User not found.",
})
}
func writeChatUsageLimitOverrideNotFound(ctx context.Context, rw http.ResponseWriter) {
httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{
Message: "Chat usage limit override not found.",
})
}
func writeChatUsageLimitGroupOverrideNotFound(ctx context.Context, rw http.ResponseWriter) {
httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{
Message: "Chat usage limit group override not found.",
})
}
func writeChatUsageLimitGroupNotFound(ctx context.Context, rw http.ResponseWriter) {
httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{
Message: "Group not found.",
})
}
func parseChatUsageLimitUserID(rw http.ResponseWriter, r *http.Request) (uuid.UUID, bool) {
userID, err := uuid.Parse(chi.URLParam(r, "user"))
if err != nil {
httpapi.Write(r.Context(), rw, http.StatusBadRequest, codersdk.Response{
Message: "Invalid chat usage limit user ID.",
Detail: err.Error(),
})
return uuid.Nil, false
}
return userID, true
}
func parseChatModelConfigID(rw http.ResponseWriter, r *http.Request) (uuid.UUID, bool) {
modelConfigID, err := uuid.Parse(chi.URLParam(r, "modelConfig"))
if err != nil {
coderd/exp_chats_test.go
+81 -275
View File
@@ -264,28 +264,6 @@ func (s *failNextUpdateChatModelConfigStore) UpdateChatModelConfig(
return s.Store.UpdateChatModelConfig(ctx, arg)
}
func enableDailyChatUsageLimit(
ctx context.Context,
t *testing.T,
db database.Store,
limitMicros int64,
) time.Time {
t.Helper()
_, err := db.UpsertChatUsageLimitConfig(
dbauthz.AsSystemRestricted(ctx),
database.UpsertChatUsageLimitConfigParams{
Enabled: true,
DefaultLimitMicros: limitMicros,
Period: string(codersdk.ChatUsageLimitPeriodDay),
},
)
require.NoError(t, err)
_, periodEnd := chatd.ComputeUsagePeriodBounds(time.Now(), codersdk.ChatUsageLimitPeriodDay)
return periodEnd
}
func insertAssistantCostMessage(
t *testing.T,
db database.Store,
@@ -9694,6 +9672,44 @@ func TestRegenerateChatTitle(t *testing.T) {
require.NoError(t, err)
require.True(t, after.UpdatedAt.Equal(before.UpdatedAt))
})
t.Run("UsageLimitExhausted", func(t *testing.T) {
t.Parallel()
ctx := testutil.Context(t, testutil.WaitLong)
client, db, api := newChatClientWithAPIAndDatabase(t)
firstUser := coderdtest.CreateFirstUser(t, client.Client)
_ = createChatModelConfigWithTitleQuotaExhausted(t, client)
chat, err := client.CreateChat(ctx, codersdk.CreateChatRequest{
OrganizationID: firstUser.OrganizationID,
Content: []codersdk.ChatInputPart{
{
Type: codersdk.ChatInputPartTypeText,
Text: "test chat",
},
},
})
require.NoError(t, err)
coderdtest.WaitForChatSettled(ctx, t, api, chat.ID)
_, err = db.UpdateChatStatus(dbauthz.AsSystemRestricted(ctx), database.UpdateChatStatusParams{
ID: chat.ID,
Status: database.ChatStatusWaiting,
WorkerID: uuid.NullUUID{},
StartedAt: sql.NullTime{},
HeartbeatAt: sql.NullTime{},
LastError: pqtype.NullRawMessage{},
})
require.NoError(t, err)
_, err = client.RegenerateChatTitle(ctx, chat.ID)
sdkErr := requireSDKError(t, err, http.StatusConflict)
require.Equal(t,
"The AI usage limit has been exceeded. Contact an administrator or check the applicable budget and quota settings.",
sdkErr.Message)
})
}
func TestProposeChatTitle(t *testing.T) {
@@ -9882,6 +9898,31 @@ func TestProposeChatTitle(t *testing.T) {
require.True(t, after.UpdatedAt.Equal(before.UpdatedAt),
"propose must not bump updated_at")
})
t.Run("UsageLimitExhausted", func(t *testing.T) {
t.Parallel()
ctx := testutil.Context(t, testutil.WaitLong)
client, _, api := newChatClientWithAPIAndDatabase(t)
firstUser := coderdtest.CreateFirstUser(t, client.Client)
_ = createChatModelConfigWithTitleQuotaExhausted(t, client)
chat, err := client.CreateChat(ctx, codersdk.CreateChatRequest{
OrganizationID: firstUser.OrganizationID,
Content: []codersdk.ChatInputPart{
{Type: codersdk.ChatInputPartTypeText, Text: "test chat"},
},
})
require.NoError(t, err)
coderdtest.WaitForChatSettled(ctx, t, api, chat.ID)
_, err = client.ProposeChatTitle(ctx, chat.ID)
sdkErr := requireSDKError(t, err, http.StatusConflict)
require.Equal(t,
"The AI usage limit has been exceeded. Contact an administrator or check the applicable budget and quota settings.",
sdkErr.Message)
})
}
func TestManualTitleEndpointsPassOwnerSyntheticAPIKeyToAIGateway(t *testing.T) {
@@ -10458,69 +10499,6 @@ func TestPromoteChatQueuedMessage(t *testing.T) {
}
})
t.Run("PromotesAlreadyQueuedMessageAfterLimitReached", func(t *testing.T) {
t.Parallel()
ctx := testutil.Context(t, testutil.WaitLong)
client, db := newChatClientWithDatabase(t)
user := coderdtest.CreateFirstUser(t, client.Client)
modelConfig := createChatModelConfig(t, client)
enableDailyChatUsageLimit(ctx, t, db, 100)
chat := dbgen.Chat(t, db, database.Chat{
OrganizationID: user.OrganizationID,
OwnerID: user.UserID,
LastModelConfigID: modelConfig.ID,
Title: "promote queued usage limit",
Status: database.ChatStatusError,
})
const queuedText = "queued message for promote route"
queuedContent, err := json.Marshal([]codersdk.ChatMessagePart{
codersdk.ChatMessageText(queuedText),
})
require.NoError(t, err)
queuedMessage := insertTestChatQueuedMessage(ctx, t, db, chat.ID, queuedContent, chat.LastModelConfigID)
insertAssistantCostMessage(t, db, chat.ID, modelConfig.ID, 100)
promoteRes, err := client.Request(
ctx,
http.MethodPost,
fmt.Sprintf("/api/experimental/chats/%s/queue/%d/promote", chat.ID, queuedMessage.ID),
nil,
)
require.NoError(t, err)
defer promoteRes.Body.Close()
require.Equal(t, http.StatusAccepted, promoteRes.StatusCode)
var resp codersdk.Response
require.NoError(t, json.NewDecoder(promoteRes.Body).Decode(&resp))
require.NotEmpty(t, resp.Message)
messagesResult, err := client.GetChatMessages(ctx, chat.ID, nil)
require.NoError(t, err)
foundPromoted := false
for _, msg := range messagesResult.Messages {
if msg.Role != codersdk.ChatMessageRoleUser {
continue
}
for _, part := range msg.Content {
if part.Type == codersdk.ChatMessagePartTypeText && part.Text == queuedText {
foundPromoted = true
}
}
}
require.True(t, foundPromoted, "promoted message must appear in chat history")
queuedMessages, err := db.GetChatQueuedMessages(dbauthz.AsSystemRestricted(ctx), chat.ID)
require.NoError(t, err)
for _, queued := range queuedMessages {
require.NotEqual(t, queuedMessage.ID, queued.ID)
}
})
t.Run("InvalidQueuedMessageID", func(t *testing.T) {
t.Parallel()
@@ -10822,196 +10800,6 @@ func TestPromoteChatQueuedMessage(t *testing.T) {
})
}
func TestChatUsageLimitOverrideRoutes(t *testing.T) {
t.Parallel()
t.Run("UpsertUserOverrideRequiresPositiveSpendLimit", func(t *testing.T) {
t.Parallel()
ctx := testutil.Context(t, testutil.WaitLong)
client, _ := newChatClientWithDatabase(t)
firstUser := coderdtest.CreateFirstUser(t, client.Client)
_, member := coderdtest.CreateAnotherUser(t, client.Client, firstUser.OrganizationID)
res, err := client.Request(
ctx,
http.MethodPut,
fmt.Sprintf("/api/experimental/chats/usage-limits/overrides/%s", member.ID),
map[string]any{},
)
require.NoError(t, err)
defer res.Body.Close()
err = codersdk.ReadBodyAsError(res)
sdkErr := requireSDKError(t, err, http.StatusBadRequest)
require.Equal(t, "Invalid chat usage limit override.", sdkErr.Message)
require.Equal(t, "Spend limit must be greater than 0.", sdkErr.Detail)
})
t.Run("UpsertUserOverrideMissingUser", func(t *testing.T) {
t.Parallel()
ctx := testutil.Context(t, testutil.WaitLong)
client := newChatClient(t)
_ = coderdtest.CreateFirstUser(t, client.Client)
_, err := client.UpsertChatUsageLimitOverride(ctx, uuid.New(), codersdk.UpsertChatUsageLimitOverrideRequest{
SpendLimitMicros: 7_000_000,
})
sdkErr := requireSDKError(t, err, http.StatusNotFound)
require.Equal(t, "User not found.", sdkErr.Message)
})
t.Run("DeleteUserOverrideMissingUser", func(t *testing.T) {
t.Parallel()
ctx := testutil.Context(t, testutil.WaitLong)
client := newChatClient(t)
_ = coderdtest.CreateFirstUser(t, client.Client)
err := client.DeleteChatUsageLimitOverride(ctx, uuid.New())
sdkErr := requireSDKError(t, err, http.StatusBadRequest)
require.Equal(t, "User not found.", sdkErr.Message)
})
t.Run("DeleteUserOverrideMissingOverride", func(t *testing.T) {
t.Parallel()
ctx := testutil.Context(t, testutil.WaitLong)
client := newChatClient(t)
firstUser := coderdtest.CreateFirstUser(t, client.Client)
_, member := coderdtest.CreateAnotherUser(t, client.Client, firstUser.OrganizationID)
err := client.DeleteChatUsageLimitOverride(ctx, member.ID)
sdkErr := requireSDKError(t, err, http.StatusBadRequest)
require.Equal(t, "Chat usage limit override not found.", sdkErr.Message)
})
t.Run("UpdateUserOverride", func(t *testing.T) {
t.Parallel()
ctx := testutil.Context(t, testutil.WaitLong)
client, _ := newChatClientWithDatabase(t)
firstUser := coderdtest.CreateFirstUser(t, client.Client)
_, member := coderdtest.CreateAnotherUser(t, client.Client, firstUser.OrganizationID)
_, err := client.UpsertChatUsageLimitOverride(ctx, member.ID, codersdk.UpsertChatUsageLimitOverrideRequest{
SpendLimitMicros: 5_000_000,
})
require.NoError(t, err)
override, err := client.UpsertChatUsageLimitOverride(ctx, member.ID, codersdk.UpsertChatUsageLimitOverrideRequest{
SpendLimitMicros: 10_000_000,
})
require.NoError(t, err)
require.Equal(t, member.ID, override.UserID)
require.NotNil(t, override.SpendLimitMicros)
require.EqualValues(t, 10_000_000, *override.SpendLimitMicros)
config, err := client.GetChatUsageLimitConfig(ctx)
require.NoError(t, err)
require.Len(t, config.Overrides, 1)
require.Equal(t, member.ID, config.Overrides[0].UserID)
require.NotNil(t, config.Overrides[0].SpendLimitMicros)
require.EqualValues(t, 10_000_000, *config.Overrides[0].SpendLimitMicros)
})
t.Run("UpsertGroupOverrideIncludesMemberCount", func(t *testing.T) {
t.Parallel()
ctx := testutil.Context(t, testutil.WaitLong)
client, db := newChatClientWithDatabase(t)
firstUser := coderdtest.CreateFirstUser(t, client.Client)
_, member := coderdtest.CreateAnotherUser(t, client.Client, firstUser.OrganizationID)
group := dbgen.Group(t, db, database.Group{OrganizationID: firstUser.OrganizationID})
dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: group.ID, UserID: member.ID})
dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: group.ID, UserID: database.PrebuildsSystemUserID})
override, err := client.UpsertChatUsageLimitGroupOverride(ctx, group.ID, codersdk.UpsertChatUsageLimitGroupOverrideRequest{
SpendLimitMicros: 7_000_000,
})
require.NoError(t, err)
require.Equal(t, group.ID, override.GroupID)
require.EqualValues(t, 1, override.MemberCount)
require.NotNil(t, override.SpendLimitMicros)
require.EqualValues(t, 7_000_000, *override.SpendLimitMicros)
config, err := client.GetChatUsageLimitConfig(ctx)
require.NoError(t, err)
var listed *codersdk.ChatUsageLimitGroupOverride
for i := range config.GroupOverrides {
if config.GroupOverrides[i].GroupID == group.ID {
listed = &config.GroupOverrides[i]
break
}
}
require.NotNil(t, listed)
require.EqualValues(t, 1, listed.MemberCount)
})
t.Run("UpdateGroupOverride", func(t *testing.T) {
t.Parallel()
ctx := testutil.Context(t, testutil.WaitLong)
client, db := newChatClientWithDatabase(t)
firstUser := coderdtest.CreateFirstUser(t, client.Client)
_, member := coderdtest.CreateAnotherUser(t, client.Client, firstUser.OrganizationID)
group := dbgen.Group(t, db, database.Group{OrganizationID: firstUser.OrganizationID})
dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: group.ID, UserID: firstUser.UserID})
dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: group.ID, UserID: member.ID})
_, err := client.UpsertChatUsageLimitGroupOverride(ctx, group.ID, codersdk.UpsertChatUsageLimitGroupOverrideRequest{
SpendLimitMicros: 5_000_000,
})
require.NoError(t, err)
override, err := client.UpsertChatUsageLimitGroupOverride(ctx, group.ID, codersdk.UpsertChatUsageLimitGroupOverrideRequest{
SpendLimitMicros: 10_000_000,
})
require.NoError(t, err)
require.Equal(t, group.ID, override.GroupID)
require.EqualValues(t, 2, override.MemberCount)
require.NotNil(t, override.SpendLimitMicros)
require.EqualValues(t, 10_000_000, *override.SpendLimitMicros)
config, err := client.GetChatUsageLimitConfig(ctx)
require.NoError(t, err)
require.Len(t, config.GroupOverrides, 1)
require.Equal(t, group.ID, config.GroupOverrides[0].GroupID)
require.EqualValues(t, 2, config.GroupOverrides[0].MemberCount)
require.NotNil(t, config.GroupOverrides[0].SpendLimitMicros)
require.EqualValues(t, 10_000_000, *config.GroupOverrides[0].SpendLimitMicros)
})
t.Run("UpsertGroupOverrideMissingGroup", func(t *testing.T) {
t.Parallel()
ctx := testutil.Context(t, testutil.WaitLong)
client := newChatClient(t)
_ = coderdtest.CreateFirstUser(t, client.Client)
_, err := client.UpsertChatUsageLimitGroupOverride(ctx, uuid.New(), codersdk.UpsertChatUsageLimitGroupOverrideRequest{
SpendLimitMicros: 7_000_000,
})
sdkErr := requireSDKError(t, err, http.StatusNotFound)
require.Equal(t, "Group not found.", sdkErr.Message)
})
t.Run("DeleteGroupOverrideMissingOverride", func(t *testing.T) {
t.Parallel()
ctx := testutil.Context(t, testutil.WaitLong)
client, db := newChatClientWithDatabase(t)
firstUser := coderdtest.CreateFirstUser(t, client.Client)
group := dbgen.Group(t, db, database.Group{OrganizationID: firstUser.OrganizationID})
err := client.DeleteChatUsageLimitGroupOverride(ctx, group.ID)
sdkErr := requireSDKError(t, err, http.StatusBadRequest)
require.Equal(t, "Chat usage limit group override not found.", sdkErr.Message)
})
}
func TestPostChatFile(t *testing.T) {
t.Parallel()
@@ -12604,6 +12392,24 @@ func createChatModelConfigWithTitleFailure(t testing.TB, client *codersdk.Experi
return createChatModelConfigWithBaseURL(t, client, baseURL)
}
// createChatModelConfigWithTitleQuotaExhausted provisions a model whose
// non-streaming responses return a provider insufficient_quota error, which
// classifies as a usage limit like an exhausted AI Gateway budget.
func createChatModelConfigWithTitleQuotaExhausted(t testing.TB, client *codersdk.ExperimentalClient) codersdk.ChatModelConfig {
t.Helper()
baseURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse {
if req.Stream {
return chattest.OpenAIStreamingResponse(chattest.OpenAITextChunks("Hello from test server.")...)
}
return chattest.OpenAIErrorResponse(
http.StatusBadRequest,
"insufficient_quota",
"You exceeded your current quota, please check your plan and billing details.",
)
})
return createChatModelConfigWithBaseURL(t, client, baseURL)
}
func createAdditionalChatModelConfig(
t *testing.T,
client *codersdk.ExperimentalClient,
coderd/x/chatd/chatd.go
-57
View File
@@ -20,7 +20,6 @@ import (
"github.com/dustin/go-humanize"
"github.com/google/uuid"
"github.com/prometheus/client_golang/prometheus"
"github.com/shopspring/decimal"
"github.com/sqlc-dev/pqtype"
"golang.org/x/sync/errgroup"
"golang.org/x/xerrors"
@@ -1133,27 +1132,6 @@ var (
ErrNothingToCompact = xerrors.New("nothing to compact")
)
// UsageLimitExceededError indicates the user has exceeded their chat spend
// limit.
type UsageLimitExceededError struct {
LimitMicros int64
ConsumedMicros int64
PeriodEnd time.Time
}
func formatMicrosAsDollars(micros int64) string {
return "$" + decimal.NewFromInt(micros).Shift(-6).StringFixed(2)
}
func (e *UsageLimitExceededError) Error() string {
return fmt.Sprintf(
"usage limit exceeded: spent %s of %s limit, resets at %s",
formatMicrosAsDollars(e.ConsumedMicros),
formatMicrosAsDollars(e.LimitMicros),
e.PeriodEnd.Format(time.RFC3339),
)
}
// CreateOptions controls chat creation in the shared chat mutation path.
type CreateOptions struct {
OrganizationID uuid.UUID
@@ -1289,11 +1267,6 @@ func (p *Server) CreateChat(ctx context.Context, opts CreateOptions) (database.C
// another pool checkout.
deploymentPrompt := p.resolveDeploymentSystemPrompt(ctx)
// Usage limits gate the create before we touch the state machine.
if limitErr := p.checkUsageLimit(ctx, p.db, opts.OwnerID, uuid.NullUUID{UUID: opts.OrganizationID, Valid: true}); limitErr != nil {
return database.Chat{}, limitErr
}
if opts.ModelConfigID != uuid.Nil {
if err := requireEnabledChatModelConfig(ctx, p.db, opts.ModelConfigID); err != nil {
return database.Chat{}, err
@@ -1456,9 +1429,6 @@ func (p *Server) SendMessage(
if chat.Archived {
return SendMessageResult{}, ErrChatArchived
}
if err := p.checkUsageLimit(ctx, p.db, chat.OwnerID, uuid.NullUUID{UUID: chat.OrganizationID, Valid: true}); err != nil {
return SendMessageResult{}, err
}
if _, err := resolveSendMessageModelConfigID(ctx, p.db, chat, opts.ModelConfigID); err != nil {
return SendMessageResult{}, err
}
@@ -1505,11 +1475,6 @@ func (p *Server) SendMessage(
return ErrChatArchived
}
// Enforce usage limits before any state-machine work.
if limitErr := p.checkUsageLimit(ctx, store, lockedChat.OwnerID, uuid.NullUUID{UUID: lockedChat.OrganizationID, Valid: true}); limitErr != nil {
return limitErr
}
if requestedPlanMode != nil {
lockedChat, err = store.UpdateChatPlanModeByID(ctx, database.UpdateChatPlanModeByIDParams{
PlanMode: *requestedPlanMode,
@@ -1599,12 +1564,6 @@ func (p *Server) SendMessage(
return result, nil
}
// checkUsageLimit is a no-op. Usage limits (a.k.a. "Budgets") are now enforced
// by AI Gateway.
func (*Server) checkUsageLimit(_ context.Context, _ database.Store, _ uuid.UUID, _ uuid.NullUUID) error {
return nil
}
func chatdModelConfigLookupContext(ctx context.Context) context.Context {
//nolint:gocritic // Chat message admission needs daemon-scoped
// deployment-config reads for model config validation.
@@ -1770,9 +1729,6 @@ func (p *Server) EditMessage(
if chat.Archived {
return EditMessageResult{}, ErrChatArchived
}
if err := p.checkUsageLimit(ctx, p.db, chat.OwnerID, uuid.NullUUID{UUID: chat.OrganizationID, Valid: true}); err != nil {
return EditMessageResult{}, err
}
if err := validateEditTarget(ctx, p.db, opts.ChatID, opts.EditedMessageID); err != nil {
return EditMessageResult{}, err
}
@@ -1815,10 +1771,6 @@ func (p *Server) EditMessage(
if lockedChat.Archived {
return ErrChatArchived
}
if limitErr := p.checkUsageLimit(ctx, store, lockedChat.OwnerID, uuid.NullUUID{UUID: lockedChat.OrganizationID, Valid: true}); limitErr != nil {
return limitErr
}
// Capture the target message for the post-commit debug
// cleanup hook below. The transition itself revalidates
// chat ownership and user-message constraints.
@@ -2364,11 +2316,6 @@ func (p *Server) CompactChat(
if _, ok := firstUncompressedAssistantAfter(messages, boundary); !ok {
return ErrNothingToCompact
}
// Usage validation runs last so rejected requests report the more
// specific state or content conflict. Its failure rolls back the marker.
if limitErr := p.checkUsageLimit(ctx, store, lockedChat.OwnerID, uuid.NullUUID{UUID: lockedChat.OrganizationID, Valid: true}); limitErr != nil {
return limitErr
}
refreshed = result.Chat
return nil
})
@@ -2523,10 +2470,6 @@ func (p *Server) generateManualTitleCandidate(
store database.Store,
chat database.Chat,
) (string, error) {
if limitErr := p.checkUsageLimit(ctx, store, chat.OwnerID, uuid.NullUUID{UUID: chat.OrganizationID, Valid: true}); limitErr != nil {
return "", limitErr
}
headMessages, err := store.GetChatMessagesByChatIDAscPaginated(
ctx,
database.GetChatMessagesByChatIDAscPaginatedParams{
coderd/x/chatd/chatd_test.go
-189
View File
@@ -64,7 +64,6 @@ import (
"github.com/coder/coder/v2/provisioner/echo"
proto "github.com/coder/coder/v2/provisionersdk/proto"
"github.com/coder/coder/v2/testutil"
"github.com/coder/quartz"
)
type recordedOpenAIRequest struct {
@@ -1988,194 +1987,6 @@ func TestAutoPromoteQueuedMessagesPreservesPerTurnModelOrder(t *testing.T) {
require.Equal(t, []uuid.UUID{modelConfigA.ID, modelConfigB.ID, modelConfigC.ID}, userModelConfigIDs)
}
func TestInterruptAutoPromotionIgnoresLaterUsageLimitIncrease(t *testing.T) {
t.Parallel()
db, ps := dbtestutil.NewDB(t)
ctx := testutil.Context(t, testutil.WaitLong)
_, err := db.UpsertChatUsageLimitConfig(ctx, database.UpsertChatUsageLimitConfigParams{
Enabled: true,
DefaultLimitMicros: 100,
Period: string(codersdk.ChatUsageLimitPeriodDay),
})
require.NoError(t, err)
clock := quartz.NewMock(t)
streamStarted := make(chan struct{})
interrupted := make(chan struct{})
secondRequestStarted := make(chan struct{}, 1)
thirdRequestStarted := make(chan struct{}, 1)
allowFinish := make(chan struct{})
allowSecondRequestFinish := make(chan struct{})
allowThirdRequestFinish := make(chan struct{})
var requestCount atomic.Int32
openAIURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse {
if !req.Stream {
return chattest.OpenAINonStreamingResponse("title")
}
switch requestCount.Add(1) {
case 1:
chunks := make(chan chattest.OpenAIChunk, 1)
go func() {
defer close(chunks)
chunks <- chattest.OpenAITextChunks("partial")[0]
select {
case <-streamStarted:
default:
close(streamStarted)
}
<-req.Context().Done()
select {
case <-interrupted:
default:
close(interrupted)
}
<-allowFinish
}()
return chattest.OpenAIResponse{StreamingChunks: chunks}
case 2:
select {
case secondRequestStarted <- struct{}{}:
default:
}
chunks := make(chan chattest.OpenAIChunk, 1)
go func() {
defer close(chunks)
chunks <- chattest.OpenAITextChunks("second run partial")[0]
select {
case <-allowSecondRequestFinish:
case <-req.Context().Done():
}
}()
return chattest.OpenAIResponse{StreamingChunks: chunks}
case 3:
select {
case thirdRequestStarted <- struct{}{}:
default:
}
chunks := make(chan chattest.OpenAIChunk, 1)
go func() {
defer close(chunks)
chunks <- chattest.OpenAITextChunks("third run partial")[0]
select {
case <-allowThirdRequestFinish:
case <-req.Context().Done():
}
}()
return chattest.OpenAIResponse{StreamingChunks: chunks}
}
return chattest.OpenAIStreamingResponse(
chattest.OpenAITextChunks("done")...,
)
})
factory := chattest.NewMockAIBridgeTransport(t, openAIURL)
server := newActiveTestServer(t, db, ps, func(cfg *chatd.Config) {
cfg.AIBridgeTransportFactory = chatAIGatewayTransportFactoryPointer(factory)
cfg.Clock = clock
// Keep periodic polling frozen so request handoff is synchronized
// through explicit mock channels.
cfg.PendingChatAcquireInterval = time.Hour
cfg.InFlightChatStaleAfter = testutil.WaitSuperLong
})
user, org, model := seedChatDependencies(t, db)
setOpenAIProviderBaseURL(ctx, t, db, openAIURL)
chat, err := server.CreateChat(ctx, chatd.CreateOptions{
OrganizationID: org.ID,
OwnerID: user.ID,
Title: "interrupt-autopromote-limit",
ModelConfigID: model.ID,
InitialUserContent: []codersdk.ChatMessagePart{codersdk.ChatMessageText("hello")},
})
require.NoError(t, err)
testutil.TryReceive(ctx, t, streamStarted)
queuedResult, err := server.SendMessage(ctx, chatd.SendMessageOptions{
ChatID: chat.ID,
Content: []codersdk.ChatMessagePart{codersdk.ChatMessageText("queued")},
BusyBehavior: chatd.SendMessageBusyBehaviorInterrupt,
})
require.NoError(t, err)
require.True(t, queuedResult.Queued)
require.NotNil(t, queuedResult.QueuedMessage)
testutil.TryReceive(ctx, t, interrupted)
close(allowFinish)
testutil.TryReceive(ctx, t, secondRequestStarted)
laterQueuedResult, err := server.SendMessage(ctx, chatd.SendMessageOptions{
ChatID: chat.ID,
Content: []codersdk.ChatMessagePart{codersdk.ChatMessageText("later queued")},
})
require.NoError(t, err)
require.True(t, laterQueuedResult.Queued)
require.NotNil(t, laterQueuedResult.QueuedMessage)
spendChat := dbgen.Chat(t, db, database.Chat{
OrganizationID: org.ID,
OwnerID: user.ID,
LastModelConfigID: model.ID,
Title: "other-spend",
})
assistantContent, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{
codersdk.ChatMessageText("spent elsewhere"),
})
require.NoError(t, err)
_ = dbgen.ChatMessage(t, db, database.ChatMessage{
ChatID: spendChat.ID,
ModelConfigID: uuid.NullUUID{UUID: model.ID, Valid: true},
Role: database.ChatMessageRoleAssistant,
ContentVersion: chatprompt.CurrentContentVersion,
Content: assistantContent,
TotalCostMicros: sql.NullInt64{Int64: 100, Valid: true},
})
close(allowSecondRequestFinish)
testutil.TryReceive(ctx, t, thirdRequestStarted)
require.GreaterOrEqual(t, requestCount.Load(), int32(3))
close(allowThirdRequestFinish)
chatd.WaitUntilIdleForTest(server)
queued, err := db.GetChatQueuedMessages(ctx, chat.ID)
require.NoError(t, err)
require.Empty(t, queued)
fromDB, err := db.GetChatByID(ctx, chat.ID)
require.NoError(t, err)
require.Equal(t, database.ChatStatusWaiting, fromDB.Status)
require.False(t, fromDB.WorkerID.Valid)
messages, err := db.GetChatMessagesByChatID(ctx, database.GetChatMessagesByChatIDParams{
ChatID: chat.ID,
AfterID: 0,
})
require.NoError(t, err)
userTexts := make([]string, 0, 3)
for _, message := range messages {
if message.Role != database.ChatMessageRoleUser {
continue
}
sdkMessage := db2sdk.ChatMessage(message)
if len(sdkMessage.Content) != 1 {
continue
}
userTexts = append(userTexts, sdkMessage.Content[0].Text)
}
require.Equal(t, []string{"hello", "queued", "later queued"}, userTexts)
}
func TestEditMessageRejectsMissingMessage(t *testing.T) {
t.Parallel()
coderd/x/chatd/chaterror/classify_test.go
+48
View File
@@ -135,6 +135,54 @@ func TestClassify(t *testing.T) {
StatusCode: 403,
},
},
{
name: "AIBudget403ClassifiesAsUsageLimit",
err: xerrors.New("status 403: AI budget of US$5.00 exceeded. Please contact an administrator for more details."),
want: chaterror.ClassifiedError{
Message: "The AI usage limit has been exceeded. Contact an administrator or check the applicable budget and quota settings.",
Detail: "status 403: AI budget of US$5.00 exceeded. Please contact an administrator for more details.",
Kind: codersdk.ChatErrorKindUsageLimit,
Provider: "",
Retryable: false,
StatusCode: 403,
},
},
{
// The SDK message reduces to a bare status line, so the
// text/plain body is the only usage-limit signal.
name: "AIBudget403PlainTextBodyClassifiesAsUsageLimit",
err: testProviderError(
`POST "http://coder-aibridge/v1/messages": 403 Forbidden`,
403,
nil,
[]byte("HTTP/1.1 403 Forbidden\r\nContent-Type: text/plain; charset=utf-8\r\n\r\nAI budget of US$10.00 exceeded. Please contact an administrator for more details.\n"),
),
want: chaterror.ClassifiedError{
Message: "The AI usage limit has been exceeded. Contact an administrator or check the applicable budget and quota settings.",
Detail: "AI budget of US$10.00 exceeded. Please contact an administrator for more details.",
Kind: codersdk.ChatErrorKindUsageLimit,
Provider: "",
Retryable: false,
StatusCode: 403,
},
},
{
name: "HTMLBodyDoesNotBecomeDetail",
err: testProviderError(
`POST "https://example.com/v1/messages": 403 Forbidden`,
403,
nil,
[]byte("HTTP/1.1 403 Forbidden\r\nContent-Type: text/html\r\n\r\n<html><body>Forbidden</body></html>"),
),
want: chaterror.ClassifiedError{
Message: "Authentication with the AI provider failed. Check the API key and permissions.",
Detail: `POST "https://example.com/v1/messages": 403 Forbidden`,
Kind: codersdk.ChatErrorKindAuth,
Provider: "",
Retryable: false,
StatusCode: 403,
},
},
{
name: "ForbiddenContextLengthClassifiesAsConfig",
err: xerrors.New("forbidden: context length exceeded"),
coderd/x/chatd/chaterror/message.go
+2 -5
View File
@@ -34,11 +34,8 @@ func terminalMessage(classified ClassifiedError) string {
))
case codersdk.ChatErrorKindUsageLimit:
return stringutil.Capitalize(fmt.Sprintf(
"The usage quota for %s has been exceeded."+
" Check the billing and quota settings for the provider account.",
subject,
))
return "The AI usage limit has been exceeded." +
" Contact an administrator or check the applicable budget and quota settings."
case codersdk.ChatErrorKindAuth:
return fmt.Sprintf(
coderd/x/chatd/chaterror/message_test.go
+2 -2
View File
@@ -81,14 +81,14 @@ func TestTerminalMessage(t *testing.T) {
kind: codersdk.ChatErrorKindUsageLimit,
provider: "openai",
retryable: false,
want: "The usage quota for OpenAI has been exceeded. Check the billing and quota settings for the provider account.",
want: "The AI usage limit has been exceeded. Contact an administrator or check the applicable budget and quota settings.",
},
{
name: "UsageLimit_UnknownProvider",
kind: codersdk.ChatErrorKindUsageLimit,
provider: "",
retryable: false,
want: "The usage quota for the AI provider has been exceeded. Check the billing and quota settings for the provider account.",
want: "The AI usage limit has been exceeded. Contact an administrator or check the applicable budget and quota settings.",
},
{
name: "MissingKey",
coderd/x/chatd/subagent.go
-4
View File
@@ -1293,10 +1293,6 @@ func (p *Server) createChildSubagentChatWithOptions(
// strip the root-only orchestration guidance from their prompt.
deploymentPrompt = strings.Replace(deploymentPrompt, subagentOrchestrationPromptBlock, "", 1)
if limitErr := p.checkUsageLimit(ctx, p.db, parent.OwnerID, uuid.NullUUID{UUID: parent.OrganizationID, Valid: true}); limitErr != nil {
return database.Chat{}, limitErr
}
// Review before persistence so spawned chats cannot bypass prompt policy.
childChatID := uuid.New()
var promptResult *chathooks.Result
coderd/x/chatd/usagelimit.go
-152
View File
@@ -1,152 +0,0 @@
package chatd
import (
"context"
"database/sql"
"errors"
"fmt"
"time"
"github.com/google/uuid"
"golang.org/x/xerrors"
"github.com/coder/coder/v2/coderd/database"
"github.com/coder/coder/v2/coderd/database/dbauthz"
"github.com/coder/coder/v2/codersdk"
)
// ComputeUsagePeriodBounds returns the UTC-aligned start and end bounds for the
// active usage-limit period containing now.
func ComputeUsagePeriodBounds(now time.Time, period codersdk.ChatUsageLimitPeriod) (start, end time.Time) {
utcNow := now.UTC()
switch period {
case codersdk.ChatUsageLimitPeriodDay:
start = time.Date(utcNow.Year(), utcNow.Month(), utcNow.Day(), 0, 0, 0, 0, time.UTC)
end = start.AddDate(0, 0, 1)
case codersdk.ChatUsageLimitPeriodWeek:
// Walk backward to Monday of the current ISO week.
// ISO 8601 weeks always start on Monday, so this never
// crosses an ISO-week boundary.
start = time.Date(utcNow.Year(), utcNow.Month(), utcNow.Day(), 0, 0, 0, 0, time.UTC)
for start.Weekday() != time.Monday {
start = start.AddDate(0, 0, -1)
}
end = start.AddDate(0, 0, 7)
case codersdk.ChatUsageLimitPeriodMonth:
start = time.Date(utcNow.Year(), utcNow.Month(), 1, 0, 0, 0, 0, time.UTC)
end = start.AddDate(0, 1, 0)
default:
panic(fmt.Sprintf("unknown chat usage limit period: %q", period))
}
return start, end
}
// 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
// parallel, allowing brief overage. This is acceptable because:
// - Cost is only known after the LLM API returns.
// - Overage is bounded by message cost × concurrency.
// - Fail-open is the deliberate design choice for this feature.
//
// Architecture note: today this path enforces one period globally
// (day/week/month) from config.
// To support simultaneous periods, add nullable
// daily/weekly/monthly_limit_micros columns on override tables, where NULL
// means no limit for that period.
// 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, 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)
config, err := db.GetChatUsageLimitConfig(authCtx)
if err != nil {
if errors.Is(err, sql.ErrNoRows) {
return nil, nil //nolint:nilnil // Nil status cleanly signals disabled limits.
}
return nil, err
}
if !config.Enabled {
return nil, nil //nolint:nilnil // Nil status cleanly signals disabled limits.
}
period, ok := mapDBPeriodToSDK(config.Period)
if !ok {
return nil, xerrors.Errorf("invalid chat usage limit period %q", config.Period)
}
// Resolve effective limit in a single query:
// individual override > group limit > global default.
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 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,
OrganizationID: spendOrgID,
StartTime: start,
EndTime: end,
})
if err != nil {
return nil, err
}
effectiveLimit := limitResult.EffectiveLimitMicros
return &codersdk.ChatUsageLimitStatus{
IsLimited: true,
Period: period,
SpendLimitMicros: &effectiveLimit,
CurrentSpend: spendTotal,
PeriodStart: start,
PeriodEnd: end,
}, 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):
return codersdk.ChatUsageLimitPeriodDay, true
case string(codersdk.ChatUsageLimitPeriodWeek):
return codersdk.ChatUsageLimitPeriodWeek, true
case string(codersdk.ChatUsageLimitPeriodMonth):
return codersdk.ChatUsageLimitPeriodMonth, true
default:
return "", false
}
}
coderd/x/chatd/usagelimit_internal_test.go
-132
View File
@@ -1,132 +0,0 @@
package chatd
import (
"testing"
"time"
"github.com/coder/coder/v2/codersdk"
)
func TestComputeUsagePeriodBounds(t *testing.T) {
t.Parallel()
newYork, err := time.LoadLocation("America/New_York")
if err != nil {
t.Fatalf("load America/New_York: %v", err)
}
tests := []struct {
name string
now time.Time
period codersdk.ChatUsageLimitPeriod
wantStart time.Time
wantEnd time.Time
}{
{
name: "day/mid_day",
now: time.Date(2025, time.June, 15, 14, 30, 0, 0, time.UTC),
period: codersdk.ChatUsageLimitPeriodDay,
wantStart: time.Date(2025, time.June, 15, 0, 0, 0, 0, time.UTC),
wantEnd: time.Date(2025, time.June, 16, 0, 0, 0, 0, time.UTC),
},
{
name: "day/midnight_exactly",
now: time.Date(2025, time.June, 15, 0, 0, 0, 0, time.UTC),
period: codersdk.ChatUsageLimitPeriodDay,
wantStart: time.Date(2025, time.June, 15, 0, 0, 0, 0, time.UTC),
wantEnd: time.Date(2025, time.June, 16, 0, 0, 0, 0, time.UTC),
},
{
name: "day/end_of_day",
now: time.Date(2025, time.June, 15, 23, 59, 59, 0, time.UTC),
period: codersdk.ChatUsageLimitPeriodDay,
wantStart: time.Date(2025, time.June, 15, 0, 0, 0, 0, time.UTC),
wantEnd: time.Date(2025, time.June, 16, 0, 0, 0, 0, time.UTC),
},
{
name: "week/wednesday",
now: time.Date(2025, time.June, 11, 10, 0, 0, 0, time.UTC),
period: codersdk.ChatUsageLimitPeriodWeek,
wantStart: time.Date(2025, time.June, 9, 0, 0, 0, 0, time.UTC),
wantEnd: time.Date(2025, time.June, 16, 0, 0, 0, 0, time.UTC),
},
{
name: "week/monday",
now: time.Date(2025, time.June, 9, 0, 0, 0, 0, time.UTC),
period: codersdk.ChatUsageLimitPeriodWeek,
wantStart: time.Date(2025, time.June, 9, 0, 0, 0, 0, time.UTC),
wantEnd: time.Date(2025, time.June, 16, 0, 0, 0, 0, time.UTC),
},
{
name: "week/sunday",
now: time.Date(2025, time.June, 15, 23, 0, 0, 0, time.UTC),
period: codersdk.ChatUsageLimitPeriodWeek,
wantStart: time.Date(2025, time.June, 9, 0, 0, 0, 0, time.UTC),
wantEnd: time.Date(2025, time.June, 16, 0, 0, 0, 0, time.UTC),
},
{
name: "week/year_boundary",
now: time.Date(2024, time.December, 31, 12, 0, 0, 0, time.UTC),
period: codersdk.ChatUsageLimitPeriodWeek,
wantStart: time.Date(2024, time.December, 30, 0, 0, 0, 0, time.UTC),
wantEnd: time.Date(2025, time.January, 6, 0, 0, 0, 0, time.UTC),
},
{
name: "month/mid_month",
now: time.Date(2025, time.June, 15, 0, 0, 0, 0, time.UTC),
period: codersdk.ChatUsageLimitPeriodMonth,
wantStart: time.Date(2025, time.June, 1, 0, 0, 0, 0, time.UTC),
wantEnd: time.Date(2025, time.July, 1, 0, 0, 0, 0, time.UTC),
},
{
name: "month/first_day",
now: time.Date(2025, time.June, 1, 0, 0, 0, 0, time.UTC),
period: codersdk.ChatUsageLimitPeriodMonth,
wantStart: time.Date(2025, time.June, 1, 0, 0, 0, 0, time.UTC),
wantEnd: time.Date(2025, time.July, 1, 0, 0, 0, 0, time.UTC),
},
{
name: "month/last_day",
now: time.Date(2025, time.June, 30, 23, 59, 59, 0, time.UTC),
period: codersdk.ChatUsageLimitPeriodMonth,
wantStart: time.Date(2025, time.June, 1, 0, 0, 0, 0, time.UTC),
wantEnd: time.Date(2025, time.July, 1, 0, 0, 0, 0, time.UTC),
},
{
name: "month/february",
now: time.Date(2025, time.February, 15, 12, 0, 0, 0, time.UTC),
period: codersdk.ChatUsageLimitPeriodMonth,
wantStart: time.Date(2025, time.February, 1, 0, 0, 0, 0, time.UTC),
wantEnd: time.Date(2025, time.March, 1, 0, 0, 0, 0, time.UTC),
},
{
name: "month/leap_year_february",
now: time.Date(2024, time.February, 29, 12, 0, 0, 0, time.UTC),
period: codersdk.ChatUsageLimitPeriodMonth,
wantStart: time.Date(2024, time.February, 1, 0, 0, 0, 0, time.UTC),
wantEnd: time.Date(2024, time.March, 1, 0, 0, 0, 0, time.UTC),
},
{
name: "day/non_utc_timezone",
now: time.Date(2025, time.June, 15, 22, 0, 0, 0, newYork),
period: codersdk.ChatUsageLimitPeriodDay,
wantStart: time.Date(2025, time.June, 16, 0, 0, 0, 0, time.UTC),
wantEnd: time.Date(2025, time.June, 17, 0, 0, 0, 0, time.UTC),
},
}
for _, tc := range tests {
tc := tc
t.Run(tc.name, func(t *testing.T) {
t.Parallel()
start, end := ComputeUsagePeriodBounds(tc.now, tc.period)
if !start.Equal(tc.wantStart) {
t.Errorf("start: got %v, want %v", start, tc.wantStart)
}
if !end.Equal(tc.wantEnd) {
t.Errorf("end: got %v, want %v", end, tc.wantEnd)
}
})
}
}