diff --git a/coderd/coderd.go b/coderd/coderd.go index 2f49c5f00f..a6121f1660 100644 --- a/coderd/coderd.go +++ b/coderd/coderd.go @@ -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) { diff --git a/coderd/database/dbauthz/dbauthz.go b/coderd/database/dbauthz/dbauthz.go index e0eab762a8..e732bf7b7d 100644 --- a/coderd/database/dbauthz/dbauthz.go +++ b/coderd/database/dbauthz/dbauthz.go @@ -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 { diff --git a/coderd/database/dbauthz/dbauthz_test.go b/coderd/database/dbauthz/dbauthz_test.go index 2cb8a82191..171015c1a2 100644 --- a/coderd/database/dbauthz/dbauthz_test.go +++ b/coderd/database/dbauthz/dbauthz_test.go @@ -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) diff --git a/coderd/database/dbmetrics/querymetrics.go b/coderd/database/dbmetrics/querymetrics.go index 7ef3f15811..716ca6da8f 100644 --- a/coderd/database/dbmetrics/querymetrics.go +++ b/coderd/database/dbmetrics/querymetrics.go @@ -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) diff --git a/coderd/database/dbmock/dbmock.go b/coderd/database/dbmock/dbmock.go index 03088ef9a8..872a47a185 100644 --- a/coderd/database/dbmock/dbmock.go +++ b/coderd/database/dbmock/dbmock.go @@ -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() diff --git a/coderd/database/querier.go b/coderd/database/querier.go index ef684c940c..2654a8328b 100644 --- a/coderd/database/querier.go +++ b/coderd/database/querier.go @@ -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. diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go index b22a2fef0d..52ace4eb12 100644 --- a/coderd/database/queries.sql.go +++ b/coderd/database/queries.sql.go @@ -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, diff --git a/coderd/database/queries/chats.sql b/coderd/database/queries/chats.sql index 9c3ed9e65a..293d4918e6 100644 --- a/coderd/database/queries/chats.sql +++ b/coderd/database/queries/chats.sql @@ -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. diff --git a/coderd/exp_chats.go b/coderd/exp_chats.go index c66ecdc71c..cd2b15c9d3 100644 --- a/coderd/exp_chats.go +++ b/coderd/exp_chats.go @@ -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 { diff --git a/coderd/exp_chats_test.go b/coderd/exp_chats_test.go index e62da540fa..ad4da69df9 100644 --- a/coderd/exp_chats_test.go +++ b/coderd/exp_chats_test.go @@ -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, diff --git a/coderd/x/chatd/chatd.go b/coderd/x/chatd/chatd.go index 3678f8e565..4a24ba6c89 100644 --- a/coderd/x/chatd/chatd.go +++ b/coderd/x/chatd/chatd.go @@ -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{ diff --git a/coderd/x/chatd/chatd_test.go b/coderd/x/chatd/chatd_test.go index b0da8283d3..cfb9c51f44 100644 --- a/coderd/x/chatd/chatd_test.go +++ b/coderd/x/chatd/chatd_test.go @@ -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() diff --git a/coderd/x/chatd/chaterror/classify_test.go b/coderd/x/chatd/chaterror/classify_test.go index f4788b488e..d1cb3d9135 100644 --- a/coderd/x/chatd/chaterror/classify_test.go +++ b/coderd/x/chatd/chaterror/classify_test.go @@ -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
Forbidden"), + ), + 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"), diff --git a/coderd/x/chatd/chaterror/message.go b/coderd/x/chatd/chaterror/message.go index a64488a8c8..847651e339 100644 --- a/coderd/x/chatd/chaterror/message.go +++ b/coderd/x/chatd/chaterror/message.go @@ -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( diff --git a/coderd/x/chatd/chaterror/message_test.go b/coderd/x/chatd/chaterror/message_test.go index ba00b595fb..784d65da7c 100644 --- a/coderd/x/chatd/chaterror/message_test.go +++ b/coderd/x/chatd/chaterror/message_test.go @@ -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", diff --git a/coderd/x/chatd/subagent.go b/coderd/x/chatd/subagent.go index a9fccfcc17..68516041cc 100644 --- a/coderd/x/chatd/subagent.go +++ b/coderd/x/chatd/subagent.go @@ -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 diff --git a/coderd/x/chatd/usagelimit.go b/coderd/x/chatd/usagelimit.go deleted file mode 100644 index cbe67f50e1..0000000000 --- a/coderd/x/chatd/usagelimit.go +++ /dev/null @@ -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 - } -} diff --git a/coderd/x/chatd/usagelimit_internal_test.go b/coderd/x/chatd/usagelimit_internal_test.go deleted file mode 100644 index 0f0dba1461..0000000000 --- a/coderd/x/chatd/usagelimit_internal_test.go +++ /dev/null @@ -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) - } - }) - } -} diff --git a/codersdk/chats.go b/codersdk/chats.go index b8124e8f98..43e91ced4c 100644 --- a/codersdk/chats.go +++ b/codersdk/chats.go @@ -1968,7 +1968,6 @@ type ChatCostSummary struct { TotalRuntimeMs int64 `json:"total_runtime_ms"` ByModel []ChatCostModelBreakdown `json:"by_model"` ByChat []ChatCostChatBreakdown `json:"by_chat"` - UsageLimit *ChatUsageLimitStatus `json:"usage_limit,omitempty"` } // ChatCostModelBreakdown contains per-model cost aggregation. @@ -2037,17 +2036,6 @@ type ChatCostUsersResponse struct { Users []ChatCostUserRollup `json:"users"` } -// ChatUsageLimitExceededResponse is the 409 response body returned when a -// chat operation exceeds the caller's usage limit. The structured fields let -// frontends render user-friendly spend, limit, and reset information without -// parsing debug text. -type ChatUsageLimitExceededResponse struct { - Response - SpentMicros int64 `json:"spent_micros"` - LimitMicros int64 `json:"limit_micros"` - ResetsAt time.Time `json:"resets_at" format:"date-time"` -} - // ChatHookDispatchFailedResponse is the error body returned when a // lifecycle hook dispatch fails during a synchronous chat operation. // Kind lets clients classify the failure without parsing message text. @@ -2063,213 +2051,6 @@ type ChatHookDeniedResponse struct { Response Kind ChatErrorKind `json:"kind"` } - -type chatUsageLimitExceededError struct { - err *Error - response ChatUsageLimitExceededResponse -} - -func (e *chatUsageLimitExceededError) Error() string { - if e.err == nil { - return e.response.Message - } - return e.err.Error() -} - -func (e *chatUsageLimitExceededError) Unwrap() error { - return e.err -} - -func readBodyAsChatUsageLimitError(res *http.Response) error { - if res == nil || res.StatusCode != http.StatusConflict { - return ReadBodyAsError(res) - } - defer res.Body.Close() - - rawBody, err := io.ReadAll(res.Body) - if err != nil { - return xerrors.Errorf("read body: %w", err) - } - - if mimeErr := ExpectJSONMime(res); mimeErr != nil { - return readRawBodyAsError(res, rawBody) - } - - var payload ChatUsageLimitExceededResponse - if err := json.NewDecoder(bytes.NewReader(rawBody)).Decode(&payload); err == nil && isChatUsageLimitExceededResponse(payload) { - return &chatUsageLimitExceededError{ - err: newResponseError(res, payload.Response), - response: payload, - } - } - - return readRawBodyAsError(res, rawBody) -} - -func isChatUsageLimitExceededResponse(resp ChatUsageLimitExceededResponse) bool { - return resp.Message != "" && !resp.ResetsAt.IsZero() -} - -func readRawBodyAsError(res *http.Response, rawBody []byte) error { - if mimeErr := ExpectJSONMime(res); mimeErr != nil { - if len(rawBody) > 2048 { - rawBody = append(rawBody[:2048], []byte("...")...) - } - if len(rawBody) == 0 { - rawBody = []byte("no response body") - } - return newResponseError(res, Response{ - Message: mimeErr.Error(), - Detail: string(rawBody), - }) - } - - var response Response - if err := json.NewDecoder(bytes.NewReader(rawBody)).Decode(&response); err != nil { - if errors.Is(err, io.EOF) { - return newResponseError(res, Response{Message: "empty response body"}) - } - return xerrors.Errorf("decode body: %w", err) - } - if response.Message == "" { - if len(rawBody) > 1024 { - rawBody = append(rawBody[:1024], []byte("...")...) - } - response.Message = fmt.Sprintf( - "unexpected status code %d, response has no message", - res.StatusCode, - ) - response.Detail = string(rawBody) - } - return newResponseError(res, response) -} - -func newResponseError(res *http.Response, response Response) *Error { - if res == nil { - return &Error{Response: response} - } - - var requestMethod, requestURL string - if res.Request != nil { - requestMethod = res.Request.Method - if res.Request.URL != nil { - requestURL = res.Request.URL.String() - } - } - - var helpMessage string - if res.StatusCode == http.StatusUnauthorized { - helpMessage = "Try logging in using 'coder login'." - } - - return &Error{ - Response: response, - statusCode: res.StatusCode, - method: requestMethod, - url: requestURL, - Helper: helpMessage, - } -} - -// ChatUsageLimitExceededFrom extracts a structured chat usage limit response -// from an SDK error returned by chat mutation methods. -func ChatUsageLimitExceededFrom(err error) *ChatUsageLimitExceededResponse { - var limitErr *chatUsageLimitExceededError - if !errors.As(err, &limitErr) { - return nil - } - return &limitErr.response -} - -// ChatUsageLimitPeriod represents the time window for usage limits. -type ChatUsageLimitPeriod string - -const ( - ChatUsageLimitPeriodDay ChatUsageLimitPeriod = "day" - ChatUsageLimitPeriodWeek ChatUsageLimitPeriod = "week" - ChatUsageLimitPeriodMonth ChatUsageLimitPeriod = "month" -) - -// Valid reports whether p is a supported chat usage limit period. -func (p ChatUsageLimitPeriod) Valid() bool { - switch p { - case ChatUsageLimitPeriodDay, ChatUsageLimitPeriodWeek, ChatUsageLimitPeriodMonth: - return true - default: - return false - } -} - -// ChatUsageLimitConfig is the deployment-wide default usage limit config. -type ChatUsageLimitConfig struct { - // Nil in the API means no default limit is set. The DB stores 0 when - // limiting is disabled. - SpendLimitMicros *int64 `json:"spend_limit_micros"` - Period ChatUsageLimitPeriod `json:"period"` - UpdatedAt time.Time `json:"updated_at" format:"date-time"` -} - -// ChatUsageLimitOverride is a per-user override of the deployment default. -type ChatUsageLimitOverride struct { - UserID uuid.UUID `json:"user_id" format:"uuid"` - Username string `json:"username"` - Name string `json:"name"` - AvatarURL string `json:"avatar_url"` - // Nil in the API means no user override is set. Persisted override rows - // store positive values. - SpendLimitMicros *int64 `json:"spend_limit_micros"` -} - -// ChatUsageLimitGroupOverride represents a group-scoped spend limit override. -type ChatUsageLimitGroupOverride struct { - GroupID uuid.UUID `json:"group_id" format:"uuid"` - GroupName string `json:"group_name"` - GroupDisplayName string `json:"group_display_name"` - GroupAvatarURL string `json:"group_avatar_url"` - MemberCount int64 `json:"member_count"` - // Nil in the API means no group override is set. Persisted override rows - // store positive values. - SpendLimitMicros *int64 `json:"spend_limit_micros"` -} - -// UpsertChatUsageLimitOverrideRequest is the body for creating/updating a -// per-user usage limit override. -type UpsertChatUsageLimitOverrideRequest struct { - SpendLimitMicros int64 `json:"spend_limit_micros"` // Must be greater than 0. -} - -// UpdateChatUsageLimitOverrideRequest is kept as a compatibility alias. -type UpdateChatUsageLimitOverrideRequest = UpsertChatUsageLimitOverrideRequest - -// UpsertChatUsageLimitGroupOverrideRequest is the request to create or update -// a group-level spend limit override. -type UpsertChatUsageLimitGroupOverrideRequest struct { - SpendLimitMicros int64 `json:"spend_limit_micros"` // Must be greater than 0. -} - -// UpdateChatUsageLimitGroupOverrideRequest is kept as a compatibility alias. -type UpdateChatUsageLimitGroupOverrideRequest = UpsertChatUsageLimitGroupOverrideRequest - -// ChatUsageLimitStatus represents the current spend status for a user -// within their active limit period. -type ChatUsageLimitStatus struct { - IsLimited bool `json:"is_limited"` - Period ChatUsageLimitPeriod `json:"period,omitempty"` - SpendLimitMicros *int64 `json:"spend_limit_micros,omitempty"` - CurrentSpend int64 `json:"current_spend"` - PeriodStart time.Time `json:"period_start,omitempty" format:"date-time"` - PeriodEnd time.Time `json:"period_end,omitempty" format:"date-time"` -} - -// ChatUsageLimitConfigResponse is returned from the admin config endpoint -// and includes the config plus a count of models without pricing. -type ChatUsageLimitConfigResponse struct { - ChatUsageLimitConfig - UnpricedModelCount int64 `json:"unpriced_model_count"` - Overrides []ChatUsageLimitOverride `json:"overrides"` - GroupOverrides []ChatUsageLimitGroupOverride `json:"group_overrides"` -} - type ChatRole string const ( @@ -3079,7 +2860,7 @@ func (c *ExperimentalClient) CreateChat(ctx context.Context, req CreateChatReque return Chat{}, err } if res.StatusCode != http.StatusCreated { - return Chat{}, readBodyAsChatUsageLimitError(res) + return Chat{}, ReadBodyAsError(res) } defer res.Body.Close() var chat Chat @@ -3466,7 +3247,7 @@ func (c *ExperimentalClient) CreateChatMessage(ctx context.Context, chatID uuid. return CreateChatMessageResponse{}, err } if res.StatusCode != http.StatusOK { - return CreateChatMessageResponse{}, readBodyAsChatUsageLimitError(res) + return CreateChatMessageResponse{}, ReadBodyAsError(res) } defer res.Body.Close() var resp CreateChatMessageResponse @@ -3490,7 +3271,7 @@ func (c *ExperimentalClient) EditChatMessage( return EditChatMessageResponse{}, err } if res.StatusCode != http.StatusOK { - return EditChatMessageResponse{}, readBodyAsChatUsageLimitError(res) + return EditChatMessageResponse{}, ReadBodyAsError(res) } defer res.Body.Close() var resp EditChatMessageResponse @@ -3522,9 +3303,7 @@ func (c *ExperimentalClient) CompactChat(ctx context.Context, chatID uuid.UUID) } defer res.Body.Close() if res.StatusCode != http.StatusOK { - // Compaction runs LLM inference, so spend-limit rejections - // carry the structured usage-limit payload. - return Chat{}, readBodyAsChatUsageLimitError(res) + return Chat{}, ReadBodyAsError(res) } var chat Chat return chat, json.NewDecoder(res.Body).Decode(&chat) @@ -3555,7 +3334,7 @@ func (c *ExperimentalClient) RegenerateChatTitle(ctx context.Context, chatID uui } defer res.Body.Close() if res.StatusCode != http.StatusOK { - return Chat{}, readBodyAsChatUsageLimitError(res) + return Chat{}, ReadBodyAsError(res) } var chat Chat return chat, json.NewDecoder(res.Body).Decode(&chat) @@ -3574,7 +3353,7 @@ func (c *ExperimentalClient) ProposeChatTitle(ctx context.Context, chatID uuid.U } defer res.Body.Close() if res.StatusCode != http.StatusOK { - return ProposeChatTitleResponse{}, readBodyAsChatUsageLimitError(res) + return ProposeChatTitleResponse{}, ReadBodyAsError(res) } var resp ProposeChatTitleResponse return resp, json.NewDecoder(res.Body).Decode(&resp) @@ -3630,120 +3409,6 @@ func (c *ExperimentalClient) GetChatFile(ctx context.Context, fileID uuid.UUID) return data, res.Header.Get("Content-Type"), nil } -// GetChatUsageLimitConfig returns the deployment-wide chat usage limit config. -func (c *ExperimentalClient) GetChatUsageLimitConfig(ctx context.Context) (ChatUsageLimitConfigResponse, error) { - res, err := c.Request(ctx, http.MethodGet, "/api/experimental/chats/usage-limits", nil) - if err != nil { - return ChatUsageLimitConfigResponse{}, err - } - defer res.Body.Close() - if res.StatusCode != http.StatusOK { - return ChatUsageLimitConfigResponse{}, ReadBodyAsError(res) - } - var resp ChatUsageLimitConfigResponse - return resp, json.NewDecoder(res.Body).Decode(&resp) -} - -// UpdateChatUsageLimitConfig updates the deployment-wide usage limit config. -func (c *ExperimentalClient) UpdateChatUsageLimitConfig(ctx context.Context, req ChatUsageLimitConfig) (ChatUsageLimitConfig, error) { - res, err := c.Request(ctx, http.MethodPut, "/api/experimental/chats/usage-limits", req) - if err != nil { - return ChatUsageLimitConfig{}, err - } - defer res.Body.Close() - if res.StatusCode != http.StatusOK { - return ChatUsageLimitConfig{}, ReadBodyAsError(res) - } - var resp ChatUsageLimitConfig - return resp, json.NewDecoder(res.Body).Decode(&resp) -} - -// UpsertChatUsageLimitOverride creates or updates a per-user usage limit override. -func (c *ExperimentalClient) UpsertChatUsageLimitOverride(ctx context.Context, userID uuid.UUID, req UpsertChatUsageLimitOverrideRequest) (ChatUsageLimitOverride, error) { - res, err := c.Request(ctx, http.MethodPut, fmt.Sprintf("/api/experimental/chats/usage-limits/overrides/%s", userID), req) - if err != nil { - return ChatUsageLimitOverride{}, err - } - defer res.Body.Close() - if res.StatusCode != http.StatusOK { - return ChatUsageLimitOverride{}, ReadBodyAsError(res) - } - var resp ChatUsageLimitOverride - return resp, json.NewDecoder(res.Body).Decode(&resp) -} - -// UpdateChatUserUsageLimitOverride creates or updates a per-user usage limit override. -func (c *ExperimentalClient) UpdateChatUserUsageLimitOverride(ctx context.Context, userID uuid.UUID, req UpdateChatUsageLimitOverrideRequest) (ChatUsageLimitOverride, error) { - return c.UpsertChatUsageLimitOverride(ctx, userID, req) -} - -// DeleteChatUsageLimitOverride removes a per-user usage limit override. -func (c *ExperimentalClient) DeleteChatUsageLimitOverride(ctx context.Context, userID uuid.UUID) error { - res, err := c.Request(ctx, http.MethodDelete, fmt.Sprintf("/api/experimental/chats/usage-limits/overrides/%s", userID), nil) - if err != nil { - return err - } - defer res.Body.Close() - if res.StatusCode != http.StatusNoContent { - return ReadBodyAsError(res) - } - return nil -} - -// DeleteChatUserUsageLimitOverride removes a per-user usage limit override. -func (c *ExperimentalClient) DeleteChatUserUsageLimitOverride(ctx context.Context, userID uuid.UUID) error { - return c.DeleteChatUsageLimitOverride(ctx, userID) -} - -// UpsertChatUsageLimitGroupOverride creates or updates a group-level -// spend limit override. EXPERIMENTAL: This API is subject to change. -func (c *ExperimentalClient) UpsertChatUsageLimitGroupOverride(ctx context.Context, groupID uuid.UUID, req UpsertChatUsageLimitGroupOverrideRequest) (ChatUsageLimitGroupOverride, error) { - res, err := c.Request(ctx, http.MethodPut, - fmt.Sprintf("/api/experimental/chats/usage-limits/group-overrides/%s", groupID), - req, - ) - if err != nil { - return ChatUsageLimitGroupOverride{}, err - } - defer res.Body.Close() - if res.StatusCode != http.StatusOK { - return ChatUsageLimitGroupOverride{}, ReadBodyAsError(res) - } - var override ChatUsageLimitGroupOverride - return override, json.NewDecoder(res.Body).Decode(&override) -} - -// DeleteChatUsageLimitGroupOverride removes a group-level spend limit -// override. EXPERIMENTAL: This API is subject to change. -func (c *ExperimentalClient) DeleteChatUsageLimitGroupOverride(ctx context.Context, groupID uuid.UUID) error { - res, err := c.Request(ctx, http.MethodDelete, - fmt.Sprintf("/api/experimental/chats/usage-limits/group-overrides/%s", groupID), - nil, - ) - if err != nil { - return err - } - defer res.Body.Close() - if res.StatusCode != http.StatusNoContent { - return ReadBodyAsError(res) - } - return nil -} - -// GetMyChatUsageLimitStatus returns the current user's chat usage limit status. -func (c *ExperimentalClient) GetMyChatUsageLimitStatus(ctx context.Context) (ChatUsageLimitStatus, error) { - res, err := c.Request(ctx, http.MethodGet, "/api/experimental/chats/usage-limits/status", nil) - if err != nil { - return ChatUsageLimitStatus{}, err - } - defer res.Body.Close() - if res.StatusCode != http.StatusOK { - return ChatUsageLimitStatus{}, ReadBodyAsError(res) - } - var resp ChatUsageLimitStatus - return resp, json.NewDecoder(res.Body).Decode(&resp) -} - // SubmitToolResults submits the results of dynamic tool calls for a chat // that is in requires_action status. func (c *ExperimentalClient) SubmitToolResults(ctx context.Context, chatID uuid.UUID, req SubmitToolResultsRequest) error { diff --git a/codersdk/chats_test.go b/codersdk/chats_test.go index 461467eef3..fb95b07a06 100644 --- a/codersdk/chats_test.go +++ b/codersdk/chats_test.go @@ -3,9 +3,6 @@ package codersdk_test import ( "context" "encoding/json" - "net/http" - "net/http/httptest" - "net/url" "reflect" "strings" "testing" @@ -58,81 +55,6 @@ func TestChatModelProviderOptions_UnmarshalJSON_ParsesPlainProviderPayloads(t *t require.Equal(t, "summarized", *decoded.Anthropic.ThinkingDisplay) } -func TestChatUsageLimitExceededFrom(t *testing.T) { - t.Parallel() - - t.Run("ExtractsTyped409", func(t *testing.T) { - t.Parallel() - - want := codersdk.ChatUsageLimitExceededResponse{ - Response: codersdk.Response{Message: "Chat usage limit exceeded."}, - SpentMicros: 123, - LimitMicros: 456, - ResetsAt: time.Date(2026, time.March, 16, 12, 0, 0, 0, time.UTC), - } - - srv := httptest.NewServer(http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) { - require.Equal(t, http.MethodPost, r.Method) - require.Equal(t, "/api/experimental/chats", r.URL.Path) - rw.Header().Set("Content-Type", "application/json") - rw.WriteHeader(http.StatusConflict) - require.NoError(t, json.NewEncoder(rw).Encode(want)) - })) - defer srv.Close() - - serverURL, err := url.Parse(srv.URL) - require.NoError(t, err) - - client := codersdk.NewExperimentalClient(codersdk.New(serverURL)) - _, err = client.CreateChat(context.Background(), codersdk.CreateChatRequest{ - Content: []codersdk.ChatInputPart{{ - Type: codersdk.ChatInputPartTypeText, - Text: "hello", - }}, - }) - require.Error(t, err) - - sdkErr, ok := codersdk.AsError(err) - require.True(t, ok) - require.Equal(t, http.StatusConflict, sdkErr.StatusCode()) - require.Equal(t, want.Message, sdkErr.Message) - - limitErr := codersdk.ChatUsageLimitExceededFrom(err) - require.NotNil(t, limitErr) - require.Equal(t, want, *limitErr) - }) - - t.Run("ReturnsNilForNonLimitErrors", func(t *testing.T) { - t.Parallel() - - require.Nil(t, codersdk.ChatUsageLimitExceededFrom(codersdk.NewError(http.StatusConflict, codersdk.Response{Message: "plain conflict"}))) - - srv := httptest.NewServer(http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) { - rw.Header().Set("Content-Type", "application/json") - rw.WriteHeader(http.StatusBadRequest) - require.NoError(t, json.NewEncoder(rw).Encode(codersdk.Response{Message: "Invalid request."})) - })) - defer srv.Close() - - serverURL, err := url.Parse(srv.URL) - require.NoError(t, err) - - client := codersdk.NewExperimentalClient(codersdk.New(serverURL)) - _, err = client.CreateChat(context.Background(), codersdk.CreateChatRequest{ - Content: []codersdk.ChatInputPart{{ - Type: codersdk.ChatInputPartTypeText, - Text: "hello", - }}, - }) - require.Error(t, err) - - sdkErr, ok := codersdk.AsError(err) - require.True(t, ok) - require.Equal(t, http.StatusBadRequest, sdkErr.StatusCode()) - require.Nil(t, codersdk.ChatUsageLimitExceededFrom(err)) - }) -} - func TestChatErrorKind_JSONRoundTrip(t *testing.T) { t.Parallel() diff --git a/docs/ai-coder/agents/platform-controls/advisor.md b/docs/ai-coder/agents/platform-controls/advisor.md index 7bff121530..7854db884d 100644 --- a/docs/ai-coder/agents/platform-controls/advisor.md +++ b/docs/ai-coder/agents/platform-controls/advisor.md @@ -39,9 +39,8 @@ Once the experiment is enabled, configure the advisor under **AI Settings** > | Advisor model | Use chat model | Optional dedicated chat model config for the advisor. When unset, the advisor reuses the root agent's model. | | Reasoning effort | Model default | Overrides the selected advisor model's reasoning effort. Available only when the model supports selectable effort. | -The advisor is not available in plan mode or to subagents. Failed advisor -invocations refund the per-turn budget, and advisor calls are not metered -against the root chat's usage limit. +The advisor is not available in plan mode or to subagents. +Failed advisor invocations refund the per-turn budget. The same configuration is available at: diff --git a/docs/ai-coder/agents/platform-controls/index.md b/docs/ai-coder/agents/platform-controls/index.md index d8b7469c47..da502c63f8 100644 --- a/docs/ai-coder/agents/platform-controls/index.md +++ b/docs/ai-coder/agents/platform-controls/index.md @@ -116,12 +116,15 @@ none, if the template does not define any). ### Spend management -Administrators can set spend limits to cap LLM usage per user within a rolling -time period, with per-user and per-group overrides. The cost tracking dashboard -provides visibility into per-user spending, token consumption, and per-model -breakdowns. +AI Gateway budgets cap each user's AI spend, including Coder Agents chats, over a monthly period. +Coder sets budgets per group, and the deployment policy selects the group with the largest spend limit when a user belongs to several budgeted groups. +A per-user override takes priority over all group budgets. -See [Spend Management](./usage-insights.md) for details. +Budgets are the only spend cap for Coder Agents chats. +Chats no longer enforce a separate limit of their own, and existing native limit values are not migrated to budgets. +Budget controls in the Coder UI, the group budget endpoints, and the AI spend status endpoints all require a license that includes AI Gateway. + +Refer to [Spend Management](./usage-insights.md) for details. ### Git providers @@ -158,10 +161,8 @@ For chat debug logging (not experiment-gated), see [Chat debug logging](./chat-d ## Where we are headed -The controls above cover providers, models, system prompts, templates, MCP -servers, usage limits, and data retention. We are continuing to invest in platform controls -based on what we hear from customers deploying agents in regulated and -enterprise environments. +The controls above cover providers, models, system prompts, templates, MCP servers, AI Gateway budgets, and data retention. +We are continuing to invest in platform controls based on what feedback we get from customers deploying agents in regulated and enterprise environments. ### Infrastructure-level enforcement diff --git a/docs/ai-coder/agents/platform-controls/usage-insights.md b/docs/ai-coder/agents/platform-controls/usage-insights.md index b6b2d1e5db..90292f1638 100644 --- a/docs/ai-coder/agents/platform-controls/usage-insights.md +++ b/docs/ai-coder/agents/platform-controls/usage-insights.md @@ -1,90 +1,55 @@ # Spend Management -Coder provides admin-only controls for monitoring and controlling agent -spend: usage limits and cost tracking. +Coder controls agent spend with AI Gateway budgets, and surfaces the resulting spend to both admins and users. -## Usage limits +## Budgets -Navigate to **Agents** > **Settings** > **Manage Agents** > **Spend**. +Coder Agents spend is controlled by AI Gateway budgets, which cap all AI Gateway usage (including Coder Agents chats) per user over the budget period. -Usage limits cap how much each user can spend on LLM usage within a rolling -time period. When enabled, the system checks the user's current spend before -processing each chat message. +- **Group budgets**: set a budget for a group from the group's settings page. + The deployment budget policy resolves which group budget applies when a user belongs to multiple budgeted groups, which defaults to the group with the highest budget. +- **Per-user overrides**: set a custom budget for an individual user, attributed to one of their groups. + Per-user overrides take priority over group budgets. -### Configuration +The deployment flags `--ai-budget-policy` and `--ai-budget-period` currently +support only `highest` and `month`. `highest` selects the group with the +largest spend limit, and `month` resets spend at the start of each UTC +calendar month. A user who belongs to no budgeted group falls back to the +Everyone group, which has no budget unless one is set for it. There is no +deployment-wide budget amount, and a configured spend limit cannot exceed +$1,000,000 per member per period. -- **Enable/disable toggle** — master on/off for the entire limit system. -- **Period** — `day`, `week`, or `month`. Periods are UTC-aligned: midnight - UTC for daily, Monday start for weekly, first of the month for monthly. -- **Default limit** — deployment-wide default in dollars. Applies to all - users who do not have a more specific override. Leave unset for no limit. -- **Per-user overrides** — set a custom dollar limit for an individual user. - Takes highest priority. -- **Per-group overrides** — set a limit for a group. When a user belongs to - multiple groups, the lowest group limit applies. +> [!IMPORTANT] +> Budget controls in the Coder UI, the group budget endpoints (`/api/v2/groups/{group}/ai/budget`), and the AI spend status and reporting endpoints all require the AI Gateway entitlement. +> No experiment is needed. +> +> Native chat usage limits are removed from the application. +> Existing native limit values are not migrated to AI Gateway budgets and are no longer enforced. +> Configure AI Gateway budgets separately. -### Priority hierarchy - -The system resolves a user's effective limit in this order: - -1. Individual user override (highest priority) -1. Minimum group limit across all of the user's groups -1. Global default limit -1. No limit (if limits are disabled or no value is configured) +The API reference documents how to [get](../../../reference/api/enterprise.md#get-group-ai-budget), [upsert](../../../reference/api/enterprise.md#upsert-group-ai-budget), and [delete](../../../reference/api/enterprise.md#delete-group-ai-budget) a group budget. ### Enforcement -- Checked before each chat message is processed. -- When current spend meets or exceeds the limit, the chat returns a - **409 Conflict** response and the message is blocked. -- Fail-open: if the limit query itself fails, the message is allowed - through. -- Brief overage is possible when concurrent messages are in flight, because - cost is determined only after the LLM returns. +- The AI Gateway checks the user's current spend before forwarding each request. + When spend meets or exceeds the budget, the request is rejected and the chat shows a terminal error explaining that the budget was exceeded. +- Brief overage is possible when concurrent requests are in flight, because cost is recorded after the LLM responds. ### User-facing status -Users can view their own spend status, including whether a limit is active, -their effective limit, current spend, and when the current period resets. +The usage indicator on the Agents page and the summary in the user menu both show the signed-in user's current AI spend, their budget, and the period reset date. +Both appear only when the deployment has the AI Gateway entitlement. + +## Spend visibility + +Spend is shown where it is actionable: + +- **Agents page and user menu**: the signed-in user's spend against their budget, as described previously. +- **Group settings**: each member's spend against the group's budget, for admins who can manage the group. +- **Chat summary panel**: the cost of one chat tree, on a chat's Summary tab. + A subagent reports the total for its whole tree, including the chat that started it. +- **Agents** > **Settings** > **Manage Agents** > **Spend**: deployment-wide chat cost per user, with per-user drill-down. > [!NOTE] -> The admin configuration page shows the count of models without pricing -> data. Models missing pricing cannot be tracked accurately against limits. - -## Cost tracking - -Navigate to **Agents** > **Settings** > **Manage Agents** > **Spend**. - -This view shows deployment-wide LLM chat costs with per-user drill-down. - -### Top-level view - -A per-user rollup table with the following columns: - -| Column | Description | -|--------------------|-------------------------------------| -| Total cost | Aggregate dollar spend for the user | -| Messages | Number of chat messages sent | -| Chats | Number of distinct chat sessions | -| Input tokens | Total input tokens consumed | -| Output tokens | Total output tokens consumed | -| Cache read tokens | Tokens served from cache | -| Cache write tokens | Tokens written to cache | - -The table supports date range filtering (default: last 30 days), search by -name or username, and pagination. - -### Per-user detail view - -Select a user to see: - -- **Summary cards** — total cost, token breakdowns, and message counts. -- **Usage limit progress** — if a limit is active, a color-coded progress - bar shows current spend relative to the limit. -- **Per-model breakdown** — table of costs and token usage by model. -- **Per-chat breakdown** — table of costs and token usage by chat session. - -> [!NOTE] -> Automatic title generation uses lightweight models, such as Claude Haiku or GPT-4o -> Mini. Its token usage is not counted towards usage limits or shown in usage -> summaries. +> Per-chat cost comes from AI Gateway records, which are pruned according to `--ai-gateway-retention` (60 days by default). +> A chat for which gateway records have been pruned reports no cost. diff --git a/docs/manifest.json b/docs/manifest.json index 81a1994bca..3c6e4fd84a 100644 --- a/docs/manifest.json +++ b/docs/manifest.json @@ -1085,7 +1085,7 @@ }, { "title": "Spend Management", - "description": "Set spend limits and track AI costs for Coder Agents with admin-only controls.", + "description": "Cap Coder Agents spend with AI Gateway budgets and track the resulting spend.", "path": "./ai-coder/agents/platform-controls/usage-insights.md", "state": ["beta"] }, diff --git a/enterprise/coderd/x/chatd/usagelimit_test.go b/enterprise/coderd/x/chatd/usagelimit_test.go deleted file mode 100644 index 9f44bfa07c..0000000000 --- a/enterprise/coderd/x/chatd/usagelimit_test.go +++ /dev/null @@ -1,324 +0,0 @@ -package chatd_test - -import ( - "database/sql" - "encoding/json" - "testing" - "time" - - "github.com/google/uuid" - "github.com/sqlc-dev/pqtype" - "github.com/stretchr/testify/require" - - "github.com/coder/coder/v2/coderd/database" - "github.com/coder/coder/v2/coderd/database/dbgen" - "github.com/coder/coder/v2/coderd/database/dbtestutil" - "github.com/coder/coder/v2/coderd/x/chatd" - "github.com/coder/coder/v2/codersdk" - "github.com/coder/coder/v2/testutil" -) - -func TestResolveUsageLimitStatus_OrgScoped(t *testing.T) { - t.Parallel() - - db, _ := dbtestutil.NewDB(t) - ctx := testutil.Context(t, testutil.WaitLong) - - // Create two orgs and a user in both. - orgA := dbgen.Organization(t, db, database.Organization{}) - orgB := dbgen.Organization(t, db, database.Organization{}) - user := dbgen.User(t, db, database.User{}) - dbgen.OrganizationMember(t, db, database.OrganizationMember{ - UserID: user.ID, - OrganizationID: orgA.ID, - }) - dbgen.OrganizationMember(t, db, database.OrganizationMember{ - UserID: user.ID, - OrganizationID: orgB.ID, - }) - - // Create groups with different spend limits. - // groupA ($5) and groupA2 ($20) are both in orgA to exercise - // MIN aggregation within a single org. - groupA := dbgen.Group(t, db, database.Group{ - OrganizationID: orgA.ID, - }) - groupA2 := dbgen.Group(t, db, database.Group{ - OrganizationID: orgA.ID, - }) - groupB := dbgen.Group(t, db, database.Group{ - OrganizationID: orgB.ID, - }) - dbgen.GroupMember(t, db, database.GroupMemberTable{ - UserID: user.ID, - GroupID: groupA.ID, - }) - dbgen.GroupMember(t, db, database.GroupMemberTable{ - UserID: user.ID, - GroupID: groupA2.ID, - }) - dbgen.GroupMember(t, db, database.GroupMemberTable{ - UserID: user.ID, - GroupID: groupB.ID, - }) - - // Set group spend limits: groupA=$5, groupA2=$20, groupB=$50. - _, err := db.UpsertChatUsageLimitGroupOverride(ctx, database.UpsertChatUsageLimitGroupOverrideParams{ - GroupID: groupA.ID, - SpendLimitMicros: 5_000_000, - }) - require.NoError(t, err) - _, err = db.UpsertChatUsageLimitGroupOverride(ctx, database.UpsertChatUsageLimitGroupOverrideParams{ - GroupID: groupA2.ID, - SpendLimitMicros: 20_000_000, - }) - require.NoError(t, err) - _, err = db.UpsertChatUsageLimitGroupOverride(ctx, database.UpsertChatUsageLimitGroupOverrideParams{ - GroupID: groupB.ID, - SpendLimitMicros: 50_000_000, - }) - require.NoError(t, err) - - // Enable usage limits with a high default so group limits win. - _, err = db.UpsertChatUsageLimitConfig(ctx, database.UpsertChatUsageLimitConfigParams{ - Enabled: true, - DefaultLimitMicros: 100_000_000, - Period: string(codersdk.ChatUsageLimitPeriodMonth), - }) - require.NoError(t, err) - - // We need a chat provider + model config for inserting chats. - _ = dbgen.ChatProvider(t, db, database.ChatProvider{ - CreatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, - }) - modelConfig := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{ - CreatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, - UpdatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, - IsDefault: true, - }) - - now := time.Now().UTC() - - // insertChatWithSpend is a test helper that creates a chat in the - // given org and inserts a single message with the specified cost. - insertChatWithSpend := func(t *testing.T, ownerID, orgID, modelCfgID uuid.UUID, costMicros int64) { - t.Helper() - c := dbgen.Chat(t, db, database.Chat{ - OrganizationID: orgID, - OwnerID: ownerID, - LastModelConfigID: modelCfgID, - Title: "test chat", - }) - _ = dbgen.ChatMessage(t, db, database.ChatMessage{ - ChatID: c.ID, - ModelConfigID: uuid.NullUUID{UUID: modelCfgID, Valid: true}, - Role: database.ChatMessageRoleAssistant, - Content: pqtype.NullRawMessage{RawMessage: json.RawMessage(`[{"type":"text","text":"hello"}]`), Valid: true}, - InputTokens: sql.NullInt64{Int64: 100, Valid: true}, - OutputTokens: sql.NullInt64{Int64: 50, Valid: true}, - TotalTokens: sql.NullInt64{Int64: 150, Valid: true}, - ContextLimit: sql.NullInt64{Int64: 128000, Valid: true}, - TotalCostMicros: sql.NullInt64{Int64: costMicros, Valid: true}, - RuntimeMs: sql.NullInt64{Int64: 500, Valid: true}, - ProviderResponseID: sql.NullString{String: uuid.NewString(), Valid: true}, - }) - } - - t.Run("OrgA_gets_orgA_limit", func(t *testing.T) { - t.Parallel() - ctx := testutil.Context(t, testutil.WaitLong) - // orgA has groupA ($5) and groupA2 ($20). MIN($5, $20) = $5. - status, err := chatd.ResolveUsageLimitStatus(ctx, db, user.ID, uuid.NullUUID{UUID: orgA.ID, Valid: true}, now) - require.NoError(t, err) - require.NotNil(t, status) - require.NotNil(t, status.SpendLimitMicros) - require.Equal(t, int64(5_000_000), *status.SpendLimitMicros, - "orgA should resolve to MIN of both groups ($5, $20) = $5") - }) - - t.Run("OrgB_gets_orgB_limit", func(t *testing.T) { - t.Parallel() - ctx := testutil.Context(t, testutil.WaitLong) - status, err := chatd.ResolveUsageLimitStatus(ctx, db, user.ID, uuid.NullUUID{UUID: orgB.ID, Valid: true}, now) - require.NoError(t, err) - require.NotNil(t, status) - require.NotNil(t, status.SpendLimitMicros) - require.Equal(t, int64(50_000_000), *status.SpendLimitMicros, - "orgB should resolve to groupB's $50 limit, not global MIN") - }) - - t.Run("UnknownOrg_gets_global_default", func(t *testing.T) { - t.Parallel() - ctx := testutil.Context(t, testutil.WaitLong) - // When the org ID does not match any group the user belongs - // to, MIN() over an empty set returns NULL, the CASE sees - // gl.limit_micros IS NOT NULL as false, and falls through - // to the global default. This subtest guards that contract: - // if someone changes the NULL-handling in - // ResolveUserChatSpendLimit, this will catch it. - randomOrg := uuid.NullUUID{UUID: uuid.New(), Valid: true} - status, err := chatd.ResolveUsageLimitStatus(ctx, db, user.ID, randomOrg, now) - require.NoError(t, err) - require.NotNil(t, status) - require.NotNil(t, status.SpendLimitMicros) - require.Equal(t, int64(100_000_000), *status.SpendLimitMicros, - "org with no matching groups should fall through to global default ($100)") - }) - - t.Run("NilOrg_gets_global_min", func(t *testing.T) { - t.Parallel() - ctx := testutil.Context(t, testutil.WaitLong) - // NULL org = global behavior: MIN across all groups. - status, err := chatd.ResolveUsageLimitStatus(ctx, db, user.ID, uuid.NullUUID{}, now) - require.NoError(t, err) - require.NotNil(t, status) - require.NotNil(t, status.SpendLimitMicros) - require.Equal(t, int64(5_000_000), *status.SpendLimitMicros, - "nil org should fall back to global MIN($5, $20, $50) = $5") - }) - - t.Run("Spend_scoped_to_org", func(t *testing.T) { - t.Parallel() - ctx := testutil.Context(t, testutil.WaitLong) - - // Dedicated user so spend insertion doesn't affect sibling subtests. - spendUser := dbgen.User(t, db, database.User{}) - dbgen.OrganizationMember(t, db, database.OrganizationMember{ - UserID: spendUser.ID, - OrganizationID: orgA.ID, - }) - dbgen.OrganizationMember(t, db, database.OrganizationMember{ - UserID: spendUser.ID, - OrganizationID: orgB.ID, - }) - dbgen.GroupMember(t, db, database.GroupMemberTable{ - UserID: spendUser.ID, - GroupID: groupA.ID, - }) - dbgen.GroupMember(t, db, database.GroupMemberTable{ - UserID: spendUser.ID, - GroupID: groupB.ID, - }) - - insertChatWithSpend(t, spendUser.ID, orgA.ID, modelConfig.ID, 3_000_000) - - // Resolve for orgB: should see zero spend (orgA's $3 not counted). - statusB, err := chatd.ResolveUsageLimitStatus(ctx, db, spendUser.ID, uuid.NullUUID{UUID: orgB.ID, Valid: true}, now) - require.NoError(t, err) - require.NotNil(t, statusB) - require.Equal(t, int64(0), statusB.CurrentSpend, - "orgB should not include orgA's spend") - - // Resolve for orgA: should see $3 spend. - statusA, err := chatd.ResolveUsageLimitStatus(ctx, db, spendUser.ID, uuid.NullUUID{UUID: orgA.ID, Valid: true}, now) - require.NoError(t, err) - require.NotNil(t, statusA) - require.Equal(t, int64(3_000_000), statusA.CurrentSpend, - "orgA should include its own spend") - - // Nil org: should see $3 (global). - statusNil, err := chatd.ResolveUsageLimitStatus(ctx, db, spendUser.ID, uuid.NullUUID{}, now) - require.NoError(t, err) - require.NotNil(t, statusNil) - require.Equal(t, int64(3_000_000), statusNil.CurrentSpend, - "nil org should include all spend globally") - }) - - t.Run("User_override_beats_group", func(t *testing.T) { - t.Parallel() - ctx := testutil.Context(t, testutil.WaitLong) - // Create a separate user with a personal override. - user2 := dbgen.User(t, db, database.User{}) - dbgen.OrganizationMember(t, db, database.OrganizationMember{ - UserID: user2.ID, - OrganizationID: orgA.ID, - }) - dbgen.GroupMember(t, db, database.GroupMemberTable{ - UserID: user2.ID, - GroupID: groupA.ID, - }) - - // Set $10 user override (beats groupA's $5 limit). - _, err := db.UpsertChatUsageLimitUserOverride(ctx, database.UpsertChatUsageLimitUserOverrideParams{ - UserID: user2.ID, - SpendLimitMicros: 10_000_000, - }) - require.NoError(t, err) - - status, err := chatd.ResolveUsageLimitStatus(ctx, db, user2.ID, uuid.NullUUID{UUID: orgA.ID, Valid: true}, now) - require.NoError(t, err) - require.NotNil(t, status) - require.NotNil(t, status.SpendLimitMicros) - require.Equal(t, int64(10_000_000), *status.SpendLimitMicros, - "user override should take priority over group limit") - }) - - t.Run("UserOverride_spend_is_global", func(t *testing.T) { - t.Parallel() - // When user override wins, spend should be checked globally, - // not per-org. Otherwise a user in N orgs can spend limit*N. - user3 := dbgen.User(t, db, database.User{}) - dbgen.OrganizationMember(t, db, database.OrganizationMember{ - UserID: user3.ID, - OrganizationID: orgA.ID, - }) - dbgen.OrganizationMember(t, db, database.OrganizationMember{ - UserID: user3.ID, - OrganizationID: orgB.ID, - }) - - // Set $10 user override. - _, err := db.UpsertChatUsageLimitUserOverride(testutil.Context(t, testutil.WaitLong), database.UpsertChatUsageLimitUserOverrideParams{ - UserID: user3.ID, - SpendLimitMicros: 10_000_000, - }) - require.NoError(t, err) - - // $6 in orgA + $6 in orgB = $12 total. - insertChatWithSpend(t, user3.ID, orgA.ID, modelConfig.ID, 6_000_000) - insertChatWithSpend(t, user3.ID, orgB.ID, modelConfig.ID, 6_000_000) - - ctx := testutil.Context(t, testutil.WaitLong) - status, err := chatd.ResolveUsageLimitStatus(ctx, db, user3.ID, uuid.NullUUID{UUID: orgA.ID, Valid: true}, now) - require.NoError(t, err) - require.NotNil(t, status) - require.NotNil(t, status.SpendLimitMicros) - require.Equal(t, int64(10_000_000), *status.SpendLimitMicros) - // Spend should be global ($12), not org-scoped ($6). - require.Equal(t, int64(12_000_000), status.CurrentSpend, - "user override should check global spend to prevent cross-org evasion") - }) - - t.Run("GlobalDefault_spend_is_global", func(t *testing.T) { - t.Parallel() - // When global default wins (no groups in the target org, - // no user override), spend should also be checked globally. - user4 := dbgen.User(t, db, database.User{}) - orgC := dbgen.Organization(t, db, database.Organization{}) - dbgen.OrganizationMember(t, db, database.OrganizationMember{ - UserID: user4.ID, - OrganizationID: orgA.ID, - }) - dbgen.OrganizationMember(t, db, database.OrganizationMember{ - UserID: user4.ID, - OrganizationID: orgC.ID, - }) - - // $30 in orgA + $40 in orgC = $70 total. - insertChatWithSpend(t, user4.ID, orgA.ID, modelConfig.ID, 30_000_000) - insertChatWithSpend(t, user4.ID, orgC.ID, modelConfig.ID, 40_000_000) - - ctx := testutil.Context(t, testutil.WaitLong) - // user4 has no groups in orgC, no override: falls through - // to global default ($100). - status, err := chatd.ResolveUsageLimitStatus(ctx, db, user4.ID, uuid.NullUUID{UUID: orgC.ID, Valid: true}, now) - require.NoError(t, err) - require.NotNil(t, status) - require.NotNil(t, status.SpendLimitMicros) - require.Equal(t, int64(100_000_000), *status.SpendLimitMicros, - "should fall through to global default ($100)") - // Spend should be global ($70), not org-scoped ($40). - require.Equal(t, int64(70_000_000), status.CurrentSpend, - "global default should check global spend") - }) -} diff --git a/site/src/api/api.ts b/site/src/api/api.ts index c7de79f98a..3683399dcb 100644 --- a/site/src/api/api.ts +++ b/site/src/api/api.ts @@ -3966,71 +3966,6 @@ class ExperimentalApiMethods { const response = await this.axios.get- {usageLimitPeriodLabel} spend limit -
- {usageLimitCurrentPeriod && ( -- {usageLimitCurrentPeriod} -
- )} -- {formatCostMicros(usageLimitCurrentSpend)} /{" "} - {formatCostMicros(usageLimitSpendMicros)} -
-- {Math.round(usageProgressPercentage)}% used -
-- {usageLimitStatusText} -
-- Resets {usageLimitResetAt} -
-
{emptyMessage}
diff --git a/site/src/pages/AISettingsPage/SpendPage/components/LimitsTab/DefaultLimitController.tsx b/site/src/pages/AISettingsPage/SpendPage/components/LimitsTab/DefaultLimitController.tsx
deleted file mode 100644
index 855a25e0da..0000000000
--- a/site/src/pages/AISettingsPage/SpendPage/components/LimitsTab/DefaultLimitController.tsx
+++ /dev/null
@@ -1,65 +0,0 @@
-import { type FC, type ReactNode, useState } from "react";
-
-import type { ChatUsageLimitPeriod } from "#/api/typesGenerated";
-import { isPositiveFiniteDollarAmount } from "#/utils/currency";
-
-export interface DefaultLimitFormValues {
- enabled: boolean;
- period: ChatUsageLimitPeriod;
- amountDollars: string;
-}
-
-interface DefaultLimitControllerProps {
- initialValues: DefaultLimitFormValues;
- onSave: (values: DefaultLimitFormValues) => void;
- children: (props: {
- enabled: boolean;
- onEnabledChange: (enabled: boolean) => void;
- period: ChatUsageLimitPeriod;
- onPeriodChange: (period: ChatUsageLimitPeriod) => void;
- amountDollars: string;
- onAmountDollarsChange: (amount: string) => void;
- isAmountValid: boolean;
- isDirty: boolean;
- saveDefault: () => void;
- }) => ReactNode;
-}
-
-export const DefaultLimitController: FC
- Set a deployment-wide spend cap that applies to all users by
- default. When disabled, users have unlimited spending.
-
- {getErrorMessage(deleteError, "Failed to delete group override.")}
-
- {getErrorMessage(groupsError, "Failed to load groups.")}
-
- {getErrorMessage(deleteError, "Failed to delete override.")}
-
- Spend limit
-
-
-
-
-
-