diff --git a/coderd/exp_chats_test.go b/coderd/exp_chats_test.go index d4762af6c37..f37f082cca0 100644 --- a/coderd/exp_chats_test.go +++ b/coderd/exp_chats_test.go @@ -253,36 +253,6 @@ func (s *failNextUpdateChatModelConfigStore) UpdateChatModelConfig( return s.Store.UpdateChatModelConfig(ctx, arg) } -func requireChatUsageLimitExceededError( - t *testing.T, - err error, - wantSpentMicros int64, - wantLimitMicros int64, - wantResetsAt time.Time, -) *codersdk.ChatUsageLimitExceededResponse { - t.Helper() - - sdkErr, ok := codersdk.AsError(err) - require.True(t, ok) - require.Equal(t, http.StatusConflict, sdkErr.StatusCode()) - require.Equal(t, "Chat usage limit exceeded.", sdkErr.Message) - - limitErr := codersdk.ChatUsageLimitExceededFrom(err) - require.NotNil(t, limitErr) - require.Equal(t, "Chat usage limit exceeded.", limitErr.Message) - require.Equal(t, wantSpentMicros, limitErr.SpentMicros) - require.Equal(t, wantLimitMicros, limitErr.LimitMicros) - require.True( - t, - limitErr.ResetsAt.Equal(wantResetsAt), - "expected resets_at %s, got %s", - wantResetsAt.UTC().Format(time.RFC3339), - limitErr.ResetsAt.UTC().Format(time.RFC3339), - ) - - return limitErr -} - func enableDailyChatUsageLimit( ctx context.Context, t *testing.T, @@ -893,33 +863,6 @@ func TestPostChats(t *testing.T) { require.Equal(t, `content[0].type "image" is not supported.`, sdkErr.Detail) }) - t.Run("UsageLimitExceeded", 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) - wantResetsAt := enableDailyChatUsageLimit(ctx, t, db, 100) - - existingChat := dbgen.Chat(t, db, database.Chat{ - OrganizationID: user.OrganizationID, - OwnerID: user.UserID, - LastModelConfigID: modelConfig.ID, - Title: "existing-limit-chat", - }) - insertAssistantCostMessage(t, db, existingChat.ID, modelConfig.ID, 100) - - _, err := client.CreateChat(ctx, codersdk.CreateChatRequest{ - OrganizationID: user.OrganizationID, - Content: []codersdk.ChatInputPart{{ - Type: codersdk.ChatInputPartTypeText, - Text: "over limit", - }}, - }) - requireChatUsageLimitExceededError(t, err, 100, 100, wantResetsAt) - }) - t.Run("NilOrganizationID", func(t *testing.T) { t.Parallel() @@ -7150,35 +7093,6 @@ func TestPostChatMessages(t *testing.T) { require.Equal(t, "content[0].text cannot be empty.", sdkErr.Detail) }) - t.Run("UsageLimitExceeded", func(t *testing.T) { - t.Parallel() - - ctx := testutil.Context(t, testutil.WaitLong) - client, db := newChatClientWithDatabase(t) - firstUser := coderdtest.CreateFirstUser(t, client.Client) - modelConfig := createChatModelConfig(t, client) - - chat, err := client.CreateChat(ctx, codersdk.CreateChatRequest{ - OrganizationID: firstUser.OrganizationID, - Content: []codersdk.ChatInputPart{{ - Type: codersdk.ChatInputPartTypeText, - Text: "initial message for usage-limit test", - }}, - }) - require.NoError(t, err) - - wantResetsAt := enableDailyChatUsageLimit(ctx, t, db, 100) - insertAssistantCostMessage(t, db, chat.ID, modelConfig.ID, 100) - - _, err = client.CreateChatMessage(ctx, chat.ID, codersdk.CreateChatMessageRequest{ - Content: []codersdk.ChatInputPart{{ - Type: codersdk.ChatInputPartTypeText, - Text: "over limit", - }}, - }) - requireChatUsageLimitExceededError(t, err, 100, 100, wantResetsAt) - }) - t.Run("ChatNotFound", func(t *testing.T) { t.Parallel() @@ -8784,47 +8698,6 @@ func TestPatchChatMessage(t *testing.T) { require.True(t, foundFileInChat, "chat should preserve file_id after edit") }) - t.Run("UsageLimitExceeded", func(t *testing.T) { - t.Parallel() - - ctx := testutil.Context(t, testutil.WaitLong) - client, db := newChatClientWithDatabase(t) - firstUser := coderdtest.CreateFirstUser(t, client.Client) - modelConfig := createChatModelConfig(t, client) - - chat, err := client.CreateChat(ctx, codersdk.CreateChatRequest{ - OrganizationID: firstUser.OrganizationID, - Content: []codersdk.ChatInputPart{{ - Type: codersdk.ChatInputPartTypeText, - Text: "hello before edit", - }}, - }) - require.NoError(t, err) - - messagesResult, err := client.GetChatMessages(ctx, chat.ID, nil) - require.NoError(t, err) - - var userMessageID int64 - for _, message := range messagesResult.Messages { - if message.Role == codersdk.ChatMessageRoleUser { - userMessageID = message.ID - break - } - } - require.NotZero(t, userMessageID) - - wantResetsAt := enableDailyChatUsageLimit(ctx, t, db, 100) - insertAssistantCostMessage(t, db, chat.ID, modelConfig.ID, 100) - - _, err = client.EditChatMessage(ctx, chat.ID, userMessageID, codersdk.EditChatMessageRequest{ - Content: []codersdk.ChatInputPart{{ - Type: codersdk.ChatInputPartTypeText, - Text: "edited over limit", - }}, - }) - requireChatUsageLimitExceededError(t, err, 100, 100, wantResetsAt) - }) - t.Run("MessageNotFound", func(t *testing.T) { t.Parallel() @@ -9665,51 +9538,6 @@ func TestRegenerateChatTitle(t *testing.T) { requireSDKError(t, err, http.StatusUnauthorized) }) - t.Run("UsageLimitExceeded", func(t *testing.T) { - t.Parallel() - - ctx := testutil.Context(t, testutil.WaitLong) - client, db := newChatClientWithDatabase(t) - firstUser := coderdtest.CreateFirstUser(t, client.Client) - modelConfig := createChatModelConfig(t, client) - - chat, err := client.CreateChat(ctx, codersdk.CreateChatRequest{ - OrganizationID: firstUser.OrganizationID, - Content: []codersdk.ChatInputPart{{ - Type: codersdk.ChatInputPartTypeText, - Text: "chat over usage limit", - }}, - }) - require.NoError(t, err) - - wantResetsAt := enableDailyChatUsageLimit(ctx, t, db, 100) - insertAssistantCostMessage(t, db, chat.ID, modelConfig.ID, 100) - - _, 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) - limitErr := codersdk.ChatUsageLimitExceededFrom(err) - require.NotNil(t, limitErr) - require.Equal(t, "Chat usage limit exceeded.", limitErr.Message) - require.Equal(t, int64(100), limitErr.SpentMicros) - require.Equal(t, int64(100), limitErr.LimitMicros) - require.True( - t, - limitErr.ResetsAt.Equal(wantResetsAt), - "expected resets_at %s, got %s", - wantResetsAt.UTC().Format(time.RFC3339), - limitErr.ResetsAt.UTC().Format(time.RFC3339), - ) - }) - t.Run("PasteOnlyChat", func(t *testing.T) { t.Parallel() diff --git a/coderd/x/chatd/chatd.go b/coderd/x/chatd/chatd.go index 59072a6e0a2..5daff17e443 100644 --- a/coderd/x/chatd/chatd.go +++ b/coderd/x/chatd/chatd.go @@ -1518,28 +1518,9 @@ func (p *Server) SendMessage( return result, nil } -func (p *Server) checkUsageLimit(ctx context.Context, store database.Store, ownerID uuid.UUID, organizationID uuid.NullUUID) error { - status, err := ResolveUsageLimitStatus(ctx, store, ownerID, organizationID, time.Now()) - if err != nil { - // Fail open: never block chat due to a limit-resolution failure. - p.logger.Warn(ctx, "usage limit check failed, allowing message", - slog.F("owner_id", ownerID), - slog.Error(err), - ) - return nil - } - if status == nil { - return nil - } - // Block when current spend reaches or exceeds limit (>= ensures - // the user cannot start new conversations once the limit is hit). - if status.SpendLimitMicros != nil && status.CurrentSpend >= *status.SpendLimitMicros { - return &UsageLimitExceededError{ - LimitMicros: *status.SpendLimitMicros, - ConsumedMicros: status.CurrentSpend, - PeriodEnd: status.PeriodEnd, - } - } +// 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 } diff --git a/coderd/x/chatd/chatd_internal_test.go b/coderd/x/chatd/chatd_internal_test.go index 15bb8df0f4e..0ed9206878a 100644 --- a/coderd/x/chatd/chatd_internal_test.go +++ b/coderd/x/chatd/chatd_internal_test.go @@ -1098,7 +1098,6 @@ func TestRegenerateChatTitle_PersistsAndBroadcasts(t *testing.T) { }}, nil).AnyTimes() db.EXPECT().GetAIProviderKeysByProviderID(gomock.Any(), providerID).Return([]database.AIProviderKey{{ProviderID: providerID, APIKey: "test-key"}}, nil).AnyTimes() db.EXPECT().GetAIProviderKeysByProviderIDs(gomock.Any(), gomock.Any()).Return([]database.AIProviderKey{{ProviderID: providerID, APIKey: "test-key"}}, nil).AnyTimes() - db.EXPECT().GetChatUsageLimitConfig(gomock.Any()).Return(database.ChatUsageLimitConfig{}, sql.ErrNoRows) db.EXPECT().GetChatGatewayAPIKey(gomock.Any(), database.GetChatGatewayAPIKeyParams{UserID: ownerID, TokenName: GatewayTokenName(ownerID)}).Return(database.APIKey{ID: activeAPIKeyID, UserID: ownerID, ExpiresAt: time.Now().Add(48 * time.Hour)}, nil) db.EXPECT().GetChatMessagesByChatIDAscPaginated( gomock.Any(), @@ -1250,7 +1249,6 @@ func TestRegenerateChatTitle_SkipsPersistWhenTitleChangedConcurrently(t *testing }}, nil).AnyTimes() db.EXPECT().GetAIProviderKeysByProviderID(gomock.Any(), providerID).Return([]database.AIProviderKey{{ProviderID: providerID, APIKey: "test-key"}}, nil).AnyTimes() db.EXPECT().GetAIProviderKeysByProviderIDs(gomock.Any(), gomock.Any()).Return([]database.AIProviderKey{{ProviderID: providerID, APIKey: "test-key"}}, nil).AnyTimes() - db.EXPECT().GetChatUsageLimitConfig(gomock.Any()).Return(database.ChatUsageLimitConfig{}, sql.ErrNoRows) db.EXPECT().GetChatGatewayAPIKey(gomock.Any(), database.GetChatGatewayAPIKeyParams{UserID: ownerID, TokenName: GatewayTokenName(ownerID)}).Return(database.APIKey{ID: activeAPIKeyID, UserID: ownerID, ExpiresAt: time.Now().Add(48 * time.Hour)}, nil) db.EXPECT().GetChatMessagesByChatIDAscPaginated( gomock.Any(), diff --git a/coderd/x/chatd/chatd_test.go b/coderd/x/chatd/chatd_test.go index 7accb04b464..baf97534623 100644 --- a/coderd/x/chatd/chatd_test.go +++ b/coderd/x/chatd/chatd_test.go @@ -1849,78 +1849,6 @@ func TestCreateChatInsertsWorkspaceAwarenessMessage(t *testing.T) { }) } -func TestCreateChatRejectsWhenUsageLimitReached(t *testing.T) { - t.Parallel() - - db, ps := dbtestutil.NewDB(t) - replica := newTestServer(t, db, ps, uuid.New()) - - ctx := testutil.Context(t, testutil.WaitLong) - user, org, model := seedChatDependencies(t, db) - - _, err := db.UpsertChatUsageLimitConfig(ctx, database.UpsertChatUsageLimitConfigParams{ - Enabled: true, - DefaultLimitMicros: 100, - Period: string(codersdk.ChatUsageLimitPeriodDay), - }) - require.NoError(t, err) - - existingChat := dbgen.Chat(t, db, database.Chat{ - OrganizationID: org.ID, - OwnerID: user.ID, - Title: "existing-limit-chat", - LastModelConfigID: model.ID, - }) - - assistantContent, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{ - codersdk.ChatMessageText("assistant"), - }) - require.NoError(t, err) - - _ = dbgen.ChatMessage(t, db, database.ChatMessage{ - ChatID: existingChat.ID, - ModelConfigID: uuid.NullUUID{UUID: model.ID, Valid: true}, - Role: database.ChatMessageRoleAssistant, - ContentVersion: chatprompt.CurrentContentVersion, - Content: assistantContent, - TotalCostMicros: sql.NullInt64{Int64: 100, Valid: true}, - }) - - beforeChats, err := db.GetChats(ctx, database.GetChatsParams{ - OwnedOnly: true, - ViewerID: user.ID, - AfterID: uuid.Nil, - OffsetOpt: 0, - LimitOpt: 100, - }) - require.NoError(t, err) - require.Len(t, beforeChats, 1) - - _, err = replica.CreateChat(ctx, chatd.CreateOptions{ - OrganizationID: org.ID, - OwnerID: user.ID, - Title: "over-limit", - ModelConfigID: model.ID, - InitialUserContent: []codersdk.ChatMessagePart{codersdk.ChatMessageText("hello")}, - }) - require.Error(t, err) - - var limitErr *chatd.UsageLimitExceededError - require.ErrorAs(t, err, &limitErr) - require.Equal(t, int64(100), limitErr.LimitMicros) - require.Equal(t, int64(100), limitErr.ConsumedMicros) - - afterChats, err := db.GetChats(ctx, database.GetChatsParams{ - OwnedOnly: true, - ViewerID: user.ID, - AfterID: uuid.Nil, - OffsetOpt: 0, - LimitOpt: 100, - }) - require.NoError(t, err) - require.Len(t, afterChats, len(beforeChats)) -} - func TestAutoPromoteQueuedMessagesPreservesPerTurnModelOrder(t *testing.T) { t.Parallel() @@ -2248,76 +2176,6 @@ func TestInterruptAutoPromotionIgnoresLaterUsageLimitIncrease(t *testing.T) { require.Equal(t, []string{"hello", "queued", "later queued"}, userTexts) } -func TestEditMessageRejectsWhenUsageLimitReached(t *testing.T) { - t.Parallel() - - db, ps := dbtestutil.NewDB(t) - replica := newTestServer(t, db, ps, uuid.New()) - - ctx := testutil.Context(t, testutil.WaitLong) - user, org, model := seedChatDependencies(t, db) - - _, err := db.UpsertChatUsageLimitConfig(ctx, database.UpsertChatUsageLimitConfigParams{ - Enabled: true, - DefaultLimitMicros: 100, - Period: string(codersdk.ChatUsageLimitPeriodDay), - }) - require.NoError(t, err) - - chat, err := replica.CreateChat(ctx, chatd.CreateOptions{ - OrganizationID: org.ID, - OwnerID: user.ID, - Title: "edit-limit-reached", - ModelConfigID: model.ID, - InitialUserContent: []codersdk.ChatMessagePart{codersdk.ChatMessageText("original")}, - }) - require.NoError(t, err) - - messages, err := db.GetChatMessagesByChatID(ctx, database.GetChatMessagesByChatIDParams{ - ChatID: chat.ID, - AfterID: 0, - }) - require.NoError(t, err) - require.Len(t, messages, 1) - editedMessageID := messages[0].ID - - assistantContent, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{ - codersdk.ChatMessageText("assistant"), - }) - require.NoError(t, err) - - _ = dbgen.ChatMessage(t, db, database.ChatMessage{ - ChatID: chat.ID, - ModelConfigID: uuid.NullUUID{UUID: model.ID, Valid: true}, - Role: database.ChatMessageRoleAssistant, - ContentVersion: chatprompt.CurrentContentVersion, - Content: assistantContent, - TotalCostMicros: sql.NullInt64{Int64: 100, Valid: true}, - }) - - _, err = replica.EditMessage(ctx, chatd.EditMessageOptions{ - ChatID: chat.ID, - EditedMessageID: editedMessageID, - Content: []codersdk.ChatMessagePart{codersdk.ChatMessageText("edited")}, - }) - require.Error(t, err) - - var limitErr *chatd.UsageLimitExceededError - require.ErrorAs(t, err, &limitErr) - require.Equal(t, int64(100), limitErr.LimitMicros) - require.Equal(t, int64(100), limitErr.ConsumedMicros) - - messages, err = db.GetChatMessagesByChatID(ctx, database.GetChatMessagesByChatIDParams{ - ChatID: chat.ID, - AfterID: 0, - }) - require.NoError(t, err) - require.Len(t, messages, 2) - originalMessage := db2sdk.ChatMessage(messages[0]) - require.Len(t, originalMessage.Content, 1) - require.Equal(t, "original", originalMessage.Content[0].Text) -} - func TestEditMessageRejectsMissingMessage(t *testing.T) { t.Parallel() @@ -5660,106 +5518,6 @@ func TestActiveServer_Compaction(t *testing.T) { }) } -func TestCompactChat(t *testing.T) { - t.Parallel() - - setupAtLimitChat := func( - t *testing.T, - status database.ChatStatus, - compactable bool, - ) (context.Context, database.Store, *chatd.Server, database.Chat) { - t.Helper() - - db, ps := dbtestutil.NewDB(t) - ctx := testutil.Context(t, testutil.WaitLong) - user, org, model := seedChatDependencies(t, db) - server := newTestServer(t, db, ps, uuid.New()) - - _, err := db.UpsertChatUsageLimitConfig(ctx, database.UpsertChatUsageLimitConfigParams{ - Enabled: true, - DefaultLimitMicros: 100, - Period: string(codersdk.ChatUsageLimitPeriodDay), - }) - require.NoError(t, err) - - spendChat := dbgen.Chat(t, db, database.Chat{ - OrganizationID: org.ID, - OwnerID: user.ID, - LastModelConfigID: model.ID, - }) - spendContent, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{ - codersdk.ChatMessageText("spent usage"), - }) - 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: spendContent, - TotalCostMicros: sql.NullInt64{Int64: 100, Valid: true}, - }) - - chat := dbgen.Chat(t, db, database.Chat{ - OrganizationID: org.ID, - OwnerID: user.ID, - LastModelConfigID: model.ID, - Status: status, - }) - if compactable { - assistantContent, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{ - codersdk.ChatMessageText("assistant response"), - }) - require.NoError(t, err) - dbgen.ChatMessage(t, db, database.ChatMessage{ - ChatID: chat.ID, - ModelConfigID: uuid.NullUUID{UUID: model.ID, Valid: true}, - Role: database.ChatMessageRoleAssistant, - ContentVersion: chatprompt.CurrentContentVersion, - Content: assistantContent, - }) - } - - return ctx, db, server, chat - } - - t.Run("usage limit rejects compactable idle chat", func(t *testing.T) { - t.Parallel() - - ctx, db, server, chat := setupAtLimitChat(t, database.ChatStatusWaiting, true) - _, err := server.CompactChat(ctx, chat) - - var limitErr *chatd.UsageLimitExceededError - require.ErrorAs(t, err, &limitErr) - refreshed, err := db.GetChatByID(ctx, chat.ID) - require.NoError(t, err) - require.Equal(t, database.ChatStatusWaiting, refreshed.Status) - require.False(t, refreshed.CompactionRequestedAt.Valid) - }) - - t.Run("nothing to compact precedes usage limit", func(t *testing.T) { - t.Parallel() - - ctx, _, server, chat := setupAtLimitChat(t, database.ChatStatusWaiting, false) - _, err := server.CompactChat(ctx, chat) - - require.ErrorIs(t, err, chatd.ErrNothingToCompact) - var limitErr *chatd.UsageLimitExceededError - require.False(t, errors.As(err, &limitErr)) - }) - - t.Run("state conflict precedes usage limit", func(t *testing.T) { - t.Parallel() - - ctx, _, server, chat := setupAtLimitChat(t, database.ChatStatusRunning, true) - _, err := server.CompactChat(ctx, chat) - - require.ErrorIs(t, err, chatstate.ErrTransitionNotAllowed) - var limitErr *chatd.UsageLimitExceededError - require.False(t, errors.As(err, &limitErr)) - }) -} - func TestActiveServer_ManualCompaction(t *testing.T) { t.Parallel() diff --git a/coderd/x/chatd/title_override_internal_test.go b/coderd/x/chatd/title_override_internal_test.go index 9ed08a47e56..204639bc2a0 100644 --- a/coderd/x/chatd/title_override_internal_test.go +++ b/coderd/x/chatd/title_override_internal_test.go @@ -586,7 +586,6 @@ func TestGenerateManualTitleCandidate_UsesSyntheticAPIKey(t *testing.T) { }, nil })} - db.EXPECT().GetChatUsageLimitConfig(gomock.Any()).Return(database.ChatUsageLimitConfig{}, sql.ErrNoRows) db.EXPECT().GetChatMessagesByChatIDAscPaginated(gomock.Any(), database.GetChatMessagesByChatIDAscPaginatedParams{ ChatID: chat.ID, AfterID: 0,