Add ClickHouse-backed metered credit cap enforcement (#19586)
## Summary Implements a ClickHouse-backed polling system to enforce metered-credit caps for workflow executions, replacing reliance on Stripe billing alerts. The system re-evaluates tier caps against live pricing on every poll cycle, allowing price/tier changes to propagate immediately without recreating Stripe alert objects. ## Key Changes - **BillingUsageCapService**: New service that queries ClickHouse for current-period credit usage and evaluates whether a subscription has reached its metered-credit allowance (tier cap + credit balance) - `isClickHouseEnabled()`: Checks if ClickHouse is configured - `getCurrentPeriodCreditsUsed()`: Sums creditsUsedMicro from usageEvent table for a workspace within a billing period - `evaluateCap()`: Determines if usage has reached the allowance by reading live pricing from the subscription - **EnforceUsageCapJob**: Cron job that polls all active subscriptions and updates `hasReachedCurrentPeriodCap` on metered items - Runs every 2 minutes to keep cap enforcement in sync with live usage - Supports shadow mode (log-only) via `BILLING_USAGE_CAP_CLICKHOUSE_ENABLED` flag for safe rollout - Continues processing after per-subscription errors with detailed logging - **EnforceUsageCapCronCommand**: CLI command to register the enforcement cron job - **MeteredCreditService**: Extracted `extractMeteredPricingInfoFromSubscription()` as a pure function for callers that already hold the subscription with pricing loaded, avoiding redundant DB queries - **Configuration**: Added `BILLING_USAGE_CAP_CLICKHOUSE_ENABLED` flag to control enforcement mode (active vs. shadow) - **Constants**: Added `METERED_OPERATION_TYPES` to define which operation types count toward the metered product's credit cap ## Implementation Details - The service queries ClickHouse for the sum of `creditsUsedMicro` in the current billing period, matching Stripe meter semantics - Pricing is re-read on every evaluation, so tier changes propagate within one poll cycle without Stripe alert recreation - The cron job only updates the database when the cap state actually changes (no-op if already in the correct state) - Shadow mode allows safe validation before enabling enforcement; transitions are logged but not persisted - Comprehensive test coverage for both the service and cron job, including error handling and state transitions https://claude.ai/code/session_01VksTSrYLXJVCPVBQhQdBTe --------- Co-authored-by: Claude <noreply@anthropic.com>
This commit is contained in:
Félix Malfait
@@ -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
|
||||
|
||||
@@ -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 {}
|
||||
|
||||
+29
@@ -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<void> {
|
||||
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<void> {
|
||||
await queryRunner.query(
|
||||
`ALTER TABLE "core"."billingCustomer" DROP COLUMN IF EXISTS "creditBalanceMicro"`,
|
||||
);
|
||||
}
|
||||
}
|
||||
+2
@@ -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,
|
||||
];
|
||||
|
||||
+21
@@ -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<BillingCustomerEntity>,
|
||||
) {}
|
||||
|
||||
async processStripeEvent(stripeCustomerId: string): Promise<void> {
|
||||
@@ -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,
|
||||
);
|
||||
|
||||
@@ -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 {}
|
||||
|
||||
+365
@@ -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<BillingUsageCapService>;
|
||||
let twentyConfigService: jest.Mocked<TwentyConfigService>;
|
||||
|
||||
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>(EnforceUsageCapJob);
|
||||
billingSubscriptionItemRepository = module.get(
|
||||
getRepositoryToken(BillingSubscriptionItemEntity),
|
||||
);
|
||||
billingUsageCapService = module.get(BillingUsageCapService);
|
||||
twentyConfigService = module.get(TwentyConfigService);
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
jest.clearAllMocks();
|
||||
});
|
||||
|
||||
const mockConfig = (overrides: Record<string, unknown> = {}) => {
|
||||
const values: Record<string, unknown> = {
|
||||
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();
|
||||
});
|
||||
});
|
||||
+33
@@ -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<void> {
|
||||
await this.messageQueueService.addCron<undefined>({
|
||||
jobName: EnforceUsageCapJob.name,
|
||||
data: undefined,
|
||||
options: {
|
||||
repeat: { pattern: enforceUsageCapCronPattern },
|
||||
},
|
||||
});
|
||||
}
|
||||
}
|
||||
+9
@@ -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 * * * *';
|
||||
+251
@@ -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<BillingSubscriptionEntity>,
|
||||
@InjectRepository(BillingSubscriptionItemEntity)
|
||||
private readonly billingSubscriptionItemRepository: Repository<BillingSubscriptionItemEntity>,
|
||||
private readonly billingUsageCapService: BillingUsageCapService,
|
||||
private readonly twentyConfigService: TwentyConfigService,
|
||||
) {}
|
||||
|
||||
@Process(EnforceUsageCapJob.name)
|
||||
@SentryCronMonitor(EnforceUsageCapJob.name, enforceUsageCapCronPattern)
|
||||
async handle(): Promise<void> {
|
||||
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<string, number>();
|
||||
const failedWorkspaceIds = new Set<string>();
|
||||
|
||||
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<string, number>();
|
||||
|
||||
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<string, BillingSubscriptionEntity[]> {
|
||||
const groups = new Map<string, BillingSubscriptionEntity[]>();
|
||||
|
||||
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;
|
||||
}
|
||||
}
|
||||
+12
@@ -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,
|
||||
|
||||
+42
@@ -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<StripeCreditGrantService>;
|
||||
let stripeBillingMeterEventService: jest.Mocked<StripeBillingMeterEventService>;
|
||||
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 },
|
||||
);
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
+449
@@ -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<ClickHouseService>;
|
||||
let meteredCreditService: jest.Mocked<MeteredCreditService>;
|
||||
let twentyConfigService: jest.Mocked<TwentyConfigService>;
|
||||
|
||||
const buildSubscription = (
|
||||
overrides: Partial<BillingSubscriptionEntity> = {},
|
||||
): 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>(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<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(
|
||||
{ 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',
|
||||
});
|
||||
});
|
||||
});
|
||||
});
|
||||
+26
@@ -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<BillingCustomerEntity>,
|
||||
) {}
|
||||
|
||||
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<void> {
|
||||
const creditBalanceMicro =
|
||||
await this.stripeCreditGrantService.getCustomerCreditBalance(
|
||||
stripeCustomerId,
|
||||
unitPriceCents,
|
||||
);
|
||||
|
||||
await this.billingCustomerRepository.update(
|
||||
{ stripeCustomerId },
|
||||
{ creditBalanceMicro },
|
||||
);
|
||||
}
|
||||
|
||||
private async voidExistingRolloverGrants(
|
||||
|
||||
+189
@@ -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<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>();
|
||||
|
||||
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<BatchUsageSumRow>(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<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>,
|
||||
creditBalanceByCustomer: Map<string, number>,
|
||||
): Map<string, BillingCapEvaluation> {
|
||||
const results = new Map<string, BillingCapEvaluation>();
|
||||
|
||||
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;
|
||||
}
|
||||
}
|
||||
+9
-3
@@ -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,
|
||||
);
|
||||
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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',
|
||||
|
||||
Reference in New Issue
Block a user