fix: remove chatd usage limit enforcement (#27535)

This PR surgically removes enforcement of Agents spend limits:
- Adjusts the relevant function that checks usage to always return nil
- Deletes tests that expect a usage limit error.

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
 }
 

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(),

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()
 

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,