From f0e6ac64b377928d35008896d67051e4972257f7 Mon Sep 17 00:00:00 2001 From: Michael Suchacz <203725896+ibetitsmike@users.noreply.github.com> Date: Tue, 4 Aug 2026 11:36:49 +0200 Subject: [PATCH] feat: remove native chat usage limits in favor of AI Gateway budgets (#27329) ## Stack Context This stack makes AI Gateway data and budgets the source of truth for AI spend controls. 1. Re-back the per-chat cost endpoint with AI Gateway data (#27328, merged). 2. **This PR:** remove native chat usage limits. 3. Remove native chat cost tracking and its dedicated admin UI (#27330). ## Summary Removes the native usage-limit API, SDK types, SQL, and chat enforcement for deployment, user, and group chat limits. Compact AI Gateway budget indicators remain in the Agents sidebar, user menu, and group settings. Gateway budget rejections and provider quota failures continue to classify as usage-limit errors, including a 409 response for synchronous title generation. Budget-period labels now use the API's UTC boundaries, so users see the same dates in every browser timezone. The documentation explains the AI Gateway replacement, its licensing requirements, and the differences from native limits. No schema is dropped in this release. The usage-limit table, index, user and group columns, constraints, audit mappings, and generated scan fields remain for mixed-version rolling upgrades. #27600 tracks their removal after the compatibility window. ## Breaking change Native day, week, and month chat spend limits are removed and are not migrated. AI Gateway budgets are month-based, group-scoped with per-user overrides, and require the AI Gateway entitlement. Deployments without that entitlement no longer have chat spend enforcement. > Mux prepared this PR on Mike's behalf. --- coderd/coderd.go | 13 - coderd/database/dbauthz/dbauthz.go | 91 --- coderd/database/dbauthz/dbauthz_test.go | 146 ----- coderd/database/dbmetrics/querymetrics.go | 104 ---- coderd/database/dbmock/dbmock.go | 193 ------ coderd/database/querier.go | 33 -- coderd/database/queries.sql.go | 390 ------------ coderd/database/queries/chats.sql | 130 ---- coderd/exp_chats.go | 553 +----------------- coderd/exp_chats_test.go | 356 +++-------- coderd/x/chatd/chatd.go | 57 -- coderd/x/chatd/chatd_test.go | 189 ------ coderd/x/chatd/chaterror/classify_test.go | 48 ++ coderd/x/chatd/chaterror/message.go | 7 +- coderd/x/chatd/chaterror/message_test.go | 4 +- coderd/x/chatd/subagent.go | 4 - coderd/x/chatd/usagelimit.go | 152 ----- coderd/x/chatd/usagelimit_internal_test.go | 132 ----- codersdk/chats.go | 347 +---------- codersdk/chats_test.go | 78 --- .../agents/platform-controls/advisor.md | 5 +- .../agents/platform-controls/index.md | 19 +- .../platform-controls/usage-insights.md | 113 ++-- docs/manifest.json | 2 +- enterprise/coderd/x/chatd/usagelimit_test.go | 324 ---------- site/src/api/api.ts | 65 -- site/src/api/queries/chats.ts | 85 --- site/src/api/typesGenerated.ts | 129 ---- .../UserDropdown/UserDropdown.stories.tsx | 50 +- .../Navbar/UserDropdown/UserDropdown.tsx | 2 + .../UserDropdown/UserDropdownAISpend.tsx | 12 +- .../AISettingsPage/SpendPage/SpendPage.tsx | 100 +--- .../SpendPage/SpendPageView.stories.tsx | 205 +------ .../SpendPage/SpendPageView.tsx | 406 ++----------- .../components/ChatCostSummaryView.tsx | 118 +--- .../LimitsTab/DefaultLimitController.tsx | 65 -- .../LimitsTab/DefaultLimitSection.tsx | 137 ----- .../components/LimitsTab/GroupLimitDialog.tsx | 183 ------ .../LimitsTab/GroupLimitsSection.stories.tsx | 182 ------ .../LimitsTab/GroupLimitsSection.tsx | 192 ------ .../components/LimitsTab/GroupLimitsTable.tsx | 167 ------ .../LimitsTab/GroupOverrideController.tsx | 141 ----- .../LimitsTab/UserOverrideController.tsx | 131 ----- .../LimitsTab/UserOverrideDialog.tsx | 157 ----- .../UserOverridesSection.stories.tsx | 155 ----- .../LimitsTab/UserOverridesSection.tsx | 173 ------ .../LimitsTab/UserOverridesTable.tsx | 144 ----- .../LimitsTab/limitsFormLogic.test.ts | 10 - .../components/LimitsTab/limitsFormLogic.ts | 14 - .../components/UsageTab/UsageTab.tsx | 57 +- .../pages/AgentsPage/AgentChatPage.test.ts | 10 +- site/src/pages/AgentsPage/AgentChatPage.tsx | 82 +-- .../AgentsPage/AgentChatPageView.stories.tsx | 2 +- .../pages/AgentsPage/AgentChatPageView.tsx | 2 +- site/src/pages/AgentsPage/AgentEmbedPage.tsx | 8 +- .../AgentsPage/AgentsPageLayout.stories.tsx | 9 +- .../src/pages/AgentsPage/AgentsPageLayout.tsx | 8 +- .../components/AgentCreateForm.stories.tsx | 26 - .../AgentsPage/components/AgentCreateForm.tsx | 31 +- .../LiveStreamTail.stories.tsx | 21 +- .../ChatConversation/LiveStreamTail.tsx | 22 +- .../ChatConversation/chatError.test.ts | 31 + .../components/ChatConversation/chatError.ts | 52 +- .../components/ChatConversation/chatStore.ts | 5 +- .../ChatConversation/liveStatusModel.test.ts | 2 +- .../ChatConversation/liveStatusModel.ts | 2 +- .../ChatConversation/useChatStore.ts | 3 +- .../AgentsPage/components/ChatPageContent.tsx | 2 +- .../DiffViewer/DiffViewer.stories.tsx | 6 +- .../components/UsageIndicator.stories.tsx | 140 +++-- .../AgentsPage/components/UsageIndicator.tsx | 57 +- .../utils/usageLimitMessage.test.ts | 187 ------ .../AgentsPage/utils/usageLimitMessage.ts | 137 ----- site/src/pages/GroupsPage/AIBudgetPeriod.tsx | 8 +- .../pages/GroupsPage/GroupPage.stories.tsx | 3 +- site/src/testHelpers/entities.ts | 2 +- site/src/utils/budget.test.ts | 15 + site/src/utils/budget.ts | 18 + 78 files changed, 611 insertions(+), 6850 deletions(-) delete mode 100644 coderd/x/chatd/usagelimit.go delete mode 100644 coderd/x/chatd/usagelimit_internal_test.go delete mode 100644 enterprise/coderd/x/chatd/usagelimit_test.go delete mode 100644 site/src/pages/AISettingsPage/SpendPage/components/LimitsTab/DefaultLimitController.tsx delete mode 100644 site/src/pages/AISettingsPage/SpendPage/components/LimitsTab/DefaultLimitSection.tsx delete mode 100644 site/src/pages/AISettingsPage/SpendPage/components/LimitsTab/GroupLimitDialog.tsx delete mode 100644 site/src/pages/AISettingsPage/SpendPage/components/LimitsTab/GroupLimitsSection.stories.tsx delete mode 100644 site/src/pages/AISettingsPage/SpendPage/components/LimitsTab/GroupLimitsSection.tsx delete mode 100644 site/src/pages/AISettingsPage/SpendPage/components/LimitsTab/GroupLimitsTable.tsx delete mode 100644 site/src/pages/AISettingsPage/SpendPage/components/LimitsTab/GroupOverrideController.tsx delete mode 100644 site/src/pages/AISettingsPage/SpendPage/components/LimitsTab/UserOverrideController.tsx delete mode 100644 site/src/pages/AISettingsPage/SpendPage/components/LimitsTab/UserOverrideDialog.tsx delete mode 100644 site/src/pages/AISettingsPage/SpendPage/components/LimitsTab/UserOverridesSection.stories.tsx delete mode 100644 site/src/pages/AISettingsPage/SpendPage/components/LimitsTab/UserOverridesSection.tsx delete mode 100644 site/src/pages/AISettingsPage/SpendPage/components/LimitsTab/UserOverridesTable.tsx delete mode 100644 site/src/pages/AISettingsPage/SpendPage/components/LimitsTab/limitsFormLogic.test.ts delete mode 100644 site/src/pages/AISettingsPage/SpendPage/components/LimitsTab/limitsFormLogic.ts create mode 100644 site/src/pages/AgentsPage/components/ChatConversation/chatError.test.ts delete mode 100644 site/src/pages/AgentsPage/utils/usageLimitMessage.test.ts delete mode 100644 site/src/pages/AgentsPage/utils/usageLimitMessage.ts 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