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
2026-04-28 16:06:49 +02:00
parent 2a3b8adcfb
commit fbaea0639a
33 changed files with 601 additions and 379 deletions
@@ -34,6 +34,7 @@ import { BillingUsageCapService } from 'src/engine/core-modules/billing/services
 import { BillingUsageService } from 'src/engine/core-modules/billing/services/billing-usage.service';
 import { BillingService } from 'src/engine/core-modules/billing/services/billing.service';
 import { MeteredCreditService } from 'src/engine/core-modules/billing/services/metered-credit.service';
+import { WorkspaceBillingSubscriptionCacheService } from 'src/engine/core-modules/billing/services/workspace-billing-subscription-cache.service';
 import { StripeModule } from 'src/engine/core-modules/billing/stripe/stripe.module';
 import { WorkspaceDomainsModule } from 'src/engine/core-modules/domain/workspace-domains/workspace-domains.module';
 import { EnterpriseModule } from 'src/engine/core-modules/enterprise/enterprise.module';
@@ -43,9 +44,8 @@ import { MessageQueueModule } from 'src/engine/core-modules/message-queue/messag
 import { MetricsModule } from 'src/engine/core-modules/metrics/metrics.module';
 import { UserWorkspaceEntity } from 'src/engine/core-modules/user-workspace/user-workspace.entity';
 import { WorkspaceEntity } from 'src/engine/core-modules/workspace/workspace.entity';
-import { AiBillingModule } from 'src/engine/metadata-modules/ai/ai-billing/ai-billing.module';
-import { AiModelsModule } from 'src/engine/metadata-modules/ai/ai-models/ai-models.module';
 import { PermissionsModule } from 'src/engine/metadata-modules/permissions/permissions.module';
+import { WorkspaceCacheModule } from 'src/engine/workspace-cache/workspace-cache.module';

@Module({
  imports: [
@@ -54,8 +54,7 @@ import { PermissionsModule } from 'src/engine/metadata-modules/permissions/permi
    StripeModule,
    MessageQueueModule,
    PermissionsModule,
-    AiBillingModule,
-    AiModelsModule,
+    WorkspaceCacheModule,
    WorkspaceDomainsModule,
    TypeOrmModule.forFeature([
      BillingSubscriptionEntity,
@@ -96,6 +95,7 @@ import { PermissionsModule } from 'src/engine/metadata-modules/permissions/permi
    MeteredCreditService,
    BillingGaugeService,
    EnforceUsageCapCronCommand,
+    WorkspaceBillingSubscriptionCacheService,
  ],
  exports: [
    BillingSubscriptionService,
@@ -5,6 +5,7 @@ import { InjectRepository } from '@nestjs/typeorm';

 import { In, IsNull, Repository } from 'typeorm';

+import { formatDateForClickHouse } from 'src/database/clickHouse/clickHouse.util';
 import { enforceUsageCapCronPattern } from 'src/engine/core-modules/billing/crons/enforce-usage-cap.cron.pattern';
 import { BillingProductEntity } from 'src/engine/core-modules/billing/entities/billing-product.entity';
 import { BillingSubscriptionItemEntity } from 'src/engine/core-modules/billing/entities/billing-subscription-item.entity';
@@ -118,7 +119,6 @@ export class EnforceUsageCapJob {
            await this.billingUsageCapService.getBatchPeriodCreditsUsed(
              workspaceIds,
              group[0].currentPeriodStart,
-              group[0].currentPeriodEnd,
            );

          for (const [id, usage] of batchUsage) {
@@ -249,7 +249,7 @@ export class EnforceUsageCapJob {
    const groups = new Map<string, BillingSubscriptionEntity[]>();

    for (const subscription of subscriptions) {
-      const key = `${subscription.currentPeriodStart.toISOString()}|${subscription.currentPeriodEnd.toISOString()}`;
+      const key = `${formatDateForClickHouse(subscription.currentPeriodStart)}`;
      const group = groups.get(key);

      if (group) {
@@ -5,10 +5,10 @@ import { Injectable } from '@nestjs/common';
 import { isDefined } from 'twenty-shared/utils';

 import { OnCustomBatchEvent } from 'src/engine/api/graphql/graphql-query-runner/decorators/on-custom-batch-event.decorator';
-import { USAGE_RECORDED } from 'src/engine/core-modules/usage/constants/usage-recorded.constant';
 import { BillingUsageService } from 'src/engine/core-modules/billing/services/billing-usage.service';
-import { type UsageEvent } from 'src/engine/core-modules/usage/types/usage-event.type';
 import { TwentyConfigService } from 'src/engine/core-modules/twenty-config/twenty-config.service';
+import { USAGE_RECORDED } from 'src/engine/core-modules/usage/constants/usage-recorded.constant';
+import { type UsageEvent } from 'src/engine/core-modules/usage/types/usage-event.type';
 import { CustomWorkspaceEventBatch } from 'src/engine/workspace-event-emitter/types/custom-workspace-batch-event.type';

@Injectable()
@@ -38,6 +38,7 @@ export class BillingUsageEventListener {
      return;
    }

+    //TODO: To be removed
    await this.billingUsageService.billUsage({
      workspaceId: payload.workspaceId,
      usageEvents: payload.events,
@@ -2,14 +2,16 @@

 import { Test, type TestingModule } from '@nestjs/testing';

+import { getRepositoryToken } from '@nestjs/typeorm';
+
 import { ClickHouseService } from 'src/database/clickHouse/clickHouse.service';
-import { type BillingSubscriptionEntity } from 'src/engine/core-modules/billing/entities/billing-subscription.entity';
-import {
-  type BillingCapEvaluation,
-  BillingUsageCapService,
-} from 'src/engine/core-modules/billing/services/billing-usage-cap.service';
+import { BillingSubscriptionItemEntity } from 'src/engine/core-modules/billing/entities/billing-subscription-item.entity';
+import { BillingSubscriptionEntity } from 'src/engine/core-modules/billing/entities/billing-subscription.entity';
+import { BillingUsageCapService } from 'src/engine/core-modules/billing/services/billing-usage-cap.service';
 import { MeteredCreditService } from 'src/engine/core-modules/billing/services/metered-credit.service';
+import { CacheStorageNamespace } from 'src/engine/core-modules/cache-storage/types/cache-storage-namespace.enum';
 import { TwentyConfigService } from 'src/engine/core-modules/twenty-config/twenty-config.service';
+import { WorkspaceCacheService } from 'src/engine/workspace-cache/services/workspace-cache.service';

 describe('BillingUsageCapService', () => {
  let service: BillingUsageCapService;
@@ -52,6 +54,34 @@ describe('BillingUsageCapService', () => {
            get: jest.fn(),
          },
        },
+        {
+          provide: CacheStorageNamespace.EngineBillingUsage,
+          useValue: {
+            get: jest.fn(),
+            set: jest.fn(),
+            del: jest.fn(),
+            incrBy: jest.fn(),
+          },
+        },
+        {
+          provide: getRepositoryToken(BillingSubscriptionEntity),
+          useValue: {
+            findOne: jest.fn(),
+          },
+        },
+        {
+          provide: getRepositoryToken(BillingSubscriptionItemEntity),
+          useValue: {
+            find: jest.fn(),
+            update: jest.fn(),
+          },
+        },
+        {
+          provide: WorkspaceCacheService,
+          useValue: {
+            getOrRecompute: jest.fn(),
+          },
+        },
      ],
    }).compile();

@@ -79,83 +109,6 @@ describe('BillingUsageCapService', () => {
    });
  });

-  describe('getCurrentPeriodCreditsUsed', () => {
-    beforeEach(() => {
-      twentyConfigService.get.mockReturnValue('http://clickhouse:8123');
-    });
-
-    it('returns 0 when ClickHouse is disabled', async () => {
-      twentyConfigService.get.mockReturnValue('');
-
-      const result = await service.getCurrentPeriodCreditsUsed(
-        'workspace_123',
-        new Date('2026-04-01T00:00:00Z'),
-        new Date('2026-05-01T00:00:00Z'),
-      );
-
-      expect(result).toBe(0);
-      expect(clickHouseService.select).not.toHaveBeenCalled();
-    });
-
-    it('sums creditsUsedMicro for the workspace in the given period', async () => {
-      clickHouseService.select.mockResolvedValue([{ total: 12345 }]);
-
-      const result = await service.getCurrentPeriodCreditsUsed(
-        'workspace_123',
-        new Date('2026-04-01T00:00:00Z'),
-        new Date('2026-05-01T00:00:00Z'),
-      );
-
-      expect(result).toBe(12345);
-      expect(clickHouseService.select).toHaveBeenCalledTimes(1);
-
-      const [query, params] = clickHouseService.select.mock.calls[0];
-
-      expect(query).toContain('sum(creditsUsedMicro)');
-      expect(query).toContain('FROM usageEvent');
-      expect(params).toMatchObject({
-        workspaceId: 'workspace_123',
-      });
-      expect(query).not.toContain('operationType');
-    });
-
-    it('coerces string totals returned by ClickHouse to a number', async () => {
-      clickHouseService.select.mockResolvedValue([{ total: '9876543210' }]);
-
-      const result = await service.getCurrentPeriodCreditsUsed(
-        'workspace_123',
-        new Date('2026-04-01T00:00:00Z'),
-        new Date('2026-05-01T00:00:00Z'),
-      );
-
-      expect(result).toBe(9876543210);
-    });
-
-    it('returns 0 when no rows are returned', async () => {
-      clickHouseService.select.mockResolvedValue([]);
-
-      const result = await service.getCurrentPeriodCreditsUsed(
-        'workspace_123',
-        new Date('2026-04-01T00:00:00Z'),
-        new Date('2026-05-01T00:00:00Z'),
-      );
-
-      expect(result).toBe(0);
-    });
-
-    it('returns 0 when total is null', async () => {
-      clickHouseService.select.mockResolvedValue([{ total: null }]);
-
-      const result = await service.getCurrentPeriodCreditsUsed(
-        'workspace_123',
-        new Date('2026-04-01T00:00:00Z'),
-        new Date('2026-05-01T00:00:00Z'),
-      );
-
-      expect(result).toBe(0);
-    });
-  });
-
  describe('getBatchPeriodCreditsUsed', () => {
    beforeEach(() => {
      twentyConfigService.get.mockReturnValue('http://clickhouse:8123');
@@ -167,7 +120,6 @@ describe('BillingUsageCapService', () => {
      const result = await service.getBatchPeriodCreditsUsed(
        ['ws_1', 'ws_2'],
        new Date('2026-04-01T00:00:00Z'),
-        new Date('2026-05-01T00:00:00Z'),
      );

      expect(result.size).toBe(0);
@@ -178,7 +130,6 @@ describe('BillingUsageCapService', () => {
      const result = await service.getBatchPeriodCreditsUsed(
        [],
        new Date('2026-04-01T00:00:00Z'),
-        new Date('2026-05-01T00:00:00Z'),
      );

      expect(result.size).toBe(0);
@@ -194,7 +145,6 @@ describe('BillingUsageCapService', () => {
      const result = await service.getBatchPeriodCreditsUsed(
        ['ws_1', 'ws_2', 'ws_3'],
        new Date('2026-04-01T00:00:00Z'),
-        new Date('2026-05-01T00:00:00Z'),
      );

      expect(result.get('ws_1')).toBe(500_000);
@@ -218,122 +168,12 @@ describe('BillingUsageCapService', () => {
      const result = await service.getBatchPeriodCreditsUsed(
        ['ws_1'],
        new Date('2026-04-01T00:00:00Z'),
-        new Date('2026-05-01T00:00:00Z'),
      );

      expect(result.get('ws_1')).toBe(9876543210);
    });
  });

-  describe('evaluateCap', () => {
-    beforeEach(() => {
-      twentyConfigService.get.mockReturnValue('http://clickhouse:8123');
-    });
-
-    it('returns skipped when ClickHouse is disabled', async () => {
-      twentyConfigService.get.mockReturnValue('');
-
-      const result = await service.evaluateCap(buildSubscription());
-
-      expect(result).toEqual<BillingCapEvaluation>({
-        skipped: true,
-        reason: 'clickhouse-disabled',
-      });
-      expect(
-        meteredCreditService.extractMeteredPricingInfoFromSubscription,
-      ).not.toHaveBeenCalled();
-    });
-
-    it('returns skipped when subscription has no metered item', async () => {
-      meteredCreditService.extractMeteredPricingInfoFromSubscription.mockReturnValue(
-        null,
-      );
-
-      const result = await service.evaluateCap(buildSubscription());
-
-      expect(result).toEqual<BillingCapEvaluation>({
-        skipped: true,
-        reason: 'no-metered-item',
-      });
-      expect(clickHouseService.select).not.toHaveBeenCalled();
-    });
-
-    it('reports hasReachedCap=false when usage is below tierCap + creditBalance', async () => {
-      meteredCreditService.extractMeteredPricingInfoFromSubscription.mockReturnValue(
-        { tierCap: 1_000_000, unitPriceCents: 10 },
-      );
-      meteredCreditService.getCreditBalance.mockResolvedValue(250_000);
-      clickHouseService.select.mockResolvedValue([{ total: 800_000 }]);
-
-      const result = await service.evaluateCap(buildSubscription());
-
-      expect(result).toEqual<BillingCapEvaluation>({
-        skipped: false,
-        hasReachedCap: false,
-        usage: 800_000,
-        allowance: 1_250_000,
-        tierCap: 1_000_000,
-        creditBalance: 250_000,
-      });
-    });
-
-    it('reports hasReachedCap=true when usage meets tierCap + creditBalance', async () => {
-      meteredCreditService.extractMeteredPricingInfoFromSubscription.mockReturnValue(
-        { tierCap: 1_000_000, unitPriceCents: 10 },
-      );
-      meteredCreditService.getCreditBalance.mockResolvedValue(0);
-      clickHouseService.select.mockResolvedValue([{ total: 1_000_000 }]);
-
-      const result = await service.evaluateCap(buildSubscription());
-
-      expect(result).toMatchObject({
-        skipped: false,
-        hasReachedCap: true,
-        usage: 1_000_000,
-        allowance: 1_000_000,
-      });
-    });
-
-    it('reports hasReachedCap=true when usage exceeds tierCap + creditBalance', async () => {
-      meteredCreditService.extractMeteredPricingInfoFromSubscription.mockReturnValue(
-        { tierCap: 1_000_000, unitPriceCents: 10 },
-      );
-      meteredCreditService.getCreditBalance.mockResolvedValue(100_000);
-      clickHouseService.select.mockResolvedValue([{ total: 5_000_000 }]);
-
-      const result = await service.evaluateCap(buildSubscription());
-
-      expect(result).toMatchObject({
-        skipped: false,
-        hasReachedCap: true,
-        usage: 5_000_000,
-        allowance: 1_100_000,
-      });
-    });
-
-    it('re-reads pricing on each call so that tier changes apply immediately', async () => {
-      meteredCreditService.extractMeteredPricingInfoFromSubscription
-        .mockReturnValueOnce({ tierCap: 2_000_000, unitPriceCents: 10 })
-        .mockReturnValueOnce({ tierCap: 500_000, unitPriceCents: 10 });
-      meteredCreditService.getCreditBalance.mockResolvedValue(0);
-      clickHouseService.select.mockResolvedValue([{ total: 1_000_000 }]);
-
-      const first = await service.evaluateCap(buildSubscription());
-      const second = await service.evaluateCap(buildSubscription());
-
-      expect(first).toMatchObject({
-        skipped: false,
-        hasReachedCap: false,
-        allowance: 2_000_000,
-      });
-      expect(second).toMatchObject({
-        skipped: false,
-        hasReachedCap: true,
-        allowance: 500_000,
-      });
-    });
-  });
-
  describe('evaluateCapBatch', () => {
    it('returns evaluations keyed by subscription id', () => {
      meteredCreditService.extractMeteredPricingInfoFromSubscription.mockReturnValue(
@@ -2,11 +2,20 @@

 import { Injectable } from '@nestjs/common';

+import { InjectRepository } from '@nestjs/typeorm';
 import { ClickHouseService } from 'src/database/clickHouse/clickHouse.service';
 import { formatDateForClickHouse } from 'src/database/clickHouse/clickHouse.util';
-import { type BillingSubscriptionEntity } from 'src/engine/core-modules/billing/entities/billing-subscription.entity';
+import {
+  BillingException,
+  BillingExceptionCode,
+} from 'src/engine/core-modules/billing/billing.exception';
+import { BillingSubscriptionItemEntity } from 'src/engine/core-modules/billing/entities/billing-subscription-item.entity';
+import { BillingSubscriptionEntity } from 'src/engine/core-modules/billing/entities/billing-subscription.entity';
+import { BillingProductKey } from 'src/engine/core-modules/billing/enums/billing-product-key.enum';
+import { SubscriptionStatus } from 'src/engine/core-modules/billing/enums/billing-subscription-status.enum';
 import { MeteredCreditService } from 'src/engine/core-modules/billing/services/metered-credit.service';
 import { TwentyConfigService } from 'src/engine/core-modules/twenty-config/twenty-config.service';
+import { Not, Raw, Repository } from 'typeorm';

 export type BillingCapEvaluation =
  | {
@@ -22,10 +31,6 @@ export type BillingCapEvaluation =
      reason: 'no-metered-item' | 'clickhouse-disabled';
    };

-type UsageSumRow = {
-  total: string | number | null;
-};
-
 type BatchUsageSumRow = {
  workspaceId: string;
  total: string | number | null;
@@ -37,45 +42,17 @@ export class BillingUsageCapService {
    private readonly clickHouseService: ClickHouseService,
    private readonly meteredCreditService: MeteredCreditService,
    private readonly twentyConfigService: TwentyConfigService,
+    @InjectRepository(BillingSubscriptionItemEntity)
+    private readonly billingSubscriptionItemRepository: Repository<BillingSubscriptionItemEntity>,
  ) {}

  isClickHouseEnabled(): boolean {
    return Boolean(this.twentyConfigService.get('CLICKHOUSE_URL'));
  }

-  async getCurrentPeriodCreditsUsed(
-    workspaceId: string,
-    periodStart: Date,
-    periodEnd: Date,
-  ): Promise<number> {
-    if (!this.isClickHouseEnabled()) {
-      return 0;
-    }
-
-    const query = `
-      SELECT sum(creditsUsedMicro) AS total
-      FROM usageEvent
-      WHERE workspaceId = {workspaceId:String}
-        AND timestamp >= {periodStart:String}
-        AND timestamp < {periodEnd:String}
-    `;
-
-    const rows = await this.clickHouseService.select<UsageSumRow>(query, {
-      workspaceId,
-      periodStart: formatDateForClickHouse(periodStart),
-      periodEnd: formatDateForClickHouse(periodEnd),
-    });
-
-    const rawTotal = rows[0]?.total ?? 0;
-    const total = typeof rawTotal === 'string' ? Number(rawTotal) : rawTotal;
-
-    return Number.isFinite(total) ? total : 0;
-  }
-
  async getBatchPeriodCreditsUsed(
    workspaceIds: string[],
    periodStart: Date,
-    periodEnd: Date,
  ): Promise<Map<string, number>> {
    const result = new Map<string, number>();

@@ -87,15 +64,13 @@ export class BillingUsageCapService {
      SELECT workspaceId, sum(creditsUsedMicro) AS total
      FROM usageEvent
      WHERE workspaceId IN {workspaceIds:Array(String)}
-        AND timestamp >= {periodStart:String}
-        AND timestamp < {periodEnd:String}
+        AND periodStart = {periodStart:Date}
      GROUP BY workspaceId
    `;

    const rows = await this.clickHouseService.select<BatchUsageSumRow>(query, {
      workspaceIds,
      periodStart: formatDateForClickHouse(periodStart),
-      periodEnd: formatDateForClickHouse(periodEnd),
    });

    for (const row of rows) {
@@ -108,46 +83,6 @@ export class BillingUsageCapService {
    return result;
  }

-  async evaluateCap(
-    subscription: BillingSubscriptionEntity,
-  ): Promise<BillingCapEvaluation> {
-    if (!this.isClickHouseEnabled()) {
-      return { skipped: true, reason: 'clickhouse-disabled' };
-    }
-
-    const meteredPricingInfo =
-      this.meteredCreditService.extractMeteredPricingInfoFromSubscription(
-        subscription,
-      );
-
-    if (!meteredPricingInfo) {
-      return { skipped: true, reason: 'no-metered-item' };
-    }
-
-    const [creditBalance, usage] = await Promise.all([
-      this.meteredCreditService.getCreditBalance(
-        subscription.stripeCustomerId,
-        meteredPricingInfo.unitPriceCents,
-      ),
-      this.getCurrentPeriodCreditsUsed(
-        subscription.workspaceId,
-        subscription.currentPeriodStart,
-        subscription.currentPeriodEnd,
-      ),
-    ]);
-
-    const allowance = meteredPricingInfo.tierCap + creditBalance;
-
-    return {
-      skipped: false,
-      hasReachedCap: usage >= allowance,
-      usage,
-      allowance,
-      tierCap: meteredPricingInfo.tierCap,
-      creditBalance,
-    };
-  }
-
  evaluateCapBatch(
    subscriptions: BillingSubscriptionEntity[],
    usageByWorkspace: Map<string, number>,
@@ -186,4 +121,38 @@ export class BillingUsageCapService {

    return results;
  }
+
+  async setSubscriptionItemHasReachedCap(
+    workspaceId: string,
+    hasReachedCap: boolean,
+  ): Promise<void> {
+    const billingSubscriptionItems =
+      await this.billingSubscriptionItemRepository.find({
+        where: {
+          billingSubscription: {
+            workspaceId,
+            status: Not(SubscriptionStatus.Canceled),
+          },
+          billingProduct: {
+            metadata: Raw((alias) => `${alias} @> :metadata::jsonb`, {
+              metadata: JSON.stringify({
+                productKey: BillingProductKey.WORKFLOW_NODE_EXECUTION,
+              }),
+            }),
+          },
+        },
+      });
+
+    if (billingSubscriptionItems.length !== 1) {
+      throw new BillingException(
+        `Expected 1 metered billing subscription item for workspace ${workspaceId}, but got ${billingSubscriptionItems.length}`,
+        BillingExceptionCode.BILLING_SUBSCRIPTION_ITEM_NOT_FOUND,
+      );
+    }
+
+    await this.billingSubscriptionItemRepository.update(
+      { id: billingSubscriptionItems[0].id },
+      { hasReachedCurrentPeriodCap: hasReachedCap },
+    );
+  }
 }
@@ -6,6 +6,8 @@ import { InjectRepository } from '@nestjs/typeorm';
 import { isDefined } from 'twenty-shared/utils';
 import { type Repository } from 'typeorm';

+import { ClickHouseService } from 'src/database/clickHouse/clickHouse.service';
+import { formatDateTimeForClickHouse } from 'src/database/clickHouse/clickHouse.util';
 import {
  BillingException,
  BillingExceptionCode,
@@ -16,11 +18,22 @@ import { BillingSubscriptionEntity } from 'src/engine/core-modules/billing/entit
 import { SubscriptionStatus } from 'src/engine/core-modules/billing/enums/billing-subscription-status.enum';
 import { BillingSubscriptionItemService } from 'src/engine/core-modules/billing/services/billing-subscription-item.service';
 import { BillingSubscriptionService } from 'src/engine/core-modules/billing/services/billing-subscription.service';
+import { BillingUsageCapService } from 'src/engine/core-modules/billing/services/billing-usage-cap.service';
+import { MeteredCreditService } from 'src/engine/core-modules/billing/services/metered-credit.service';
 import { StripeBillingMeterEventService } from 'src/engine/core-modules/billing/stripe/services/stripe-billing-meter-event.service';
 import { StripeCreditGrantService } from 'src/engine/core-modules/billing/stripe/services/stripe-credit-grant.service';
-import { type UsageEvent } from 'src/engine/core-modules/usage/types/usage-event.type';
+import { buildBillingUsageAvailableCreditsCacheKey } from 'src/engine/core-modules/billing/utils/build-billing-usage-available-credits-cache-key.util';
+import { InjectCacheStorage } from 'src/engine/core-modules/cache-storage/decorators/cache-storage.decorator';
+import { CacheStorageService } from 'src/engine/core-modules/cache-storage/services/cache-storage.service';
+import { CacheStorageNamespace } from 'src/engine/core-modules/cache-storage/types/cache-storage-namespace.enum';
 import { TwentyConfigService } from 'src/engine/core-modules/twenty-config/twenty-config.service';
+import { type UsageEvent } from 'src/engine/core-modules/usage/types/usage-event.type';
 import { type WorkspaceEntity } from 'src/engine/core-modules/workspace/workspace.entity';
+import { WorkspaceCacheService } from 'src/engine/workspace-cache/services/workspace-cache.service';
+
+type UsageSumRow = {
+  total: string | number | null;
+};

@Injectable()
 export class BillingUsageService {
@@ -33,6 +46,14 @@ export class BillingUsageService {
    private readonly twentyConfigService: TwentyConfigService,
    private readonly billingSubscriptionItemService: BillingSubscriptionItemService,
    private readonly stripeCreditGrantService: StripeCreditGrantService,
+    @InjectCacheStorage(CacheStorageNamespace.EngineBillingUsage)
+    private readonly billingUsageCacheStorage: CacheStorageService,
+    @InjectRepository(BillingSubscriptionEntity)
+    private readonly billingSubscriptionRepository: Repository<BillingSubscriptionEntity>,
+    private readonly meteredCreditService: MeteredCreditService,
+    private readonly workspaceCacheService: WorkspaceCacheService,
+    private readonly clickHouseService: ClickHouseService,
+    private readonly billingUsageCapService: BillingUsageCapService,
  ) {}

  async canFeatureBeUsed(workspaceId: string): Promise<boolean> {
@@ -48,6 +69,7 @@ export class BillingUsageService {
    return !!billingSubscription;
  }

+  //TODO: TO be deprecated
  async billUsage({
    workspaceId,
    usageEvents,
@@ -86,6 +108,7 @@ export class BillingUsageService {
    }
  }

+  //TODO: TO be deprecated
  async getMeteredProductsUsage(
    workspace: WorkspaceEntity,
  ): Promise<BillingMeteredProductUsageDTO[]> {
@@ -113,6 +136,7 @@ export class BillingUsageService {
    );
  }

+  //TODO: TO be deprecated
  private getSubscriptionPeriod(subscription: BillingSubscriptionEntity): {
    periodStart: Date;
    periodEnd: Date;
@@ -135,6 +159,7 @@ export class BillingUsageService {
    };
  }

+  //TODO: TO be deprecated
  private async buildMeteredProductUsage(
    subscription: BillingSubscriptionEntity,
    item: Awaited<
@@ -175,4 +200,194 @@ export class BillingUsageService {
      unitPriceCents: item.unitPriceCents,
    };
  }
+
+  async flushAvailableCreditsFromCache(workspaceId: string): Promise<void> {
+    await this.billingUsageCacheStorage.flushByPattern(
+      `available-credits:${workspaceId}:*`,
+    );
+  }
+
+  private async warmAvailableCredits(
+    workspaceId: string,
+    periodStart: Date | string,
+    periodEnd: Date | string,
+    availableCredits: number,
+  ): Promise<void> {
+    const ttlMs = Math.max(new Date(periodEnd).getTime() - Date.now(), 0);
+
+    await this.billingUsageCacheStorage.set(
+      buildBillingUsageAvailableCreditsCacheKey(workspaceId, periodStart),
+      availableCredits,
+      ttlMs,
+    );
+  }
+
+  private async getAvailableCreditsFromCache(
+    workspaceId: string,
+    periodStart: Date | string,
+  ): Promise<number | undefined> {
+    return this.billingUsageCacheStorage.get<number>(
+      buildBillingUsageAvailableCreditsCacheKey(workspaceId, periodStart),
+    );
+  }
+
+  private async getAvailableCreditsFromClickHouse({
+    workspaceId,
+    currentPeriodStart,
+  }: {
+    workspaceId: string;
+    currentPeriodStart: Date | string;
+  }): Promise<number> {
+    const subscription = await this.billingSubscriptionRepository.findOne({
+      where: { workspaceId, currentPeriodStart: new Date(currentPeriodStart) },
+      relations: [
+        'billingSubscriptionItems',
+        'billingSubscriptionItems.billingProduct',
+        'billingSubscriptionItems.billingProduct.billingPrices',
+      ],
+    });
+
+    if (!isDefined(subscription)) {
+      throw new BillingException(
+        `Subscription not found for workspace ${workspaceId}`,
+        BillingExceptionCode.BILLING_SUBSCRIPTION_NOT_FOUND,
+      );
+    }
+
+    const meteredPricingInfo =
+      this.meteredCreditService.extractMeteredPricingInfoFromSubscription(
+        subscription,
+      );
+
+    if (!meteredPricingInfo) {
+      throw new BillingException(
+        `No metered item found for workspace ${workspaceId}`,
+        BillingExceptionCode.BILLING_SUBSCRIPTION_ITEM_NOT_FOUND,
+      );
+    }
+
+    const [creditBalance, usage] = await Promise.all([
+      this.meteredCreditService.getCreditBalance(
+        subscription.stripeCustomerId,
+        meteredPricingInfo.unitPriceCents,
+      ),
+      this.getCurrentPeriodCreditsUsed(
+        subscription.workspaceId,
+        subscription.currentPeriodStart,
+      ),
+    ]);
+
+    return meteredPricingInfo.tierCap + creditBalance - usage;
+  }
+
+  async decrementAvailableCredits({
+    workspaceId,
+    usedCredits,
+  }: {
+    workspaceId: string;
+    usedCredits: number;
+  }): Promise<void> {
+    const {
+      billingSubscription: { currentPeriodStart, currentPeriodEnd },
+    } = await this.workspaceCacheService.getOrRecompute(workspaceId, [
+      'billingSubscription',
+    ]);
+
+    const cachedAvailableCredits = await this.getAvailableCreditsFromCache(
+      workspaceId,
+      currentPeriodStart,
+    );
+
+    const availableCredits = isDefined(cachedAvailableCredits)
+      ? cachedAvailableCredits
+      : await this.getAvailableCreditsFromClickHouse({
+          workspaceId,
+          currentPeriodStart,
+        });
+
+    if (!isDefined(cachedAvailableCredits)) {
+      await this.warmAvailableCredits(
+        workspaceId,
+        currentPeriodStart,
+        currentPeriodEnd,
+        availableCredits,
+      );
+    }
+
+    const decrementedAvailableCredits =
+      await this.billingUsageCacheStorage.incrBy(
+        buildBillingUsageAvailableCreditsCacheKey(
+          workspaceId,
+          currentPeriodStart,
+        ),
+        -usedCredits,
+      );
+
+    if (decrementedAvailableCredits <= 0) {
+      await this.billingUsageCapService.setSubscriptionItemHasReachedCap(
+        workspaceId,
+        true,
+      );
+    }
+  }
+
+  async invalidateAvailableCredits(
+    workspaceId: string,
+    periodStart: Date,
+  ): Promise<void> {
+    await this.billingUsageCacheStorage.del(
+      buildBillingUsageAvailableCreditsCacheKey(workspaceId, periodStart),
+    );
+  }
+
+  async hasAvailableCredits(workspaceId: string): Promise<boolean> {
+    const { billingSubscription: subscription } =
+      await this.workspaceCacheService.getOrRecompute(workspaceId, [
+        'billingSubscription',
+      ]);
+    const cached = await this.getAvailableCreditsFromCache(
+      subscription.workspaceId,
+      subscription.currentPeriodStart,
+    );
+
+    if (isDefined(cached)) {
+      return cached > 0;
+    }
+
+    const availableCredits = await this.getAvailableCreditsFromClickHouse({
+      workspaceId: subscription.workspaceId,
+      currentPeriodStart: subscription.currentPeriodStart,
+    });
+
+    await this.warmAvailableCredits(
+      subscription.workspaceId,
+      subscription.currentPeriodStart,
+      subscription.currentPeriodEnd,
+      availableCredits,
+    );
+
+    return availableCredits > 0;
+  }
+
+  async getCurrentPeriodCreditsUsed(
+    workspaceId: string,
+    periodStart: Date,
+  ): Promise<number> {
+    const query = `
+      SELECT sum(creditsUsedMicro) AS total
+      FROM usageEvent
+      WHERE workspaceId = {workspaceId:String}
+        AND periodStart = {periodStart:DateTime64(3)}
+    `;
+
+    const rows = await this.clickHouseService.select<UsageSumRow>(query, {
+      workspaceId,
+      periodStart: formatDateTimeForClickHouse(periodStart),
+    });
+
+    const rawTotal = rows[0]?.total ?? 0;
+    const total = typeof rawTotal === 'string' ? Number(rawTotal) : rawTotal;
+
+    return Number.isFinite(total) ? total : 0;
+  }
 }
@@ -8,11 +8,8 @@ import { type Repository } from 'typeorm';

 import { BillingSubscriptionEntity } from 'src/engine/core-modules/billing/entities/billing-subscription.entity';
 import { type BillingEntitlementKey } from 'src/engine/core-modules/billing/enums/billing-entitlement-key.enum';
-import { type BillingProductKey } from 'src/engine/core-modules/billing/enums/billing-product-key.enum';
-import { SubscriptionStatus } from 'src/engine/core-modules/billing/enums/billing-subscription-status.enum';
 import { BillingProductService } from 'src/engine/core-modules/billing/services/billing-product.service';
 import { BillingSubscriptionService } from 'src/engine/core-modules/billing/services/billing-subscription.service';
-import { getPlanKeyFromSubscription } from 'src/engine/core-modules/billing/utils/get-plan-key-from-subscription.util';
 import { TwentyConfigService } from 'src/engine/core-modules/twenty-config/twenty-config.service';

@Injectable()
@@ -66,41 +63,4 @@ export class BillingService {

    return !hasAnySubscription;
  }
-
-  async canBillMeteredProduct(
-    workspaceId: string,
-    productKey: BillingProductKey,
-  ): Promise<boolean> {
-    const subscription =
-      await this.billingSubscriptionService.getCurrentBillingSubscriptionOrThrow(
-        { workspaceId },
-      );
-
-    const billableStatuses = [
-      SubscriptionStatus.Active,
-      SubscriptionStatus.Trialing,
-    ];
-
-    if (!billableStatuses.includes(subscription.status)) {
-      return false;
-    }
-
-    const planKey = getPlanKeyFromSubscription(subscription);
-    const products =
-      await this.billingProductService.getProductsByPlan(planKey);
-
-    const targetProduct = products.find(
-      ({ metadata }) => metadata.productKey === productKey,
-    );
-
-    if (!targetProduct) {
-      return false;
-    }
-
-    const subscriptionItem = subscription.billingSubscriptionItems.find(
-      (item) => item.stripeProductId === targetProduct.stripeProductId,
-    );
-
-    return subscriptionItem?.hasReachedCurrentPeriodCap === false;
-  }
 }
@@ -0,0 +1,46 @@
+/* @license Enterprise */
+
+import { Injectable } from '@nestjs/common';
+import { BillingSubscriptionService } from 'src/engine/core-modules/billing/services/billing-subscription.service';
+
+import { type FlatBillingSubscription } from 'src/engine/core-modules/billing/types/flat-billing-subscription.type';
+import { WorkspaceCache } from 'src/engine/workspace-cache/decorators/workspace-cache.decorator';
+import { WorkspaceCacheProvider } from 'src/engine/workspace-cache/interfaces/workspace-cache-provider.service';
+
+@Injectable()
+@WorkspaceCache('billingSubscription')
+export class WorkspaceBillingSubscriptionCacheService extends WorkspaceCacheProvider<FlatBillingSubscription> {
+  constructor(
+    private readonly billingSubscriptionService: BillingSubscriptionService,
+  ) {
+    super();
+  }
+
+  async computeForCache(workspaceId: string): Promise<FlatBillingSubscription> {
+    const subscription =
+      await this.billingSubscriptionService.getCurrentBillingSubscriptionOrThrow(
+        {
+          workspaceId,
+        },
+      );
+
+    return {
+      id: subscription.id,
+      workspaceId: subscription.workspaceId,
+      stripeCustomerId: subscription.stripeCustomerId,
+      stripeSubscriptionId: subscription.stripeSubscriptionId,
+      status: subscription.status,
+      interval: subscription.interval,
+      currency: subscription.currency,
+      currentPeriodStart: subscription.currentPeriodStart,
+      currentPeriodEnd: subscription.currentPeriodEnd,
+      cancelAtPeriodEnd: subscription.cancelAtPeriodEnd,
+      cancelAt: subscription.cancelAt,
+      canceledAt: subscription.canceledAt,
+      endedAt: subscription.endedAt,
+      trialStart: subscription.trialStart,
+      trialEnd: subscription.trialEnd,
+      collectionMethod: subscription.collectionMethod,
+    };
+  }
+}
@@ -0,0 +1,24 @@
+/* @license Enterprise */
+
+import { type BillingSubscriptionCollectionMethod } from 'src/engine/core-modules/billing/enums/billing-subscription-collection-method.enum';
+import { type SubscriptionInterval } from 'src/engine/core-modules/billing/enums/billing-subscription-interval.enum';
+import { type SubscriptionStatus } from 'src/engine/core-modules/billing/enums/billing-subscription-status.enum';
+
+export type FlatBillingSubscription = {
+  id: string;
+  workspaceId: string;
+  stripeCustomerId: string;
+  stripeSubscriptionId: string;
+  status: SubscriptionStatus;
+  interval: SubscriptionInterval;
+  currency: string;
+  currentPeriodStart: Date;
+  currentPeriodEnd: Date;
+  cancelAtPeriodEnd: boolean;
+  cancelAt: Date | null;
+  canceledAt: Date | null;
+  endedAt: Date | null;
+  trialStart: Date | null;
+  trialEnd: Date | null;
+  collectionMethod: BillingSubscriptionCollectionMethod;
+};
@@ -0,0 +1,6 @@
+export const buildBillingUsageAvailableCreditsCacheKey = (
+  workspaceId: string,
+  periodStart: Date | string,
+): string => {
+  return `available-credits:${workspaceId}:${new Date(periodStart).getTime()}`;
+};