Billing - optimize usageEvent CH table (#20019)

- Update usageEvent clickhouse table, partitioning, indexing and
projection (auto materialized view) to optimize credit usage queries
- Add caching for available credits and billing subscription


To do in next PR: deprecate enforceCapUsage cron. Bonus : real-time on
billingSubscription

packages/twenty-server/src/engine/metadata-modules/ai/ai-billing/ai-billing.module.ts

@@ -1,11 +1,12 @@
 import { Module } from '@nestjs/common';
 
+import { BillingModule } from 'src/engine/core-modules/billing/billing.module';
 import { AiBillingService } from 'src/engine/metadata-modules/ai/ai-billing/services/ai-billing.service';
 import { AiModelsModule } from 'src/engine/metadata-modules/ai/ai-models/ai-models.module';
 import { WorkspaceEventEmitterModule } from 'src/engine/workspace-event-emitter/workspace-event-emitter.module';
 
 @Module({
-  imports: [WorkspaceEventEmitterModule, AiModelsModule],
+  imports: [WorkspaceEventEmitterModule, AiModelsModule, BillingModule],
   providers: [AiBillingService],
   exports: [AiBillingService],
 })

packages/twenty-server/src/engine/metadata-modules/ai/ai-billing/services/__tests__/ai-billing.service.spec.ts

@@ -4,6 +4,8 @@ import { USAGE_RECORDED } from 'src/engine/core-modules/usage/constants/usage-re
 import { UsageOperationType } from 'src/engine/core-modules/usage/enums/usage-operation-type.enum';
 import { UsageResourceType } from 'src/engine/core-modules/usage/enums/usage-resource-type.enum';
 import { UsageUnit } from 'src/engine/core-modules/usage/enums/usage-unit.enum';
+import { BillingService } from 'src/engine/core-modules/billing/services/billing.service';
+import { BillingUsageService } from 'src/engine/core-modules/billing/services/billing-usage.service';
 import { AiBillingService } from 'src/engine/metadata-modules/ai/ai-billing/services/ai-billing.service';
 import { ModelFamily } from 'src/engine/metadata-modules/ai/ai-models/types/model-family.enum';
 import { AiModelRegistryService } from 'src/engine/metadata-modules/ai/ai-models/services/ai-model-registry.service';
@@ -73,6 +75,18 @@ describe('AiBillingService', () => {
           provide: AiModelRegistryService,
           useValue: mockAiModelRegistryMethods,
         },
+        {
+          provide: BillingService,
+          useValue: {
+            isBillingEnabled: jest.fn().mockReturnValue(false),
+          },
+        },
+        {
+          provide: BillingUsageService,
+          useValue: {
+            decrementAvailableCredits: jest.fn().mockResolvedValue(undefined),
+          },
+        },
       ],
     }).compile();
 

packages/twenty-server/src/engine/metadata-modules/ai/ai-billing/services/ai-billing.service.ts

@@ -1,6 +1,8 @@
 import { Injectable, Logger } from '@nestjs/common';
 
 import { type LanguageModelUsage } from 'ai';
+import { BillingUsageService } from 'src/engine/core-modules/billing/services/billing-usage.service';
+import { BillingService } from 'src/engine/core-modules/billing/services/billing.service';
 
 import { USAGE_RECORDED } from 'src/engine/core-modules/usage/constants/usage-recorded.constant';
 import { UsageOperationType } from 'src/engine/core-modules/usage/enums/usage-operation-type.enum';
@@ -10,8 +12,8 @@ import { type UsageEvent } from 'src/engine/core-modules/usage/types/usage-event
 import { NATIVE_WEB_SEARCH_COST_PER_CALL_DOLLARS } from 'src/engine/metadata-modules/ai/ai-billing/constants/native-web-search-cost-per-call-dollars';
 import { computeCostBreakdown } from 'src/engine/metadata-modules/ai/ai-billing/utils/compute-cost-breakdown.util';
 import { convertDollarsToBillingCredits } from 'src/engine/metadata-modules/ai/ai-billing/utils/convert-dollars-to-billing-credits.util';
-import { type ModelId } from 'src/engine/metadata-modules/ai/ai-models/types/model-id.type';
 import { AiModelRegistryService } from 'src/engine/metadata-modules/ai/ai-models/services/ai-model-registry.service';
+import { type ModelId } from 'src/engine/metadata-modules/ai/ai-models/types/model-id.type';
 import { WorkspaceEventEmitter } from 'src/engine/workspace-event-emitter/workspace-event-emitter';
 
 export type BillingUsageInput = {
@@ -26,6 +28,8 @@ export class AiBillingService {
   constructor(
     private readonly workspaceEventEmitter: WorkspaceEventEmitter,
     private readonly aiModelRegistryService: AiModelRegistryService,
+    private readonly billingService: BillingService,
+    private readonly billingUsageService: BillingUsageService,
   ) {}
 
   calculateCost(modelId: ModelId, billingInput: BillingUsageInput): number {
@@ -52,14 +56,14 @@ export class AiBillingService {
     return breakdown.totalCostInDollars;
   }
 
-  calculateAndBillUsage(
+  async calculateAndBillUsage(
     modelId: ModelId,
     billingInput: BillingUsageInput,
     workspaceId: string,
     operationType: UsageOperationType,
     agentId?: string | null,
     userWorkspaceId?: string | null,
-  ): void {
+  ): Promise<void> {
     const costInDollars = this.calculateCost(modelId, billingInput);
     const creditsUsedMicro = Math.round(
       convertDollarsToBillingCredits(costInDollars),
@@ -70,7 +74,7 @@ export class AiBillingService {
       (billingInput.usage.outputTokens ?? 0) +
       (billingInput.cacheCreationTokens ?? 0);
 
-    this.emitAiTokenUsageEvent(
+    await this.emitAiTokenUsageEvent(
       workspaceId,
       creditsUsedMicro,
       totalTokens,
@@ -81,11 +85,11 @@ export class AiBillingService {
     );
   }
 
-  billNativeWebSearchUsage(
+  async billNativeWebSearchUsage(
     nativeWebSearchCallCount: number,
     workspaceId: string,
     userWorkspaceId?: string | null,
-  ): void {
+  ): Promise<void> {
     if (nativeWebSearchCallCount <= 0) {
       return;
     }
@@ -114,9 +118,16 @@ export class AiBillingService {
       ],
       workspaceId,
     );
+
+    if (this.billingService.isBillingEnabled()) {
+      await this.billingUsageService.decrementAvailableCredits({
+        workspaceId,
+        usedCredits: creditsUsedMicro,
+      });
+    }
   }
 
-  private emitAiTokenUsageEvent(
+  private async emitAiTokenUsageEvent(
     workspaceId: string,
     creditsUsedMicro: number,
     totalTokens: number,
@@ -124,7 +135,7 @@ export class AiBillingService {
     operationType: UsageOperationType,
     agentId?: string | null,
     userWorkspaceId?: string | null,
-  ): void {
+  ): Promise<void> {
     this.workspaceEventEmitter.emitCustomBatchEvent<UsageEvent>(
       USAGE_RECORDED,
       [
@@ -141,5 +152,11 @@ export class AiBillingService {
       ],
       workspaceId,
     );
+    if (this.billingService.isBillingEnabled()) {
+      await this.billingUsageService.decrementAvailableCredits({
+        workspaceId,
+        usedCredits: creditsUsedMicro,
+      });
+    }
   }
 }

packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/resolvers/agent-chat.resolver.ts

@@ -9,10 +9,10 @@ import {
 } from '@nestjs/graphql';
 
 import { InjectRepository } from '@nestjs/typeorm';
+import GraphQLJSON from 'graphql-type-json';
 import { PermissionFlagType } from 'twenty-shared/constants';
 import { isDefined } from 'twenty-shared/utils';
 import { Repository } from 'typeorm';
-import GraphQLJSON from 'graphql-type-json';
 
 import { MetadataResolver } from 'src/engine/api/graphql/graphql-config/decorators/metadata-resolver.decorator';
 import { UUIDScalarType } from 'src/engine/api/graphql/workspace-schema-builder/graphql-types/scalars';
@@ -20,8 +20,7 @@ import {
   BillingException,
   BillingExceptionCode,
 } from 'src/engine/core-modules/billing/billing.exception';
-import { BillingProductKey } from 'src/engine/core-modules/billing/enums/billing-product-key.enum';
-import { BillingService } from 'src/engine/core-modules/billing/services/billing.service';
+import { BillingUsageService } from 'src/engine/core-modules/billing/services/billing-usage.service';
 import { RedisClientService } from 'src/engine/core-modules/redis-client/redis-client.service';
 import { TwentyConfigService } from 'src/engine/core-modules/twenty-config/twenty-config.service';
 import { toDisplayCredits } from 'src/engine/core-modules/usage/utils/to-display-credits.util';
@@ -30,13 +29,8 @@ import { AuthUserWorkspaceId } from 'src/engine/decorators/auth/auth-user-worksp
 import { AuthWorkspace } from 'src/engine/decorators/auth/auth-workspace.decorator';
 import { SettingsPermissionGuard } from 'src/engine/guards/settings-permission.guard';
 import { WorkspaceAuthGuard } from 'src/engine/guards/workspace-auth.guard';
-import {
-  AiException,
-  AiExceptionCode,
-} from 'src/engine/metadata-modules/ai/ai.exception';
-import { AiGraphqlApiExceptionInterceptor } from 'src/engine/metadata-modules/ai/interceptors/ai-graphql-api-exception.interceptor';
-import { type BrowsingContextType } from 'src/engine/metadata-modules/ai/ai-agent/types/browsingContext.type';
 import { AgentMessageDTO } from 'src/engine/metadata-modules/ai/ai-agent-execution/dtos/agent-message.dto';
+import { type BrowsingContextType } from 'src/engine/metadata-modules/ai/ai-agent/types/browsingContext.type';
 import { AgentChatThreadDTO } from 'src/engine/metadata-modules/ai/ai-chat/dtos/agent-chat-thread.dto';
 import { AiSystemPromptPreviewDTO } from 'src/engine/metadata-modules/ai/ai-chat/dtos/ai-system-prompt-preview.dto';
 import { ChatStreamCatchupChunksDTO } from 'src/engine/metadata-modules/ai/ai-chat/dtos/chat-stream-catchup-chunks.dto';
@@ -45,9 +39,14 @@ import { AgentChatThreadEntity } from 'src/engine/metadata-modules/ai/ai-chat/en
 import { AgentChatEventPublisherService } from 'src/engine/metadata-modules/ai/ai-chat/services/agent-chat-event-publisher.service';
 import { AgentChatStreamingService } from 'src/engine/metadata-modules/ai/ai-chat/services/agent-chat-streaming.service';
 import { AgentChatService } from 'src/engine/metadata-modules/ai/ai-chat/services/agent-chat.service';
-import { getCancelChannel } from 'src/engine/metadata-modules/ai/ai-chat/utils/get-cancel-channel.util';
 import { SystemPromptBuilderService } from 'src/engine/metadata-modules/ai/ai-chat/services/system-prompt-builder.service';
+import { getCancelChannel } from 'src/engine/metadata-modules/ai/ai-chat/utils/get-cancel-channel.util';
 import { AiModelRegistryService } from 'src/engine/metadata-modules/ai/ai-models/services/ai-model-registry.service';
+import {
+  AiException,
+  AiExceptionCode,
+} from 'src/engine/metadata-modules/ai/ai.exception';
+import { AiGraphqlApiExceptionInterceptor } from 'src/engine/metadata-modules/ai/interceptors/ai-graphql-api-exception.interceptor';
 
 @UseGuards(WorkspaceAuthGuard, SettingsPermissionGuard(PermissionFlagType.AI))
 @UseInterceptors(AiGraphqlApiExceptionInterceptor)
@@ -58,7 +57,7 @@ export class AgentChatResolver {
     private readonly agentChatStreamingService: AgentChatStreamingService,
     private readonly eventPublisherService: AgentChatEventPublisherService,
     private readonly systemPromptBuilderService: SystemPromptBuilderService,
-    private readonly billingService: BillingService,
+    private readonly billingUsageService: BillingUsageService,
     private readonly twentyConfigService: TwentyConfigService,
     private readonly aiModelRegistryService: AiModelRegistryService,
     private readonly redisClientService: RedisClientService,
@@ -135,9 +134,8 @@ export class AgentChatResolver {
     );
 
     if (this.twentyConfigService.get('IS_BILLING_ENABLED')) {
-      const canBill = await this.billingService.canBillMeteredProduct(
+      const canBill = await this.billingUsageService.hasAvailableCredits(
         workspace.id,
-        BillingProductKey.WORKFLOW_NODE_EXECUTION,
       );
 
       if (!canBill) {

packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/services/chat-execution.service.ts

@@ -266,7 +266,7 @@ export class ChatExecutionService {
 
     const modelMessages = pruningResult.messages;
 
-    const billUsageFromSteps = (steps: StepResult<ToolSet>[]) => {
+    const billUsageFromSteps = async (steps: StepResult<ToolSet>[]) => {
       const usage = steps.reduce<LanguageModelUsage>(
         (acc, step) => ({
           inputTokens: (acc.inputTokens ?? 0) + (step.usage.inputTokens ?? 0),
@@ -308,7 +308,7 @@ export class ChatExecutionService {
 
       const cacheCreationTokens = extractCacheCreationTokensFromSteps(steps);
 
-      this.aiBillingService.calculateAndBillUsage(
+      await this.aiBillingService.calculateAndBillUsage(
         registeredModel.modelId,
         { usage, cacheCreationTokens },
         workspace.id,
@@ -333,8 +333,8 @@ export class ChatExecutionService {
       abortSignal,
       stopWhen: stepCountIs(AGENT_CONFIG.MAX_STEPS),
       experimental_telemetry: AI_TELEMETRY_CONFIG,
-      onAbort: ({ steps }) => {
-        billUsageFromSteps(steps);
+      onAbort: async ({ steps }) => {
+        await billUsageFromSteps(steps);
       },
       experimental_repairToolCall: async ({
         toolCall,
@@ -360,8 +360,8 @@ export class ChatExecutionService {
     });
 
     Promise.all([stream.usage, stream.steps])
-      .then(([, steps]) => {
-        billUsageFromSteps(steps);
+      .then(async ([, steps]) => {
+        await billUsageFromSteps(steps);
       })
       .catch((error) => {
         if (error?.name === 'AbortError') {