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\nForbidden"), + ), + 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(url); return response.data; }; - - getChatUsageLimitConfig = - async (): Promise => { - const response = - await this.axios.get( - "/api/experimental/chats/usage-limits", - ); - return response.data; - }; - - getChatUsageLimitStatus = - async (): Promise => { - const response = await this.axios.get( - "/api/experimental/chats/usage-limits/status", - ); - return response.data; - }; - - updateChatUsageLimitConfig = async ( - req: TypesGen.ChatUsageLimitConfig, - ): Promise => { - const response = await this.axios.put( - "/api/experimental/chats/usage-limits", - req, - ); - return response.data; - }; - - upsertChatUsageLimitOverride = async ( - userID: string, - req: TypesGen.UpsertChatUsageLimitOverrideRequest, - ): Promise => { - const response = await this.axios.put( - `/api/experimental/chats/usage-limits/overrides/${encodeURIComponent(userID)}`, - req, - ); - return response.data; - }; - - deleteChatUsageLimitOverride = async (userID: string): Promise => { - const response = await this.axios.delete( - `/api/experimental/chats/usage-limits/overrides/${encodeURIComponent(userID)}`, - ); - return response.data; - }; - - upsertChatUsageLimitGroupOverride = async ( - groupID: string, - req: TypesGen.UpsertChatUsageLimitGroupOverrideRequest, - ): Promise => { - const response = await this.axios.put( - `/api/experimental/chats/usage-limits/group-overrides/${encodeURIComponent(groupID)}`, - req, - ); - return response.data; - }; - - deleteChatUsageLimitGroupOverride = async ( - groupID: string, - ): Promise => { - const response = await this.axios.delete( - `/api/experimental/chats/usage-limits/group-overrides/${encodeURIComponent(groupID)}`, - ); - return response.data; - }; } // This is a hard coded CSRF token/cookie pair for local development. In prod, diff --git a/site/src/api/queries/chats.ts b/site/src/api/queries/chats.ts index 6d33afcfb1..0d69200e63 100644 --- a/site/src/api/queries/chats.ts +++ b/site/src/api/queries/chats.ts @@ -2005,91 +2005,6 @@ export function paginatedChatCostUsers( }; } -export const chatUsageLimitStatusKey = [ - ...chatsKey, - "usageLimitStatus", -] as const; - -export const chatUsageLimitStatus = () => ({ - queryKey: chatUsageLimitStatusKey, - queryFn: () => API.experimental.getChatUsageLimitStatus(), - refetchInterval: 60_000, -}); - -const chatUsageLimitConfigKey = [...chatsKey, "usageLimitConfig"] as const; - -export const chatUsageLimitConfig = () => ({ - queryKey: chatUsageLimitConfigKey, - queryFn: () => API.experimental.getChatUsageLimitConfig(), -}); - -export const updateChatUsageLimitConfig = (queryClient: QueryClient) => ({ - mutationFn: (req: TypesGen.ChatUsageLimitConfig) => - API.experimental.updateChatUsageLimitConfig(req), - onSuccess: async () => { - await queryClient.invalidateQueries({ - queryKey: chatUsageLimitConfigKey, - }); - }, -}); - -type UpsertChatUsageLimitOverrideMutationArgs = { - userID: string; - req: TypesGen.UpsertChatUsageLimitOverrideRequest; -}; - -export const upsertChatUsageLimitOverride = (queryClient: QueryClient) => ({ - mutationFn: ({ userID, req }: UpsertChatUsageLimitOverrideMutationArgs) => - API.experimental.upsertChatUsageLimitOverride(userID, req), - onSuccess: async () => { - await queryClient.invalidateQueries({ - queryKey: chatUsageLimitConfigKey, - }); - }, -}); - -export const deleteChatUsageLimitOverride = (queryClient: QueryClient) => ({ - mutationFn: (userID: string) => - API.experimental.deleteChatUsageLimitOverride(userID), - onSuccess: async () => { - await queryClient.invalidateQueries({ - queryKey: chatUsageLimitConfigKey, - }); - }, -}); - -type UpsertChatUsageLimitGroupOverrideMutationArgs = { - groupID: string; - req: TypesGen.UpsertChatUsageLimitGroupOverrideRequest; -}; - -export const upsertChatUsageLimitGroupOverride = ( - queryClient: QueryClient, -) => ({ - mutationFn: ({ - groupID, - req, - }: UpsertChatUsageLimitGroupOverrideMutationArgs) => - API.experimental.upsertChatUsageLimitGroupOverride(groupID, req), - onSuccess: async () => { - await queryClient.invalidateQueries({ - queryKey: chatUsageLimitConfigKey, - }); - }, -}); - -export const deleteChatUsageLimitGroupOverride = ( - queryClient: QueryClient, -) => ({ - mutationFn: (groupID: string) => - API.experimental.deleteChatUsageLimitGroupOverride(groupID), - onSuccess: async () => { - await queryClient.invalidateQueries({ - queryKey: chatUsageLimitConfigKey, - }); - }, -}); - // ── MCP Server Configs ─────────────────────────────────────── export const mcpServerConfigsKey = ["mcp-server-configs"] as const; diff --git a/site/src/api/typesGenerated.ts b/site/src/api/typesGenerated.ts index 7c528b5785..af44c9dc73 100644 --- a/site/src/api/typesGenerated.ts +++ b/site/src/api/typesGenerated.ts @@ -2253,7 +2253,6 @@ export interface ChatCostSummary { readonly total_runtime_ms: number; readonly by_model: readonly ChatCostModelBreakdown[]; readonly by_chat: readonly ChatCostChatBreakdown[]; - readonly usage_limit?: ChatUsageLimitStatus; } // From codersdk/chats.go @@ -3629,100 +3628,6 @@ export interface ChatUnsupportedProvider { readonly display_name: string; } -// From codersdk/chats.go -/** - * ChatUsageLimitConfig is the deployment-wide default usage limit config. - */ -export interface ChatUsageLimitConfig { - /** - * Nil in the API means no default limit is set. The DB stores 0 when - * limiting is disabled. - */ - readonly spend_limit_micros: number | null; - readonly period: ChatUsageLimitPeriod; - readonly updated_at: string; -} - -// From codersdk/chats.go -/** - * ChatUsageLimitConfigResponse is returned from the admin config endpoint - * and includes the config plus a count of models without pricing. - */ -export interface ChatUsageLimitConfigResponse extends ChatUsageLimitConfig { - readonly unpriced_model_count: number; - readonly overrides: readonly ChatUsageLimitOverride[]; - readonly group_overrides: readonly ChatUsageLimitGroupOverride[]; -} - -// From codersdk/chats.go -/** - * 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. - */ -export interface ChatUsageLimitExceededResponse extends Response { - readonly spent_micros: number; - readonly limit_micros: number; - readonly resets_at: string; -} - -// From codersdk/chats.go -/** - * ChatUsageLimitGroupOverride represents a group-scoped spend limit override. - */ -export interface ChatUsageLimitGroupOverride { - readonly group_id: string; - readonly group_name: string; - readonly group_display_name: string; - readonly group_avatar_url: string; - readonly member_count: number; - /** - * Nil in the API means no group override is set. Persisted override rows - * store positive values. - */ - readonly spend_limit_micros: number | null; -} - -// From codersdk/chats.go -/** - * ChatUsageLimitOverride is a per-user override of the deployment default. - */ -export interface ChatUsageLimitOverride { - readonly user_id: string; - readonly username: string; - readonly name: string; - readonly avatar_url: string; - /** - * Nil in the API means no user override is set. Persisted override rows - * store positive values. - */ - readonly spend_limit_micros: number | null; -} - -// From codersdk/chats.go -export type ChatUsageLimitPeriod = "day" | "month" | "week"; - -export const ChatUsageLimitPeriods: ChatUsageLimitPeriod[] = [ - "day", - "month", - "week", -]; - -// From codersdk/chats.go -/** - * ChatUsageLimitStatus represents the current spend status for a user - * within their active limit period. - */ -export interface ChatUsageLimitStatus { - readonly is_limited: boolean; - readonly period?: ChatUsageLimitPeriod; - readonly spend_limit_micros?: number; - readonly current_spend: number; - readonly period_start?: string; - readonly period_end?: string; -} - // From codersdk/chats.go export interface ChatUser extends MinimalUser { readonly role: ChatRole; @@ -9766,22 +9671,6 @@ export interface UpdateChatSystemPromptRequest { readonly include_default_system_prompt?: boolean; } -// From codersdk/chats.go -/** - * UpdateChatUsageLimitGroupOverrideRequest is kept as a compatibility alias. - */ -export interface UpdateChatUsageLimitGroupOverrideRequest { - readonly spend_limit_micros: number; // Must be greater than 0. -} - -// From codersdk/chats.go -/** - * UpdateChatUsageLimitOverrideRequest is kept as a compatibility alias. - */ -export interface UpdateChatUsageLimitOverrideRequest { - readonly spend_limit_micros: number; // Must be greater than 0. -} - // From codersdk/chats.go /** * UpdateChatWorkspaceTTLRequest is the request to update the chat @@ -10250,24 +10139,6 @@ export interface UploadResponse { readonly hash: string; } -// From codersdk/chats.go -/** - * UpsertChatUsageLimitGroupOverrideRequest is the request to create or update - * a group-level spend limit override. - */ -export interface UpsertChatUsageLimitGroupOverrideRequest { - readonly spend_limit_micros: number; // Must be greater than 0. -} - -// From codersdk/chats.go -/** - * UpsertChatUsageLimitOverrideRequest is the body for creating/updating a - * per-user usage limit override. - */ -export interface UpsertChatUsageLimitOverrideRequest { - readonly spend_limit_micros: number; // Must be greater than 0. -} - // From codersdk/aibridge.go export interface UpsertGroupAIBudgetRequest { /** diff --git a/site/src/modules/dashboard/Navbar/UserDropdown/UserDropdown.stories.tsx b/site/src/modules/dashboard/Navbar/UserDropdown/UserDropdown.stories.tsx index ff48b73254..4577a78173 100644 --- a/site/src/modules/dashboard/Navbar/UserDropdown/UserDropdown.stories.tsx +++ b/site/src/modules/dashboard/Navbar/UserDropdown/UserDropdown.stories.tsx @@ -18,6 +18,8 @@ const mockAISpend: UserAISpendStatus = { period_end: "2026-07-01T00:00:00Z", }; +const spendPeriodLabel = "Estimated AI spend: June 1 - July 1, 2026"; + const aiCostControl: { features: FeatureName[] } = { features: ["aibridge"], }; @@ -45,8 +47,8 @@ type Story = StoryObj; const openDropdown = async (canvasElement: HTMLElement) => { const canvas = within(canvasElement); await userEvent.click(canvas.getByRole("button")); - await waitFor(async () => - expect(await screen.findByText(/v2\.\d+\.\d+/i)).toBeInTheDocument(), + return within( + await within(canvasElement.ownerDocument.body).findByRole("menu"), ); }; @@ -69,11 +71,11 @@ export const WithAISpend: Story = { }, play: async ({ canvasElement, step }) => { await step("shows AI spend", async () => { - await openDropdown(canvasElement); - await waitFor(() => - expect(document.body).toHaveTextContent("$819 / $1,200 USD"), - ); - expect(document.body).toHaveTextContent("(Estimated AI spend/month)"); + const menu = await openDropdown(canvasElement); + await waitFor(() => { + expect(document.body).toHaveTextContent("$819 / $1,200 USD"); + expect(menu.getByText(spendPeriodLabel)).toBeVisible(); + }); expect( screen.getByRole("progressbar", { name: "AI spend usage" }), ).toHaveAttribute("aria-valuenow", "68"); @@ -93,11 +95,11 @@ export const AISpendWarning: Story = { }, play: async ({ canvasElement, step }) => { await step("shows the warning marker near the limit", async () => { - await openDropdown(canvasElement); - await waitFor(() => - expect(document.body).toHaveTextContent("$1,080 / $1,200 USD"), - ); - expect(document.body).toHaveTextContent("(Estimated AI spend/month)"); + const menu = await openDropdown(canvasElement); + await waitFor(() => { + expect(document.body).toHaveTextContent("$1,080 / $1,200 USD"); + expect(menu.getByText(spendPeriodLabel)).toBeVisible(); + }); expect( screen.getByRole("progressbar", { name: "AI spend usage" }), ).toHaveAttribute("aria-valuenow", "90"); @@ -140,11 +142,11 @@ export const AISpendExceeded: Story = { }, play: async ({ canvasElement, step }) => { await step("shows the exceeded marker at the limit", async () => { - await openDropdown(canvasElement); - await waitFor(() => - expect(document.body).toHaveTextContent("$1,500 / $1,200 USD"), - ); - expect(document.body).toHaveTextContent("(Estimated AI spend/month)"); + const menu = await openDropdown(canvasElement); + await waitFor(() => { + expect(document.body).toHaveTextContent("$1,500 / $1,200 USD"); + expect(menu.getByText(spendPeriodLabel)).toBeVisible(); + }); expect( screen.getByRole("progressbar", { name: "AI spend usage" }), ).toHaveAttribute("aria-valuenow", "100"); @@ -161,11 +163,11 @@ export const AISpendUnlimited: Story = { }, play: async ({ canvasElement, step }) => { await step("shows unlimited spend without a bar", async () => { - await openDropdown(canvasElement); - await waitFor(() => - expect(document.body).toHaveTextContent("$819 / Unlimited USD"), - ); - expect(document.body).toHaveTextContent("(Estimated AI spend/month)"); + const menu = await openDropdown(canvasElement); + await waitFor(() => { + expect(document.body).toHaveTextContent("$819 / Unlimited USD"); + expect(menu.getByText(spendPeriodLabel)).toBeVisible(); + }); expect( screen.queryByRole("progressbar", { name: "AI spend usage" }), ).not.toBeInTheDocument(); @@ -272,7 +274,7 @@ export const AISpendHiddenOnInvalidData: Story = { play: async ({ canvasElement, step }) => { await step("hides AI spend on invalid data", async () => { await openDropdown(canvasElement); - expect(screen.queryByText(/spend\/month/)).not.toBeInTheDocument(); + expect(screen.queryByText(spendPeriodLabel)).not.toBeInTheDocument(); }); }, }; @@ -293,7 +295,7 @@ export const AISpendHiddenOnNegativeLimit: Story = { play: async ({ canvasElement, step }) => { await step("hides AI spend on a negative limit", async () => { await openDropdown(canvasElement); - expect(screen.queryByText(/spend\/month/)).not.toBeInTheDocument(); + expect(screen.queryByText(spendPeriodLabel)).not.toBeInTheDocument(); }); }, }; diff --git a/site/src/modules/dashboard/Navbar/UserDropdown/UserDropdown.tsx b/site/src/modules/dashboard/Navbar/UserDropdown/UserDropdown.tsx index 2cc33a8fa0..36d08b3284 100644 --- a/site/src/modules/dashboard/Navbar/UserDropdown/UserDropdown.tsx +++ b/site/src/modules/dashboard/Navbar/UserDropdown/UserDropdown.tsx @@ -50,6 +50,8 @@ export const UserDropdown: FC = ({ ? { currentSpend: data.current_spend_micros, spendLimit: data.effective_budget?.spend_limit_micros ?? null, + periodStart: data.period_start, + periodEnd: data.period_end, } : null; const severity = diff --git a/site/src/modules/dashboard/Navbar/UserDropdown/UserDropdownAISpend.tsx b/site/src/modules/dashboard/Navbar/UserDropdown/UserDropdownAISpend.tsx index 17e9816b7b..8060e5efd1 100644 --- a/site/src/modules/dashboard/Navbar/UserDropdown/UserDropdownAISpend.tsx +++ b/site/src/modules/dashboard/Navbar/UserDropdown/UserDropdownAISpend.tsx @@ -1,17 +1,25 @@ import type { FC } from "react"; import { UsageBar } from "#/components/UsageBar/UsageBar"; -import { getSeverity, usageProgressPercentage } from "#/utils/budget"; +import { + formatSpendPeriodLabel, + getSeverity, + usageProgressPercentage, +} from "#/utils/budget"; import { formatBudgetUSD } from "#/utils/currency"; interface UserDropdownAISpendProps { currentSpend: number; /** A null limit means unlimited. */ spendLimit: number | null; + periodStart: string; + periodEnd: string; } export const UserDropdownAISpend: FC = ({ currentSpend, spendLimit, + periodStart, + periodEnd, }) => { return (
@@ -31,7 +39,7 @@ export const UserDropdownAISpend: FC = ({ /> )}
- (Estimated AI spend/month) + Estimated AI spend: {formatSpendPeriodLabel(periodStart, periodEnd)}
); diff --git a/site/src/pages/AISettingsPage/SpendPage/SpendPage.tsx b/site/src/pages/AISettingsPage/SpendPage/SpendPage.tsx index 19306a7cc9..1acc4dade0 100644 --- a/site/src/pages/AISettingsPage/SpendPage/SpendPage.tsx +++ b/site/src/pages/AISettingsPage/SpendPage/SpendPage.tsx @@ -1,18 +1,8 @@ import dayjs from "dayjs"; import { type FC, useState } from "react"; -import { useMutation, useQuery, useQueryClient } from "react-query"; +import { useQuery } from "react-query"; import { useSearchParams } from "react-router"; -import { - chatCostSummary, - chatUsageLimitConfig, - deleteChatUsageLimitGroupOverride, - deleteChatUsageLimitOverride, - paginatedChatCostUsers, - updateChatUsageLimitConfig, - upsertChatUsageLimitGroupOverride, - upsertChatUsageLimitOverride, -} from "#/api/queries/chats"; -import { groups } from "#/api/queries/groups"; +import { chatCostSummary, paginatedChatCostUsers } from "#/api/queries/chats"; import { user } from "#/api/queries/users"; import type { ChatCostUserRollup } from "#/api/typesGenerated"; import type { DateRangeValue } from "#/components/DateRangePicker/DateRangePicker"; @@ -25,7 +15,6 @@ import { toExclusiveEndOfDayDateRange } from "./utils/dateRange"; const startDateSearchParam = "startDate"; const endDateSearchParam = "endDate"; -const tabSearchParam = "tab"; const DEFAULT_DATE_RANGE_DAYS = 30; const SEARCH_DEBOUNCE_MS = 300; const USAGE_USERS_PAGE_SIZE = 10; @@ -44,33 +33,10 @@ interface SpendPageProps { const SpendPage: FC = ({ now }) => { const { permissions } = useAuthenticated(); - const queryClient = useQueryClient(); - - const configQuery = useQuery(chatUsageLimitConfig()); - const groupsQuery = useQuery(groups()); - - const updateConfigMutation = useMutation( - updateChatUsageLimitConfig(queryClient), - ); - const upsertOverrideMutation = useMutation( - upsertChatUsageLimitOverride(queryClient), - ); - const deleteOverrideMutation = useMutation( - deleteChatUsageLimitOverride(queryClient), - ); - const upsertGroupOverrideMutation = useMutation( - upsertChatUsageLimitGroupOverride(queryClient), - ); - const deleteGroupOverrideMutation = useMutation( - deleteChatUsageLimitGroupOverride(queryClient), - ); - const [searchParams, setSearchParams] = useSearchParams(); const searchFilter = searchParams.get("search") ?? ""; const debouncedSearch = useDebouncedValue(searchFilter, SEARCH_DEBOUNCE_MS); - const tabParam = searchParams.get(tabSearchParam); - const activeTab = tabParam === "usage" ? "usage" : "limits"; const setSearchFilter = (value: string) => { setSearchParams( @@ -117,21 +83,6 @@ const SpendPage: FC = ({ now }) => { end_date: dateRange.endDate.toISOString(), }; - const onActiveTabChange = (tab: "limits" | "usage") => { - setSearchParams( - (prev) => { - const next = new URLSearchParams(prev); - if (tab === "usage") { - next.set(tabSearchParam, tab); - } else { - next.delete(tabSearchParam); - } - return next; - }, - { replace: true }, - ); - }; - const onDateRangeChange = (value: DateRangeValue) => { const nextDateRange = toExclusiveEndOfDayDateRange(value); @@ -170,51 +121,6 @@ const SpendPage: FC = ({ now }) => { return ( void configQuery.refetch()} - groupsData={groupsQuery.data} - isLoadingGroups={groupsQuery.isLoading} - groupsError={groupsQuery.isError ? groupsQuery.error : null} - onUpdateConfig={(req, options) => { - updateConfigMutation.mutate(req, { - onSuccess: options?.onSuccess, - }); - }} - isUpdatingConfig={updateConfigMutation.isPending} - updateConfigError={ - updateConfigMutation.isError ? updateConfigMutation.error : null - } - resetUpdateConfig={updateConfigMutation.reset} - onUpsertOverride={({ userID, req, onSuccess }) => - upsertOverrideMutation.mutate({ userID, req }, { onSuccess }) - } - isUpsertingOverride={upsertOverrideMutation.isPending} - upsertOverrideError={ - upsertOverrideMutation.isError ? upsertOverrideMutation.error : null - } - onDeleteOverride={deleteOverrideMutation.mutate} - isDeletingOverride={deleteOverrideMutation.isPending} - deleteOverrideError={ - deleteOverrideMutation.isError ? deleteOverrideMutation.error : null - } - onUpsertGroupOverride={({ groupID, req, onSuccess }) => - upsertGroupOverrideMutation.mutate({ groupID, req }, { onSuccess }) - } - isUpsertingGroupOverride={upsertGroupOverrideMutation.isPending} - upsertGroupOverrideError={ - upsertGroupOverrideMutation.isError - ? upsertGroupOverrideMutation.error - : null - } - onDeleteGroupOverride={deleteGroupOverrideMutation.mutate} - isDeletingGroupOverride={deleteGroupOverrideMutation.isPending} - deleteGroupOverrideError={ - deleteGroupOverrideMutation.isError - ? deleteGroupOverrideMutation.error - : null - } dateRange={dateRange} endDateIsExclusive={endDateIsExclusive} onDateRangeChange={onDateRangeChange} @@ -245,8 +151,6 @@ const SpendPage: FC = ({ now }) => { isSummaryLoading={summaryQuery.isLoading} summaryError={summaryQuery.error} onSummaryRetry={() => void summaryQuery.refetch()} - activeTab={activeTab} - onActiveTabChange={onActiveTabChange} /> ); diff --git a/site/src/pages/AISettingsPage/SpendPage/SpendPageView.stories.tsx b/site/src/pages/AISettingsPage/SpendPage/SpendPageView.stories.tsx index 91e63fb29b..f3c1b5c880 100644 --- a/site/src/pages/AISettingsPage/SpendPage/SpendPageView.stories.tsx +++ b/site/src/pages/AISettingsPage/SpendPage/SpendPageView.stories.tsx @@ -1,5 +1,4 @@ import type { Meta, StoryObj } from "@storybook/react-vite"; -import { useState } from "react"; import { expect, fn, userEvent, within } from "storybook/test"; import type * as TypesGen from "#/api/typesGenerated"; import type { PaginationResult } from "#/components/PaginationWidget/PaginationContainer"; @@ -100,68 +99,6 @@ const mockCostSummary = { ], } as TypesGen.ChatCostSummary; -const mockConfigData = { - spend_limit_micros: 50_000_000, - period: "month", - updated_at: "2026-03-01T00:00:00Z", - unpriced_model_count: 0, - overrides: [ - { - user_id: "user-3", - username: "dave", - name: "Dave Grohl", - avatar_url: "", - spend_limit_micros: 100_000_000, - }, - { - user_id: "user-4", - username: "charlie", - name: "Charlie Chaplin", - avatar_url: "", - spend_limit_micros: 25_000_000, - }, - ], - group_overrides: [ - { - group_id: "group-1", - group_name: "engineering", - group_display_name: "Engineering", - group_avatar_url: "", - member_count: 12, - spend_limit_micros: 75_000_000, - }, - ], -} as TypesGen.ChatUsageLimitConfigResponse; - -const mockGroups = [ - { - id: "group-1", - name: "engineering", - display_name: "Engineering", - organization_id: "org-1", - members: [], - total_member_count: 12, - avatar_url: "", - quota_allowance: 0, - source: "user", - organization_name: "default", - organization_display_name: "Default", - }, - { - id: "group-2", - name: "design", - display_name: "Design", - organization_id: "org-1", - members: [], - total_member_count: 5, - avatar_url: "", - quota_allowance: 0, - source: "user", - organization_name: "default", - organization_display_name: "Default", - }, -] as TypesGen.Group[]; - const defaultDateRange = { startDate: new Date("2026-02-10T00:00:00Z"), endDate: new Date("2026-03-12T00:00:00Z"), @@ -219,22 +156,6 @@ function mockUsersQuery( } const baseProps = { - configData: undefined as TypesGen.ChatUsageLimitConfigResponse | undefined, - isLoadingConfig: false, - configError: null as Error | null, - groupsData: undefined as TypesGen.Group[] | undefined, - isLoadingGroups: false, - groupsError: null as Error | null, - isUpdatingConfig: false, - updateConfigError: null as Error | null, - isUpsertingOverride: false, - upsertOverrideError: null as Error | null, - isDeletingOverride: false, - deleteOverrideError: null as Error | null, - isUpsertingGroupOverride: false, - upsertGroupOverrideError: null as Error | null, - isDeletingGroupOverride: false, - deleteGroupOverrideError: null as Error | null, dateRange: defaultDateRange, endDateIsExclusive: false, searchFilter: "", @@ -254,84 +175,22 @@ const meta = { component: SpendPageView, // TODO: Stories in this file fail when pixel runs their play functions. Fix them and remove the exclude. parameters: { pixel: { exclude: true } }, - render: (args) => { - const [activeTab, setActiveTab] = useState(args.activeTab); - return ( - { - setActiveTab(tab); - args.onActiveTabChange(tab); - }} - /> - ); - }, args: { ...baseProps, - refetchConfig: fn(), - onUpdateConfig: fn(), - resetUpdateConfig: fn(), - onUpsertOverride: fn(), - onDeleteOverride: fn(), - onUpsertGroupOverride: fn(), - onDeleteGroupOverride: fn(), onDateRangeChange: fn(), onSearchFilterChange: fn(), onDrillInUserRetry: fn(), onClearSelectedUser: fn(), onSelectUser: fn(), onSummaryRetry: fn(), - activeTab: "limits", - onActiveTabChange: fn(), }, } satisfies Meta; export default meta; type Story = StoryObj; -export const SpendWithLimitsAndUsers: Story = { - args: { - configData: mockConfigData, - groupsData: mockGroups, - usersQuery: mockUsersQuery({ data: mockUsersResponse }), - }, - play: async ({ canvasElement }) => { - const canvas = within(canvasElement); - - await canvas.findByText("Spend limits and usage"); - await expect( - canvas.getByText( - /Cost controls features will move to AI Governance in v2\.36\./, - ), - ).toBeInTheDocument(); - await expect( - canvas.getByRole("link", { name: /Read more here/ }), - ).toHaveAttribute( - "href", - expect.stringContaining("/ai-coder/ai-gateway/cost-controls"), - ); - await expect( - canvas.getByRole("switch", { name: "Spend limit" }), - ).toBeInTheDocument(); - await expect(canvas.getByText("Group limits")).toBeInTheDocument(); - await expect(canvas.getByText("Usage")).toBeInTheDocument(); - - await userEvent.click(canvas.getByRole("tab", { name: "Usage" })); - - await expect(await canvas.findByText("Alice Liddell")).toBeInTheDocument(); - await expect(canvas.getByText("Bob Builder")).toBeInTheDocument(); - - await expect( - canvas.getByPlaceholderText("Search by name or username"), - ).toBeInTheDocument(); - }, -}; - export const SpendUsersEmpty: Story = { args: { - configData: mockConfigData, - groupsData: mockGroups, usersQuery: mockUsersQuery({ data: { start_date: "2026-02-10T00:00:00Z", @@ -344,8 +203,16 @@ export const SpendUsersEmpty: Story = { play: async ({ canvasElement }) => { const canvas = within(canvasElement); - await canvas.findByText("Spend limits and usage"); - await userEvent.click(canvas.getByRole("tab", { name: "Usage" })); + await canvas.findByText("AI spend usage"); + await expect(canvas.getByRole("alert")).toHaveTextContent( + "As of v2.36, AI Governance Cost Control replaces Coder Agents Cost Control.", + ); + await expect( + canvas.getByRole("link", { name: /Read more here/ }), + ).toHaveAttribute( + "href", + expect.stringContaining("/ai-coder/ai-gateway/cost-controls"), + ); await expect( await canvas.findByText("No usage data for this period."), ).toBeInTheDocument(); @@ -354,8 +221,6 @@ export const SpendUsersEmpty: Story = { export const SpendUserDrillIn: Story = { args: { - configData: mockConfigData, - groupsData: mockGroups, drillInUserId: "user-1", drillInUser: mockUserProfile, summaryData: mockCostSummary, @@ -373,8 +238,6 @@ export const SpendUserDrillIn: Story = { export const SpendUserDrillInAndBack: Story = { args: { - configData: mockConfigData, - groupsData: mockGroups, drillInUserId: "user-1", drillInUser: mockUserProfile, summaryData: mockCostSummary, @@ -392,8 +255,6 @@ export const SpendUserDrillInAndBack: Story = { export const SpendDrillInLoading: Story = { args: { - configData: mockConfigData, - groupsData: mockGroups, drillInUserId: "user-1", drillInUser: null, isDrillInUserLoading: true, @@ -408,8 +269,6 @@ export const SpendDrillInLoading: Story = { export const SpendDrillInError: Story = { args: { - configData: mockConfigData, - groupsData: mockGroups, drillInUserId: "user-1", drillInUser: null, isDrillInUserError: true, @@ -424,8 +283,6 @@ export const SpendDrillInError: Story = { export const SpendRefetchOverlay: Story = { args: { - configData: mockConfigData, - groupsData: mockGroups, usersQuery: mockUsersQuery({ data: mockUsersResponse, isFetching: true, @@ -434,8 +291,6 @@ export const SpendRefetchOverlay: Story = { play: async ({ canvasElement }) => { const canvas = within(canvasElement); - await userEvent.click(canvas.getByRole("tab", { name: "Usage" })); - await canvas.findByText("Alice Liddell"); await expect( @@ -444,43 +299,13 @@ export const SpendRefetchOverlay: Story = { }, }; -export const SpendConfigLoading: Story = { - args: { - isLoadingConfig: true, - }, -}; - -export const SpendConfigError: Story = { - args: { - configError: new Error("Network error: failed to fetch config"), - usersQuery: mockUsersQuery({ data: mockUsersResponse }), - }, - play: async ({ canvasElement }) => { - const canvas = within(canvasElement); - - await canvas.findByText("Network error: failed to fetch config"); - await expect(canvas.getByText("Retry")).toBeInTheDocument(); - - await userEvent.click(canvas.getByRole("tab", { name: "Usage" })); - - await expect(canvas.getByText("Alice Liddell")).toBeInTheDocument(); - await expect(canvas.getByText("Bob Builder")).toBeInTheDocument(); - }, -}; - export const SpendUsersLoading: Story = { args: { - configData: mockConfigData, - groupsData: mockGroups, usersQuery: mockUsersQuery({ isLoading: true }), }, play: async ({ canvasElement }) => { const canvas = within(canvasElement); - await canvas.findByRole("switch", { name: "Spend limit" }); - - await userEvent.click(canvas.getByRole("tab", { name: "Usage" })); - await expect( await canvas.findByRole("status", { name: "Loading usage" }), ).toBeInTheDocument(); @@ -489,8 +314,6 @@ export const SpendUsersLoading: Story = { export const SpendUsersError: Story = { args: { - configData: mockConfigData, - groupsData: mockGroups, usersQuery: mockUsersQuery({ error: new Error("Failed to load usage data"), }), @@ -498,10 +321,6 @@ export const SpendUsersError: Story = { play: async ({ canvasElement }) => { const canvas = within(canvasElement); - await canvas.findByRole("switch", { name: "Spend limit" }); - - await userEvent.click(canvas.getByRole("tab", { name: "Usage" })); - await expect( canvas.getByText("Failed to load usage data"), ).toBeInTheDocument(); @@ -511,15 +330,11 @@ export const SpendUsersError: Story = { export const SpendUserClickToDrillIn: Story = { args: { - configData: mockConfigData, - groupsData: mockGroups, usersQuery: mockUsersQuery({ data: mockUsersResponse }), }, play: async ({ canvasElement, args }) => { const canvas = within(canvasElement); - await userEvent.click(canvas.getByRole("tab", { name: "Usage" })); - const row = await canvas.findByRole("button", { name: /^View details for Alice Liddell/, }); diff --git a/site/src/pages/AISettingsPage/SpendPage/SpendPageView.tsx b/site/src/pages/AISettingsPage/SpendPage/SpendPageView.tsx index 3cf1aec74d..f836ce0c7c 100644 --- a/site/src/pages/AISettingsPage/SpendPage/SpendPageView.tsx +++ b/site/src/pages/AISettingsPage/SpendPage/SpendPageView.tsx @@ -1,9 +1,6 @@ import type { FC } from "react"; -import { getErrorMessage } from "#/api/errors"; import type * as TypesGen from "#/api/typesGenerated"; import { Alert, AlertDescription } from "#/components/Alert/Alert"; -import { ErrorAlert } from "#/components/Alert/ErrorAlert"; -import { Button } from "#/components/Button/Button"; import type { DateRangeValue } from "#/components/DateRangePicker/DateRangePicker"; import { Link } from "#/components/Link/Link"; import type { PaginationResult } from "#/components/PaginationWidget/PaginationContainer"; @@ -12,69 +9,12 @@ import { SettingsHeaderDescription, SettingsHeaderTitle, } from "#/components/SettingsHeader/SettingsHeader"; -import { Spinner } from "#/components/Spinner/Spinner"; -import { - Tabs, - TabsContent, - TabsList, - TabsTrigger, -} from "#/components/Tabs/Tabs"; -import { - TemporarySavedState, - useTemporarySavedState, -} from "#/components/TemporarySavedState/TemporarySavedState"; -import { dollarsToMicros, microsToDollars } from "#/utils/currency"; import { docs } from "#/utils/docs"; -import { - DefaultLimitController, - type DefaultLimitFormValues, -} from "./components/LimitsTab/DefaultLimitController"; -import { DefaultLimitSection } from "./components/LimitsTab/DefaultLimitSection"; -import { GroupLimitsSection } from "./components/LimitsTab/GroupLimitsSection"; -import { GroupOverrideController } from "./components/LimitsTab/GroupOverrideController"; -import { normalizeChatUsageLimitPeriod } from "./components/LimitsTab/limitsFormLogic"; -import { UserOverrideController } from "./components/LimitsTab/UserOverrideController"; -import { UserOverridesSection } from "./components/LimitsTab/UserOverridesSection"; import { SpendDrillInView } from "./components/SpendDrillInView"; -import { SpendSectionHeader } from "./components/SpendSectionHeader"; import { UsageTab } from "./components/UsageTab/UsageTab"; import { formatUsageDateRange, toInclusiveDateRange } from "./utils/dateRange"; interface SpendPageViewProps { - configData: TypesGen.ChatUsageLimitConfigResponse | undefined; - isLoadingConfig: boolean; - configError: Error | null; - refetchConfig: () => void; - groupsData: TypesGen.Group[] | undefined; - isLoadingGroups: boolean; - groupsError: Error | null; - onUpdateConfig: ( - req: TypesGen.ChatUsageLimitConfig, - options?: { onSuccess?: () => void }, - ) => void; - isUpdatingConfig: boolean; - updateConfigError: Error | null; - resetUpdateConfig: () => void; - onUpsertOverride: (args: { - userID: string; - req: TypesGen.UpsertChatUsageLimitOverrideRequest; - onSuccess: () => void; - }) => void; - isUpsertingOverride: boolean; - upsertOverrideError: Error | null; - onDeleteOverride: (userID: string) => void; - isDeletingOverride: boolean; - deleteOverrideError: Error | null; - onUpsertGroupOverride: (args: { - groupID: string; - req: TypesGen.UpsertChatUsageLimitGroupOverrideRequest; - onSuccess: () => void; - }) => void; - isUpsertingGroupOverride: boolean; - upsertGroupOverrideError: Error | null; - onDeleteGroupOverride: (groupID: string) => void; - isDeletingGroupOverride: boolean; - deleteGroupOverrideError: Error | null; dateRange: DateRangeValue; endDateIsExclusive: boolean; onDateRangeChange: (value: DateRangeValue) => void; @@ -99,34 +39,9 @@ interface SpendPageViewProps { isSummaryLoading: boolean; summaryError: unknown; onSummaryRetry: () => void; - activeTab: "limits" | "usage"; - onActiveTabChange: (tab: "limits" | "usage") => void; } export const SpendPageView: FC = ({ - configData, - isLoadingConfig, - configError, - refetchConfig, - groupsData, - isLoadingGroups, - groupsError, - onUpdateConfig, - isUpdatingConfig, - updateConfigError, - resetUpdateConfig, - onUpsertOverride, - isUpsertingOverride, - upsertOverrideError, - onDeleteOverride, - isDeletingOverride, - deleteOverrideError, - onUpsertGroupOverride, - isUpsertingGroupOverride, - upsertGroupOverrideError, - onDeleteGroupOverride, - isDeletingGroupOverride, - deleteGroupOverrideError, dateRange, endDateIsExclusive, onDateRangeChange, @@ -145,63 +60,12 @@ export const SpendPageView: FC = ({ isSummaryLoading, summaryError, onSummaryRetry, - activeTab, - onActiveTabChange, }) => { - const defaultLimitValues: DefaultLimitFormValues = (() => { - const spendLimitMicros = configData?.spend_limit_micros; - const enabled = spendLimitMicros !== null && spendLimitMicros !== undefined; - - return { - enabled, - period: normalizeChatUsageLimitPeriod(configData?.period), - amountDollars: enabled - ? microsToDollars(spendLimitMicros).toString() - : "", - }; - })(); - const defaultLimitKey = JSON.stringify({ - spend_limit_micros: configData?.spend_limit_micros ?? null, - period: defaultLimitValues.period, - }); - const displayDateRange = toInclusiveDateRange(dateRange, endDateIsExclusive); const dateRangeLabel = formatUsageDateRange(dateRange, { endDateIsExclusive, }); - const handleResetUpdateConfig = () => { - if (!isUpdatingConfig) { - resetUpdateConfig(); - } - }; - - const handleSaveDefault = ({ - enabled, - period, - amountDollars, - }: DefaultLimitFormValues) => { - const spendLimitMicros = enabled ? dollarsToMicros(amountDollars) : null; - onUpdateConfig( - { - spend_limit_micros: spendLimitMicros, - period, - updated_at: new Date().toISOString(), - }, - { onSuccess: showSavedState }, - ); - }; - - const groupOverrides = configData?.group_overrides ?? []; - const overrides = configData?.overrides ?? []; - const unpricedModelCount = configData?.unpriced_model_count ?? 0; - const groupOrganizationNames = - groupsData?.reduce>((acc, group) => { - acc[group.id] = group.organization_name; - return acc; - }, {}) ?? {}; - const { isSavedVisible, showSavedState } = useTemporarySavedState(); - if (drillInUserId) { return ( = ({ /> ); } + return ( - - {(groupCtrl) => ( - - {(userCtrl) => ( -
- - - As of v2.36, AI Governance Cost Control replaces Coder Agents - Cost Control. The limits on this page are no longer enforced - and do not carry over. Recreate each limit as an AI Governance - budget to restore enforcement.{" "} - - Read more here - - - +
+ + + As of v2.36, AI Governance Cost Control replaces Coder Agents Cost + Control. The limits on this page are no longer enforced and do not + carry over. Recreate each limit as an AI Governance budget to restore + enforcement.{" "} + + Read more here + + + - - - Spend limits and usage - - - Configure spend limits and monitor AI usage across your - deployment. - - + + AI spend usage + + Monitor AI usage across your deployment. + + - { - if (tab === "limits" || tab === "usage") { - onActiveTabChange(tab); - } - }} - > - - Spend limits - Usage - - - -
- {isLoadingConfig ? ( -
- -
- ) : configError ? ( -
- - -
- ) : ( - <> - - {({ - enabled, - onEnabledChange, - period, - onPeriodChange, - amountDollars, - onAmountDollarsChange, - isAmountValid, - isDirty, - saveDefault, - }) => ( - { - handleResetUpdateConfig(); - onEnabledChange(v); - }} - period={period} - onPeriodChange={(v) => { - handleResetUpdateConfig(); - onPeriodChange(v); - }} - amountDollars={amountDollars} - onAmountDollarsChange={(v) => { - handleResetUpdateConfig(); - onAmountDollarsChange(v); - }} - unpricedModelCount={unpricedModelCount} - isSaving={isUpdatingConfig} - isSavedVisible={isSavedVisible} - saveDisabled={ - isUpdatingConfig || !isAmountValid || !isDirty - } - onSave={isDirty ? saveDefault : undefined} - saveStatus={ - isSavedVisible ? ( - - ) : updateConfigError ? ( - getErrorMessage( - updateConfigError, - "Failed to save the default spend limit.", - ) - ) : null - } - /> - )} - - -
- - { - userCtrl.handleShowUserFormChange(false); - groupCtrl.handleEditGroupOverride(override); - }} - onAddGroupOverride={ - groupCtrl.handleAddGroupOverride - } - onDeleteGroupOverride={onDeleteGroupOverride} - upsertPending={isUpsertingGroupOverride} - upsertError={upsertGroupOverrideError} - deletePending={isDeletingGroupOverride} - deleteError={deleteGroupOverrideError} - groupsError={groupsError} - /> -
- -
- - { - groupCtrl.handleShowGroupFormChange(false); - userCtrl.handleEditUserOverride(override); - }} - onAddOverride={userCtrl.handleAddOverride} - onDeleteOverride={onDeleteOverride} - upsertPending={isUpsertingOverride} - upsertError={upsertOverrideError} - deletePending={isDeletingOverride} - deleteError={deleteOverrideError} - /> -
- - )} -
-
- - - { - groupCtrl.handleShowGroupFormChange(false); - userCtrl.handleEditUserOverride(override); - }} - /> - -
-
- )} - - )} - +
+ +
+
); }; diff --git a/site/src/pages/AISettingsPage/SpendPage/components/ChatCostSummaryView.tsx b/site/src/pages/AISettingsPage/SpendPage/components/ChatCostSummaryView.tsx index 4f1e954f55..622cf1eec2 100644 --- a/site/src/pages/AISettingsPage/SpendPage/components/ChatCostSummaryView.tsx +++ b/site/src/pages/AISettingsPage/SpendPage/components/ChatCostSummaryView.tsx @@ -1,5 +1,4 @@ -import dayjs from "dayjs"; -import { InfoIcon, TriangleAlertIcon } from "lucide-react"; +import { TriangleAlertIcon } from "lucide-react"; import { type FC, useState } from "react"; import { getErrorMessage } from "#/api/errors"; import type * as TypesGen from "#/api/typesGenerated"; @@ -32,25 +31,6 @@ interface ChatCostSummaryViewProps { emptyMessage: string; } -export const getUsageLimitPeriodLabel = ( - period: TypesGen.ChatUsageLimitPeriod | undefined, -): string => { - if (!period) { - return ""; - } - - switch (period) { - case "day": - return "Daily"; - case "week": - return "Weekly"; - case "month": - return "Monthly"; - default: - return ""; - } -}; - export const ChatCostSummaryView: FC = ({ summary, isLoading, @@ -110,41 +90,6 @@ export const ChatCostSummaryView: FC = ({ hasNextPage: hasChatNext, } = paginateItems(summary.by_chat, chatPageSize, chatPage); - const usageLimit = summary.usage_limit; - const showUsageLimitCard = usageLimit?.is_limited === true; - const usageLimitCurrentSpend = usageLimit?.current_spend ?? 0; - const usageLimitSpendMicros = usageLimit?.spend_limit_micros ?? 0; - const usageLimitPeriodLabel = usageLimit - ? getUsageLimitPeriodLabel(usageLimit.period) - : ""; - const usageProgressPercentage = - showUsageLimitCard && usageLimitSpendMicros > 0 - ? Math.min((usageLimitCurrentSpend / usageLimitSpendMicros) * 100, 100) - : 0; - const usageProgressBarClass = - usageProgressPercentage > 90 - ? "bg-content-destructive" - : usageProgressPercentage >= 75 - ? "bg-content-warning" - : "bg-content-success"; - const usageLimitExceeded = - showUsageLimitCard && usageLimitCurrentSpend >= usageLimitSpendMicros; - const usageLimitStatusText = usageLimitExceeded - ? "Limit exceeded" - : `${formatCostMicros( - Math.max(usageLimitSpendMicros - usageLimitCurrentSpend, 0), - )} remaining`; - const usageLimitCurrentPeriod = - showUsageLimitCard && usageLimit?.period_start && usageLimit?.period_end - ? `Current period: ${dayjs(usageLimit.period_start).format("MMM D")} to ${dayjs( - usageLimit.period_end, - ).format("MMM D")}` - : ""; - const usageLimitResetAt = - showUsageLimitCard && usageLimit?.period_end - ? dayjs(usageLimit.period_end).format("MMM D, YYYY h:mm A") - : ""; - return (
@@ -201,59 +146,6 @@ export const ChatCostSummaryView: FC = ({
- {showUsageLimitCard && usageLimit && ( -
-
-
-
-

- {usageLimitPeriodLabel} spend limit -

- {usageLimitCurrentPeriod && ( -

- {usageLimitCurrentPeriod} -

- )} -

- {formatCostMicros(usageLimitCurrentSpend)} /{" "} - {formatCostMicros(usageLimitSpendMicros)} -

-
-

- {Math.round(usageProgressPercentage)}% used -

-
-
-
-
-
-

- {usageLimitStatusText} -

-

- Resets {usageLimitResetAt} -

-
-
-
- )} - {summary.unpriced_messages_having_usage_count > 0 && (
@@ -266,14 +158,6 @@ export const ChatCostSummaryView: FC = ({
)} -
- - - Automatic title generation uses lightweight models and is not counted - towards usage limits. - -
- {summary.by_model.length === 0 && summary.by_chat.length === 0 ? (

{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 = ({ - initialValues, - onSave, - children, -}) => { - const [enabled, setEnabled] = useState(initialValues.enabled); - const [period, setPeriod] = useState( - initialValues.period, - ); - const [amountDollars, setAmountDollars] = useState( - initialValues.amountDollars, - ); - const isAmountValid = !enabled || isPositiveFiniteDollarAmount(amountDollars); - const isDirty = - enabled !== initialValues.enabled || - period !== initialValues.period || - amountDollars !== initialValues.amountDollars; - - const handleSave = () => { - if (enabled && !isPositiveFiniteDollarAmount(amountDollars)) { - return; - } - - onSave({ enabled, period, amountDollars }); - }; - - return children({ - enabled, - onEnabledChange: setEnabled, - period, - onPeriodChange: setPeriod, - amountDollars, - onAmountDollarsChange: setAmountDollars, - isAmountValid, - isDirty, - saveDefault: handleSave, - }); -}; diff --git a/site/src/pages/AISettingsPage/SpendPage/components/LimitsTab/DefaultLimitSection.tsx b/site/src/pages/AISettingsPage/SpendPage/components/LimitsTab/DefaultLimitSection.tsx deleted file mode 100644 index 3e837469e9..0000000000 --- a/site/src/pages/AISettingsPage/SpendPage/components/LimitsTab/DefaultLimitSection.tsx +++ /dev/null @@ -1,137 +0,0 @@ -import { TriangleAlertIcon } from "lucide-react"; -import type { FC, ReactNode } from "react"; -import type { ChatUsageLimitPeriod } from "#/api/typesGenerated"; -import { Button } from "#/components/Button/Button"; -import { - InputGroup, - InputGroupAddon, - InputGroupInput, -} from "#/components/InputGroup/InputGroup"; -import { - Select, - SelectContent, - SelectItem, - SelectTrigger, - SelectValue, -} from "#/components/Select/Select"; -import { Spinner } from "#/components/Spinner/Spinner"; -import { Switch } from "#/components/Switch/Switch"; - -interface DefaultLimitSectionProps { - enabled: boolean; - onEnabledChange: (enabled: boolean) => void; - period: ChatUsageLimitPeriod; - onPeriodChange: (period: ChatUsageLimitPeriod) => void; - amountDollars: string; - onAmountDollarsChange: (amount: string) => void; - unpricedModelCount: number; - isSaving?: boolean; - isSavedVisible?: boolean; - saveDisabled?: boolean; - onSave?: () => void; - saveStatus?: ReactNode; -} - -export const DefaultLimitSection: FC = ({ - enabled, - onEnabledChange, - period, - onPeriodChange, - amountDollars, - onAmountDollarsChange, - unpricedModelCount, - isSaving = false, - isSavedVisible = false, - saveDisabled = false, - onSave, - saveStatus, -}) => { - return ( -

-
- -
-

- Spend limit -

-

- Set a deployment-wide spend cap that applies to all users by - default. When disabled, users have unlimited spending. -

-
- - $ - onAmountDollarsChange(event.target.value)} - placeholder="0.00" - aria-label="Spend limit amount in dollars" - disabled={!enabled} - /> - - -
- {(onSave || isSavedVisible || isSaving) && - (isSavedVisible ? ( - saveStatus - ) : ( - - ))} -
-
- {saveStatus && !isSavedVisible && ( -
{saveStatus}
- )} -
-
- - {enabled && unpricedModelCount > 0 && ( -
- -
- {unpricedModelCount === 1 - ? "1 enabled model does not have pricing configured." - : `${unpricedModelCount} enabled models do not have pricing configured.`}{" "} - Usage of unpriced models cannot be tracked against the spend limit. -
-
- )} -
- ); -}; diff --git a/site/src/pages/AISettingsPage/SpendPage/components/LimitsTab/GroupLimitDialog.tsx b/site/src/pages/AISettingsPage/SpendPage/components/LimitsTab/GroupLimitDialog.tsx deleted file mode 100644 index b3c8e55521..0000000000 --- a/site/src/pages/AISettingsPage/SpendPage/components/LimitsTab/GroupLimitDialog.tsx +++ /dev/null @@ -1,183 +0,0 @@ -import { CheckIcon } from "lucide-react"; -import type { FC } from "react"; -import { getErrorMessage } from "#/api/errors"; -import type { Group } from "#/api/typesGenerated"; -import { Autocomplete } from "#/components/Autocomplete/Autocomplete"; -import { AvatarData } from "#/components/Avatar/AvatarData"; -import { Button } from "#/components/Button/Button"; -import { - Dialog, - DialogContent, - DialogDescription, - DialogFooter, - DialogHeader, - DialogTitle, -} from "#/components/Dialog/Dialog"; -import { Input } from "#/components/Input/Input"; -import { Label } from "#/components/Label/Label"; -import { Spinner } from "#/components/Spinner/Spinner"; -import { getGroupSubtitle } from "#/modules/groups"; -import { isPositiveFiniteDollarAmount } from "#/utils/currency"; -import type { GroupLimitOverrideGroup } from "./GroupLimitsSection"; - -interface GroupLimitDialogProps { - open: boolean; - onOpenChange: (open: boolean) => void; - selectedGroup: Group | null; - onSelectedGroupChange: (group: Group | null) => void; - groupAmount: string; - onGroupAmountChange: (amount: string) => void; - availableGroups: Group[]; - groupAutocompleteNoOptionsText: string; - groupsLoading: boolean; - editingGroupOverride: GroupLimitOverrideGroup | null; - upsertPending: boolean; - upsertError: Error | null; - onSave: () => void; - groupAutocompleteId: string; - groupAmountId: string; -} - -export const GroupLimitDialog: FC = ({ - open, - onOpenChange, - selectedGroup, - onSelectedGroupChange, - groupAmount, - onGroupAmountChange, - availableGroups, - groupAutocompleteNoOptionsText, - groupsLoading, - editingGroupOverride, - upsertPending, - upsertError, - onSave, - groupAutocompleteId, - groupAmountId, -}) => { - const isEditing = editingGroupOverride !== null; - const saveDisabled = isEditing - ? upsertPending || !isPositiveFiniteDollarAmount(groupAmount) - : upsertPending || - selectedGroup === null || - !isPositiveFiniteDollarAmount(groupAmount); - - return ( - { - if (upsertPending) { - return; - } - - onOpenChange(nextOpen); - if (!nextOpen) { - onSelectedGroupChange(null); - onGroupAmountChange(""); - } - }} - > - - - - {isEditing ? "Update group budget" : "Add group budget"} - - - {isEditing - ? "Update this group's spend limit override." - : "Set a spend limit override for a specific group."} - - - -
-
- {editingGroupOverride ? ( - <> - -
- -
- - ) : ( - <> - - group.id} - getOptionLabel={(group) => group.display_name || group.name} - isOptionEqualToValue={(option, optionValue) => - option.id === optionValue.id - } - renderOption={(option, isSelected) => ( -
- - {isSelected && } -
- )} - placeholder="Search groups..." - noOptionsText={groupAutocompleteNoOptionsText} - loading={groupsLoading} - disabled={groupsLoading} - className="w-full" - /> - - )} -
-
- - onGroupAmountChange(event.target.value)} - placeholder="0.00" - /> -
- {upsertError && ( -

- {getErrorMessage(upsertError, "Failed to save group override.")} -

- )} -
- - - - - -
-
- ); -}; diff --git a/site/src/pages/AISettingsPage/SpendPage/components/LimitsTab/GroupLimitsSection.stories.tsx b/site/src/pages/AISettingsPage/SpendPage/components/LimitsTab/GroupLimitsSection.stories.tsx deleted file mode 100644 index eb1963020a..0000000000 --- a/site/src/pages/AISettingsPage/SpendPage/components/LimitsTab/GroupLimitsSection.stories.tsx +++ /dev/null @@ -1,182 +0,0 @@ -import type { Meta, StoryObj } from "@storybook/react-vite"; -import { expect, fn, userEvent, waitFor, within } from "storybook/test"; -import type { Group } from "#/api/typesGenerated"; -import { GroupLimitsSection } from "./GroupLimitsSection"; - -const mockGroupOverrides = [ - { - group_id: "group-1", - group_display_name: "Engineering", - group_name: "engineering", - group_avatar_url: "", - member_count: 15, - spend_limit_micros: 10_000_000, - }, - { - group_id: "group-2", - group_display_name: "Design", - group_name: "design", - group_avatar_url: "", - member_count: 8, - spend_limit_micros: 5_000_000, - }, - { - group_id: "group-4", - group_display_name: "Support", - group_name: "support", - group_avatar_url: "", - member_count: 11, - spend_limit_micros: null, - }, -]; - -const mockAvailableGroups: Group[] = [ - { - id: "group-3", - name: "marketing", - display_name: "Marketing", - avatar_url: "", - organization_id: "org-1", - organization_name: "Acme", - organization_display_name: "Acme", - members: [], - quota_allowance: 0, - source: "user", - total_member_count: 5, - }, - { - id: "group-5", - name: "sales", - display_name: "Sales", - avatar_url: "", - organization_id: "org-1", - organization_name: "Acme", - organization_display_name: "Acme", - members: [], - quota_allowance: 0, - source: "user", - total_member_count: 12, - }, -]; - -const editingGroupOverride = { - group_id: mockGroupOverrides[0].group_id, - group_display_name: mockGroupOverrides[0].group_display_name, - group_name: mockGroupOverrides[0].group_name, - group_avatar_url: mockGroupOverrides[0].group_avatar_url, - member_count: mockGroupOverrides[0].member_count, -}; - -const meta: Meta = { - title: "pages/AISettingsPage/SpendPage/LimitsTab/GroupLimitsSection", - component: GroupLimitsSection, - args: { - groupOverrides: mockGroupOverrides, - groupOrganizationNames: { - "group-1": "Acme", - "group-2": "Acme", - }, - showGroupForm: false, - onShowGroupFormChange: fn(), - selectedGroup: null, - onSelectedGroupChange: fn(), - groupAmount: "", - onGroupAmountChange: fn(), - availableGroups: [], - groupAutocompleteNoOptionsText: "No groups available", - groupsLoading: false, - editingGroupOverride: null, - onEditGroupOverride: fn(), - onAddGroupOverride: fn(), - onDeleteGroupOverride: fn(), - upsertPending: false, - upsertError: null, - deletePending: false, - deleteError: null, - groupsError: null, - }, -}; - -export default meta; -type Story = StoryObj; - -export const Default: Story = {}; - -export const EmptyState: Story = { - args: { - groupOverrides: [], - }, -}; - -export const AddForm: Story = { - args: { - showGroupForm: true, - availableGroups: mockAvailableGroups, - }, -}; - -export const EditForm: Story = { - args: { - showGroupForm: true, - editingGroupOverride, - groupAmount: "10.00", - }, - play: async ({ canvasElement }) => { - const body = within(canvasElement.ownerDocument.body); - // Update budget button confirms edit mode is active. - await expect( - body.getByRole("button", { name: /update budget/i }), - ).toBeInTheDocument(); - // The editing group name appears in both the table row and the - // read-only edit form identity, confirming it was populated. - const nameElements = body.getAllByText( - editingGroupOverride.group_display_name, - ); - expect(nameElements.length).toBeGreaterThanOrEqual(2); - }, -}; - -export const DeleteGroupOverride: Story = { - play: async ({ canvasElement }) => { - const body = within(canvasElement.ownerDocument.body); - - const actionButtons = await body.findAllByRole("button", { - name: /actions for/i, - }); - await userEvent.click(actionButtons[0]); - await userEvent.click( - await body.findByRole("menuitem", { name: /remove group limit/i }), - ); - - // The confirmation dialog should appear. - const dialog = await body.findByRole("dialog"); - await expect(dialog).toBeInTheDocument(); - await expect( - body.getByText(/Are you sure you want to delete this group override/i), - ).toBeInTheDocument(); - }, -}; - -export const DeleteGroupOverrideCancelled: Story = { - play: async ({ canvasElement, args }) => { - const body = within(canvasElement.ownerDocument.body); - - const actionButtons = await body.findAllByRole("button", { - name: /actions for/i, - }); - await userEvent.click(actionButtons[0]); - await userEvent.click( - await body.findByRole("menuitem", { name: /remove group limit/i }), - ); - - // Cancel the dialog. - await body.findByRole("dialog"); - await userEvent.click(body.getByRole("button", { name: "Cancel" })); - - // The dialog should be closed and the callback should not have been called. - await waitFor(() => { - expect(body.queryByRole("dialog")).not.toBeInTheDocument(); - }); - expect(args.onDeleteGroupOverride).not.toHaveBeenCalled(); - }, -}; diff --git a/site/src/pages/AISettingsPage/SpendPage/components/LimitsTab/GroupLimitsSection.tsx b/site/src/pages/AISettingsPage/SpendPage/components/LimitsTab/GroupLimitsSection.tsx deleted file mode 100644 index d0d912b444..0000000000 --- a/site/src/pages/AISettingsPage/SpendPage/components/LimitsTab/GroupLimitsSection.tsx +++ /dev/null @@ -1,192 +0,0 @@ -import { PlusIcon } from "lucide-react"; -import { type FC, useId, useState } from "react"; -import { getErrorMessage } from "#/api/errors"; -import type { ChatUsageLimitGroupOverride, Group } from "#/api/typesGenerated"; -import { Button } from "#/components/Button/Button"; -import { ConfirmDialog } from "#/components/Dialog/ConfirmDialog/ConfirmDialog"; -import { SearchField } from "#/components/SearchField/SearchField"; -import { paginateItems } from "#/utils/paginateItems"; -import { SpendSectionHeader } from "../SpendSectionHeader"; -import { GroupLimitDialog } from "./GroupLimitDialog"; -import { GROUP_LIMITS_PAGE_SIZE, GroupLimitsTable } from "./GroupLimitsTable"; - -export interface GroupLimitOverrideGroup { - group_id: string; - group_display_name: string; - group_name: string; - group_avatar_url: string; - member_count: number; -} - -export type GroupLimitOverride = ChatUsageLimitGroupOverride; - -interface GroupLimitsSectionProps { - hideHeader?: boolean; - groupOverrides: readonly GroupLimitOverride[]; - groupOrganizationNames?: Record; - showGroupForm: boolean; - onShowGroupFormChange: (show: boolean) => void; - selectedGroup: Group | null; - onSelectedGroupChange: (group: Group | null) => void; - groupAmount: string; - onGroupAmountChange: (amount: string) => void; - availableGroups: Group[]; - groupAutocompleteNoOptionsText: string; - groupsLoading: boolean; - editingGroupOverride: GroupLimitOverrideGroup | null; - onEditGroupOverride: (override: GroupLimitOverride) => void; - onAddGroupOverride: () => void; - onDeleteGroupOverride: (groupID: string) => void; - upsertPending: boolean; - upsertError: Error | null; - deletePending: boolean; - deleteError: Error | null; - groupsError: Error | null; -} - -export const GroupLimitsSection: FC = ({ - hideHeader, - groupOverrides, - groupOrganizationNames, - showGroupForm, - onShowGroupFormChange, - selectedGroup, - onSelectedGroupChange, - groupAmount, - onGroupAmountChange, - availableGroups, - groupAutocompleteNoOptionsText, - groupsLoading, - editingGroupOverride, - onEditGroupOverride, - onAddGroupOverride, - onDeleteGroupOverride, - upsertPending, - upsertError, - deletePending, - deleteError, - groupsError, -}) => { - const groupAutocompleteId = useId(); - const groupAmountId = useId(); - const isEditing = editingGroupOverride !== null; - const [pendingDeleteGroupId, setPendingDeleteGroupId] = useState< - string | null - >(null); - const [page, setPage] = useState(1); - const [searchQuery, setSearchQuery] = useState(""); - const normalizedSearchQuery = searchQuery.trim().toLowerCase(); - const filteredGroupOverrides = normalizedSearchQuery - ? groupOverrides.filter((override) => - [override.group_display_name, override.group_name].some((value) => - value.toLowerCase().includes(normalizedSearchQuery), - ), - ) - : groupOverrides; - const { pagedItems, clampedPage, hasPreviousPage, hasNextPage } = - paginateItems(filteredGroupOverrides, GROUP_LIMITS_PAGE_SIZE, page); - - return ( -
- {!hideHeader && ( - - )} -
-
- { - setSearchQuery(value); - setPage(1); - }} - placeholder="Search groups..." - aria-label="Search group limits" - className="w-full sm:max-w-md" - /> - -
- - - - {deleteError && ( -

- {getErrorMessage(deleteError, "Failed to delete group override.")} -

- )} - - - {groupsError && ( -

- {getErrorMessage(groupsError, "Failed to load groups.")} -

- )} -
- {pendingDeleteGroupId && ( - { - void onDeleteGroupOverride(pendingDeleteGroupId); - setPendingDeleteGroupId(null); - }} - confirmLoading={deletePending} - open - onClose={() => setPendingDeleteGroupId(null)} - /> - )} -
- ); -}; diff --git a/site/src/pages/AISettingsPage/SpendPage/components/LimitsTab/GroupLimitsTable.tsx b/site/src/pages/AISettingsPage/SpendPage/components/LimitsTab/GroupLimitsTable.tsx deleted file mode 100644 index 56ff608434..0000000000 --- a/site/src/pages/AISettingsPage/SpendPage/components/LimitsTab/GroupLimitsTable.tsx +++ /dev/null @@ -1,167 +0,0 @@ -import { EllipsisVerticalIcon, TrashIcon } from "lucide-react"; -import type { FC } from "react"; -import { Link } from "react-router"; -import { AvatarData } from "#/components/Avatar/AvatarData"; -import { Button } from "#/components/Button/Button"; -import { - DropdownMenu, - DropdownMenuContent, - DropdownMenuItem, - DropdownMenuTrigger, -} from "#/components/DropdownMenu/DropdownMenu"; -import { PaginationAmount } from "#/components/PaginationWidget/PaginationAmount"; -import { PaginationWidgetBase } from "#/components/PaginationWidget/PaginationWidgetBase"; -import { - Table, - TableBody, - TableCell, - TableHead, - TableHeader, - TableRow, -} from "#/components/Table/Table"; -import { formatCostMicros } from "#/utils/currency"; -import type { GroupLimitOverride } from "./GroupLimitsSection"; - -export const GROUP_LIMITS_PAGE_SIZE = 10; - -interface GroupLimitsTableProps { - pagedOverrides: readonly GroupLimitOverride[]; - totalOverrides: number; - clampedPage: number; - hasPreviousPage: boolean; - hasNextPage: boolean; - onPageChange: (page: number) => void; - onEditGroupOverride: (override: GroupLimitOverride) => void; - onRequestDelete: (groupID: string) => void; - groupOrganizationNames?: Record; - deletePending: boolean; - upsertPending: boolean; - isEditing: boolean; - emptyMessage: string; -} - -export const GroupLimitsTable: FC = ({ - pagedOverrides, - totalOverrides, - clampedPage, - hasPreviousPage, - hasNextPage, - onPageChange, - onEditGroupOverride, - onRequestDelete, - groupOrganizationNames, - deletePending, - upsertPending, - isEditing, - emptyMessage, -}) => { - if (totalOverrides === 0) { - return ( -
- {emptyMessage} -
- ); - } - - return ( - <> -
- - - - Group - Members - Spend limit - Actions - - - - {pagedOverrides.map((override) => { - const orgName = groupOrganizationNames?.[override.group_id]; - const groupAvatar = ( - - ); - - return ( - - - {orgName ? ( - - {groupAvatar} - - ) : ( - groupAvatar - )} - - {override.member_count} - - {override.spend_limit_micros !== null - ? formatCostMicros(override.spend_limit_micros) - : "Unlimited"} - - - - - - - - onEditGroupOverride(override)} - > - Update budget - - onRequestDelete(override.group_id)} - > - - Remove group limit - - - - - - ); - })} - -
- -
- {totalOverrides > GROUP_LIMITS_PAGE_SIZE && ( - - )} - - ); -}; diff --git a/site/src/pages/AISettingsPage/SpendPage/components/LimitsTab/GroupOverrideController.tsx b/site/src/pages/AISettingsPage/SpendPage/components/LimitsTab/GroupOverrideController.tsx deleted file mode 100644 index 2ccd1ae44f..0000000000 --- a/site/src/pages/AISettingsPage/SpendPage/components/LimitsTab/GroupOverrideController.tsx +++ /dev/null @@ -1,141 +0,0 @@ -import { type FC, type ReactNode, useState } from "react"; - -import type { - Group, - UpsertChatUsageLimitGroupOverrideRequest, -} from "#/api/typesGenerated"; -import { - dollarsToMicros, - isPositiveFiniteDollarAmount, - microsToDollars, -} from "#/utils/currency"; - -interface EditingGroupOverride { - group_id: string; - group_display_name: string; - group_name: string; - group_avatar_url: string; - member_count: number; -} - -type GroupOverrideChildProps = { - showGroupForm: boolean; - setShowGroupForm: (show: boolean) => void; - selectedGroup: Group | null; - setSelectedGroup: (group: Group | null) => void; - groupAmount: string; - setGroupAmount: (amount: string) => void; - editingGroupOverride: EditingGroupOverride | null; - setEditingGroupOverride: (override: EditingGroupOverride | null) => void; - handleShowGroupFormChange: (show: boolean) => void; - handleEditGroupOverride: ( - override: EditingGroupOverride & { - spend_limit_micros: number | null; - }, - ) => void; - handleAddGroupOverride: () => void; - existingGroupIds: Set; - availableGroups: Group[]; - groupAutocompleteNoOptionsText: string; -}; - -interface GroupOverrideControllerProps { - groupOverrides: ReadonlyArray<{ group_id: string }>; - groups: ReadonlyArray; - isLoadingGroups: boolean; - onUpsertGroupOverride: (args: { - groupID: string; - req: UpsertChatUsageLimitGroupOverrideRequest; - onSuccess: () => void; - }) => void; - children: (props: GroupOverrideChildProps) => ReactNode; -} - -export const GroupOverrideController: FC = ({ - groupOverrides, - groups, - isLoadingGroups, - onUpsertGroupOverride, - children, -}) => { - const [showGroupForm, setShowGroupForm] = useState(false); - const [selectedGroup, setSelectedGroup] = useState(null); - const [groupAmount, setGroupAmount] = useState(""); - const [editingGroupOverride, setEditingGroupOverride] = - useState(null); - - // Derived values. - const existingGroupIds = new Set(groupOverrides.map((g) => g.group_id)); - const availableGroups = groups.filter((g) => !existingGroupIds.has(g.id)); - const groupAutocompleteNoOptionsText = isLoadingGroups - ? "Loading groups..." - : groups.length === 0 - ? "No groups configured" - : availableGroups.length === 0 - ? "All groups already have overrides" - : "No groups available"; - - // Handlers. - const handleShowGroupFormChange = (show: boolean) => { - setShowGroupForm(show); - if (!show) { - setEditingGroupOverride(null); - } - }; - - const handleEditGroupOverride = ( - override: EditingGroupOverride & { - spend_limit_micros: number | null; - }, - ) => { - setEditingGroupOverride({ - group_id: override.group_id, - group_display_name: override.group_display_name, - group_name: override.group_name, - group_avatar_url: override.group_avatar_url, - member_count: override.member_count, - }); - setSelectedGroup(null); - setGroupAmount( - override.spend_limit_micros !== null - ? microsToDollars(override.spend_limit_micros).toString() - : "", - ); - setShowGroupForm(true); - }; - - const handleAddGroupOverride = () => { - const targetGroupID = editingGroupOverride?.group_id ?? selectedGroup?.id; - - if (!targetGroupID || !isPositiveFiniteDollarAmount(groupAmount)) { - return; - } - onUpsertGroupOverride({ - groupID: targetGroupID, - req: { spend_limit_micros: dollarsToMicros(groupAmount) }, - onSuccess: () => { - setEditingGroupOverride(null); - setSelectedGroup(null); - setGroupAmount(""); - setShowGroupForm(false); - }, - }); - }; - - return children({ - showGroupForm, - setShowGroupForm, - selectedGroup, - setSelectedGroup, - groupAmount, - setGroupAmount, - editingGroupOverride, - setEditingGroupOverride, - handleShowGroupFormChange, - handleEditGroupOverride, - handleAddGroupOverride, - existingGroupIds, - availableGroups, - groupAutocompleteNoOptionsText, - }); -}; diff --git a/site/src/pages/AISettingsPage/SpendPage/components/LimitsTab/UserOverrideController.tsx b/site/src/pages/AISettingsPage/SpendPage/components/LimitsTab/UserOverrideController.tsx deleted file mode 100644 index 99f406592a..0000000000 --- a/site/src/pages/AISettingsPage/SpendPage/components/LimitsTab/UserOverrideController.tsx +++ /dev/null @@ -1,131 +0,0 @@ -import { type FC, type ReactNode, useState } from "react"; - -import type { - UpsertChatUsageLimitOverrideRequest, - User, -} from "#/api/typesGenerated"; -import { - dollarsToMicros, - isPositiveFiniteDollarAmount, - microsToDollars, -} from "#/utils/currency"; - -interface EditingUserOverride { - user_id: string; - name: string; - username: string; - avatar_url: string; -} - -type UserOverrideChildProps = { - showUserForm: boolean; - setShowUserForm: (show: boolean) => void; - selectedUserOverride: User | null; - setSelectedUserOverride: (user: User | null) => void; - userOverrideAmount: string; - setUserOverrideAmount: (amount: string) => void; - editingUserOverride: EditingUserOverride | null; - setEditingUserOverride: (override: EditingUserOverride | null) => void; - handleShowUserFormChange: (show: boolean) => void; - handleEditUserOverride: ( - override: EditingUserOverride & { - spend_limit_micros: number | null; - }, - ) => void; - handleAddOverride: () => void; - existingUserIds: Set; - selectedUserAlreadyOverridden: boolean; -}; - -interface UserOverrideControllerProps { - overrides: ReadonlyArray<{ user_id: string }>; - onUpsertOverride: (args: { - userID: string; - req: UpsertChatUsageLimitOverrideRequest; - onSuccess: () => void; - }) => void; - children: (props: UserOverrideChildProps) => ReactNode; -} - -export const UserOverrideController: FC = ({ - overrides, - onUpsertOverride, - children, -}) => { - const [showUserForm, setShowUserForm] = useState(false); - const [selectedUserOverride, setSelectedUserOverride] = useState( - null, - ); - const [userOverrideAmount, setUserOverrideAmount] = useState(""); - const [editingUserOverride, setEditingUserOverride] = - useState(null); - - // Derived values. - const existingUserIds = new Set(overrides.map((o) => o.user_id)); - const selectedUserAlreadyOverridden = selectedUserOverride - ? existingUserIds.has(selectedUserOverride.id) - : false; - - // Handlers. - const handleShowUserFormChange = (show: boolean) => { - setShowUserForm(show); - if (!show) { - setEditingUserOverride(null); - } - }; - - const handleEditUserOverride = ( - override: EditingUserOverride & { - spend_limit_micros: number | null; - }, - ) => { - setEditingUserOverride({ - user_id: override.user_id, - name: override.name, - username: override.username, - avatar_url: override.avatar_url, - }); - setSelectedUserOverride(null); - setUserOverrideAmount( - override.spend_limit_micros !== null - ? microsToDollars(override.spend_limit_micros).toString() - : "", - ); - setShowUserForm(true); - }; - - const handleAddOverride = () => { - const targetUserID = - editingUserOverride?.user_id ?? selectedUserOverride?.id; - - if (!targetUserID || !isPositiveFiniteDollarAmount(userOverrideAmount)) { - return; - } - onUpsertOverride({ - userID: targetUserID, - req: { spend_limit_micros: dollarsToMicros(userOverrideAmount) }, - onSuccess: () => { - setEditingUserOverride(null); - setSelectedUserOverride(null); - setUserOverrideAmount(""); - setShowUserForm(false); - }, - }); - }; - - return children({ - showUserForm, - setShowUserForm, - selectedUserOverride, - setSelectedUserOverride, - userOverrideAmount, - setUserOverrideAmount, - editingUserOverride, - setEditingUserOverride, - handleShowUserFormChange, - handleEditUserOverride, - handleAddOverride, - existingUserIds, - selectedUserAlreadyOverridden, - }); -}; diff --git a/site/src/pages/AISettingsPage/SpendPage/components/LimitsTab/UserOverrideDialog.tsx b/site/src/pages/AISettingsPage/SpendPage/components/LimitsTab/UserOverrideDialog.tsx deleted file mode 100644 index 5226f10687..0000000000 --- a/site/src/pages/AISettingsPage/SpendPage/components/LimitsTab/UserOverrideDialog.tsx +++ /dev/null @@ -1,157 +0,0 @@ -import type { FC } from "react"; -import { getErrorMessage } from "#/api/errors"; -import type { User } from "#/api/typesGenerated"; -import { AvatarData } from "#/components/Avatar/AvatarData"; -import { Button } from "#/components/Button/Button"; -import { - Dialog, - DialogContent, - DialogDescription, - DialogFooter, - DialogHeader, - DialogTitle, -} from "#/components/Dialog/Dialog"; -import { Input } from "#/components/Input/Input"; -import { Label } from "#/components/Label/Label"; -import { Spinner } from "#/components/Spinner/Spinner"; -import { UserAutocomplete } from "#/components/UserAutocomplete/UserAutocomplete"; -import { isPositiveFiniteDollarAmount } from "#/utils/currency"; -import type { UserOverrideUser } from "./UserOverridesSection"; - -interface UserOverrideDialogProps { - open: boolean; - onOpenChange: (open: boolean) => void; - selectedUser: User | null; - onSelectedUserChange: (user: User | null) => void; - userOverrideAmount: string; - onUserOverrideAmountChange: (amount: string) => void; - selectedUserAlreadyOverridden: boolean; - editingUserOverride: UserOverrideUser | null; - upsertPending: boolean; - upsertError: Error | null; - onSave: () => void; - amountInputId: string; -} - -export const UserOverrideDialog: FC = ({ - open, - onOpenChange, - selectedUser, - onSelectedUserChange, - userOverrideAmount, - onUserOverrideAmountChange, - selectedUserAlreadyOverridden, - editingUserOverride, - upsertPending, - upsertError, - onSave, - amountInputId, -}) => { - const isEditing = editingUserOverride !== null; - const saveDisabled = isEditing - ? upsertPending || !isPositiveFiniteDollarAmount(userOverrideAmount) - : upsertPending || - !selectedUser || - selectedUserAlreadyOverridden || - !isPositiveFiniteDollarAmount(userOverrideAmount); - - return ( - { - if (upsertPending) { - return; - } - - onOpenChange(nextOpen); - if (!nextOpen) { - onSelectedUserChange(null); - onUserOverrideAmountChange(""); - } - }} - > - - - - {isEditing ? "Update user budget" : "Add user budget"} - - - {isEditing - ? "Update this user's spend limit override." - : "Set a spend limit override for a specific user."} - - - -
-
- {editingUserOverride ? ( - <> - -
- -
- - ) : ( - - )} -
-
- - - onUserOverrideAmountChange(event.target.value) - } - placeholder="0.00" - /> -
- {!isEditing && selectedUserAlreadyOverridden && ( -

- This user already has an override. -

- )} - {upsertError && ( -

- {getErrorMessage(upsertError, "Failed to save the override.")} -

- )} -
- - - - - -
-
- ); -}; diff --git a/site/src/pages/AISettingsPage/SpendPage/components/LimitsTab/UserOverridesSection.stories.tsx b/site/src/pages/AISettingsPage/SpendPage/components/LimitsTab/UserOverridesSection.stories.tsx deleted file mode 100644 index 7486a749c9..0000000000 --- a/site/src/pages/AISettingsPage/SpendPage/components/LimitsTab/UserOverridesSection.stories.tsx +++ /dev/null @@ -1,155 +0,0 @@ -import type { Meta, StoryObj } from "@storybook/react-vite"; -import { QueryClient, QueryClientProvider } from "react-query"; -import { expect, fn, userEvent, waitFor, within } from "storybook/test"; -import { UserOverridesSection } from "./UserOverridesSection"; - -const queryClient = new QueryClient({ - defaultOptions: { - queries: { - retry: false, - gcTime: 0, - refetchOnWindowFocus: false, - }, - }, -}); - -const mockOverrides = [ - { - user_id: "user-1", - name: "Alice Johnson", - username: "alice", - avatar_url: "", - spend_limit_micros: 5_000_000, - }, - { - user_id: "user-2", - name: "Bob Smith", - username: "bob", - avatar_url: "", - spend_limit_micros: 10_000_000, - }, - { - user_id: "user-3", - name: "Charlie Davis", - username: "charlie", - avatar_url: "", - spend_limit_micros: null, - }, -]; - -const meta: Meta = { - title: "pages/AISettingsPage/SpendPage/LimitsTab/UserOverridesSection", - component: UserOverridesSection, - args: { - overrides: mockOverrides, - showUserForm: false, - onShowUserFormChange: fn(), - selectedUser: null, - onSelectedUserChange: fn(), - userOverrideAmount: "", - onUserOverrideAmountChange: fn(), - selectedUserAlreadyOverridden: false, - editingUserOverride: null, - onEditUserOverride: fn(), - onAddOverride: fn(), - onDeleteOverride: fn(), - upsertPending: false, - upsertError: null, - deletePending: false, - deleteError: null, - }, - decorators: [ - (Story) => ( - - - - ), - ], -}; - -export default meta; -type Story = StoryObj; - -export const Default: Story = {}; - -export const EmptyState: Story = { - args: { - overrides: [], - }, -}; - -export const AddForm: Story = { - args: { - showUserForm: true, - overrides: mockOverrides, - }, -}; - -export const EditForm: Story = { - args: { - showUserForm: true, - editingUserOverride: { - user_id: mockOverrides[0].user_id, - name: mockOverrides[0].name, - username: mockOverrides[0].username, - avatar_url: mockOverrides[0].avatar_url, - }, - userOverrideAmount: "5.00", - }, - play: async ({ canvasElement }) => { - const body = within(canvasElement.ownerDocument.body); - // Update budget button confirms edit mode is active. - await expect( - body.getByRole("button", { name: /update budget/i }), - ).toBeInTheDocument(); - // The editing user name appears in both the table row and the - // read-only edit form identity, confirming it was populated. - const nameElements = body.getAllByText("Alice Johnson"); - expect(nameElements.length).toBeGreaterThanOrEqual(2); - }, -}; - -export const DeleteUserOverride: Story = { - play: async ({ canvasElement }) => { - const body = within(canvasElement.ownerDocument.body); - - const actionButtons = await body.findAllByRole("button", { - name: /actions for/i, - }); - await userEvent.click(actionButtons[0]); - await userEvent.click( - await body.findByRole("menuitem", { name: /remove override/i }), - ); - - // The confirmation dialog should appear. - const dialog = await body.findByRole("dialog"); - await expect(dialog).toBeInTheDocument(); - await expect( - body.getByText(/Are you sure you want to delete this user override/i), - ).toBeInTheDocument(); - }, -}; - -export const DeleteUserOverrideCancelled: Story = { - play: async ({ canvasElement, args }) => { - const body = within(canvasElement.ownerDocument.body); - - const actionButtons = await body.findAllByRole("button", { - name: /actions for/i, - }); - await userEvent.click(actionButtons[0]); - await userEvent.click( - await body.findByRole("menuitem", { name: /remove override/i }), - ); - - // Cancel the dialog. - await body.findByRole("dialog"); - await userEvent.click(body.getByRole("button", { name: "Cancel" })); - - // The dialog should be closed and the callback should not have been called. - await waitFor(() => { - expect(body.queryByRole("dialog")).not.toBeInTheDocument(); - }); - expect(args.onDeleteOverride).not.toHaveBeenCalled(); - }, -}; diff --git a/site/src/pages/AISettingsPage/SpendPage/components/LimitsTab/UserOverridesSection.tsx b/site/src/pages/AISettingsPage/SpendPage/components/LimitsTab/UserOverridesSection.tsx deleted file mode 100644 index 882f4df2ad..0000000000 --- a/site/src/pages/AISettingsPage/SpendPage/components/LimitsTab/UserOverridesSection.tsx +++ /dev/null @@ -1,173 +0,0 @@ -import { PlusIcon } from "lucide-react"; -import { type FC, useId, useState } from "react"; -import { getErrorMessage } from "#/api/errors"; -import type { User } from "#/api/typesGenerated"; -import { Button } from "#/components/Button/Button"; -import { ConfirmDialog } from "#/components/Dialog/ConfirmDialog/ConfirmDialog"; -import { SearchField } from "#/components/SearchField/SearchField"; -import { paginateItems } from "#/utils/paginateItems"; -import { SpendSectionHeader } from "../SpendSectionHeader"; -import { UserOverrideDialog } from "./UserOverrideDialog"; -import { - USER_OVERRIDES_PAGE_SIZE, - UserOverridesTable, -} from "./UserOverridesTable"; - -export interface UserOverrideUser { - user_id: string; - name: string; - username: string; - avatar_url: string; -} - -export interface UserOverride extends UserOverrideUser { - spend_limit_micros: number | null; -} - -interface UserOverridesSectionProps { - hideHeader?: boolean; - overrides: readonly UserOverride[]; - showUserForm: boolean; - onShowUserFormChange: (show: boolean) => void; - selectedUser: User | null; - onSelectedUserChange: (user: User | null) => void; - userOverrideAmount: string; - onUserOverrideAmountChange: (amount: string) => void; - selectedUserAlreadyOverridden: boolean; - editingUserOverride: UserOverrideUser | null; - onEditUserOverride: (override: UserOverride) => void; - onAddOverride: () => void; - onDeleteOverride: (userID: string) => void; - upsertPending: boolean; - upsertError: Error | null; - deletePending: boolean; - deleteError: Error | null; -} - -export const UserOverridesSection: FC = ({ - hideHeader, - overrides, - showUserForm, - onShowUserFormChange, - selectedUser, - onSelectedUserChange, - userOverrideAmount, - onUserOverrideAmountChange, - selectedUserAlreadyOverridden, - editingUserOverride, - onEditUserOverride, - onAddOverride, - onDeleteOverride, - upsertPending, - upsertError, - deletePending, - deleteError, -}) => { - const userOverrideAmountId = useId(); - const isEditing = editingUserOverride !== null; - const [pendingDeleteUserId, setPendingDeleteUserId] = useState( - null, - ); - const [page, setPage] = useState(1); - const [searchQuery, setSearchQuery] = useState(""); - const normalizedSearchQuery = searchQuery.trim().toLowerCase(); - const filteredOverrides = normalizedSearchQuery - ? overrides.filter((override) => - [override.name, override.username].some((value) => - value.toLowerCase().includes(normalizedSearchQuery), - ), - ) - : overrides; - const { pagedItems, clampedPage, hasPreviousPage, hasNextPage } = - paginateItems(filteredOverrides, USER_OVERRIDES_PAGE_SIZE, page); - - return ( -
- {!hideHeader && ( - - )} -
-
- { - setSearchQuery(value); - setPage(1); - }} - placeholder="Search users..." - aria-label="Search user overrides" - className="w-full sm:max-w-md" - /> - -
- - - - {deleteError && ( -

- {getErrorMessage(deleteError, "Failed to delete override.")} -

- )} - - -
- {pendingDeleteUserId && ( - { - void onDeleteOverride(pendingDeleteUserId); - setPendingDeleteUserId(null); - }} - confirmLoading={deletePending} - open - onClose={() => setPendingDeleteUserId(null)} - /> - )} -
- ); -}; diff --git a/site/src/pages/AISettingsPage/SpendPage/components/LimitsTab/UserOverridesTable.tsx b/site/src/pages/AISettingsPage/SpendPage/components/LimitsTab/UserOverridesTable.tsx deleted file mode 100644 index 1ab125f2c4..0000000000 --- a/site/src/pages/AISettingsPage/SpendPage/components/LimitsTab/UserOverridesTable.tsx +++ /dev/null @@ -1,144 +0,0 @@ -import { EllipsisVerticalIcon, TrashIcon } from "lucide-react"; -import type { FC } from "react"; -import { AvatarData } from "#/components/Avatar/AvatarData"; -import { Button } from "#/components/Button/Button"; -import { - DropdownMenu, - DropdownMenuContent, - DropdownMenuItem, - DropdownMenuTrigger, -} from "#/components/DropdownMenu/DropdownMenu"; -import { PaginationAmount } from "#/components/PaginationWidget/PaginationAmount"; -import { PaginationWidgetBase } from "#/components/PaginationWidget/PaginationWidgetBase"; -import { - Table, - TableBody, - TableCell, - TableHead, - TableHeader, - TableRow, -} from "#/components/Table/Table"; -import { formatCostMicros } from "#/utils/currency"; -import type { UserOverride } from "./UserOverridesSection"; - -export const USER_OVERRIDES_PAGE_SIZE = 10; - -interface UserOverridesTableProps { - pagedOverrides: readonly UserOverride[]; - totalOverrides: number; - clampedPage: number; - hasPreviousPage: boolean; - hasNextPage: boolean; - onPageChange: (page: number) => void; - onEditUserOverride: (override: UserOverride) => void; - onRequestDelete: (userID: string) => void; - deletePending: boolean; - upsertPending: boolean; - isEditing: boolean; - emptyMessage: string; -} - -export const UserOverridesTable: FC = ({ - pagedOverrides, - totalOverrides, - clampedPage, - hasPreviousPage, - hasNextPage, - onPageChange, - onEditUserOverride, - onRequestDelete, - deletePending, - upsertPending, - isEditing, - emptyMessage, -}) => { - if (totalOverrides === 0) { - return ( -
- {emptyMessage} -
- ); - } - - return ( - <> -
- - - - User - Spend limit - Actions - - - - {pagedOverrides.map((override) => ( - - - - - - {override.spend_limit_micros !== null - ? formatCostMicros(override.spend_limit_micros) - : "Unlimited"} - - - - - - - - onEditUserOverride(override)} - > - Update budget - - onRequestDelete(override.user_id)} - > - - Remove override - - - - - - ))} - -
- -
- {totalOverrides > USER_OVERRIDES_PAGE_SIZE && ( - - )} - - ); -}; diff --git a/site/src/pages/AISettingsPage/SpendPage/components/LimitsTab/limitsFormLogic.test.ts b/site/src/pages/AISettingsPage/SpendPage/components/LimitsTab/limitsFormLogic.test.ts deleted file mode 100644 index 97307cee08..0000000000 --- a/site/src/pages/AISettingsPage/SpendPage/components/LimitsTab/limitsFormLogic.test.ts +++ /dev/null @@ -1,10 +0,0 @@ -import { normalizeChatUsageLimitPeriod } from "./limitsFormLogic"; - -describe("limitsFormLogic", () => { - describe("normalizeChatUsageLimitPeriod", () => { - it("defaults invalid periods to month", () => { - expect(normalizeChatUsageLimitPeriod("year")).toBe("month"); - expect(normalizeChatUsageLimitPeriod(undefined)).toBe("month"); - }); - }); -}); diff --git a/site/src/pages/AISettingsPage/SpendPage/components/LimitsTab/limitsFormLogic.ts b/site/src/pages/AISettingsPage/SpendPage/components/LimitsTab/limitsFormLogic.ts deleted file mode 100644 index 63da616d9c..0000000000 --- a/site/src/pages/AISettingsPage/SpendPage/components/LimitsTab/limitsFormLogic.ts +++ /dev/null @@ -1,14 +0,0 @@ -import type { ChatUsageLimitPeriod } from "#/api/typesGenerated"; - -export const normalizeChatUsageLimitPeriod = ( - period: string | null | undefined, -): ChatUsageLimitPeriod => { - switch (period) { - case "day": - case "week": - case "month": - return period; - default: - return "month"; - } -}; diff --git a/site/src/pages/AISettingsPage/SpendPage/components/UsageTab/UsageTab.tsx b/site/src/pages/AISettingsPage/SpendPage/components/UsageTab/UsageTab.tsx index c3d5f63460..676a89614e 100644 --- a/site/src/pages/AISettingsPage/SpendPage/components/UsageTab/UsageTab.tsx +++ b/site/src/pages/AISettingsPage/SpendPage/components/UsageTab/UsageTab.tsx @@ -1,4 +1,3 @@ -import { EllipsisVerticalIcon } from "lucide-react"; import type { FC } from "react"; import type * as TypesGen from "#/api/typesGenerated"; import { ErrorAlert } from "#/components/Alert/ErrorAlert"; @@ -8,12 +7,6 @@ import { DateRangePicker, type DateRangeValue, } from "#/components/DateRangePicker/DateRangePicker"; -import { - DropdownMenu, - DropdownMenuContent, - DropdownMenuItem, - DropdownMenuTrigger, -} from "#/components/DropdownMenu/DropdownMenu"; import { PaginationContainer, type PaginationResult, @@ -38,14 +31,6 @@ import { formatTokenCount } from "#/utils/analytics"; import { formatCostMicros } from "#/utils/currency"; import { SpendSectionHeader } from "../SpendSectionHeader"; -type UsageUserOverride = { - user_id: string; - name: string; - username: string; - avatar_url: string; - spend_limit_micros: number | null; -}; - interface UsageTabProps { displayDateRange: DateRangeValue; onDateRangeChange: (value: DateRangeValue) => void; @@ -58,9 +43,7 @@ interface UsageTabProps { error: unknown; refetch: () => unknown; }; - overrides: readonly UsageUserOverride[]; onSelectUser: (user: TypesGen.ChatCostUserRollup) => void; - onEditBudget: (override: UsageUserOverride) => void; } export const UsageTab: FC = ({ @@ -69,9 +52,7 @@ export const UsageTab: FC = ({ searchFilter, onSearchFilterChange, usersQuery, - overrides, onSelectUser, - onEditBudget, }) => { return (
@@ -146,7 +127,6 @@ export const UsageTab: FC = ({ Output Cache Read Cache Write - Actions @@ -155,18 +135,6 @@ export const UsageTab: FC = ({ key={user.user_id} user={user} onSelect={onSelectUser} - onEditBudget={(selectedUser) => { - const override = overrides.find( - (o) => o.user_id === selectedUser.user_id, - ) ?? { - user_id: selectedUser.user_id, - name: selectedUser.name, - username: selectedUser.username, - avatar_url: selectedUser.avatar_url, - spend_limit_micros: null, - }; - onEditBudget(override); - }} /> ))} @@ -183,8 +151,7 @@ export const UsageTab: FC = ({ const UserRow: FC<{ user: TypesGen.ChatCostUserRollup; onSelect: (user: TypesGen.ChatCostUserRollup) => void; - onEditBudget: (user: TypesGen.ChatCostUserRollup) => void; -}> = ({ user, onSelect, onEditBudget }) => { +}> = ({ user, onSelect }) => { const clickableRowProps = useClickableTableRow({ onClick: () => onSelect(user), }); @@ -237,28 +204,6 @@ const UserRow: FC<{ {formatTokenCount(user.total_cache_creation_tokens)} - event.stopPropagation()}> - - - - - - onEditBudget(user)}> - Update budget - - onSelect(user)}> - View spend details - - - - ); }; diff --git a/site/src/pages/AgentsPage/AgentChatPage.test.ts b/site/src/pages/AgentsPage/AgentChatPage.test.ts index 294678d1b0..c608a91bb2 100644 --- a/site/src/pages/AgentsPage/AgentChatPage.test.ts +++ b/site/src/pages/AgentsPage/AgentChatPage.test.ts @@ -237,7 +237,7 @@ describe("runPromoteQueuedMessage", () => { const promote = vi.fn(async (_id: number) => undefined); const clearChatErrorReason = vi.fn(); - const handleUsageLimitError = vi.fn(); + const onError = vi.fn(); await runPromoteQueuedMessage({ id: b.id, @@ -245,7 +245,7 @@ describe("runPromoteQueuedMessage", () => { promoteQueuedMessage: promote, agentId: "chat-1", clearChatErrorReason, - handleUsageLimitError, + onError, }); expect(promote).toHaveBeenCalledWith(b.id); @@ -268,7 +268,7 @@ describe("runPromoteQueuedMessage", () => { throw apiError; }); const clearChatErrorReason = vi.fn(); - const handleUsageLimitError = vi.fn(); + const onError = vi.fn(); await expect( runPromoteQueuedMessage({ @@ -277,11 +277,11 @@ describe("runPromoteQueuedMessage", () => { promoteQueuedMessage: promote, agentId: "chat-1", clearChatErrorReason, - handleUsageLimitError, + onError, }), ).rejects.toBe(apiError); - expect(handleUsageLimitError).toHaveBeenCalledWith(apiError); + expect(onError).toHaveBeenCalledWith(apiError); const snapshot = store.getSnapshot(); expect(snapshot.queuedMessages.map((m) => m.id)).toEqual([a.id, b.id]); diff --git a/site/src/pages/AgentsPage/AgentChatPage.tsx b/site/src/pages/AgentsPage/AgentChatPage.tsx index c7b526579e..7a4ef233e7 100644 --- a/site/src/pages/AgentsPage/AgentChatPage.tsx +++ b/site/src/pages/AgentsPage/AgentChatPage.tsx @@ -75,7 +75,12 @@ import { } from "./AgentChatPageView"; import type { AgentsPageOutletContext } from "./AgentsPageLayout"; import type { ChatMessageInputRef } from "./components/AgentChatInput"; -import { normalizeChatErrorPayload } from "./components/ChatConversation/chatError"; +import { + type ChatDetailError, + isChatHookDeniedResponse, + isChatHookDispatchFailedResponse, + normalizeChatErrorPayload, +} from "./components/ChatConversation/chatError"; import { getParentChatID, getWorkspaceAgent, @@ -114,13 +119,6 @@ import { chatSlashCommandTriggerText, resolveChatSlashCommandAvailability, } from "./utils/slashCommands"; -import { - type ChatDetailError, - formatUsageLimitMessage, - isChatHookDeniedResponse, - isChatHookDispatchFailedResponse, - isChatUsageLimitExceededResponse, -} from "./utils/usageLimitMessage"; /** localStorage key controlling whether the right panel is visible. */ export const RIGHT_PANEL_OPEN_KEY = "agents.right-panel-open"; @@ -203,7 +201,7 @@ export const runPromoteQueuedMessage = async (params: { promoteQueuedMessage: (id: number) => Promise; agentId: string | undefined; clearChatErrorReason: (chatID: string) => void; - handleUsageLimitError: (error: unknown) => void; + onError: (error: unknown) => void; }): Promise => { const { id, @@ -211,7 +209,7 @@ export const runPromoteQueuedMessage = async (params: { promoteQueuedMessage, agentId, clearChatErrorReason, - handleUsageLimitError, + onError, } = params; const previousSnapshot = store.getSnapshot(); store.batch(() => { @@ -231,7 +229,7 @@ export const runPromoteQueuedMessage = async (params: { } catch (error) { store.unsuppressQueuedMessageID(id); restoreOptimisticRequestSnapshot(store, previousSnapshot); - handleUsageLimitError(error); + onError(error); throw error; } }; @@ -740,9 +738,6 @@ const getPersistedDetailError = ({ chatRecord: TypesGen.Chat | undefined; cachedError: ChatDetailError | undefined; }): ChatDetailError | undefined => { - if (cachedError?.kind === "usage_limit") { - return cachedError; - } if (chatStatus !== "error") { return undefined; } @@ -1350,36 +1345,23 @@ const AgentChatPage: FC = () => { ); }; - const handleUsageLimitError = (error: unknown): void => { - if (!agentId) { + const handleRequestError = (error: unknown): void => { + if (!agentId || !isApiError(error)) { return; } - if ( - isApiError(error) && - error.response?.status === 409 && - isChatUsageLimitExceededResponse(error.response.data) - ) { - const reason: ChatDetailError = { - kind: "usage_limit", - message: formatUsageLimitMessage(error.response.data), - }; - store.setStreamError(reason); - setChatErrorReason(agentId, reason); - } else if (isApiError(error)) { - const detail = error.response?.data?.detail?.trim() || undefined; - const kind = isChatHookDeniedResponse(error.response?.data) - ? "hook_denied" - : isChatHookDispatchFailedResponse(error.response?.data) - ? "hook_dispatch_failed" - : "generic"; - const reason: ChatDetailError = { - kind, - message: getErrorMessage(error, "An unexpected error occurred."), - ...(detail ? { detail } : {}), - }; - store.setStreamError(reason); - setChatErrorReason(agentId, reason); - } + const detail = error.response?.data?.detail?.trim() || undefined; + const kind = isChatHookDeniedResponse(error.response?.data) + ? "hook_denied" + : isChatHookDispatchFailedResponse(error.response?.data) + ? "hook_dispatch_failed" + : "generic"; + const reason: ChatDetailError = { + kind, + message: getErrorMessage(error, "An unexpected error occurred."), + ...(detail ? { detail } : {}), + }; + store.setStreamError(reason); + setChatErrorReason(agentId, reason); }; const handleInterrupt = () => { @@ -1422,7 +1404,7 @@ const AgentChatPage: FC = () => { promoteQueuedMessage, agentId, clearChatErrorReason, - handleUsageLimitError, + onError: handleRequestError, }); const editing = useConversationEditingState({ @@ -1648,15 +1630,7 @@ const AgentChatPage: FC = () => { await compact(); } catch (error) { restoreOptimisticRequestSnapshot(store, previousSnapshot); - if ( - isApiError(error) && - error.response?.status === 409 && - isChatUsageLimitExceededResponse(error.response.data) - ) { - handleUsageLimitError(error); - } else { - toast.error(getErrorMessage(error, "Failed to compact chat.")); - } + toast.error(getErrorMessage(error, "Failed to compact chat.")); throw error; } return; @@ -1714,7 +1688,7 @@ const AgentChatPage: FC = () => { scrollToBottom: scrollToBottomRef.current, onError: (error) => { restoreOptimisticRequestSnapshot(store, previousSnapshot); - handleUsageLimitError(error); + handleRequestError(error); // Hook dispatch failures can park an idle chat in error before returning the request error. acceptServerChatStatus(); void queryClient.invalidateQueries({ @@ -1765,7 +1739,7 @@ const AgentChatPage: FC = () => { try { response = await sendMessage(request); } catch (error) { - handleUsageLimitError(error); + handleRequestError(error); // Hook dispatch failures can park an idle chat in error before returning the request error. acceptServerChatStatus(); void queryClient.invalidateQueries({ diff --git a/site/src/pages/AgentsPage/AgentChatPageView.stories.tsx b/site/src/pages/AgentsPage/AgentChatPageView.stories.tsx index f8affb2562..3381ad9012 100644 --- a/site/src/pages/AgentsPage/AgentChatPageView.stories.tsx +++ b/site/src/pages/AgentsPage/AgentChatPageView.stories.tsx @@ -26,13 +26,13 @@ import { AgentChatPageNotFoundView, AgentChatPageView, } from "./AgentChatPageView"; +import type { ChatDetailError } from "./components/ChatConversation/chatError"; import { createChatStore, useChatSelector, } from "./components/ChatConversation/chatStore"; import type { ModelSelectorOption } from "./components/ChatElements"; import { lastActiveSidebarTabStorageKeyPrefix } from "./utils/sidebarTabStorage"; -import type { ChatDetailError } from "./utils/usageLimitMessage"; // --------------------------------------------------------------------------- // Shared constants & helpers diff --git a/site/src/pages/AgentsPage/AgentChatPageView.tsx b/site/src/pages/AgentsPage/AgentChatPageView.tsx index 4bb0b50221..bb772d3dba 100644 --- a/site/src/pages/AgentsPage/AgentChatPageView.tsx +++ b/site/src/pages/AgentsPage/AgentChatPageView.tsx @@ -33,6 +33,7 @@ import { ChatConversationSkeleton, RightPanelSkeleton, } from "./components/AgentsSkeletons"; +import type { ChatDetailError } from "./components/ChatConversation/chatError"; import type { useChatStore } from "./components/ChatConversation/chatStore"; import type { ModelSelectorOption } from "./components/ChatElements"; import { DesktopPanelContext } from "./components/ChatElements/tools/DesktopPanelContext"; @@ -70,7 +71,6 @@ import { getPersistedSidebarTabId, savePersistedSidebarTabId, } from "./utils/sidebarTabStorage"; -import type { ChatDetailError } from "./utils/usageLimitMessage"; type ChatStoreHandle = ReturnType["store"]; diff --git a/site/src/pages/AgentsPage/AgentEmbedPage.tsx b/site/src/pages/AgentsPage/AgentEmbedPage.tsx index 37b7018271..658c363d30 100644 --- a/site/src/pages/AgentsPage/AgentEmbedPage.tsx +++ b/site/src/pages/AgentsPage/AgentEmbedPage.tsx @@ -15,14 +15,14 @@ import { isConcreteThemeName, } from "#/theme"; import type { AgentsPageOutletContext } from "./AgentsPageLayout"; +import { + type ChatDetailError, + chatDetailErrorsEqual, +} from "./components/ChatConversation/chatError"; import { bootstrapChatEmbedSession, EmbedContext, } from "./components/EmbedContext"; -import { - type ChatDetailError, - chatDetailErrorsEqual, -} from "./utils/usageLimitMessage"; type BootstrapMessage = { type: "coder:vscode-auth-bootstrap"; diff --git a/site/src/pages/AgentsPage/AgentsPageLayout.stories.tsx b/site/src/pages/AgentsPage/AgentsPageLayout.stories.tsx index 1ff2a42f2f..3a6cd0913f 100644 --- a/site/src/pages/AgentsPage/AgentsPageLayout.stories.tsx +++ b/site/src/pages/AgentsPage/AgentsPageLayout.stories.tsx @@ -500,14 +500,7 @@ const meta: Meta = { retention_days: 30, }); spyOn(API.experimental, "updateChatRetentionDays").mockResolvedValue(); - spyOn(API.experimental, "getChatUsageLimitConfig").mockResolvedValue({ - spend_limit_micros: null, - period: "month", - updated_at: "2026-02-18T00:00:00.000Z", - unpriced_model_count: 0, - overrides: [], - group_overrides: [], - }); + spyOn(API, "getGroups").mockResolvedValue([]); spyOn(API.experimental, "getChatCostUsers").mockResolvedValue({ start_date: "2026-02-10T00:00:00Z", diff --git a/site/src/pages/AgentsPage/AgentsPageLayout.tsx b/site/src/pages/AgentsPage/AgentsPageLayout.tsx index c8266076eb..c5fac2c055 100644 --- a/site/src/pages/AgentsPage/AgentsPageLayout.tsx +++ b/site/src/pages/AgentsPage/AgentsPageLayout.tsx @@ -59,6 +59,10 @@ import { cn } from "#/utils/cn"; import { pageTitle } from "#/utils/page"; import { createReconnectingWebSocket } from "#/utils/reconnectingWebSocket"; import { emptyInputStorageKey } from "./components/AgentCreateForm"; +import { + type ChatDetailError, + chatDetailErrorsEqual, +} from "./components/ChatConversation/chatError"; import { getChatCostTreeID } from "./components/ChatConversation/chatHelpers"; import { isActiveChatStatus } from "./components/ChatConversation/chatStore"; import { @@ -85,10 +89,6 @@ import { } from "./utils/modelOptions"; import { clearPersistedRightPanelState } from "./utils/rightPanelTabStorage"; import { clearPersistedSidebarTabId } from "./utils/sidebarTabStorage"; -import { - type ChatDetailError, - chatDetailErrorsEqual, -} from "./utils/usageLimitMessage"; export interface AgentsPageOutletContext { chatErrorReasons: Record; diff --git a/site/src/pages/AgentsPage/components/AgentCreateForm.stories.tsx b/site/src/pages/AgentsPage/components/AgentCreateForm.stories.tsx index 886cd3475d..5f31532e46 100644 --- a/site/src/pages/AgentsPage/components/AgentCreateForm.stories.tsx +++ b/site/src/pages/AgentsPage/components/AgentCreateForm.stories.tsx @@ -790,32 +790,6 @@ export const PreservesAttachmentsOnFailedSend: Story = { }, }; -export const UsageLimitExceeded: Story = { - args: { - ...defaultArgs, - createError: Object.assign( - new Error("Request failed with status code 409"), - { - isAxiosError: true, - response: { - status: 409, - statusText: "Conflict", - data: { - message: "Chat usage limit exceeded.", - spent_micros: 900_000, - limit_micros: 500_000, - resets_at: "2026-03-16T00:00:00Z", - }, - headers: {}, - config: {}, - }, - config: {}, - toJSON: () => ({}), - }, - ), - }, -}; - export const HookDispatchFailed: Story = { args: { ...defaultArgs, diff --git a/site/src/pages/AgentsPage/components/AgentCreateForm.tsx b/site/src/pages/AgentsPage/components/AgentCreateForm.tsx index cf8b15c4b6..da796a80b5 100644 --- a/site/src/pages/AgentsPage/components/AgentCreateForm.tsx +++ b/site/src/pages/AgentsPage/components/AgentCreateForm.tsx @@ -1,6 +1,5 @@ import { type FC, useEffect, useEffectEvent, useRef, useState } from "react"; import { useQuery } from "react-query"; -import { Link } from "react-router"; import { toast } from "sonner"; import { isApiError } from "#/api/errors"; import { permittedOrganizations } from "#/api/queries/organizations"; @@ -8,7 +7,6 @@ import type * as TypesGen from "#/api/typesGenerated"; import type { AgentChatSendShortcut } from "#/api/typesGenerated"; import { Alert, AlertDescription, AlertTitle } from "#/components/Alert/Alert"; import { ErrorAlert } from "#/components/Alert/ErrorAlert"; -import { Button } from "#/components/Button/Button"; import { ConfirmDialog } from "#/components/Dialog/ConfirmDialog/ConfirmDialog"; import { useDashboard } from "#/modules/dashboard/useDashboard"; import { docs } from "#/utils/docs"; @@ -25,14 +23,12 @@ import { pickReasoningEffort, saveReasoningEffortForModel, } from "../utils/reasoningEffort"; -import { - formatUsageLimitMessage, - isChatHookDeniedResponse, - isChatHookDispatchFailedResponse, - isChatUsageLimitExceededResponse, -} from "../utils/usageLimitMessage"; import { AgentChatInput } from "./AgentChatInput"; import { ChatAccessDeniedAlert } from "./ChatAccessDeniedAlert"; +import { + isChatHookDeniedResponse, + isChatHookDispatchFailedResponse, +} from "./ChatConversation/chatError"; import { getErrorTitle } from "./ChatConversation/chatStatusHelpers"; import type { ModelSelectorOption } from "./ChatElements"; import { CompactOrgSelector } from "./ChatElements"; @@ -516,23 +512,8 @@ export const AgentCreateForm: FC = ({ ) : createError ? ( isApiError(createError) && - createError.response?.status === 409 && - isChatUsageLimitExceededResponse(createError.response.data) ? ( - - View usage - - } - > - - {formatUsageLimitMessage(createError.response.data)} - - - ) : isApiError(createError) && - createError.response.status === 502 && - isChatHookDispatchFailedResponse(createError.response.data) ? ( + createError.response.status === 502 && + isChatHookDispatchFailedResponse(createError.response.data) ? ( {getErrorTitle("hook_dispatch_failed", "error")} diff --git a/site/src/pages/AgentsPage/components/ChatConversation/LiveStreamTail.stories.tsx b/site/src/pages/AgentsPage/components/ChatConversation/LiveStreamTail.stories.tsx index 0cf96976d4..269b5d0e5e 100644 --- a/site/src/pages/AgentsPage/components/ChatConversation/LiveStreamTail.stories.tsx +++ b/site/src/pages/AgentsPage/components/ChatConversation/LiveStreamTail.stories.tsx @@ -40,32 +40,28 @@ export const EmptyConversationPrompt: Story = { }, }; -/** Usage-limit failures replace the idle prompt with the analytics CTA. */ export const UsageLimitExceeded: Story = { args: { ...defaultArgs, liveStatus: buildLiveStatus({ persistedError: { kind: "usage_limit", - message: - "You've used $50.00 of your $50.00 spend limit. Your limit resets on July 1, 2025.", + message: "Your AI spend budget has been reached.", }, }), }, play: async ({ canvasElement }) => { const canvas = within(canvasElement); - expect(canvas.getByText(/spend limit/i)).toBeVisible(); - const link = canvas.getByRole("link", { name: /view usage/i }); - expect(link).toBeVisible(); - expect(link).toHaveAttribute("href", "/agents/analytics"); + expect( + canvas.getByRole("heading", { name: /usage limit reached/i }), + ).toBeVisible(); + expect(canvas.getByText(/ai spend budget has been reached/i)).toBeVisible(); + expect( + canvas.queryByRole("link", { name: /view usage/i }), + ).not.toBeInTheDocument(); }, }; -/** - * Provider quota errors use the standard ChatStatusCallout instead of the - * "View usage" CTA (which links to Coder's analytics, not the provider's - * billing page). - */ export const ProviderQuotaExceeded: Story = { args: { ...defaultArgs, @@ -87,7 +83,6 @@ export const ProviderQuotaExceeded: Story = { expect( canvas.queryByRole("link", { name: /view usage/i }), ).not.toBeInTheDocument(); - // Should render ChatStatusCallout instead. expect( canvas.getByRole("heading", { name: /usage limit reached/i }), ).toBeVisible(); diff --git a/site/src/pages/AgentsPage/components/ChatConversation/LiveStreamTail.tsx b/site/src/pages/AgentsPage/components/ChatConversation/LiveStreamTail.tsx index b4b732a21b..ac396eefda 100644 --- a/site/src/pages/AgentsPage/components/ChatConversation/LiveStreamTail.tsx +++ b/site/src/pages/AgentsPage/components/ChatConversation/LiveStreamTail.tsx @@ -1,11 +1,8 @@ -import { Link } from "react-router"; import type { UrlTransform } from "streamdown"; import type * as TypesGen from "#/api/typesGenerated"; -import { Alert, AlertDescription } from "#/components/Alert/Alert"; -import { Button } from "#/components/Button/Button"; -import type { ChatDetailError } from "../../utils/usageLimitMessage"; import type { SubagentVariant } from "../ChatElements/tools/subagentDescriptor"; import { ChatStatusCallout } from "./ChatStatusCallout"; +import type { ChatDetailError } from "./chatError"; import { selectIsAwaitingFirstStreamChunk, selectReconnectState, @@ -55,8 +52,6 @@ export const LiveStreamTailContent = ({ }: LiveStreamTailContentProps) => { const shouldRenderStreamSection = shouldRenderStreamingSection(liveStatus); const terminalStatus = liveStatus.phase === "failed" ? liveStatus : null; - const usageLimitStatus = - terminalStatus?.kind === "usage_limit" ? terminalStatus : null; const shouldRenderEmptyState = isTranscriptEmpty && liveStatus.phase === "idle"; @@ -93,20 +88,7 @@ export const LiveStreamTailContent = ({ mcpServers={mcpServers} /> )} - {usageLimitStatus && !usageLimitStatus.provider ? ( - - View usage - - } - > - {usageLimitStatus.message} - - ) : terminalStatus ? ( - - ) : null} + {terminalStatus && }
); }; diff --git a/site/src/pages/AgentsPage/components/ChatConversation/chatError.test.ts b/site/src/pages/AgentsPage/components/ChatConversation/chatError.test.ts new file mode 100644 index 0000000000..28bd338bf2 --- /dev/null +++ b/site/src/pages/AgentsPage/components/ChatConversation/chatError.test.ts @@ -0,0 +1,31 @@ +import { describe, expect, it } from "vitest"; +import { type ChatDetailError, chatDetailErrorsEqual } from "./chatError"; + +describe("chatDetailErrorsEqual", () => { + it("compares matching errors by value", () => { + const left: ChatDetailError = { + kind: "rate_limit", + message: "Slow down.", + provider: "anthropic", + retryable: true, + statusCode: 429, + }; + + expect(chatDetailErrorsEqual(left, { ...left })).toBe(true); + }); + + it("treats missing and mismatched errors as different", () => { + const error: ChatDetailError = { + kind: "generic", + message: "Provider request failed.", + }; + + expect(chatDetailErrorsEqual(error, null)).toBe(false); + expect(chatDetailErrorsEqual(error, { ...error, statusCode: 500 })).toBe( + false, + ); + expect( + chatDetailErrorsEqual(error, { ...error, detail: "Bad image." }), + ).toBe(false); + }); +}); diff --git a/site/src/pages/AgentsPage/components/ChatConversation/chatError.ts b/site/src/pages/AgentsPage/components/ChatConversation/chatError.ts index 0f24cf65fc..fef0ff268c 100644 --- a/site/src/pages/AgentsPage/components/ChatConversation/chatError.ts +++ b/site/src/pages/AgentsPage/components/ChatConversation/chatError.ts @@ -1,5 +1,55 @@ import type * as TypesGen from "#/api/typesGenerated"; -import type { ChatDetailError } from "../../utils/usageLimitMessage"; + +export type ChatDetailError = { + message: string; + detail?: string; + kind: TypesGen.ChatErrorKind; + provider?: string; + retryable?: boolean; + statusCode?: number; +}; + +export const chatDetailErrorsEqual = ( + left: ChatDetailError | null | undefined, + right: ChatDetailError | null | undefined, +): boolean => { + if (left === right) { + return true; + } + if (!left || !right) { + return false; + } + return ( + left.kind === right.kind && + left.message === right.message && + left.detail === right.detail && + left.provider === right.provider && + left.retryable === right.retryable && + left.statusCode === right.statusCode + ); +}; + +export function isChatHookDispatchFailedResponse( + value: unknown, +): value is TypesGen.ChatHookDispatchFailedResponse { + return ( + typeof value === "object" && + value !== null && + "kind" in value && + value.kind === "hook_dispatch_failed" + ); +} + +export function isChatHookDeniedResponse( + value: unknown, +): value is TypesGen.ChatHookDeniedResponse { + return ( + typeof value === "object" && + value !== null && + "kind" in value && + value.kind === "hook_denied" + ); +} export const normalizeChatErrorPayload = ( error: TypesGen.ChatError | undefined, diff --git a/site/src/pages/AgentsPage/components/ChatConversation/chatStore.ts b/site/src/pages/AgentsPage/components/ChatConversation/chatStore.ts index 3993af14f2..99d5da05f4 100644 --- a/site/src/pages/AgentsPage/components/ChatConversation/chatStore.ts +++ b/site/src/pages/AgentsPage/components/ChatConversation/chatStore.ts @@ -1,9 +1,6 @@ import { useSyncExternalStore } from "react"; import type * as TypesGen from "#/api/typesGenerated"; -import { - type ChatDetailError, - chatDetailErrorsEqual, -} from "../../utils/usageLimitMessage"; +import { type ChatDetailError, chatDetailErrorsEqual } from "./chatError"; import { applyMessagePartToStreamState } from "./streamState"; import type { ReconnectState, RetryState, StreamState } from "./types"; diff --git a/site/src/pages/AgentsPage/components/ChatConversation/liveStatusModel.test.ts b/site/src/pages/AgentsPage/components/ChatConversation/liveStatusModel.test.ts index eedfc3464e..f1a19d1f99 100644 --- a/site/src/pages/AgentsPage/components/ChatConversation/liveStatusModel.test.ts +++ b/site/src/pages/AgentsPage/components/ChatConversation/liveStatusModel.test.ts @@ -1,5 +1,5 @@ import { describe, expect, it } from "vitest"; -import type { ChatDetailError } from "../../utils/usageLimitMessage"; +import type { ChatDetailError } from "./chatError"; import { deriveLiveStatus } from "./liveStatusModel"; import { buildReconnectState, buildRetryState } from "./storyFixtures"; import type { StreamState } from "./types"; diff --git a/site/src/pages/AgentsPage/components/ChatConversation/liveStatusModel.ts b/site/src/pages/AgentsPage/components/ChatConversation/liveStatusModel.ts index 9672c0486a..77f3ed6cab 100644 --- a/site/src/pages/AgentsPage/components/ChatConversation/liveStatusModel.ts +++ b/site/src/pages/AgentsPage/components/ChatConversation/liveStatusModel.ts @@ -1,5 +1,5 @@ import type * as TypesGen from "#/api/typesGenerated"; -import type { ChatDetailError } from "../../utils/usageLimitMessage"; +import type { ChatDetailError } from "./chatError"; import { getErrorTitle } from "./chatStatusHelpers"; import type { ReconnectState, RetryState, StreamState } from "./types"; diff --git a/site/src/pages/AgentsPage/components/ChatConversation/useChatStore.ts b/site/src/pages/AgentsPage/components/ChatConversation/useChatStore.ts index 51810a2292..367430a5c8 100644 --- a/site/src/pages/AgentsPage/components/ChatConversation/useChatStore.ts +++ b/site/src/pages/AgentsPage/components/ChatConversation/useChatStore.ts @@ -19,8 +19,7 @@ import { import type * as TypesGen from "#/api/typesGenerated"; import type { OneWayMessageEvent } from "#/utils/OneWayWebSocket"; import { createReconnectingWebSocket } from "#/utils/reconnectingWebSocket"; -import type { ChatDetailError } from "../../utils/usageLimitMessage"; -import { normalizeChatErrorPayload } from "./chatError"; +import { type ChatDetailError, normalizeChatErrorPayload } from "./chatError"; import { type ChatStore, type ChatStoreState, diff --git a/site/src/pages/AgentsPage/components/ChatPageContent.tsx b/site/src/pages/AgentsPage/components/ChatPageContent.tsx index 830e6aa329..fb08c2e303 100644 --- a/site/src/pages/AgentsPage/components/ChatPageContent.tsx +++ b/site/src/pages/AgentsPage/components/ChatPageContent.tsx @@ -12,7 +12,6 @@ import { useFileAttachments } from "../hooks/useFileAttachments"; import { getChatFileURL } from "../utils/chatAttachments"; import { getProviderForModelOption } from "../utils/modelOptions"; import { CHAT_SLASH_COMMANDS } from "../utils/slashCommands"; -import type { ChatDetailError } from "../utils/usageLimitMessage"; import { AgentChatInput, type AttachedWorkspaceInfo, @@ -21,6 +20,7 @@ import { type UploadState, } from "./AgentChatInput"; import { ConversationTimeline } from "./ChatConversation/ConversationTimeline"; +import type { ChatDetailError } from "./ChatConversation/chatError"; import { getLatestContextUsage } from "./ChatConversation/chatHelpers"; import { selectChatStatus, diff --git a/site/src/pages/AgentsPage/components/DiffViewer/DiffViewer.stories.tsx b/site/src/pages/AgentsPage/components/DiffViewer/DiffViewer.stories.tsx index 140711ccae..08cc63ab53 100644 --- a/site/src/pages/AgentsPage/components/DiffViewer/DiffViewer.stories.tsx +++ b/site/src/pages/AgentsPage/components/DiffViewer/DiffViewer.stories.tsx @@ -410,13 +410,13 @@ export const CrossSideAdditionsToDeletions: Story = { // 3. File names truncate with ellipsis // biome-ignore format: raw diff string must preserve exact whitespace const renameDiff = [ -"diff --git a/site/src/pages/AgentsPage/components/LimitsTab/DefaultLimitSection.tsx b/site/src/pages/AgentsPage/components/UsageLimitsTab/DefaultLimitSection.tsx", +"diff --git a/site/src/pages/AgentsPage/components/LimitsTab/DefaultLimitSection.tsx b/site/src/pages/AgentsPage/components/SpendingTab/DefaultLimitSection.tsx", "similarity index 95%", "rename from site/src/pages/AgentsPage/components/LimitsTab/DefaultLimitSection.tsx", -"rename to site/src/pages/AgentsPage/components/UsageLimitsTab/DefaultLimitSection.tsx", +"rename to site/src/pages/AgentsPage/components/SpendingTab/DefaultLimitSection.tsx", "index abc1234..def5678 100644", "--- a/site/src/pages/AgentsPage/components/LimitsTab/DefaultLimitSection.tsx", -"+++ b/site/src/pages/AgentsPage/components/UsageLimitsTab/DefaultLimitSection.tsx", +"+++ b/site/src/pages/AgentsPage/components/SpendingTab/DefaultLimitSection.tsx", "@@ -1,3 +1,3 @@", " export function DefaultLimitSection() {", "- return null;", diff --git a/site/src/pages/AgentsPage/components/UsageIndicator.stories.tsx b/site/src/pages/AgentsPage/components/UsageIndicator.stories.tsx index 1d177bfbb2..2b449c7e41 100644 --- a/site/src/pages/AgentsPage/components/UsageIndicator.stories.tsx +++ b/site/src/pages/AgentsPage/components/UsageIndicator.stories.tsx @@ -1,12 +1,12 @@ import type { Decorator, Meta, StoryObj } from "@storybook/react-vite"; import type { FC } from "react"; import { useQueryClient } from "react-query"; -import { expect, userEvent, within } from "storybook/test"; -import { chatUsageLimitStatusKey } from "#/api/queries/chats"; +import { expect, userEvent, waitFor, within } from "storybook/test"; +import { meAISpendKey } from "#/api/queries/users"; import { getWorkspaceQuotaQueryKey } from "#/api/queries/workspaceQuota"; import { workspacesKey } from "#/api/queries/workspaces"; import type { - ChatUsageLimitStatus, + UserAISpendStatus, WorkspaceQuota, WorkspacesResponse, } from "#/api/typesGenerated"; @@ -21,9 +21,9 @@ import { } from "#/testHelpers/storybook"; import { UsageIndicator } from "./UsageIndicator"; -const withUsageLimitStatus = (status: ChatUsageLimitStatus) => (Story: FC) => { +const withAISpend = (status: UserAISpendStatus) => (Story: FC) => { const queryClient = useQueryClient(); - queryClient.setQueryData(chatUsageLimitStatusKey, status); + queryClient.setQueryData(meAISpendKey, status); return ; }; @@ -79,22 +79,23 @@ const openUsageMenu = async (canvasElement: HTMLElement) => { await userEvent.click(canvas.getByRole("button")); }; -const limitedUsageStatus = ( - overrides: Partial = {}, -): ChatUsageLimitStatus => ({ - is_limited: true, - period: "month", - spend_limit_micros: 50_000_000, - current_spend: 12_500_000, - period_start: "2026-02-10T00:00:00Z", - period_end: "2026-03-12T00:00:00Z", +const aiSpendStatus = ( + overrides: Partial = {}, +): UserAISpendStatus => ({ + user_id: MockUserOwner.id, + effective_group_id: "group-1", + effective_budget: { spend_limit_micros: 50_000_000, limit_source: "group" }, + current_spend_micros: 12_500_000, + period_start: "2026-07-01T00:00:00Z", + period_end: "2026-08-01T00:00:00Z", ...overrides, }); -const unlimitedUsageStatus = { - is_limited: false, - current_spend: 0, -} satisfies ChatUsageLimitStatus; +const noBudgetStatus = aiSpendStatus({ + effective_group_id: null, + effective_budget: null, + current_spend_micros: 0, +}); const userWorkspacesRequest = { q: `owner:me organization:${MockDefaultOrganization.name}`, @@ -120,6 +121,7 @@ const meta: Meta = { parameters: { user: MockUserOwner, permissions: MockPermissions, + features: ["aibridge"], }, }; @@ -128,18 +130,20 @@ type Story = StoryObj; export const LowUsage: Story = { decorators: [ - withUsageLimitStatus(limitedUsageStatus()), + withAISpend(aiSpendStatus()), withWorkspaceQuota(noWorkspaceQuota), ], }; export const MediumUsage: Story = { decorators: [ - withUsageLimitStatus( - limitedUsageStatus({ - period: "week", - spend_limit_micros: 20_000_000, - current_spend: 16_000_000, + withAISpend( + aiSpendStatus({ + effective_budget: { + spend_limit_micros: 20_000_000, + limit_source: "group", + }, + current_spend_micros: 16_000_000, }), ), withWorkspaceQuota(noWorkspaceQuota), @@ -148,11 +152,13 @@ export const MediumUsage: Story = { export const HighUsage: Story = { decorators: [ - withUsageLimitStatus( - limitedUsageStatus({ - period: "day", - spend_limit_micros: 10_000_000, - current_spend: 9_500_000, + withAISpend( + aiSpendStatus({ + effective_budget: { + spend_limit_micros: 10_000_000, + limit_source: "group", + }, + current_spend_micros: 9_500_000, }), ), withWorkspaceQuota(noWorkspaceQuota), @@ -161,10 +167,13 @@ export const HighUsage: Story = { export const LimitExceeded: Story = { decorators: [ - withUsageLimitStatus( - limitedUsageStatus({ - spend_limit_micros: 30_000_000, - current_spend: 32_000_000, + withAISpend( + aiSpendStatus({ + effective_budget: { + spend_limit_micros: 30_000_000, + limit_source: "group", + }, + current_spend_micros: 32_000_000, }), ), withWorkspaceQuota(noWorkspaceQuota), @@ -173,7 +182,7 @@ export const LimitExceeded: Story = { export const WorkspaceQuotaOnly: Story = { decorators: [ - withUsageLimitStatus(unlimitedUsageStatus), + withAISpend(noBudgetStatus), withWorkspaceQuota(defaultWorkspaceQuota), withWorkspaceCount(3), ], @@ -189,7 +198,7 @@ export const WorkspaceQuotaOnly: Story = { export const UsageAndWorkspaceQuota: Story = { decorators: [ - withUsageLimitStatus(limitedUsageStatus()), + withAISpend(aiSpendStatus()), withWorkspaceQuota(defaultWorkspaceQuota), withWorkspaceCount(3), ], @@ -199,18 +208,23 @@ export const UsageAndWorkspaceQuota: Story = { expect(canvas.getByRole("button", { name: "Usage" })).toBeVisible(); expect(progressBars.map((bar) => bar.getAttribute("aria-label"))).toEqual([ - "Monthly spend usage", + "AI spend usage", "Workspace quota usage", ]); await openUsageMenu(canvasElement); + const menu = within(await within(document.body).findByRole("menu")); + await waitFor(() => { + expect(menu.getByText("$12.50 of $50.00 used")).toBeVisible(); + expect(menu.getByText("July 1 - August 1, 2026")).toBeVisible(); + }); }, }; export const TriggerTiny: Story = { decorators: [ withUsageIndicatorFrame("w-[240px]", "usage-indicator-frame"), - withUsageLimitStatus(limitedUsageStatus()), + withAISpend(aiSpendStatus()), withWorkspaceQuota(defaultWorkspaceQuota), withWorkspaceCount(3), ], @@ -218,7 +232,7 @@ export const TriggerTiny: Story = { export const WorkspaceQuotaUnused: Story = { decorators: [ - withUsageLimitStatus(unlimitedUsageStatus), + withAISpend(noBudgetStatus), withWorkspaceQuota({ credits_consumed: 0, budget: 100, @@ -233,7 +247,7 @@ export const WorkspaceQuotaUnused: Story = { export const WorkspaceQuotaWithoutBudget: Story = { decorators: [ - withUsageLimitStatus(unlimitedUsageStatus), + withAISpend(noBudgetStatus), withWorkspaceQuota({ credits_consumed: 20, budget: 0, @@ -258,7 +272,7 @@ export const WorkspaceQuotaWithoutBudget: Story = { export const WorkspaceQuotaExceeded: Story = { decorators: [ - withUsageLimitStatus(unlimitedUsageStatus), + withAISpend(noBudgetStatus), withWorkspaceQuota({ credits_consumed: 125, budget: 100, @@ -272,7 +286,7 @@ export const WorkspaceQuotaExceeded: Story = { export const WorkspaceQuotaWithoutWorkspaceCount: Story = { decorators: [ - withUsageLimitStatus(unlimitedUsageStatus), + withAISpend(noBudgetStatus), withWorkspaceQuota(defaultWorkspaceQuota), withUnavailableWorkspaceCount, ], @@ -283,7 +297,49 @@ export const WorkspaceQuotaWithoutWorkspaceCount: Story = { export const NotLimited: Story = { decorators: [ - withUsageLimitStatus(unlimitedUsageStatus), + withAISpend(noBudgetStatus), withWorkspaceQuota(noWorkspaceQuota), ], }; + +export const ZeroBudget: Story = { + decorators: [ + withAISpend( + aiSpendStatus({ + effective_budget: { spend_limit_micros: 0, limit_source: "group" }, + current_spend_micros: 0, + }), + ), + withWorkspaceQuota(noWorkspaceQuota), + ], + play: async ({ canvasElement }) => { + const canvas = within(canvasElement); + const trigger = await canvas.findByRole("button"); + await userEvent.click(trigger); + const menu = within( + await within(canvasElement.ownerDocument.body).findByRole("menu"), + ); + await waitFor(() => { + expect(menu.getByText(/limit exceeded/)).toBeVisible(); + }); + }, +}; + +export const GatewayUnavailable: Story = { + parameters: { features: [], experiments: [] }, + decorators: [ + withAISpend(aiSpendStatus()), + withWorkspaceQuota(defaultWorkspaceQuota), + withWorkspaceCount(3), + ], + play: async ({ canvasElement }) => { + await openUsageMenu(canvasElement); + const progressBars = within(canvasElement.ownerDocument.body).getAllByRole( + "progressbar", + ); + + expect(progressBars.map((bar) => bar.getAttribute("aria-label"))).toEqual([ + "Workspace quota usage", + ]); + }, +}; diff --git a/site/src/pages/AgentsPage/components/UsageIndicator.tsx b/site/src/pages/AgentsPage/components/UsageIndicator.tsx index 3e582864da..872e20b76a 100644 --- a/site/src/pages/AgentsPage/components/UsageIndicator.tsx +++ b/site/src/pages/AgentsPage/components/UsageIndicator.tsx @@ -1,15 +1,12 @@ -import dayjs from "dayjs"; import { CoinsIcon, InfoIcon, ServerIcon } from "lucide-react"; import { type FC, Fragment, type ReactNode } from "react"; import { useQuery } from "react-query"; -import { Link } from "react-router"; -import { chatUsageLimitStatus } from "#/api/queries/chats"; +import { meAISpend } from "#/api/queries/users"; import { workspaceQuota } from "#/api/queries/workspaceQuota"; import { workspaces } from "#/api/queries/workspaces"; import { DropdownMenu, DropdownMenuContent, - DropdownMenuItem, DropdownMenuSeparator, DropdownMenuTrigger, } from "#/components/DropdownMenu/DropdownMenu"; @@ -25,9 +22,10 @@ import { getDefaultOrganizationName, useDashboard, } from "#/modules/dashboard/useDashboard"; -import { getUsageLimitPeriodLabel } from "#/pages/AISettingsPage/SpendPage/components/ChatCostSummaryView"; +import { useFeatureVisibility } from "#/modules/dashboard/useFeatureVisibility"; import { clampPercentage, + formatSpendPeriodLabel, getSeverity, type UsageSeverity, usageProgressPercentage, @@ -52,11 +50,13 @@ type UsageSectionData = { const numberFormatter = new Intl.NumberFormat("en-US"); export const UsageIndicator: FC = () => { - const { data: chatUsage, isError: isChatUsageError } = useQuery( - chatUsageLimitStatus(), - ); const { user } = useAuthenticated(); const { organizations } = useDashboard(); + const aiSpendAvailable = Boolean(useFeatureVisibility().aibridge); + const { data: aiSpend, isError: isAISpendError } = useQuery({ + ...meAISpend(), + enabled: aiSpendAvailable, + }); const organizationName = getDefaultOrganizationName(organizations); const username = user.username; const { data: quota, isError: isQuotaError } = useQuery({ @@ -75,18 +75,26 @@ export const UsageIndicator: FC = () => { }); const sections: UsageSectionData[] = []; - if (!isChatUsageError && chatUsage?.is_limited) { - const spendLimit = chatUsage.spend_limit_micros ?? 0; - const currentSpend = chatUsage.current_spend; - const periodLabel = getUsageLimitPeriodLabel(chatUsage.period); - const exceeded = spendLimit > 0 && currentSpend >= spendLimit; + if ( + aiSpendAvailable && + !isAISpendError && + aiSpend && + aiSpend.effective_budget !== null + ) { + const spendLimit = aiSpend.effective_budget.spend_limit_micros; + const currentSpend = aiSpend.current_spend_micros; + // The gateway blocks once spend >= limit, so a zero budget is + // always exhausted. + const exceeded = currentSpend >= spendLimit; sections.push({ - id: "ai-usage", - title: `${periodLabel} usage`, - progressLabel: `${periodLabel} spend usage`, - percent: usageProgressPercentage(currentSpend, spendLimit), - severity: getSeverity(currentSpend, spendLimit), + id: "ai-spend", + title: "AI spend", + progressLabel: "AI spend usage", + percent: exceeded + ? 100 + : usageProgressPercentage(currentSpend, spendLimit), + severity: exceeded ? "exceeded" : getSeverity(currentSpend, spendLimit), icon: , hoverLabel: `Spend ${formatCostMicros(currentSpend)}`, detail: ( @@ -100,9 +108,10 @@ export const UsageIndicator: FC = () => { )} ), - secondaryDetail: chatUsage.period_end - ? `Resets ${dayjs(chatUsage.period_end).format("MMM D, YYYY")}` - : undefined, + secondaryDetail: formatSpendPeriodLabel( + aiSpend.period_start, + aiSpend.period_end, + ), }); } @@ -167,12 +176,6 @@ const UsageMenu: FC<{ sections: readonly UsageSectionData[] }> = ({ ))} - - - - - View usage - ); diff --git a/site/src/pages/AgentsPage/utils/usageLimitMessage.test.ts b/site/src/pages/AgentsPage/utils/usageLimitMessage.test.ts deleted file mode 100644 index 91ef15e437..0000000000 --- a/site/src/pages/AgentsPage/utils/usageLimitMessage.test.ts +++ /dev/null @@ -1,187 +0,0 @@ -import { - type ChatDetailError, - chatDetailErrorsEqual, - formatUsageLimitMessage, - isChatUsageLimitExceededResponse, -} from "./usageLimitMessage"; - -describe("formatUsageLimitMessage", () => { - it("formats a full structured message", () => { - const result = formatUsageLimitMessage({ - spent_micros: 900_000, - limit_micros: 500_000, - resets_at: "2026-03-16T00:00:00Z", - }); - expect(result).toContain("$0.90"); - expect(result).toContain("$0.50"); - expect(result).toContain("Mar"); - expect(result).toContain("2026"); - }); - - it("returns fallback when fields are missing", () => { - expect(formatUsageLimitMessage({})).toBe( - "Your usage limit has been reached.", - ); - expect(formatUsageLimitMessage({ spent_micros: 100 })).toBe( - "Your usage limit has been reached.", - ); - }); - - it("returns fallback for custom fallback message", () => { - expect(formatUsageLimitMessage({}, "Custom fallback.")).toBe( - "Custom fallback.", - ); - }); - - it("formats zero-value amounts", () => { - const result = formatUsageLimitMessage({ - spent_micros: 0, - limit_micros: 0, - resets_at: "2026-03-16T00:00:00Z", - }); - expect(result).toContain("$0.00"); - }); - - it("formats high-value amounts with locale grouping", () => { - const result = formatUsageLimitMessage({ - spent_micros: 1_234_560_000, - limit_micros: 5_000_000_000, - resets_at: "2026-03-16T00:00:00Z", - }); - expect(result).toContain("$1,234.56"); - expect(result).toContain("$5,000.00"); - }); - - it("formats sub-cent values with four decimal places", () => { - const result = formatUsageLimitMessage({ - spent_micros: 500, - limit_micros: 1_000, - resets_at: "2026-03-16T00:00:00Z", - }); - expect(result).toContain("$0.0005"); - expect(result).toContain("$0.0010"); - }); - - it("handles invalid resets_at gracefully", () => { - const result = formatUsageLimitMessage({ - spent_micros: 900_000, - limit_micros: 500_000, - resets_at: "not-a-date", - }); - expect(result).toContain("$0.90"); - expect(result).toContain("$0.50"); - expect(result).not.toContain("Resets"); - }); -}); - -describe("chatDetailErrorsEqual", () => { - it("compares matching errors by value", () => { - const left: ChatDetailError = { - kind: "rate_limit", - message: "Slow down.", - provider: "anthropic", - retryable: true, - statusCode: 429, - }; - - expect(chatDetailErrorsEqual(left, { ...left })).toBe(true); - }); - - it("treats missing and mismatched errors as different", () => { - const error: ChatDetailError = { - kind: "generic", - message: "Provider request failed.", - }; - - expect(chatDetailErrorsEqual(error, null)).toBe(false); - expect(chatDetailErrorsEqual(error, { ...error, statusCode: 500 })).toBe( - false, - ); - expect( - chatDetailErrorsEqual(error, { ...error, detail: "Bad image." }), - ).toBe(false); - }); -}); - -describe("isChatUsageLimitExceededResponse", () => { - it("accepts a payload with usage fields", () => { - const error: ChatDetailError = { - message: "Your usage limit has been reached.", - kind: "usage_limit", - }; - - expect(error.kind).toBe("usage_limit"); - expect( - isChatUsageLimitExceededResponse({ - spent_micros: 900_000, - limit_micros: 500_000, - resets_at: "2026-03-16T00:00:00Z", - }), - ).toBe(true); - }); - - it("rejects null", () => { - expect(isChatUsageLimitExceededResponse(null)).toBe(false); - }); - - it("rejects undefined", () => { - expect(isChatUsageLimitExceededResponse(undefined)).toBe(false); - }); - - it("rejects an empty object (missing all fields)", () => { - expect(isChatUsageLimitExceededResponse({})).toBe(false); - }); - - it("rejects when spent_micros is missing", () => { - expect( - isChatUsageLimitExceededResponse({ - message: "Chat usage limit exceeded.", - limit_micros: 500_000, - resets_at: "2026-03-16T00:00:00Z", - }), - ).toBe(false); - }); - - it("rejects when limit_micros is missing", () => { - expect( - isChatUsageLimitExceededResponse({ - message: "Chat usage limit exceeded.", - spent_micros: 900_000, - resets_at: "2026-03-16T00:00:00Z", - }), - ).toBe(false); - }); - - it("rejects when resets_at is missing", () => { - expect( - isChatUsageLimitExceededResponse({ - message: "Chat usage limit exceeded.", - spent_micros: 900_000, - limit_micros: 500_000, - }), - ).toBe(false); - }); - - it("rejects wrong field types (string for spent_micros)", () => { - expect( - isChatUsageLimitExceededResponse({ - message: "Chat usage limit exceeded.", - spent_micros: "900000", - limit_micros: 500_000, - resets_at: "2026-03-16T00:00:00Z", - }), - ).toBe(false); - }); - - it("accepts payload with extra fields", () => { - expect( - isChatUsageLimitExceededResponse({ - message: "Chat usage limit exceeded.", - spent_micros: 900_000, - limit_micros: 500_000, - resets_at: "2026-03-16T00:00:00Z", - extra_field: "ignored", - }), - ).toBe(true); - }); -}); diff --git a/site/src/pages/AgentsPage/utils/usageLimitMessage.ts b/site/src/pages/AgentsPage/utils/usageLimitMessage.ts deleted file mode 100644 index 0f714ad331..0000000000 --- a/site/src/pages/AgentsPage/utils/usageLimitMessage.ts +++ /dev/null @@ -1,137 +0,0 @@ -import type * as TypesGen from "#/api/typesGenerated"; -import { formatCostMicros } from "#/utils/currency"; - -type UsageLimitData = Partial< - Pick< - TypesGen.ChatUsageLimitExceededResponse, - "spent_micros" | "limit_micros" | "resets_at" - > ->; - -/** - * Typed classification for errors surfaced in the agent detail view. - * - "usage_limit": the user hit a spending cap (409 + valid usage data). - * - other kinds come from normalized stream/provider failures. - * See ChatErrorKind for the full set. - */ -export type ChatDetailError = { - message: string; - detail?: string; - kind: TypesGen.ChatErrorKind; - provider?: string; - retryable?: boolean; - statusCode?: number; -}; - -/** - * Compare two chat-detail errors by their user-visible fields. - */ -export const chatDetailErrorsEqual = ( - left: ChatDetailError | null | undefined, - right: ChatDetailError | null | undefined, -): boolean => { - if (left === right) { - return true; - } - if (!left || !right) { - return false; - } - return ( - left.kind === right.kind && - left.message === right.message && - left.detail === right.detail && - left.provider === right.provider && - left.retryable === right.retryable && - left.statusCode === right.statusCode - ); -}; - -/** - * Format a resets_at RFC3339 timestamp into a user-friendly string. - * Example: "2026-03-16T00:00:00Z" → "Mar 16, 2026 at 12:00 AM" - */ -function formatResetDate(isoString: string): string { - const date = new Date(isoString); - if (Number.isNaN(date.getTime())) { - return ""; - } - return date.toLocaleDateString("en-US", { - month: "short", - day: "numeric", - year: "numeric", - hour: "numeric", - minute: "2-digit", - }); -} - -/** - * Runtime guard for the structured 409 usage-limit response. - */ -export function isChatUsageLimitExceededResponse( - value: unknown, -): value is TypesGen.ChatUsageLimitExceededResponse { - if (value == null || typeof value !== "object") { - return false; - } - const obj = value as Record; - return ( - typeof obj.spent_micros === "number" && - typeof obj.limit_micros === "number" && - typeof obj.resets_at === "string" - ); -} - -export function isChatHookDispatchFailedResponse( - value: unknown, -): value is TypesGen.ChatHookDispatchFailedResponse { - return ( - typeof value === "object" && - value !== null && - "kind" in value && - value.kind === "hook_dispatch_failed" - ); -} - -export function isChatHookDeniedResponse( - value: unknown, -): value is TypesGen.ChatHookDeniedResponse { - return ( - typeof value === "object" && - value !== null && - "kind" in value && - value.kind === "hook_denied" - ); -} - -/** - * Build a user-friendly usage-limit message from structured 409 - * response data. Falls back to a generic message if structured - * fields are missing or invalid. - */ -export function formatUsageLimitMessage( - data: UsageLimitData, - fallback = "Your usage limit has been reached.", -): string { - const { spent_micros, limit_micros, resets_at } = data; - - // All structured fields must be present and valid for the - // detailed message. - if ( - typeof spent_micros !== "number" || - typeof limit_micros !== "number" || - typeof resets_at !== "string" || - !resets_at - ) { - return fallback; - } - - const spent = formatCostMicros(spent_micros); - const limit = formatCostMicros(limit_micros); - const resetDate = formatResetDate(resets_at); - - if (!resetDate) { - return `You've used ${spent} of your ${limit} limit.`; - } - - return `You've used ${spent} of your ${limit} limit. Resets ${resetDate}.`; -} diff --git a/site/src/pages/GroupsPage/AIBudgetPeriod.tsx b/site/src/pages/GroupsPage/AIBudgetPeriod.tsx index b716da6313..3ab592efa3 100644 --- a/site/src/pages/GroupsPage/AIBudgetPeriod.tsx +++ b/site/src/pages/GroupsPage/AIBudgetPeriod.tsx @@ -1,8 +1,8 @@ -import dayjs from "dayjs"; import type { FC } from "react"; import { useQuery } from "react-query"; import { meAISpend } from "#/api/queries/users"; import { useFeatureVisibility } from "#/modules/dashboard/useFeatureVisibility"; +import { formatSpendPeriodLabel } from "#/utils/budget"; /** The current AI budget window, e.g. "June 1 - July 1, 2026". */ export const AIBudgetPeriod: FC = () => { @@ -13,12 +13,10 @@ export const AIBudgetPeriod: FC = () => { return null; } - // Local time and raw exclusive period_end, matching the spend page. - const start = dayjs(aiSpend.period_start).format("MMMM D"); - const end = dayjs(aiSpend.period_end).format("MMMM D, YYYY"); return ( - AI budget period: {start} - {end} + AI budget period:{" "} + {formatSpendPeriodLabel(aiSpend.period_start, aiSpend.period_end)} ); }; diff --git a/site/src/pages/GroupsPage/GroupPage.stories.tsx b/site/src/pages/GroupsPage/GroupPage.stories.tsx index c8a5e62231..a761ec8598 100644 --- a/site/src/pages/GroupsPage/GroupPage.stories.tsx +++ b/site/src/pages/GroupsPage/GroupPage.stories.tsx @@ -348,9 +348,8 @@ export const WithMemberAIBudget: Story = { const canvas = within(canvasElement); await expect(await canvas.findByText("AI budget")).toBeInTheDocument(); await expect(await canvas.findByText("Budget group")).toBeInTheDocument(); - // Dates depend on the runner's timezone; match loosely. await expect( - await canvas.findByText(/^AI budget period: \w+ \d+ - \w+ \d+, 2026$/), + await canvas.findByText("AI budget period: June 1 - July 1, 2026"), ).toBeInTheDocument(); await expect( diff --git a/site/src/testHelpers/entities.ts b/site/src/testHelpers/entities.ts index 87f42e08a7..2b3892a6b8 100644 --- a/site/src/testHelpers/entities.ts +++ b/site/src/testHelpers/entities.ts @@ -3381,7 +3381,7 @@ export const MockNoPermissions: Permissions = { editAnySettings: false, viewAnyIdpSyncSettings: false, viewAnyMembers: false, - viewAnyAIBridgeInterception: true, + viewAnyAIBridgeInterception: false, viewAnyAIProvider: false, viewAIGatewayKeys: false, createOAuth2App: false, diff --git a/site/src/utils/budget.test.ts b/site/src/utils/budget.test.ts index 7ee182c8fa..a475eaa615 100644 --- a/site/src/utils/budget.test.ts +++ b/site/src/utils/budget.test.ts @@ -1,10 +1,25 @@ import { describe, expect, it } from "vitest"; import { clampPercentage, + formatSpendPeriodLabel, getSeverity, usageProgressPercentage, } from "./budget"; +describe("formatSpendPeriodLabel", () => { + it("renders API timestamps in UTC with the year on the exclusive end", () => { + expect( + formatSpendPeriodLabel("2026-06-01T00:00:00Z", "2026-07-01T00:00:00Z"), + ).toBe("June 1 - July 1, 2026"); + }); + + it("uses the end year when the window crosses into the next year", () => { + expect( + formatSpendPeriodLabel("2026-12-01T00:00:00Z", "2027-01-01T00:00:00Z"), + ).toBe("December 1 - January 1, 2027"); + }); +}); + describe("getSeverity", () => { it("returns normal below the warning threshold", () => { expect(getSeverity(0, 50)).toBe("normal"); diff --git a/site/src/utils/budget.ts b/site/src/utils/budget.ts index 6b14b76970..78d537a244 100644 --- a/site/src/utils/budget.ts +++ b/site/src/utils/budget.ts @@ -1,5 +1,23 @@ +import dayjs from "dayjs"; +import utc from "dayjs/plugin/utc"; + +dayjs.extend(utc); + export type UsageSeverity = "normal" | "warning" | "exceeded"; +/** + * Formats an AI budget window in UTC, e.g. "June 1 - July 1, 2026". Renders + * the exclusive period_end as-is. + */ +export function formatSpendPeriodLabel( + periodStart: string, + periodEnd: string, +): string { + const start = dayjs.utc(periodStart).format("MMMM D"); + const end = dayjs.utc(periodEnd).format("MMMM D, YYYY"); + return `${start} - ${end}`; +} + /** * Classifies usage against a budget. Returns "warning" once usage reaches 85% * of the budget and "exceeded" once it meets or passes the budget. A budget of