diff --git a/packages/twenty-server/src/database/commands/cron-register-all.command.ts b/packages/twenty-server/src/database/commands/cron-register-all.command.ts index f19e9d2806..0381719d4d 100644 --- a/packages/twenty-server/src/database/commands/cron-register-all.command.ts +++ b/packages/twenty-server/src/database/commands/cron-register-all.command.ts @@ -5,6 +5,7 @@ import { Command, CommandRunner, Option } from 'nest-commander'; import { MarketplaceCatalogSyncCronCommand } from 'src/engine/core-modules/application/application-marketplace/crons/commands/marketplace-catalog-sync.cron.command'; import { StaleRegistrationCleanupCronCommand } from 'src/engine/core-modules/application/application-oauth/stale-registration-cleanup/commands/stale-registration-cleanup.cron.command'; import { ApplicationVersionCheckCronCommand } from 'src/engine/core-modules/application/application-upgrade/crons/commands/application-version-check.cron.command'; +import { EnforceUsageCapCronCommand } from 'src/engine/core-modules/billing/crons/commands/enforce-usage-cap.cron.command'; import { EnterpriseKeyValidationCronCommand } from 'src/engine/core-modules/enterprise/cron/command/enterprise-key-validation.cron.command'; import { EventLogCleanupCronCommand } from 'src/engine/core-modules/event-logs/cleanup/commands/event-log-cleanup.cron.command'; import { CronTriggerCronCommand } from 'src/engine/core-modules/logic-function/logic-function-trigger/triggers/cron/cron-trigger.cron.command'; @@ -62,6 +63,7 @@ export class CronRegisterAllCommand extends CommandRunner { private readonly marketplaceCatalogSyncCronCommand: MarketplaceCatalogSyncCronCommand, private readonly applicationVersionCheckCronCommand: ApplicationVersionCheckCronCommand, private readonly staleRegistrationCleanupCronCommand: StaleRegistrationCleanupCronCommand, + private readonly enforceUsageCapCronCommand: EnforceUsageCapCronCommand, ) { super(); } @@ -187,6 +189,10 @@ export class CronRegisterAllCommand extends CommandRunner { name: 'StaleRegistrationCleanup', command: this.staleRegistrationCleanupCronCommand, }, + { + name: 'EnforceUsageCap', + command: this.enforceUsageCapCronCommand, + }, ]; const commands = this.devMode diff --git a/packages/twenty-server/src/database/commands/database-command.module.ts b/packages/twenty-server/src/database/commands/database-command.module.ts index d6a6eb0422..c36907285a 100644 --- a/packages/twenty-server/src/database/commands/database-command.module.ts +++ b/packages/twenty-server/src/database/commands/database-command.module.ts @@ -17,6 +17,7 @@ import { GenerateApiKeyCommand } from 'src/engine/core-modules/api-key/commands/ import { MarketplaceModule } from 'src/engine/core-modules/application/application-marketplace/marketplace.module'; import { StaleRegistrationCleanupModule } from 'src/engine/core-modules/application/application-oauth/stale-registration-cleanup/stale-registration-cleanup.module'; import { ApplicationUpgradeModule } from 'src/engine/core-modules/application/application-upgrade/application-upgrade.module'; +import { EnforceUsageCapCronCommand } from 'src/engine/core-modules/billing/crons/commands/enforce-usage-cap.cron.command'; import { EnterpriseKeyValidationCronCommand } from 'src/engine/core-modules/enterprise/cron/command/enterprise-key-validation.cron.command'; import { EnterpriseModule } from 'src/engine/core-modules/enterprise/enterprise.module'; import { EventLogCleanupModule } from 'src/engine/core-modules/event-logs/cleanup/event-log-cleanup.module'; @@ -86,6 +87,7 @@ import { AutomatedTriggerModule } from 'src/modules/workflow/workflow-trigger/au ListOrphanedWorkspaceEntitiesCommand, EnterpriseKeyValidationCronCommand, GenerateApiKeyCommand, + EnforceUsageCapCronCommand, ], }) export class DatabaseCommandModule {} diff --git a/packages/twenty-server/src/database/commands/upgrade-version-command/1-22/1-22-instance-command-fast-1776078919203-add-credit-balance-to-billing-customer.ts b/packages/twenty-server/src/database/commands/upgrade-version-command/1-22/1-22-instance-command-fast-1776078919203-add-credit-balance-to-billing-customer.ts new file mode 100644 index 0000000000..a5b8f55768 --- /dev/null +++ b/packages/twenty-server/src/database/commands/upgrade-version-command/1-22/1-22-instance-command-fast-1776078919203-add-credit-balance-to-billing-customer.ts @@ -0,0 +1,29 @@ +import { QueryRunner } from 'typeorm'; + +import { RegisteredInstanceCommand } from 'src/engine/core-modules/upgrade/decorators/registered-instance-command.decorator'; +import { FastInstanceCommand } from 'src/engine/core-modules/upgrade/interfaces/fast-instance-command.interface'; + +@RegisteredInstanceCommand('1.22.0', 1776078919203) +export class AddCreditBalanceToBillingCustomerFastInstanceCommand + implements FastInstanceCommand +{ + public async up(queryRunner: QueryRunner): Promise { + const tableExists = await queryRunner.query( + `SELECT 1 FROM pg_tables WHERE schemaname = 'core' AND tablename = 'billingCustomer'`, + ); + + if (tableExists.length === 0) { + return; + } + + await queryRunner.query( + `ALTER TABLE "core"."billingCustomer" ADD COLUMN IF NOT EXISTS "creditBalanceMicro" bigint NOT NULL DEFAULT 0`, + ); + } + + public async down(queryRunner: QueryRunner): Promise { + await queryRunner.query( + `ALTER TABLE "core"."billingCustomer" DROP COLUMN IF EXISTS "creditBalanceMicro"`, + ); + } +} diff --git a/packages/twenty-server/src/database/commands/upgrade-version-command/instance-commands.constant.ts b/packages/twenty-server/src/database/commands/upgrade-version-command/instance-commands.constant.ts index cb2459c3b7..2f63cd90e8 100644 --- a/packages/twenty-server/src/database/commands/upgrade-version-command/instance-commands.constant.ts +++ b/packages/twenty-server/src/database/commands/upgrade-version-command/instance-commands.constant.ts @@ -7,6 +7,7 @@ import { AddPermissionFlagRoleIdIndexFastInstanceCommand } from 'src/database/co import { AddWorkspaceIdToIndirectEntitiesFastInstanceCommand } from 'src/database/commands/upgrade-version-command/1-22/1-22-instance-command-fast-1775758621017-add-workspace-id-to-indirect-entities'; import { AddWorkspaceIdIndexesAndFksFastInstanceCommand } from 'src/database/commands/upgrade-version-command/1-22/1-22-instance-command-fast-1775761294897-add-workspace-id-indexes-and-fks-to-indirect-entities'; import { DropObjectMetadataDataSourceFkFastInstanceCommand } from 'src/database/commands/upgrade-version-command/1-22/1-22-instance-command-fast-1775804361516-drop-object-metadata-data-source-fk'; +import { AddCreditBalanceToBillingCustomerFastInstanceCommand } from 'src/database/commands/upgrade-version-command/1-22/1-22-instance-command-fast-1776078919203-add-credit-balance-to-billing-customer'; import { BackfillWorkspaceIdOnIndirectEntitiesSlowInstanceCommand } from 'src/database/commands/upgrade-version-command/1-22/1-22-instance-command-slow-1775758621018-backfill-workspace-id-on-indirect-entities'; import { DropWorkspaceVersionColumnFastInstanceCommand } from 'src/database/commands/upgrade-version-command/1-23/1-23-instance-command-fast-1785000000000-drop-workspace-version-column'; @@ -19,5 +20,6 @@ export const INSTANCE_COMMANDS = [ BackfillWorkspaceIdOnIndirectEntitiesSlowInstanceCommand, AddWorkspaceIdIndexesAndFksFastInstanceCommand, DropObjectMetadataDataSourceFkFastInstanceCommand, + AddCreditBalanceToBillingCustomerFastInstanceCommand, DropWorkspaceVersionColumnFastInstanceCommand, ]; diff --git a/packages/twenty-server/src/engine/core-modules/billing-webhook/services/billing-webhook-credit-grant.service.ts b/packages/twenty-server/src/engine/core-modules/billing-webhook/services/billing-webhook-credit-grant.service.ts index e9f2474613..976796d70e 100644 --- a/packages/twenty-server/src/engine/core-modules/billing-webhook/services/billing-webhook-credit-grant.service.ts +++ b/packages/twenty-server/src/engine/core-modules/billing-webhook/services/billing-webhook-credit-grant.service.ts @@ -1,9 +1,12 @@ /* @license Enterprise */ import { Injectable } from '@nestjs/common'; +import { InjectRepository } from '@nestjs/typeorm'; import { isDefined } from 'twenty-shared/utils'; +import { Repository } from 'typeorm'; +import { BillingCustomerEntity } from 'src/engine/core-modules/billing/entities/billing-customer.entity'; import { BillingSubscriptionService } from 'src/engine/core-modules/billing/services/billing-subscription.service'; import { MeteredCreditService } from 'src/engine/core-modules/billing/services/metered-credit.service'; @@ -12,6 +15,8 @@ export class BillingWebhookCreditGrantService { constructor( private readonly meteredCreditService: MeteredCreditService, private readonly billingSubscriptionService: BillingSubscriptionService, + @InjectRepository(BillingCustomerEntity) + private readonly billingCustomerRepository: Repository, ) {} async processStripeEvent(stripeCustomerId: string): Promise { @@ -24,6 +29,22 @@ export class BillingWebhookCreditGrantService { return; } + const meteredPricingInfo = + await this.meteredCreditService.getMeteredPricingInfo(subscription.id); + + if (isDefined(meteredPricingInfo)) { + const creditBalanceMicro = + await this.meteredCreditService.getCreditBalance( + stripeCustomerId, + meteredPricingInfo.unitPriceCents, + ); + + await this.billingCustomerRepository.update( + { stripeCustomerId }, + { creditBalanceMicro }, + ); + } + await this.meteredCreditService.recreateBillingAlertForSubscription( subscription, ); diff --git a/packages/twenty-server/src/engine/core-modules/billing/billing.module.ts b/packages/twenty-server/src/engine/core-modules/billing/billing.module.ts index 8675b75982..ff0205c08d 100644 --- a/packages/twenty-server/src/engine/core-modules/billing/billing.module.ts +++ b/packages/twenty-server/src/engine/core-modules/billing/billing.module.ts @@ -3,12 +3,14 @@ import { Module } from '@nestjs/common'; import { TypeOrmModule } from '@nestjs/typeorm'; +import { ClickHouseModule } from 'src/database/clickHouse/clickHouse.module'; import { WorkspaceIteratorModule } from 'src/database/commands/command-runners/workspace-iterator.module'; import { BillingGaugeService } from 'src/engine/core-modules/billing/billing-gauge.service'; import { BillingResolver } from 'src/engine/core-modules/billing/billing.resolver'; import { BillingSyncCustomerDataCommand } from 'src/engine/core-modules/billing/commands/billing-sync-customer-data.command'; import { BillingSyncPlansDataCommand } from 'src/engine/core-modules/billing/commands/billing-sync-plans-data.command'; import { BillingUpdateSubscriptionPriceCommand } from 'src/engine/core-modules/billing/commands/billing-update-subscription-price.command'; +import { EnforceUsageCapCronCommand } from 'src/engine/core-modules/billing/crons/commands/enforce-usage-cap.cron.command'; import { BillingCustomerEntity } from 'src/engine/core-modules/billing/entities/billing-customer.entity'; import { BillingEntitlementEntity } from 'src/engine/core-modules/billing/entities/billing-entitlement.entity'; import { BillingMeterEntity } from 'src/engine/core-modules/billing/entities/billing-meter.entity'; @@ -28,6 +30,7 @@ import { BillingSubscriptionItemService } from 'src/engine/core-modules/billing/ import { BillingSubscriptionPhaseService } from 'src/engine/core-modules/billing/services/billing-subscription-phase.service'; import { BillingSubscriptionUpdateService } from 'src/engine/core-modules/billing/services/billing-subscription-update.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 { 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'; @@ -46,6 +49,7 @@ import { PermissionsModule } from 'src/engine/metadata-modules/permissions/permi @Module({ imports: [ + ClickHouseModule, FeatureFlagModule, StripeModule, MessageQueueModule, @@ -86,10 +90,12 @@ import { PermissionsModule } from 'src/engine/metadata-modules/permissions/permi BillingUpdateSubscriptionPriceCommand, BillingSyncPlansDataCommand, BillingUsageService, + BillingUsageCapService, BillingPriceService, BillingCreditRolloverService, MeteredCreditService, BillingGaugeService, + EnforceUsageCapCronCommand, ], exports: [ BillingSubscriptionService, @@ -98,8 +104,10 @@ import { PermissionsModule } from 'src/engine/metadata-modules/permissions/permi BillingPortalWorkspaceService, BillingService, BillingUsageService, + BillingUsageCapService, BillingCreditRolloverService, MeteredCreditService, + EnforceUsageCapCronCommand, ], }) export class BillingModule {} diff --git a/packages/twenty-server/src/engine/core-modules/billing/crons/__test__/enforce-usage-cap.job.spec.ts b/packages/twenty-server/src/engine/core-modules/billing/crons/__test__/enforce-usage-cap.job.spec.ts new file mode 100644 index 0000000000..4f37f1ac1c --- /dev/null +++ b/packages/twenty-server/src/engine/core-modules/billing/crons/__test__/enforce-usage-cap.job.spec.ts @@ -0,0 +1,365 @@ +/* @license Enterprise */ + +import { Test, type TestingModule } from '@nestjs/testing'; +import { getRepositoryToken } from '@nestjs/typeorm'; + +import { In } from 'typeorm'; + +import { EnforceUsageCapJob } from 'src/engine/core-modules/billing/crons/enforce-usage-cap.job'; +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 { BillingUsageCapService } from 'src/engine/core-modules/billing/services/billing-usage-cap.service'; +import { TwentyConfigService } from 'src/engine/core-modules/twenty-config/twenty-config.service'; + +describe('EnforceUsageCapJob', () => { + let job: EnforceUsageCapJob; + let idQueryMock: jest.Mock; + let fullQueryMock: jest.Mock; + let billingSubscriptionItemRepository: jest.Mocked<{ + update: jest.Mock; + }>; + let billingUsageCapService: jest.Mocked; + let twentyConfigService: jest.Mocked; + + const buildSubscription = ({ + id = 'sub_123', + workspaceId = 'workspace_123', + itemId = 'item_123', + hasReachedCurrentPeriodCap = false, + stripeCustomerId = 'cus_123', + creditBalanceMicro = 0, + } = {}) => + ({ + id, + workspaceId, + stripeCustomerId, + currentPeriodStart: new Date('2026-04-01T00:00:00Z'), + currentPeriodEnd: new Date('2026-05-01T00:00:00Z'), + billingCustomer: { + stripeCustomerId, + creditBalanceMicro, + }, + billingSubscriptionItems: [ + { + id: itemId, + hasReachedCurrentPeriodCap, + billingProduct: { + metadata: { + productKey: BillingProductKey.WORKFLOW_NODE_EXECUTION, + }, + }, + }, + ], + }) as unknown as BillingSubscriptionEntity; + + beforeEach(async () => { + idQueryMock = jest.fn().mockResolvedValue([]); + fullQueryMock = jest.fn().mockResolvedValue([]); + + const idQueryBuilderMock = { + select: jest.fn().mockReturnThis(), + innerJoin: jest.fn().mockReturnThis(), + where: jest.fn().mockReturnThis(), + andWhere: jest.fn().mockReturnThis(), + orderBy: jest.fn().mockReturnThis(), + limit: jest.fn().mockReturnThis(), + offset: jest.fn().mockReturnThis(), + getRawMany: idQueryMock, + }; + + const fullQueryBuilderMock = { + innerJoinAndSelect: jest.fn().mockReturnThis(), + leftJoinAndSelect: jest.fn().mockReturnThis(), + where: jest.fn().mockReturnThis(), + orderBy: jest.fn().mockReturnThis(), + getMany: fullQueryMock, + }; + + const createQueryBuilder = jest + .fn() + .mockReturnValueOnce(idQueryBuilderMock) + .mockReturnValueOnce(fullQueryBuilderMock); + + const module: TestingModule = await Test.createTestingModule({ + providers: [ + EnforceUsageCapJob, + { + provide: getRepositoryToken(BillingSubscriptionEntity), + useValue: { createQueryBuilder }, + }, + { + provide: getRepositoryToken(BillingSubscriptionItemEntity), + useValue: { update: jest.fn() }, + }, + { + provide: BillingUsageCapService, + useValue: { + isClickHouseEnabled: jest.fn().mockReturnValue(true), + getBatchPeriodCreditsUsed: jest.fn().mockResolvedValue(new Map()), + evaluateCapBatch: jest.fn().mockReturnValue(new Map()), + }, + }, + { + provide: TwentyConfigService, + useValue: { get: jest.fn() }, + }, + ], + }).compile(); + + job = module.get(EnforceUsageCapJob); + billingSubscriptionItemRepository = module.get( + getRepositoryToken(BillingSubscriptionItemEntity), + ); + billingUsageCapService = module.get(BillingUsageCapService); + twentyConfigService = module.get(TwentyConfigService); + }); + + afterEach(() => { + jest.clearAllMocks(); + }); + + const mockConfig = (overrides: Record = {}) => { + const values: Record = { + IS_BILLING_ENABLED: true, + BILLING_USAGE_CAP_CLICKHOUSE_ENABLED: false, + ...overrides, + }; + + twentyConfigService.get.mockImplementation( + (key: string) => values[key] as never, + ); + }; + + it('no-ops when billing is disabled', async () => { + mockConfig({ IS_BILLING_ENABLED: false }); + + await job.handle(); + + expect(idQueryMock).not.toHaveBeenCalled(); + }); + + it('no-ops when ClickHouse is not configured', async () => { + mockConfig(); + billingUsageCapService.isClickHouseEnabled.mockReturnValue(false); + + await job.handle(); + + expect(idQueryMock).not.toHaveBeenCalled(); + }); + + it('skips transitions in shadow mode (flag off)', async () => { + mockConfig({ BILLING_USAGE_CAP_CLICKHOUSE_ENABLED: false }); + const sub = buildSubscription({ hasReachedCurrentPeriodCap: false }); + + idQueryMock.mockResolvedValueOnce([{ subscription_id: 'sub_123' }]); + fullQueryMock.mockResolvedValueOnce([sub]); + + billingUsageCapService.getBatchPeriodCreditsUsed.mockResolvedValue( + new Map([['workspace_123', 2_000_000]]), + ); + billingUsageCapService.evaluateCapBatch.mockReturnValue( + new Map([ + [ + 'sub_123', + { + skipped: false as const, + hasReachedCap: true, + usage: 2_000_000, + allowance: 1_000_000, + tierCap: 1_000_000, + creditBalance: 0, + }, + ], + ]), + ); + + await job.handle(); + + expect(billingSubscriptionItemRepository.update).not.toHaveBeenCalled(); + }); + + it('batch-updates hasReachedCurrentPeriodCap=true in active mode when usage exceeds allowance', async () => { + mockConfig({ BILLING_USAGE_CAP_CLICKHOUSE_ENABLED: true }); + const sub = buildSubscription({ + itemId: 'item_123', + hasReachedCurrentPeriodCap: false, + }); + + idQueryMock.mockResolvedValueOnce([{ subscription_id: 'sub_123' }]); + fullQueryMock.mockResolvedValueOnce([sub]); + + billingUsageCapService.getBatchPeriodCreditsUsed.mockResolvedValue( + new Map([['workspace_123', 2_000_000]]), + ); + billingUsageCapService.evaluateCapBatch.mockReturnValue( + new Map([ + [ + 'sub_123', + { + skipped: false as const, + hasReachedCap: true, + usage: 2_000_000, + allowance: 1_000_000, + tierCap: 1_000_000, + creditBalance: 0, + }, + ], + ]), + ); + + await job.handle(); + + expect(billingSubscriptionItemRepository.update).toHaveBeenCalledWith( + { id: In(['item_123']) }, + { hasReachedCurrentPeriodCap: true }, + ); + }); + + it('batch-updates hasReachedCurrentPeriodCap=false in active mode when usage drops below allowance', async () => { + mockConfig({ BILLING_USAGE_CAP_CLICKHOUSE_ENABLED: true }); + const sub = buildSubscription({ + itemId: 'item_123', + hasReachedCurrentPeriodCap: true, + }); + + idQueryMock.mockResolvedValueOnce([{ subscription_id: 'sub_123' }]); + fullQueryMock.mockResolvedValueOnce([sub]); + + billingUsageCapService.getBatchPeriodCreditsUsed.mockResolvedValue( + new Map([['workspace_123', 500_000]]), + ); + billingUsageCapService.evaluateCapBatch.mockReturnValue( + new Map([ + [ + 'sub_123', + { + skipped: false as const, + hasReachedCap: false, + usage: 500_000, + allowance: 1_000_000, + tierCap: 1_000_000, + creditBalance: 0, + }, + ], + ]), + ); + + await job.handle(); + + expect(billingSubscriptionItemRepository.update).toHaveBeenCalledWith( + { id: In(['item_123']) }, + { hasReachedCurrentPeriodCap: false }, + ); + }); + + it('does not update when state already matches', async () => { + mockConfig({ BILLING_USAGE_CAP_CLICKHOUSE_ENABLED: true }); + const sub = buildSubscription({ hasReachedCurrentPeriodCap: true }); + + idQueryMock.mockResolvedValueOnce([{ subscription_id: 'sub_123' }]); + fullQueryMock.mockResolvedValueOnce([sub]); + + billingUsageCapService.getBatchPeriodCreditsUsed.mockResolvedValue( + new Map([['workspace_123', 2_000_000]]), + ); + billingUsageCapService.evaluateCapBatch.mockReturnValue( + new Map([ + [ + 'sub_123', + { + skipped: false as const, + hasReachedCap: true, + usage: 2_000_000, + allowance: 1_000_000, + tierCap: 1_000_000, + creditBalance: 0, + }, + ], + ]), + ); + + await job.handle(); + + expect(billingSubscriptionItemRepository.update).not.toHaveBeenCalled(); + }); + + it('skips subscriptions without a metered item', async () => { + mockConfig({ BILLING_USAGE_CAP_CLICKHOUSE_ENABLED: true }); + const sub = buildSubscription(); + + idQueryMock.mockResolvedValueOnce([{ subscription_id: 'sub_123' }]); + fullQueryMock.mockResolvedValueOnce([sub]); + + billingUsageCapService.getBatchPeriodCreditsUsed.mockResolvedValue( + new Map(), + ); + billingUsageCapService.evaluateCapBatch.mockReturnValue( + new Map([ + [ + 'sub_123', + { skipped: true as const, reason: 'no-metered-item' as const }, + ], + ]), + ); + + await job.handle(); + + expect(billingSubscriptionItemRepository.update).not.toHaveBeenCalled(); + }); + + it('passes credit balance from billingCustomer to evaluateCapBatch', async () => { + mockConfig({ BILLING_USAGE_CAP_CLICKHOUSE_ENABLED: true }); + const sub = buildSubscription({ + stripeCustomerId: 'cus_456', + creditBalanceMicro: 300_000, + }); + + idQueryMock.mockResolvedValueOnce([{ subscription_id: 'sub_123' }]); + fullQueryMock.mockResolvedValueOnce([sub]); + + billingUsageCapService.getBatchPeriodCreditsUsed.mockResolvedValue( + new Map(), + ); + billingUsageCapService.evaluateCapBatch.mockReturnValue(new Map()); + + await job.handle(); + + expect(billingUsageCapService.evaluateCapBatch).toHaveBeenCalledWith( + [sub], + expect.any(Map), + new Map([['cus_456', 300_000]]), + ); + }); + + it('skips subscriptions whose ClickHouse query failed', async () => { + mockConfig({ BILLING_USAGE_CAP_CLICKHOUSE_ENABLED: true }); + const sub = buildSubscription({ hasReachedCurrentPeriodCap: true }); + + idQueryMock.mockResolvedValueOnce([{ subscription_id: 'sub_123' }]); + fullQueryMock.mockResolvedValueOnce([sub]); + + billingUsageCapService.getBatchPeriodCreditsUsed.mockRejectedValue( + new Error('clickhouse exploded'), + ); + billingUsageCapService.evaluateCapBatch.mockReturnValue( + new Map([ + [ + 'sub_123', + { + skipped: false as const, + hasReachedCap: false, + usage: 0, + allowance: 1_000_000, + tierCap: 1_000_000, + creditBalance: 0, + }, + ], + ]), + ); + + await expect(job.handle()).resolves.not.toThrow(); + + expect(billingSubscriptionItemRepository.update).not.toHaveBeenCalled(); + }); +}); diff --git a/packages/twenty-server/src/engine/core-modules/billing/crons/commands/enforce-usage-cap.cron.command.ts b/packages/twenty-server/src/engine/core-modules/billing/crons/commands/enforce-usage-cap.cron.command.ts new file mode 100644 index 0000000000..0b2a39b0d5 --- /dev/null +++ b/packages/twenty-server/src/engine/core-modules/billing/crons/commands/enforce-usage-cap.cron.command.ts @@ -0,0 +1,33 @@ +/* @license Enterprise */ + +import { Command, CommandRunner } from 'nest-commander'; + +import { enforceUsageCapCronPattern } from 'src/engine/core-modules/billing/crons/enforce-usage-cap.cron.pattern'; +import { EnforceUsageCapJob } from 'src/engine/core-modules/billing/crons/enforce-usage-cap.job'; +import { InjectMessageQueue } from 'src/engine/core-modules/message-queue/decorators/message-queue.decorator'; +import { MessageQueue } from 'src/engine/core-modules/message-queue/message-queue.constants'; +import { MessageQueueService } from 'src/engine/core-modules/message-queue/services/message-queue.service'; + +@Command({ + name: 'cron:billing:enforce-usage-cap', + description: + 'Starts the cron that re-evaluates metered-credit caps from ClickHouse usage', +}) +export class EnforceUsageCapCronCommand extends CommandRunner { + constructor( + @InjectMessageQueue(MessageQueue.cronQueue) + private readonly messageQueueService: MessageQueueService, + ) { + super(); + } + + async run(): Promise { + await this.messageQueueService.addCron({ + jobName: EnforceUsageCapJob.name, + data: undefined, + options: { + repeat: { pattern: enforceUsageCapCronPattern }, + }, + }); + } +} diff --git a/packages/twenty-server/src/engine/core-modules/billing/crons/enforce-usage-cap.cron.pattern.ts b/packages/twenty-server/src/engine/core-modules/billing/crons/enforce-usage-cap.cron.pattern.ts new file mode 100644 index 0000000000..b52c669cb2 --- /dev/null +++ b/packages/twenty-server/src/engine/core-modules/billing/crons/enforce-usage-cap.cron.pattern.ts @@ -0,0 +1,9 @@ +/* @license Enterprise */ + +// Poll every 2 minutes. +// hasReachedCurrentPeriodCap is reset to false at the start of each billing +// period by BillingWebhookInvoiceService (on subscription_cycle invoice), +// on trial end (BillingSubscriptionService), and on plan changes +// (BillingSubscriptionUpdateService). This cron flips it true/false based +// on live ClickHouse usage within the current period. +export const enforceUsageCapCronPattern = '*/2 * * * *'; diff --git a/packages/twenty-server/src/engine/core-modules/billing/crons/enforce-usage-cap.job.ts b/packages/twenty-server/src/engine/core-modules/billing/crons/enforce-usage-cap.job.ts new file mode 100644 index 0000000000..53ece0ed54 --- /dev/null +++ b/packages/twenty-server/src/engine/core-modules/billing/crons/enforce-usage-cap.job.ts @@ -0,0 +1,251 @@ +/* @license Enterprise */ + +import { Injectable, Logger } from '@nestjs/common'; +import { InjectRepository } from '@nestjs/typeorm'; + +import { In, Repository } from 'typeorm'; + +import { enforceUsageCapCronPattern } from 'src/engine/core-modules/billing/crons/enforce-usage-cap.cron.pattern'; +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 { BillingUsageCapService } from 'src/engine/core-modules/billing/services/billing-usage-cap.service'; +import { SentryCronMonitor } from 'src/engine/core-modules/cron/sentry-cron-monitor.decorator'; +import { Process } from 'src/engine/core-modules/message-queue/decorators/process.decorator'; +import { Processor } from 'src/engine/core-modules/message-queue/decorators/processor.decorator'; +import { MessageQueue } from 'src/engine/core-modules/message-queue/message-queue.constants'; +import { TwentyConfigService } from 'src/engine/core-modules/twenty-config/twenty-config.service'; + +const BATCH_SIZE = 100; + +@Injectable() +@Processor(MessageQueue.cronQueue) +export class EnforceUsageCapJob { + private readonly logger = new Logger(EnforceUsageCapJob.name); + + constructor( + @InjectRepository(BillingSubscriptionEntity) + private readonly billingSubscriptionRepository: Repository, + @InjectRepository(BillingSubscriptionItemEntity) + private readonly billingSubscriptionItemRepository: Repository, + private readonly billingUsageCapService: BillingUsageCapService, + private readonly twentyConfigService: TwentyConfigService, + ) {} + + @Process(EnforceUsageCapJob.name) + @SentryCronMonitor(EnforceUsageCapJob.name, enforceUsageCapCronPattern) + async handle(): Promise { + if (!this.twentyConfigService.get('IS_BILLING_ENABLED')) { + return; + } + + if (!this.billingUsageCapService.isClickHouseEnabled()) { + this.logger.debug( + 'ClickHouse is not configured; skipping usage cap enforcement', + ); + + return; + } + + const isEnforcementActive = this.twentyConfigService.get( + 'BILLING_USAGE_CAP_CLICKHOUSE_ENABLED', + ); + + let evaluated = 0; + let transitioned = 0; + let errors = 0; + let offset = 0; + + let batch: BillingSubscriptionEntity[]; + let idRows: { subscription_id: string }[]; + + do { + idRows = await this.billingSubscriptionRepository + .createQueryBuilder('subscription') + .select('subscription.id') + .innerJoin('subscription.workspace', 'workspace') + .where('subscription.status IN (:...statuses)', { + statuses: [ + SubscriptionStatus.Active, + SubscriptionStatus.Trialing, + SubscriptionStatus.PastDue, + ], + }) + .andWhere('workspace.suspendedAt IS NULL') + .orderBy('subscription.id', 'ASC') + .limit(BATCH_SIZE) + .offset(offset) + .getRawMany<{ subscription_id: string }>(); + + if (idRows.length === 0) { + break; + } + + const ids = idRows.map((row) => row.subscription_id); + + batch = await this.billingSubscriptionRepository + .createQueryBuilder('subscription') + .innerJoinAndSelect( + 'subscription.billingSubscriptionItems', + 'item', + 'item.billingSubscriptionId = subscription.id', + ) + .innerJoinAndSelect('item.billingProduct', 'product') + .leftJoinAndSelect('product.billingPrices', 'price') + .innerJoinAndSelect('subscription.billingCustomer', 'customer') + .where('subscription.id IN (:...ids)', { ids }) + .orderBy('subscription.id', 'ASC') + .getMany(); + + if (batch.length === 0) { + break; + } + + const periodGroups = this.groupByPeriod(batch); + const usageByWorkspace = new Map(); + const failedWorkspaceIds = new Set(); + + for (const [, group] of periodGroups) { + try { + const workspaceIds = group.map((s) => s.workspaceId); + const batchUsage = + await this.billingUsageCapService.getBatchPeriodCreditsUsed( + workspaceIds, + group[0].currentPeriodStart, + group[0].currentPeriodEnd, + ); + + for (const [id, usage] of batchUsage) { + usageByWorkspace.set(id, usage); + } + } catch (error) { + for (const sub of group) { + failedWorkspaceIds.add(sub.workspaceId); + } + errors += group.length; + this.logger.error( + `Failed to fetch batch usage from ClickHouse for ${group.length} subscriptions`, + error instanceof Error ? error.stack : String(error), + ); + } + } + + const creditBalanceByCustomer = new Map(); + + for (const subscription of batch) { + if (subscription.billingCustomer) { + creditBalanceByCustomer.set( + subscription.stripeCustomerId, + subscription.billingCustomer.creditBalanceMicro, + ); + } + } + + const evaluations = this.billingUsageCapService.evaluateCapBatch( + batch, + usageByWorkspace, + creditBalanceByCustomer, + ); + + const idsToCapTrue: string[] = []; + const idsToCapFalse: string[] = []; + + for (const subscription of batch) { + if (failedWorkspaceIds.has(subscription.workspaceId)) { + continue; + } + + const evaluation = evaluations.get(subscription.id); + + if (!evaluation || evaluation.skipped) { + continue; + } + + evaluated += 1; + + const meteredItem = subscription.billingSubscriptionItems.find( + (item) => + item.billingProduct?.metadata?.productKey === + BillingProductKey.WORKFLOW_NODE_EXECUTION, + ); + + if (!meteredItem) { + continue; + } + + const shouldBeCapped = evaluation.hasReachedCap; + + if (meteredItem.hasReachedCurrentPeriodCap === shouldBeCapped) { + continue; + } + + if (!isEnforcementActive) { + this.logger.log( + `[shadow] would set hasReachedCurrentPeriodCap=${shouldBeCapped} ` + + `for subscription=${subscription.id} workspace=${subscription.workspaceId} ` + + `usage=${evaluation.usage} allowance=${evaluation.allowance} ` + + `tierCap=${evaluation.tierCap} creditBalance=${evaluation.creditBalance}`, + ); + continue; + } + + if (shouldBeCapped) { + idsToCapTrue.push(meteredItem.id); + } else { + idsToCapFalse.push(meteredItem.id); + } + + this.logger.log( + `Set hasReachedCurrentPeriodCap=${shouldBeCapped} ` + + `for subscription=${subscription.id} workspace=${subscription.workspaceId} ` + + `usage=${evaluation.usage} allowance=${evaluation.allowance} ` + + `tierCap=${evaluation.tierCap} creditBalance=${evaluation.creditBalance}`, + ); + } + + if (idsToCapTrue.length > 0) { + await this.billingSubscriptionItemRepository.update( + { id: In(idsToCapTrue) }, + { hasReachedCurrentPeriodCap: true }, + ); + transitioned += idsToCapTrue.length; + } + + if (idsToCapFalse.length > 0) { + await this.billingSubscriptionItemRepository.update( + { id: In(idsToCapFalse) }, + { hasReachedCurrentPeriodCap: false }, + ); + transitioned += idsToCapFalse.length; + } + + offset += idRows.length; + } while (idRows.length === BATCH_SIZE); + + this.logger.log( + `Usage cap enforcement run complete: evaluated=${evaluated} ` + + `transitioned=${transitioned} errors=${errors} ` + + `mode=${isEnforcementActive ? 'active' : 'shadow'}`, + ); + } + + private groupByPeriod( + subscriptions: BillingSubscriptionEntity[], + ): Map { + const groups = new Map(); + + for (const subscription of subscriptions) { + const key = `${subscription.currentPeriodStart.toISOString()}|${subscription.currentPeriodEnd.toISOString()}`; + const group = groups.get(key); + + if (group) { + group.push(subscription); + } else { + groups.set(key, [subscription]); + } + } + + return groups; + } +} diff --git a/packages/twenty-server/src/engine/core-modules/billing/entities/billing-customer.entity.ts b/packages/twenty-server/src/engine/core-modules/billing/entities/billing-customer.entity.ts index 71da4bd93f..d7f6514500 100644 --- a/packages/twenty-server/src/engine/core-modules/billing/entities/billing-customer.entity.ts +++ b/packages/twenty-server/src/engine/core-modules/billing/entities/billing-customer.entity.ts @@ -41,6 +41,18 @@ export class BillingCustomerEntity extends WorkspaceRelatedEntity { @Column({ nullable: false, unique: true }) stripeCustomerId: string; + @Column({ + type: 'bigint', + nullable: false, + default: 0, + transformer: { + to: (value: number) => value, + from: (value: string | number | null) => + typeof value === 'string' ? Number(value) : (value ?? 0), + }, + }) + creditBalanceMicro: number; + @OneToMany( () => BillingSubscriptionEntity, (billingSubscription) => billingSubscription.billingCustomer, diff --git a/packages/twenty-server/src/engine/core-modules/billing/services/__test__/billing-credit-rollover.service.spec.ts b/packages/twenty-server/src/engine/core-modules/billing/services/__test__/billing-credit-rollover.service.spec.ts index a975d5d73b..ce57d02588 100644 --- a/packages/twenty-server/src/engine/core-modules/billing/services/__test__/billing-credit-rollover.service.spec.ts +++ b/packages/twenty-server/src/engine/core-modules/billing/services/__test__/billing-credit-rollover.service.spec.ts @@ -1,9 +1,11 @@ /* @license Enterprise */ import { Test, type TestingModule } from '@nestjs/testing'; +import { getRepositoryToken } from '@nestjs/typeorm'; import type Stripe from 'stripe'; +import { BillingCustomerEntity } from 'src/engine/core-modules/billing/entities/billing-customer.entity'; import { BillingCreditRolloverService } from 'src/engine/core-modules/billing/services/billing-credit-rollover.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'; @@ -12,6 +14,7 @@ describe('BillingCreditRolloverService', () => { let service: BillingCreditRolloverService; let stripeCreditGrantService: jest.Mocked; let stripeBillingMeterEventService: jest.Mocked; + let billingCustomerRepository: jest.Mocked<{ update: jest.Mock }>; beforeEach(async () => { const module: TestingModule = await Test.createTestingModule({ @@ -23,6 +26,7 @@ describe('BillingCreditRolloverService', () => { createCreditGrant: jest.fn(), listCreditGrants: jest.fn().mockResolvedValue([]), voidCreditGrant: jest.fn(), + getCustomerCreditBalance: jest.fn().mockResolvedValue(0), }, }, { @@ -31,6 +35,12 @@ describe('BillingCreditRolloverService', () => { sumMeterEvents: jest.fn(), }, }, + { + provide: getRepositoryToken(BillingCustomerEntity), + useValue: { + update: jest.fn(), + }, + }, ], }).compile(); @@ -39,6 +49,9 @@ describe('BillingCreditRolloverService', () => { ); stripeCreditGrantService = module.get(StripeCreditGrantService); stripeBillingMeterEventService = module.get(StripeBillingMeterEventService); + billingCustomerRepository = module.get( + getRepositoryToken(BillingCustomerEntity), + ); }); describe('processRolloverOnPeriodTransition', () => { @@ -139,5 +152,34 @@ describe('BillingCreditRolloverService', () => { }), ); }); + + it('should persist credit balance to Postgres after rollover', async () => { + stripeBillingMeterEventService.sumMeterEvents.mockResolvedValue(300); + stripeCreditGrantService.getCustomerCreditBalance.mockResolvedValue( + 700_000, + ); + + await service.processRolloverOnPeriodTransition(baseParams); + + expect( + stripeCreditGrantService.getCustomerCreditBalance, + ).toHaveBeenCalledWith('cus_123', 10); + expect(billingCustomerRepository.update).toHaveBeenCalledWith( + { stripeCustomerId: 'cus_123' }, + { creditBalanceMicro: 700_000 }, + ); + }); + + it('should persist credit balance even when no grant is created', async () => { + stripeBillingMeterEventService.sumMeterEvents.mockResolvedValue(1000); + stripeCreditGrantService.getCustomerCreditBalance.mockResolvedValue(0); + + await service.processRolloverOnPeriodTransition(baseParams); + + expect(billingCustomerRepository.update).toHaveBeenCalledWith( + { stripeCustomerId: 'cus_123' }, + { creditBalanceMicro: 0 }, + ); + }); }); }); diff --git a/packages/twenty-server/src/engine/core-modules/billing/services/__test__/billing-usage-cap.service.spec.ts b/packages/twenty-server/src/engine/core-modules/billing/services/__test__/billing-usage-cap.service.spec.ts new file mode 100644 index 0000000000..541ca2c9e1 --- /dev/null +++ b/packages/twenty-server/src/engine/core-modules/billing/services/__test__/billing-usage-cap.service.spec.ts @@ -0,0 +1,449 @@ +/* @license Enterprise */ + +import { Test, type TestingModule } from '@nestjs/testing'; + +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 { MeteredCreditService } from 'src/engine/core-modules/billing/services/metered-credit.service'; +import { TwentyConfigService } from 'src/engine/core-modules/twenty-config/twenty-config.service'; + +describe('BillingUsageCapService', () => { + let service: BillingUsageCapService; + let clickHouseService: jest.Mocked; + let meteredCreditService: jest.Mocked; + let twentyConfigService: jest.Mocked; + + const buildSubscription = ( + overrides: Partial = {}, + ): BillingSubscriptionEntity => + ({ + id: 'sub_123', + workspaceId: 'workspace_123', + stripeCustomerId: 'cus_123', + currentPeriodStart: new Date('2026-04-01T00:00:00Z'), + currentPeriodEnd: new Date('2026-05-01T00:00:00Z'), + ...overrides, + }) as BillingSubscriptionEntity; + + beforeEach(async () => { + const module: TestingModule = await Test.createTestingModule({ + providers: [ + BillingUsageCapService, + { + provide: ClickHouseService, + useValue: { + select: jest.fn(), + }, + }, + { + provide: MeteredCreditService, + useValue: { + extractMeteredPricingInfoFromSubscription: jest.fn(), + getCreditBalance: jest.fn(), + }, + }, + { + provide: TwentyConfigService, + useValue: { + get: jest.fn(), + }, + }, + ], + }).compile(); + + service = module.get(BillingUsageCapService); + clickHouseService = module.get(ClickHouseService); + meteredCreditService = module.get(MeteredCreditService); + twentyConfigService = module.get(TwentyConfigService); + }); + + afterEach(() => { + jest.clearAllMocks(); + }); + + describe('isClickHouseEnabled', () => { + it('returns true when CLICKHOUSE_URL is configured', () => { + twentyConfigService.get.mockReturnValue('http://clickhouse:8123'); + + expect(service.isClickHouseEnabled()).toBe(true); + }); + + it('returns false when CLICKHOUSE_URL is empty', () => { + twentyConfigService.get.mockReturnValue(''); + + expect(service.isClickHouseEnabled()).toBe(false); + }); + }); + + 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'); + }); + + it('returns empty map when ClickHouse is disabled', async () => { + twentyConfigService.get.mockReturnValue(''); + + 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); + expect(clickHouseService.select).not.toHaveBeenCalled(); + }); + + it('returns empty map when workspaceIds is empty', async () => { + const result = await service.getBatchPeriodCreditsUsed( + [], + new Date('2026-04-01T00:00:00Z'), + new Date('2026-05-01T00:00:00Z'), + ); + + expect(result.size).toBe(0); + expect(clickHouseService.select).not.toHaveBeenCalled(); + }); + + it('returns usage grouped by workspaceId', async () => { + clickHouseService.select.mockResolvedValue([ + { workspaceId: 'ws_1', total: 500_000 }, + { workspaceId: 'ws_2', total: 1_200_000 }, + ]); + + 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); + expect(result.get('ws_2')).toBe(1_200_000); + expect(result.has('ws_3')).toBe(false); + + const [query, params] = clickHouseService.select.mock.calls[0]; + + expect(query).toContain('GROUP BY workspaceId'); + expect(query).toContain('IN {workspaceIds:Array(String)}'); + expect(params).toMatchObject({ + workspaceIds: ['ws_1', 'ws_2', 'ws_3'], + }); + }); + + it('coerces string totals from ClickHouse', async () => { + clickHouseService.select.mockResolvedValue([ + { workspaceId: 'ws_1', total: '9876543210' }, + ]); + + 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({ + 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({ + 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({ + 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( + { tierCap: 1_000_000, unitPriceCents: 10 }, + ); + + const sub1 = buildSubscription({ + id: 'sub_1', + workspaceId: 'ws_1', + stripeCustomerId: 'cus_1', + }); + const sub2 = buildSubscription({ + id: 'sub_2', + workspaceId: 'ws_2', + stripeCustomerId: 'cus_2', + }); + + const usageByWorkspace = new Map([ + ['ws_1', 500_000], + ['ws_2', 1_500_000], + ]); + const creditBalanceByCustomer = new Map([ + ['cus_1', 0], + ['cus_2', 200_000], + ]); + + const results = service.evaluateCapBatch( + [sub1, sub2], + usageByWorkspace, + creditBalanceByCustomer, + ); + + expect(results.get('sub_1')).toMatchObject({ + skipped: false, + hasReachedCap: false, + usage: 500_000, + allowance: 1_000_000, + }); + expect(results.get('sub_2')).toMatchObject({ + skipped: false, + hasReachedCap: true, + usage: 1_500_000, + allowance: 1_200_000, + }); + }); + + it('defaults usage to 0 for workspaces not in the map', () => { + meteredCreditService.extractMeteredPricingInfoFromSubscription.mockReturnValue( + { tierCap: 1_000_000, unitPriceCents: 10 }, + ); + + const sub = buildSubscription({ + id: 'sub_1', + workspaceId: 'ws_unknown', + stripeCustomerId: 'cus_1', + }); + + const results = service.evaluateCapBatch( + [sub], + new Map(), + new Map([['cus_1', 0]]), + ); + + expect(results.get('sub_1')).toMatchObject({ + skipped: false, + hasReachedCap: false, + usage: 0, + allowance: 1_000_000, + }); + }); + + it('defaults credit balance to 0 for unknown customers', () => { + meteredCreditService.extractMeteredPricingInfoFromSubscription.mockReturnValue( + { tierCap: 1_000_000, unitPriceCents: 10 }, + ); + + const sub = buildSubscription({ + id: 'sub_1', + workspaceId: 'ws_1', + stripeCustomerId: 'cus_unknown', + }); + + const results = service.evaluateCapBatch( + [sub], + new Map([['ws_1', 500_000]]), + new Map(), + ); + + expect(results.get('sub_1')).toMatchObject({ + skipped: false, + hasReachedCap: false, + usage: 500_000, + creditBalance: 0, + allowance: 1_000_000, + }); + }); + + it('returns skipped for subscriptions without metered pricing', () => { + meteredCreditService.extractMeteredPricingInfoFromSubscription.mockReturnValue( + null, + ); + + const sub = buildSubscription({ id: 'sub_1' }); + + const results = service.evaluateCapBatch([sub], new Map(), new Map()); + + expect(results.get('sub_1')).toEqual({ + skipped: true, + reason: 'no-metered-item', + }); + }); + }); +}); diff --git a/packages/twenty-server/src/engine/core-modules/billing/services/billing-credit-rollover.service.ts b/packages/twenty-server/src/engine/core-modules/billing/services/billing-credit-rollover.service.ts index 69c57c08d9..c889209c01 100644 --- a/packages/twenty-server/src/engine/core-modules/billing/services/billing-credit-rollover.service.ts +++ b/packages/twenty-server/src/engine/core-modules/billing/services/billing-credit-rollover.service.ts @@ -1,7 +1,11 @@ /* @license Enterprise */ import { Injectable } from '@nestjs/common'; +import { InjectRepository } from '@nestjs/typeorm'; +import { Repository } from 'typeorm'; + +import { BillingCustomerEntity } from 'src/engine/core-modules/billing/entities/billing-customer.entity'; 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'; @@ -10,6 +14,8 @@ export class BillingCreditRolloverService { constructor( private readonly stripeCreditGrantService: StripeCreditGrantService, private readonly stripeBillingMeterEventService: StripeBillingMeterEventService, + @InjectRepository(BillingCustomerEntity) + private readonly billingCustomerRepository: Repository, ) {} async processRolloverOnPeriodTransition({ @@ -46,6 +52,8 @@ export class BillingCreditRolloverService { const unusedCredits = Math.max(0, tierQuantity - usedCredits); if (unusedCredits <= 0) { + await this.refreshCreditBalance(stripeCustomerId, unitPriceCents); + return; } @@ -63,6 +71,24 @@ export class BillingCreditRolloverService { subscriptionId, }, }); + + await this.refreshCreditBalance(stripeCustomerId, unitPriceCents); + } + + private async refreshCreditBalance( + stripeCustomerId: string, + unitPriceCents: number, + ): Promise { + const creditBalanceMicro = + await this.stripeCreditGrantService.getCustomerCreditBalance( + stripeCustomerId, + unitPriceCents, + ); + + await this.billingCustomerRepository.update( + { stripeCustomerId }, + { creditBalanceMicro }, + ); } private async voidExistingRolloverGrants( diff --git a/packages/twenty-server/src/engine/core-modules/billing/services/billing-usage-cap.service.ts b/packages/twenty-server/src/engine/core-modules/billing/services/billing-usage-cap.service.ts new file mode 100644 index 0000000000..0f19e092c1 --- /dev/null +++ b/packages/twenty-server/src/engine/core-modules/billing/services/billing-usage-cap.service.ts @@ -0,0 +1,189 @@ +/* @license Enterprise */ + +import { Injectable } from '@nestjs/common'; + +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 { MeteredCreditService } from 'src/engine/core-modules/billing/services/metered-credit.service'; +import { TwentyConfigService } from 'src/engine/core-modules/twenty-config/twenty-config.service'; + +export type BillingCapEvaluation = + | { + skipped: false; + hasReachedCap: boolean; + usage: number; + allowance: number; + tierCap: number; + creditBalance: number; + } + | { + skipped: true; + reason: 'no-metered-item' | 'clickhouse-disabled'; + }; + +type UsageSumRow = { + total: string | number | null; +}; + +type BatchUsageSumRow = { + workspaceId: string; + total: string | number | null; +}; + +@Injectable() +export class BillingUsageCapService { + constructor( + private readonly clickHouseService: ClickHouseService, + private readonly meteredCreditService: MeteredCreditService, + private readonly twentyConfigService: TwentyConfigService, + ) {} + + isClickHouseEnabled(): boolean { + return Boolean(this.twentyConfigService.get('CLICKHOUSE_URL')); + } + + async getCurrentPeriodCreditsUsed( + workspaceId: string, + periodStart: Date, + periodEnd: Date, + ): Promise { + 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(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> { + const result = new Map(); + + if (!this.isClickHouseEnabled() || workspaceIds.length === 0) { + return result; + } + + const query = ` + SELECT workspaceId, sum(creditsUsedMicro) AS total + FROM usageEvent + WHERE workspaceId IN {workspaceIds:Array(String)} + AND timestamp >= {periodStart:String} + AND timestamp < {periodEnd:String} + GROUP BY workspaceId + `; + + const rows = await this.clickHouseService.select(query, { + workspaceIds, + periodStart: formatDateForClickHouse(periodStart), + periodEnd: formatDateForClickHouse(periodEnd), + }); + + for (const row of rows) { + const rawTotal = row.total ?? 0; + const total = typeof rawTotal === 'string' ? Number(rawTotal) : rawTotal; + + result.set(row.workspaceId, Number.isFinite(total) ? total : 0); + } + + return result; + } + + async evaluateCap( + subscription: BillingSubscriptionEntity, + ): Promise { + 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, + creditBalanceByCustomer: Map, + ): Map { + const results = new Map(); + + for (const subscription of subscriptions) { + const meteredPricingInfo = + this.meteredCreditService.extractMeteredPricingInfoFromSubscription( + subscription, + ); + + if (!meteredPricingInfo) { + results.set(subscription.id, { + skipped: true, + reason: 'no-metered-item', + }); + continue; + } + + const usage = usageByWorkspace.get(subscription.workspaceId) ?? 0; + const creditBalance = + creditBalanceByCustomer.get(subscription.stripeCustomerId) ?? 0; + const allowance = meteredPricingInfo.tierCap + creditBalance; + + results.set(subscription.id, { + skipped: false, + hasReachedCap: usage >= allowance, + usage, + allowance, + tierCap: meteredPricingInfo.tierCap, + creditBalance, + }); + } + + return results; + } +} diff --git a/packages/twenty-server/src/engine/core-modules/billing/services/metered-credit.service.ts b/packages/twenty-server/src/engine/core-modules/billing/services/metered-credit.service.ts index 0bb62fd663..90ba18dd7f 100644 --- a/packages/twenty-server/src/engine/core-modules/billing/services/metered-credit.service.ts +++ b/packages/twenty-server/src/engine/core-modules/billing/services/metered-credit.service.ts @@ -54,9 +54,15 @@ export class MeteredCreditService { return null; } - const meteredItem = subscription.billingSubscriptionItems.find( + return this.extractMeteredPricingInfoFromSubscription(subscription); + } + + extractMeteredPricingInfoFromSubscription( + subscription: BillingSubscriptionEntity, + ): MeteredPricingInfo | null { + const meteredItem = subscription.billingSubscriptionItems?.find( (item) => - item.billingProduct.metadata.productKey === + item.billingProduct?.metadata?.productKey === BillingProductKey.WORKFLOW_NODE_EXECUTION, ); @@ -64,7 +70,7 @@ export class MeteredCreditService { return null; } - const matchingPrice = meteredItem.billingProduct.billingPrices.find( + const matchingPrice = meteredItem.billingProduct.billingPrices?.find( (price) => price.stripePriceId === meteredItem.stripePriceId, ); diff --git a/packages/twenty-server/src/engine/core-modules/message-queue/jobs.module.ts b/packages/twenty-server/src/engine/core-modules/message-queue/jobs.module.ts index 709fa30706..2ef8f03a95 100644 --- a/packages/twenty-server/src/engine/core-modules/message-queue/jobs.module.ts +++ b/packages/twenty-server/src/engine/core-modules/message-queue/jobs.module.ts @@ -6,6 +6,8 @@ import { TypeORMModule } from 'src/database/typeorm/typeorm.module'; import { AuditJobModule } from 'src/engine/core-modules/audit/jobs/audit-job.module'; import { AuthModule } from 'src/engine/core-modules/auth/auth.module'; import { BillingModule } from 'src/engine/core-modules/billing/billing.module'; +import { EnforceUsageCapJob } from 'src/engine/core-modules/billing/crons/enforce-usage-cap.job'; +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 { UpdateSubscriptionQuantityJob } from 'src/engine/core-modules/billing/jobs/update-subscription-quantity.job'; import { StripeModule } from 'src/engine/core-modules/billing/stripe/stripe.module'; @@ -42,7 +44,11 @@ import { WorkflowModule } from 'src/modules/workflow/workflow.module'; @Module({ imports: [ - TypeOrmModule.forFeature([WorkspaceEntity, BillingSubscriptionEntity]), + TypeOrmModule.forFeature([ + WorkspaceEntity, + BillingSubscriptionEntity, + BillingSubscriptionItemEntity, + ]), ObjectMetadataModule, TypeORMModule, UserModule, @@ -75,6 +81,7 @@ import { WorkflowModule } from 'src/modules/workflow/workflow.module'; CleanSuspendedWorkspacesJob, CleanOnboardingWorkspacesJob, EmailSenderJob, + EnforceUsageCapJob, UpdateSubscriptionQuantityJob, HandleWorkspaceMemberDeletedJob, CleanWorkspaceDeletionWarningUserVarsJob, diff --git a/packages/twenty-server/src/engine/core-modules/twenty-config/config-variables.ts b/packages/twenty-server/src/engine/core-modules/twenty-config/config-variables.ts index 3b62dcf6a0..6b2429851f 100644 --- a/packages/twenty-server/src/engine/core-modules/twenty-config/config-variables.ts +++ b/packages/twenty-server/src/engine/core-modules/twenty-config/config-variables.ts @@ -815,6 +815,15 @@ export class ConfigVariables { @ValidateIf((env) => env.IS_BILLING_ENABLED === true) BILLING_STRIPE_WEBHOOK_SECRET: string; + @ConfigVariablesMetadata({ + group: ConfigVariablesGroup.BILLING_CONFIG, + description: + 'Use the ClickHouse-backed poller (instead of Stripe billing alerts) as the source of truth for metered-credit cap enforcement', + type: ConfigVariableType.BOOLEAN, + }) + @IsOptional() + BILLING_USAGE_CAP_CLICKHOUSE_ENABLED = false; + @ConfigVariablesMetadata({ group: ConfigVariablesGroup.SERVER_CONFIG, description: 'Url for the frontend application',