diff --git a/packages/twenty-server/src/database/clickHouse/clickHouse.util.ts b/packages/twenty-server/src/database/clickHouse/clickHouse.util.ts index 56964cbb3e2..78dcb038e7f 100644 --- a/packages/twenty-server/src/database/clickHouse/clickHouse.util.ts +++ b/packages/twenty-server/src/database/clickHouse/clickHouse.util.ts @@ -3,9 +3,12 @@ * ClickHouse expects: YYYY-MM-DD HH:mm:ss.SSS (no 'T' separator, no 'Z' suffix) * JavaScript toISOString() returns: YYYY-MM-DDTHH:mm:ss.SSSZ */ -export const formatDateForClickHouse = (date: Date | string): string => { +export const formatDateTimeForClickHouse = (date: Date | string): string => { const iso = typeof date === 'string' ? date : date.toISOString(); // Extract date (YYYY-MM-DD) and time with milliseconds (HH:mm:ss.SSS) return `${iso.slice(0, 10)} ${iso.slice(11, 23)}`; }; + +export const formatDateForClickHouse = (date: Date): string => + date.toISOString().slice(0, 10); diff --git a/packages/twenty-server/src/database/clickHouse/migrations/006-update-usage-event-table.sql b/packages/twenty-server/src/database/clickHouse/migrations/006-update-usage-event-table.sql new file mode 100644 index 00000000000..e1eac19aafd --- /dev/null +++ b/packages/twenty-server/src/database/clickHouse/migrations/006-update-usage-event-table.sql @@ -0,0 +1,53 @@ +CREATE TABLE + IF NOT EXISTS usageEvent_v2 ( + `timestamp` DateTime64 (3) NOT NULL, + `workspaceId` String NOT NULL, + `periodStart` DateTime64 (3), + `userWorkspaceId` String DEFAULT '', + `resourceType` LowCardinality (String) NOT NULL, + `operationType` LowCardinality (String) NOT NULL, + `quantity` Int64 NOT NULL DEFAULT 0, + `unit` LowCardinality (String) NOT NULL DEFAULT 'CREDIT', + `creditsUsedMicro` Int64 NOT NULL DEFAULT 0, + `resourceId` String DEFAULT '', + `resourceContext` String DEFAULT '', + `metadata` JSON + ) ENGINE = MergeTree +PARTITION BY + toYYYYMM (timestamp) +ORDER BY + (workspaceId, timestamp, userWorkspaceId) +PRIMARY KEY (workspaceId, timestamp) +TTL toDateTime(timestamp) + INTERVAL 3 YEAR DELETE; + +INSERT INTO + usageEvent_v2 (timestamp, workspaceId, userWorkspaceId, resourceType, operationType, quantity, unit, creditsUsedMicro, resourceId, resourceContext, metadata) +SELECT + timestamp, + workspaceId, + userWorkspaceId, + resourceType, + operationType, + quantity, + unit, + creditsUsedMicro, + resourceId, + resourceContext, + metadata +FROM + usageEvent; + +RENAME TABLE usageEvent TO usageEvent_old, +usageEvent_v2 TO usageEvent; + +DROP TABLE IF EXISTS usageEvent_old; + +ALTER TABLE usageEvent ADD PROJECTION IF NOT EXISTS billing_by_workspace_period ( + SELECT + workspaceId, + periodStart, + sum(creditsUsedMicro) AS totalCreditsUsedMicro + GROUP BY + workspaceId, periodStart +); + diff --git a/packages/twenty-server/src/database/clickHouse/seeds/fixtures.ts b/packages/twenty-server/src/database/clickHouse/seeds/fixtures.ts index 16a3dece868..c4ca8c4e2fa 100644 --- a/packages/twenty-server/src/database/clickHouse/seeds/fixtures.ts +++ b/packages/twenty-server/src/database/clickHouse/seeds/fixtures.ts @@ -4,7 +4,7 @@ import { OBJECT_RECORD_UPDATED_EVENT } from 'src/engine/core-modules/audit/utils import { CUSTOM_DOMAIN_ACTIVATED_EVENT } from 'src/engine/core-modules/audit/utils/events/workspace-event/custom-domain/custom-domain-activated'; import { CUSTOM_DOMAIN_DEACTIVATED_EVENT } from 'src/engine/core-modules/audit/utils/events/workspace-event/custom-domain/custom-domain-deactivated'; import { type GenericTrackEvent } from 'src/engine/core-modules/audit/utils/events/workspace-event/track'; -import { formatDateForClickHouse } from 'src/database/clickHouse/clickHouse.util'; +import { formatDateTimeForClickHouse } from 'src/database/clickHouse/clickHouse.util'; import { SEED_APPLE_WORKSPACE_ID } from 'src/engine/workspace-manager/dev-seeder/core/constants/seeder-workspaces.constant'; import { USER_WORKSPACE_DATA_SEED_IDS } from 'src/engine/workspace-manager/dev-seeder/core/utils/seed-user-workspaces.util'; @@ -17,6 +17,7 @@ export type ObjectEventFixture = GenericTrackEvent & { export type UsageEventFixture = { timestamp: string; workspaceId: string; + periodStart: string; userWorkspaceId: string; resourceType: string; operationType: string; @@ -215,8 +216,17 @@ const buildUsageEventFixtures = (): UsageEventFixture[] => { : ''; fixtures.push({ - timestamp: formatDateForClickHouse(eventDate), + timestamp: formatDateTimeForClickHouse(eventDate), workspaceId: SEED_APPLE_WORKSPACE_ID, + periodStart: formatDateTimeForClickHouse( + new Date( + Date.UTC( + eventDate.getUTCFullYear(), + eventDate.getUTCMonth(), + 1, + ), + ), + ), userWorkspaceId: users[userIdx], resourceType: op.resourceType, operationType: op.operationType, diff --git a/packages/twenty-server/src/engine/core-modules/application-logs/drivers/clickhouse.driver.ts b/packages/twenty-server/src/engine/core-modules/application-logs/drivers/clickhouse.driver.ts index 48c87f6f204..968e125278a 100644 --- a/packages/twenty-server/src/engine/core-modules/application-logs/drivers/clickhouse.driver.ts +++ b/packages/twenty-server/src/engine/core-modules/application-logs/drivers/clickhouse.driver.ts @@ -1,7 +1,7 @@ import { Logger } from '@nestjs/common'; import { ClickHouseService } from 'src/database/clickHouse/clickHouse.service'; -import { formatDateForClickHouse } from 'src/database/clickHouse/clickHouse.util'; +import { formatDateTimeForClickHouse } from 'src/database/clickHouse/clickHouse.util'; import { type ApplicationLogEntry } from 'src/engine/core-modules/application-logs/interfaces/application-log-entry.interface'; import { type ApplicationLogDriverInterface } from 'src/engine/core-modules/application-logs/interfaces/application-log-driver.interface'; @@ -18,7 +18,7 @@ export class ClickHouseApplicationLogDriver } const rows = entries.map((entry) => ({ - timestamp: formatDateForClickHouse(entry.timestamp), + timestamp: formatDateTimeForClickHouse(entry.timestamp), workspaceId: entry.workspaceId, applicationId: entry.applicationId, logicFunctionId: entry.logicFunctionId, diff --git a/packages/twenty-server/src/engine/core-modules/billing-webhook/billing-webhook.module.ts b/packages/twenty-server/src/engine/core-modules/billing-webhook/billing-webhook.module.ts index 8dfec800c39..883543e0a85 100644 --- a/packages/twenty-server/src/engine/core-modules/billing-webhook/billing-webhook.module.ts +++ b/packages/twenty-server/src/engine/core-modules/billing-webhook/billing-webhook.module.ts @@ -28,6 +28,7 @@ import { WorkspaceEntity } from 'src/engine/core-modules/workspace/workspace.ent import { WorkspaceModule } from 'src/engine/core-modules/workspace/workspace.module'; import { PermissionsModule } from 'src/engine/metadata-modules/permissions/permissions.module'; import { RowLevelPermissionModule } from 'src/engine/metadata-modules/row-level-permission-predicate/row-level-permission.module'; +import { WorkspaceCacheModule } from 'src/engine/workspace-cache/workspace-cache.module'; @Module({ imports: [ @@ -37,6 +38,7 @@ import { RowLevelPermissionModule } from 'src/engine/metadata-modules/row-level- PermissionsModule, WorkspaceModule, BillingModule, + WorkspaceCacheModule, TypeOrmModule.forFeature([ BillingSubscriptionEntity, BillingSubscriptionItemEntity, diff --git a/packages/twenty-server/src/engine/core-modules/billing-webhook/services/billing-webhook-subscription.service.ts b/packages/twenty-server/src/engine/core-modules/billing-webhook/services/billing-webhook-subscription.service.ts index 52d39a9e420..6e449d031bf 100644 --- a/packages/twenty-server/src/engine/core-modules/billing-webhook/services/billing-webhook-subscription.service.ts +++ b/packages/twenty-server/src/engine/core-modules/billing-webhook/services/billing-webhook-subscription.service.ts @@ -24,6 +24,7 @@ import { BillingSubscriptionEntity } from 'src/engine/core-modules/billing/entit import { SubscriptionStatus } from 'src/engine/core-modules/billing/enums/billing-subscription-status.enum'; import { BillingWebhookEvent } from 'src/engine/core-modules/billing/enums/billing-webhook-events.enum'; import { BillingSubscriptionService } from 'src/engine/core-modules/billing/services/billing-subscription.service'; +import { BillingUsageService } from 'src/engine/core-modules/billing/services/billing-usage.service'; import { StripeBillingAlertService } from 'src/engine/core-modules/billing/stripe/services/stripe-billing-alert.service'; import { StripeCustomerService } from 'src/engine/core-modules/billing/stripe/services/stripe-customer.service'; import { StripeSubscriptionScheduleService } from 'src/engine/core-modules/billing/stripe/services/stripe-subscription-schedule.service'; @@ -32,6 +33,7 @@ import { MessageQueue } from 'src/engine/core-modules/message-queue/message-queu import { MessageQueueService } from 'src/engine/core-modules/message-queue/services/message-queue.service'; import { WorkspaceService } from 'src/engine/core-modules/workspace/services/workspace.service'; import { WorkspaceEntity } from 'src/engine/core-modules/workspace/workspace.entity'; +import { WorkspaceCacheService } from 'src/engine/workspace-cache/services/workspace-cache.service'; import { CleanWorkspaceDeletionWarningUserVarsJob, type CleanWorkspaceDeletionWarningUserVarsJobData, @@ -59,6 +61,8 @@ export class BillingWebhookSubscriptionService { private readonly workspaceService: WorkspaceService, private readonly stripeSubscriptionScheduleService: StripeSubscriptionScheduleService, private readonly stripeBillingAlertService: StripeBillingAlertService, + private readonly billingUsageService: BillingUsageService, + private readonly workspaceCacheService: WorkspaceCacheService, ) {} async processStripeEvent( @@ -140,6 +144,11 @@ export class BillingWebhookSubscriptionService { workspaceId, ); + await this.billingUsageService.flushAvailableCreditsFromCache(workspace.id); + await this.workspaceCacheService.invalidateAndRecompute(workspace.id, [ + 'billingSubscription', + ]); + const shouldSuspend = this.shouldSuspendWorkspace(data); if (shouldSuspend) { 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 ff0205c08d4..5b0e5fdea37 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 @@ -34,6 +34,7 @@ import { BillingUsageCapService } from 'src/engine/core-modules/billing/services import { BillingUsageService } from 'src/engine/core-modules/billing/services/billing-usage.service'; import { BillingService } from 'src/engine/core-modules/billing/services/billing.service'; import { MeteredCreditService } from 'src/engine/core-modules/billing/services/metered-credit.service'; +import { WorkspaceBillingSubscriptionCacheService } from 'src/engine/core-modules/billing/services/workspace-billing-subscription-cache.service'; import { StripeModule } from 'src/engine/core-modules/billing/stripe/stripe.module'; import { WorkspaceDomainsModule } from 'src/engine/core-modules/domain/workspace-domains/workspace-domains.module'; import { EnterpriseModule } from 'src/engine/core-modules/enterprise/enterprise.module'; @@ -43,9 +44,8 @@ import { MessageQueueModule } from 'src/engine/core-modules/message-queue/messag import { MetricsModule } from 'src/engine/core-modules/metrics/metrics.module'; import { UserWorkspaceEntity } from 'src/engine/core-modules/user-workspace/user-workspace.entity'; import { WorkspaceEntity } from 'src/engine/core-modules/workspace/workspace.entity'; -import { AiBillingModule } from 'src/engine/metadata-modules/ai/ai-billing/ai-billing.module'; -import { AiModelsModule } from 'src/engine/metadata-modules/ai/ai-models/ai-models.module'; import { PermissionsModule } from 'src/engine/metadata-modules/permissions/permissions.module'; +import { WorkspaceCacheModule } from 'src/engine/workspace-cache/workspace-cache.module'; @Module({ imports: [ @@ -54,8 +54,7 @@ import { PermissionsModule } from 'src/engine/metadata-modules/permissions/permi StripeModule, MessageQueueModule, PermissionsModule, - AiBillingModule, - AiModelsModule, + WorkspaceCacheModule, WorkspaceDomainsModule, TypeOrmModule.forFeature([ BillingSubscriptionEntity, @@ -96,6 +95,7 @@ import { PermissionsModule } from 'src/engine/metadata-modules/permissions/permi MeteredCreditService, BillingGaugeService, EnforceUsageCapCronCommand, + WorkspaceBillingSubscriptionCacheService, ], exports: [ BillingSubscriptionService, 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 index 0b5bfd5ab9b..a99e825021c 100644 --- 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 @@ -5,6 +5,7 @@ import { InjectRepository } from '@nestjs/typeorm'; import { In, IsNull, Repository } from 'typeorm'; +import { formatDateForClickHouse } from 'src/database/clickHouse/clickHouse.util'; import { enforceUsageCapCronPattern } from 'src/engine/core-modules/billing/crons/enforce-usage-cap.cron.pattern'; import { BillingProductEntity } from 'src/engine/core-modules/billing/entities/billing-product.entity'; import { BillingSubscriptionItemEntity } from 'src/engine/core-modules/billing/entities/billing-subscription-item.entity'; @@ -118,7 +119,6 @@ export class EnforceUsageCapJob { await this.billingUsageCapService.getBatchPeriodCreditsUsed( workspaceIds, group[0].currentPeriodStart, - group[0].currentPeriodEnd, ); for (const [id, usage] of batchUsage) { @@ -249,7 +249,7 @@ export class EnforceUsageCapJob { const groups = new Map(); for (const subscription of subscriptions) { - const key = `${subscription.currentPeriodStart.toISOString()}|${subscription.currentPeriodEnd.toISOString()}`; + const key = `${formatDateForClickHouse(subscription.currentPeriodStart)}`; const group = groups.get(key); if (group) { diff --git a/packages/twenty-server/src/engine/core-modules/billing/listeners/billing-usage-event.listener.ts b/packages/twenty-server/src/engine/core-modules/billing/listeners/billing-usage-event.listener.ts index 0053743b3dc..d9668705d33 100644 --- a/packages/twenty-server/src/engine/core-modules/billing/listeners/billing-usage-event.listener.ts +++ b/packages/twenty-server/src/engine/core-modules/billing/listeners/billing-usage-event.listener.ts @@ -5,10 +5,10 @@ import { Injectable } from '@nestjs/common'; import { isDefined } from 'twenty-shared/utils'; import { OnCustomBatchEvent } from 'src/engine/api/graphql/graphql-query-runner/decorators/on-custom-batch-event.decorator'; -import { USAGE_RECORDED } from 'src/engine/core-modules/usage/constants/usage-recorded.constant'; import { BillingUsageService } from 'src/engine/core-modules/billing/services/billing-usage.service'; -import { type UsageEvent } from 'src/engine/core-modules/usage/types/usage-event.type'; import { TwentyConfigService } from 'src/engine/core-modules/twenty-config/twenty-config.service'; +import { USAGE_RECORDED } from 'src/engine/core-modules/usage/constants/usage-recorded.constant'; +import { type UsageEvent } from 'src/engine/core-modules/usage/types/usage-event.type'; import { CustomWorkspaceEventBatch } from 'src/engine/workspace-event-emitter/types/custom-workspace-batch-event.type'; @Injectable() @@ -38,6 +38,7 @@ export class BillingUsageEventListener { return; } + //TODO: To be removed await this.billingUsageService.billUsage({ workspaceId: payload.workspaceId, usageEvents: payload.events, 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 index 541ca2c9e19..caa25f03caf 100644 --- 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 @@ -2,14 +2,16 @@ import { Test, type TestingModule } from '@nestjs/testing'; +import { getRepositoryToken } from '@nestjs/typeorm'; + import { ClickHouseService } from 'src/database/clickHouse/clickHouse.service'; -import { type BillingSubscriptionEntity } from 'src/engine/core-modules/billing/entities/billing-subscription.entity'; -import { - type BillingCapEvaluation, - BillingUsageCapService, -} from 'src/engine/core-modules/billing/services/billing-usage-cap.service'; +import { BillingSubscriptionItemEntity } from 'src/engine/core-modules/billing/entities/billing-subscription-item.entity'; +import { BillingSubscriptionEntity } from 'src/engine/core-modules/billing/entities/billing-subscription.entity'; +import { BillingUsageCapService } from 'src/engine/core-modules/billing/services/billing-usage-cap.service'; import { MeteredCreditService } from 'src/engine/core-modules/billing/services/metered-credit.service'; +import { CacheStorageNamespace } from 'src/engine/core-modules/cache-storage/types/cache-storage-namespace.enum'; import { TwentyConfigService } from 'src/engine/core-modules/twenty-config/twenty-config.service'; +import { WorkspaceCacheService } from 'src/engine/workspace-cache/services/workspace-cache.service'; describe('BillingUsageCapService', () => { let service: BillingUsageCapService; @@ -52,6 +54,34 @@ describe('BillingUsageCapService', () => { get: jest.fn(), }, }, + { + provide: CacheStorageNamespace.EngineBillingUsage, + useValue: { + get: jest.fn(), + set: jest.fn(), + del: jest.fn(), + incrBy: jest.fn(), + }, + }, + { + provide: getRepositoryToken(BillingSubscriptionEntity), + useValue: { + findOne: jest.fn(), + }, + }, + { + provide: getRepositoryToken(BillingSubscriptionItemEntity), + useValue: { + find: jest.fn(), + update: jest.fn(), + }, + }, + { + provide: WorkspaceCacheService, + useValue: { + getOrRecompute: jest.fn(), + }, + }, ], }).compile(); @@ -79,83 +109,6 @@ describe('BillingUsageCapService', () => { }); }); - describe('getCurrentPeriodCreditsUsed', () => { - beforeEach(() => { - twentyConfigService.get.mockReturnValue('http://clickhouse:8123'); - }); - - it('returns 0 when ClickHouse is disabled', async () => { - twentyConfigService.get.mockReturnValue(''); - - const result = await service.getCurrentPeriodCreditsUsed( - 'workspace_123', - new Date('2026-04-01T00:00:00Z'), - new Date('2026-05-01T00:00:00Z'), - ); - - expect(result).toBe(0); - expect(clickHouseService.select).not.toHaveBeenCalled(); - }); - - it('sums creditsUsedMicro for the workspace in the given period', async () => { - clickHouseService.select.mockResolvedValue([{ total: 12345 }]); - - const result = await service.getCurrentPeriodCreditsUsed( - 'workspace_123', - new Date('2026-04-01T00:00:00Z'), - new Date('2026-05-01T00:00:00Z'), - ); - - expect(result).toBe(12345); - expect(clickHouseService.select).toHaveBeenCalledTimes(1); - - const [query, params] = clickHouseService.select.mock.calls[0]; - - expect(query).toContain('sum(creditsUsedMicro)'); - expect(query).toContain('FROM usageEvent'); - expect(params).toMatchObject({ - workspaceId: 'workspace_123', - }); - expect(query).not.toContain('operationType'); - }); - - it('coerces string totals returned by ClickHouse to a number', async () => { - clickHouseService.select.mockResolvedValue([{ total: '9876543210' }]); - - const result = await service.getCurrentPeriodCreditsUsed( - 'workspace_123', - new Date('2026-04-01T00:00:00Z'), - new Date('2026-05-01T00:00:00Z'), - ); - - expect(result).toBe(9876543210); - }); - - it('returns 0 when no rows are returned', async () => { - clickHouseService.select.mockResolvedValue([]); - - const result = await service.getCurrentPeriodCreditsUsed( - 'workspace_123', - new Date('2026-04-01T00:00:00Z'), - new Date('2026-05-01T00:00:00Z'), - ); - - expect(result).toBe(0); - }); - - it('returns 0 when total is null', async () => { - clickHouseService.select.mockResolvedValue([{ total: null }]); - - const result = await service.getCurrentPeriodCreditsUsed( - 'workspace_123', - new Date('2026-04-01T00:00:00Z'), - new Date('2026-05-01T00:00:00Z'), - ); - - expect(result).toBe(0); - }); - }); - describe('getBatchPeriodCreditsUsed', () => { beforeEach(() => { twentyConfigService.get.mockReturnValue('http://clickhouse:8123'); @@ -167,7 +120,6 @@ describe('BillingUsageCapService', () => { const result = await service.getBatchPeriodCreditsUsed( ['ws_1', 'ws_2'], new Date('2026-04-01T00:00:00Z'), - new Date('2026-05-01T00:00:00Z'), ); expect(result.size).toBe(0); @@ -178,7 +130,6 @@ describe('BillingUsageCapService', () => { const result = await service.getBatchPeriodCreditsUsed( [], new Date('2026-04-01T00:00:00Z'), - new Date('2026-05-01T00:00:00Z'), ); expect(result.size).toBe(0); @@ -194,7 +145,6 @@ describe('BillingUsageCapService', () => { const result = await service.getBatchPeriodCreditsUsed( ['ws_1', 'ws_2', 'ws_3'], new Date('2026-04-01T00:00:00Z'), - new Date('2026-05-01T00:00:00Z'), ); expect(result.get('ws_1')).toBe(500_000); @@ -218,122 +168,12 @@ describe('BillingUsageCapService', () => { const result = await service.getBatchPeriodCreditsUsed( ['ws_1'], new Date('2026-04-01T00:00:00Z'), - new Date('2026-05-01T00:00:00Z'), ); expect(result.get('ws_1')).toBe(9876543210); }); }); - describe('evaluateCap', () => { - beforeEach(() => { - twentyConfigService.get.mockReturnValue('http://clickhouse:8123'); - }); - - it('returns skipped when ClickHouse is disabled', async () => { - twentyConfigService.get.mockReturnValue(''); - - const result = await service.evaluateCap(buildSubscription()); - - expect(result).toEqual({ - 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( 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 index 0f19e092c1f..ba703281ef3 100644 --- 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 @@ -2,11 +2,20 @@ import { Injectable } from '@nestjs/common'; +import { InjectRepository } from '@nestjs/typeorm'; import { ClickHouseService } from 'src/database/clickHouse/clickHouse.service'; import { formatDateForClickHouse } from 'src/database/clickHouse/clickHouse.util'; -import { type BillingSubscriptionEntity } from 'src/engine/core-modules/billing/entities/billing-subscription.entity'; +import { + BillingException, + BillingExceptionCode, +} from 'src/engine/core-modules/billing/billing.exception'; +import { BillingSubscriptionItemEntity } from 'src/engine/core-modules/billing/entities/billing-subscription-item.entity'; +import { BillingSubscriptionEntity } from 'src/engine/core-modules/billing/entities/billing-subscription.entity'; +import { BillingProductKey } from 'src/engine/core-modules/billing/enums/billing-product-key.enum'; +import { SubscriptionStatus } from 'src/engine/core-modules/billing/enums/billing-subscription-status.enum'; import { MeteredCreditService } from 'src/engine/core-modules/billing/services/metered-credit.service'; import { TwentyConfigService } from 'src/engine/core-modules/twenty-config/twenty-config.service'; +import { Not, Raw, Repository } from 'typeorm'; export type BillingCapEvaluation = | { @@ -22,10 +31,6 @@ export type BillingCapEvaluation = reason: 'no-metered-item' | 'clickhouse-disabled'; }; -type UsageSumRow = { - total: string | number | null; -}; - type BatchUsageSumRow = { workspaceId: string; total: string | number | null; @@ -37,45 +42,17 @@ export class BillingUsageCapService { private readonly clickHouseService: ClickHouseService, private readonly meteredCreditService: MeteredCreditService, private readonly twentyConfigService: TwentyConfigService, + @InjectRepository(BillingSubscriptionItemEntity) + private readonly billingSubscriptionItemRepository: Repository, ) {} 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(); @@ -87,15 +64,13 @@ export class BillingUsageCapService { SELECT workspaceId, sum(creditsUsedMicro) AS total FROM usageEvent WHERE workspaceId IN {workspaceIds:Array(String)} - AND timestamp >= {periodStart:String} - AND timestamp < {periodEnd:String} + AND periodStart = {periodStart:Date} GROUP BY workspaceId `; const rows = await this.clickHouseService.select(query, { workspaceIds, periodStart: formatDateForClickHouse(periodStart), - periodEnd: formatDateForClickHouse(periodEnd), }); for (const row of rows) { @@ -108,46 +83,6 @@ export class BillingUsageCapService { return result; } - async evaluateCap( - subscription: BillingSubscriptionEntity, - ): Promise { - 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, @@ -186,4 +121,38 @@ export class BillingUsageCapService { return results; } + + async setSubscriptionItemHasReachedCap( + workspaceId: string, + hasReachedCap: boolean, + ): Promise { + const billingSubscriptionItems = + await this.billingSubscriptionItemRepository.find({ + where: { + billingSubscription: { + workspaceId, + status: Not(SubscriptionStatus.Canceled), + }, + billingProduct: { + metadata: Raw((alias) => `${alias} @> :metadata::jsonb`, { + metadata: JSON.stringify({ + productKey: BillingProductKey.WORKFLOW_NODE_EXECUTION, + }), + }), + }, + }, + }); + + if (billingSubscriptionItems.length !== 1) { + throw new BillingException( + `Expected 1 metered billing subscription item for workspace ${workspaceId}, but got ${billingSubscriptionItems.length}`, + BillingExceptionCode.BILLING_SUBSCRIPTION_ITEM_NOT_FOUND, + ); + } + + await this.billingSubscriptionItemRepository.update( + { id: billingSubscriptionItems[0].id }, + { hasReachedCurrentPeriodCap: hasReachedCap }, + ); + } } diff --git a/packages/twenty-server/src/engine/core-modules/billing/services/billing-usage.service.ts b/packages/twenty-server/src/engine/core-modules/billing/services/billing-usage.service.ts index a0053148396..1f4ac7bf09e 100644 --- a/packages/twenty-server/src/engine/core-modules/billing/services/billing-usage.service.ts +++ b/packages/twenty-server/src/engine/core-modules/billing/services/billing-usage.service.ts @@ -6,6 +6,8 @@ import { InjectRepository } from '@nestjs/typeorm'; import { isDefined } from 'twenty-shared/utils'; import { type Repository } from 'typeorm'; +import { ClickHouseService } from 'src/database/clickHouse/clickHouse.service'; +import { formatDateTimeForClickHouse } from 'src/database/clickHouse/clickHouse.util'; import { BillingException, BillingExceptionCode, @@ -16,11 +18,22 @@ import { BillingSubscriptionEntity } from 'src/engine/core-modules/billing/entit import { SubscriptionStatus } from 'src/engine/core-modules/billing/enums/billing-subscription-status.enum'; import { BillingSubscriptionItemService } from 'src/engine/core-modules/billing/services/billing-subscription-item.service'; import { BillingSubscriptionService } from 'src/engine/core-modules/billing/services/billing-subscription.service'; +import { BillingUsageCapService } from 'src/engine/core-modules/billing/services/billing-usage-cap.service'; +import { MeteredCreditService } from 'src/engine/core-modules/billing/services/metered-credit.service'; import { StripeBillingMeterEventService } from 'src/engine/core-modules/billing/stripe/services/stripe-billing-meter-event.service'; import { StripeCreditGrantService } from 'src/engine/core-modules/billing/stripe/services/stripe-credit-grant.service'; -import { type UsageEvent } from 'src/engine/core-modules/usage/types/usage-event.type'; +import { buildBillingUsageAvailableCreditsCacheKey } from 'src/engine/core-modules/billing/utils/build-billing-usage-available-credits-cache-key.util'; +import { InjectCacheStorage } from 'src/engine/core-modules/cache-storage/decorators/cache-storage.decorator'; +import { CacheStorageService } from 'src/engine/core-modules/cache-storage/services/cache-storage.service'; +import { CacheStorageNamespace } from 'src/engine/core-modules/cache-storage/types/cache-storage-namespace.enum'; import { TwentyConfigService } from 'src/engine/core-modules/twenty-config/twenty-config.service'; +import { type UsageEvent } from 'src/engine/core-modules/usage/types/usage-event.type'; import { type WorkspaceEntity } from 'src/engine/core-modules/workspace/workspace.entity'; +import { WorkspaceCacheService } from 'src/engine/workspace-cache/services/workspace-cache.service'; + +type UsageSumRow = { + total: string | number | null; +}; @Injectable() export class BillingUsageService { @@ -33,6 +46,14 @@ export class BillingUsageService { private readonly twentyConfigService: TwentyConfigService, private readonly billingSubscriptionItemService: BillingSubscriptionItemService, private readonly stripeCreditGrantService: StripeCreditGrantService, + @InjectCacheStorage(CacheStorageNamespace.EngineBillingUsage) + private readonly billingUsageCacheStorage: CacheStorageService, + @InjectRepository(BillingSubscriptionEntity) + private readonly billingSubscriptionRepository: Repository, + private readonly meteredCreditService: MeteredCreditService, + private readonly workspaceCacheService: WorkspaceCacheService, + private readonly clickHouseService: ClickHouseService, + private readonly billingUsageCapService: BillingUsageCapService, ) {} async canFeatureBeUsed(workspaceId: string): Promise { @@ -48,6 +69,7 @@ export class BillingUsageService { return !!billingSubscription; } + //TODO: TO be deprecated async billUsage({ workspaceId, usageEvents, @@ -86,6 +108,7 @@ export class BillingUsageService { } } + //TODO: TO be deprecated async getMeteredProductsUsage( workspace: WorkspaceEntity, ): Promise { @@ -113,6 +136,7 @@ export class BillingUsageService { ); } + //TODO: TO be deprecated private getSubscriptionPeriod(subscription: BillingSubscriptionEntity): { periodStart: Date; periodEnd: Date; @@ -135,6 +159,7 @@ export class BillingUsageService { }; } + //TODO: TO be deprecated private async buildMeteredProductUsage( subscription: BillingSubscriptionEntity, item: Awaited< @@ -175,4 +200,194 @@ export class BillingUsageService { unitPriceCents: item.unitPriceCents, }; } + + async flushAvailableCreditsFromCache(workspaceId: string): Promise { + await this.billingUsageCacheStorage.flushByPattern( + `available-credits:${workspaceId}:*`, + ); + } + + private async warmAvailableCredits( + workspaceId: string, + periodStart: Date | string, + periodEnd: Date | string, + availableCredits: number, + ): Promise { + const ttlMs = Math.max(new Date(periodEnd).getTime() - Date.now(), 0); + + await this.billingUsageCacheStorage.set( + buildBillingUsageAvailableCreditsCacheKey(workspaceId, periodStart), + availableCredits, + ttlMs, + ); + } + + private async getAvailableCreditsFromCache( + workspaceId: string, + periodStart: Date | string, + ): Promise { + return this.billingUsageCacheStorage.get( + buildBillingUsageAvailableCreditsCacheKey(workspaceId, periodStart), + ); + } + + private async getAvailableCreditsFromClickHouse({ + workspaceId, + currentPeriodStart, + }: { + workspaceId: string; + currentPeriodStart: Date | string; + }): Promise { + const subscription = await this.billingSubscriptionRepository.findOne({ + where: { workspaceId, currentPeriodStart: new Date(currentPeriodStart) }, + relations: [ + 'billingSubscriptionItems', + 'billingSubscriptionItems.billingProduct', + 'billingSubscriptionItems.billingProduct.billingPrices', + ], + }); + + if (!isDefined(subscription)) { + throw new BillingException( + `Subscription not found for workspace ${workspaceId}`, + BillingExceptionCode.BILLING_SUBSCRIPTION_NOT_FOUND, + ); + } + + const meteredPricingInfo = + this.meteredCreditService.extractMeteredPricingInfoFromSubscription( + subscription, + ); + + if (!meteredPricingInfo) { + throw new BillingException( + `No metered item found for workspace ${workspaceId}`, + BillingExceptionCode.BILLING_SUBSCRIPTION_ITEM_NOT_FOUND, + ); + } + + const [creditBalance, usage] = await Promise.all([ + this.meteredCreditService.getCreditBalance( + subscription.stripeCustomerId, + meteredPricingInfo.unitPriceCents, + ), + this.getCurrentPeriodCreditsUsed( + subscription.workspaceId, + subscription.currentPeriodStart, + ), + ]); + + return meteredPricingInfo.tierCap + creditBalance - usage; + } + + async decrementAvailableCredits({ + workspaceId, + usedCredits, + }: { + workspaceId: string; + usedCredits: number; + }): Promise { + const { + billingSubscription: { currentPeriodStart, currentPeriodEnd }, + } = await this.workspaceCacheService.getOrRecompute(workspaceId, [ + 'billingSubscription', + ]); + + const cachedAvailableCredits = await this.getAvailableCreditsFromCache( + workspaceId, + currentPeriodStart, + ); + + const availableCredits = isDefined(cachedAvailableCredits) + ? cachedAvailableCredits + : await this.getAvailableCreditsFromClickHouse({ + workspaceId, + currentPeriodStart, + }); + + if (!isDefined(cachedAvailableCredits)) { + await this.warmAvailableCredits( + workspaceId, + currentPeriodStart, + currentPeriodEnd, + availableCredits, + ); + } + + const decrementedAvailableCredits = + await this.billingUsageCacheStorage.incrBy( + buildBillingUsageAvailableCreditsCacheKey( + workspaceId, + currentPeriodStart, + ), + -usedCredits, + ); + + if (decrementedAvailableCredits <= 0) { + await this.billingUsageCapService.setSubscriptionItemHasReachedCap( + workspaceId, + true, + ); + } + } + + async invalidateAvailableCredits( + workspaceId: string, + periodStart: Date, + ): Promise { + await this.billingUsageCacheStorage.del( + buildBillingUsageAvailableCreditsCacheKey(workspaceId, periodStart), + ); + } + + async hasAvailableCredits(workspaceId: string): Promise { + const { billingSubscription: subscription } = + await this.workspaceCacheService.getOrRecompute(workspaceId, [ + 'billingSubscription', + ]); + const cached = await this.getAvailableCreditsFromCache( + subscription.workspaceId, + subscription.currentPeriodStart, + ); + + if (isDefined(cached)) { + return cached > 0; + } + + const availableCredits = await this.getAvailableCreditsFromClickHouse({ + workspaceId: subscription.workspaceId, + currentPeriodStart: subscription.currentPeriodStart, + }); + + await this.warmAvailableCredits( + subscription.workspaceId, + subscription.currentPeriodStart, + subscription.currentPeriodEnd, + availableCredits, + ); + + return availableCredits > 0; + } + + async getCurrentPeriodCreditsUsed( + workspaceId: string, + periodStart: Date, + ): Promise { + const query = ` + SELECT sum(creditsUsedMicro) AS total + FROM usageEvent + WHERE workspaceId = {workspaceId:String} + AND periodStart = {periodStart:DateTime64(3)} + `; + + const rows = await this.clickHouseService.select(query, { + workspaceId, + periodStart: formatDateTimeForClickHouse(periodStart), + }); + + const rawTotal = rows[0]?.total ?? 0; + const total = typeof rawTotal === 'string' ? Number(rawTotal) : rawTotal; + + return Number.isFinite(total) ? total : 0; + } } diff --git a/packages/twenty-server/src/engine/core-modules/billing/services/billing.service.ts b/packages/twenty-server/src/engine/core-modules/billing/services/billing.service.ts index 3db15b1a740..de12ec91b83 100644 --- a/packages/twenty-server/src/engine/core-modules/billing/services/billing.service.ts +++ b/packages/twenty-server/src/engine/core-modules/billing/services/billing.service.ts @@ -8,11 +8,8 @@ import { type Repository } from 'typeorm'; import { BillingSubscriptionEntity } from 'src/engine/core-modules/billing/entities/billing-subscription.entity'; import { type BillingEntitlementKey } from 'src/engine/core-modules/billing/enums/billing-entitlement-key.enum'; -import { type BillingProductKey } from 'src/engine/core-modules/billing/enums/billing-product-key.enum'; -import { SubscriptionStatus } from 'src/engine/core-modules/billing/enums/billing-subscription-status.enum'; import { BillingProductService } from 'src/engine/core-modules/billing/services/billing-product.service'; import { BillingSubscriptionService } from 'src/engine/core-modules/billing/services/billing-subscription.service'; -import { getPlanKeyFromSubscription } from 'src/engine/core-modules/billing/utils/get-plan-key-from-subscription.util'; import { TwentyConfigService } from 'src/engine/core-modules/twenty-config/twenty-config.service'; @Injectable() @@ -66,41 +63,4 @@ export class BillingService { return !hasAnySubscription; } - - async canBillMeteredProduct( - workspaceId: string, - productKey: BillingProductKey, - ): Promise { - const subscription = - await this.billingSubscriptionService.getCurrentBillingSubscriptionOrThrow( - { workspaceId }, - ); - - const billableStatuses = [ - SubscriptionStatus.Active, - SubscriptionStatus.Trialing, - ]; - - if (!billableStatuses.includes(subscription.status)) { - return false; - } - - const planKey = getPlanKeyFromSubscription(subscription); - const products = - await this.billingProductService.getProductsByPlan(planKey); - - const targetProduct = products.find( - ({ metadata }) => metadata.productKey === productKey, - ); - - if (!targetProduct) { - return false; - } - - const subscriptionItem = subscription.billingSubscriptionItems.find( - (item) => item.stripeProductId === targetProduct.stripeProductId, - ); - - return subscriptionItem?.hasReachedCurrentPeriodCap === false; - } } diff --git a/packages/twenty-server/src/engine/core-modules/billing/services/workspace-billing-subscription-cache.service.ts b/packages/twenty-server/src/engine/core-modules/billing/services/workspace-billing-subscription-cache.service.ts new file mode 100644 index 00000000000..9f9703b8be1 --- /dev/null +++ b/packages/twenty-server/src/engine/core-modules/billing/services/workspace-billing-subscription-cache.service.ts @@ -0,0 +1,46 @@ +/* @license Enterprise */ + +import { Injectable } from '@nestjs/common'; +import { BillingSubscriptionService } from 'src/engine/core-modules/billing/services/billing-subscription.service'; + +import { type FlatBillingSubscription } from 'src/engine/core-modules/billing/types/flat-billing-subscription.type'; +import { WorkspaceCache } from 'src/engine/workspace-cache/decorators/workspace-cache.decorator'; +import { WorkspaceCacheProvider } from 'src/engine/workspace-cache/interfaces/workspace-cache-provider.service'; + +@Injectable() +@WorkspaceCache('billingSubscription') +export class WorkspaceBillingSubscriptionCacheService extends WorkspaceCacheProvider { + constructor( + private readonly billingSubscriptionService: BillingSubscriptionService, + ) { + super(); + } + + async computeForCache(workspaceId: string): Promise { + const subscription = + await this.billingSubscriptionService.getCurrentBillingSubscriptionOrThrow( + { + workspaceId, + }, + ); + + return { + id: subscription.id, + workspaceId: subscription.workspaceId, + stripeCustomerId: subscription.stripeCustomerId, + stripeSubscriptionId: subscription.stripeSubscriptionId, + status: subscription.status, + interval: subscription.interval, + currency: subscription.currency, + currentPeriodStart: subscription.currentPeriodStart, + currentPeriodEnd: subscription.currentPeriodEnd, + cancelAtPeriodEnd: subscription.cancelAtPeriodEnd, + cancelAt: subscription.cancelAt, + canceledAt: subscription.canceledAt, + endedAt: subscription.endedAt, + trialStart: subscription.trialStart, + trialEnd: subscription.trialEnd, + collectionMethod: subscription.collectionMethod, + }; + } +} diff --git a/packages/twenty-server/src/engine/core-modules/billing/types/flat-billing-subscription.type.ts b/packages/twenty-server/src/engine/core-modules/billing/types/flat-billing-subscription.type.ts new file mode 100644 index 00000000000..940d7602bbf --- /dev/null +++ b/packages/twenty-server/src/engine/core-modules/billing/types/flat-billing-subscription.type.ts @@ -0,0 +1,24 @@ +/* @license Enterprise */ + +import { type BillingSubscriptionCollectionMethod } from 'src/engine/core-modules/billing/enums/billing-subscription-collection-method.enum'; +import { type SubscriptionInterval } from 'src/engine/core-modules/billing/enums/billing-subscription-interval.enum'; +import { type SubscriptionStatus } from 'src/engine/core-modules/billing/enums/billing-subscription-status.enum'; + +export type FlatBillingSubscription = { + id: string; + workspaceId: string; + stripeCustomerId: string; + stripeSubscriptionId: string; + status: SubscriptionStatus; + interval: SubscriptionInterval; + currency: string; + currentPeriodStart: Date; + currentPeriodEnd: Date; + cancelAtPeriodEnd: boolean; + cancelAt: Date | null; + canceledAt: Date | null; + endedAt: Date | null; + trialStart: Date | null; + trialEnd: Date | null; + collectionMethod: BillingSubscriptionCollectionMethod; +}; diff --git a/packages/twenty-server/src/engine/core-modules/billing/utils/build-billing-usage-available-credits-cache-key.util.ts b/packages/twenty-server/src/engine/core-modules/billing/utils/build-billing-usage-available-credits-cache-key.util.ts new file mode 100644 index 00000000000..3f82adfacab --- /dev/null +++ b/packages/twenty-server/src/engine/core-modules/billing/utils/build-billing-usage-available-credits-cache-key.util.ts @@ -0,0 +1,6 @@ +export const buildBillingUsageAvailableCreditsCacheKey = ( + workspaceId: string, + periodStart: Date | string, +): string => { + return `available-credits:${workspaceId}:${new Date(periodStart).getTime()}`; +}; diff --git a/packages/twenty-server/src/engine/core-modules/cache-storage/types/cache-storage-namespace.enum.ts b/packages/twenty-server/src/engine/core-modules/cache-storage/types/cache-storage-namespace.enum.ts index 8ac71bf71ed..5cca874475f 100644 --- a/packages/twenty-server/src/engine/core-modules/cache-storage/types/cache-storage-namespace.enum.ts +++ b/packages/twenty-server/src/engine/core-modules/cache-storage/types/cache-storage-namespace.enum.ts @@ -8,5 +8,6 @@ export enum CacheStorageNamespace { EngineHealth = 'engine:health', EngineMetrics = 'engine:metrics', EngineSubscriptions = 'engine:subscriptions', + EngineBillingUsage = 'engine:billing-usage', IntegrationTests = 'integration-tests', } diff --git a/packages/twenty-server/src/engine/core-modules/event-logs/cleanup/services/event-log-cleanup.service.ts b/packages/twenty-server/src/engine/core-modules/event-logs/cleanup/services/event-log-cleanup.service.ts index 2d0255ca041..a7cda8eb47f 100644 --- a/packages/twenty-server/src/engine/core-modules/event-logs/cleanup/services/event-log-cleanup.service.ts +++ b/packages/twenty-server/src/engine/core-modules/event-logs/cleanup/services/event-log-cleanup.service.ts @@ -5,7 +5,7 @@ import { Injectable, Logger } from '@nestjs/common'; import { EventLogTable } from 'twenty-shared/types'; import { ClickHouseService } from 'src/database/clickHouse/clickHouse.service'; -import { formatDateForClickHouse } from 'src/database/clickHouse/clickHouse.util'; +import { formatDateTimeForClickHouse } from 'src/database/clickHouse/clickHouse.util'; const CLICKHOUSE_TABLE_NAMES: Record = { [EventLogTable.WORKSPACE_EVENT]: 'workspaceEvent', @@ -52,7 +52,7 @@ export class EventLogCleanupService { `ALTER TABLE ${tableName} DELETE WHERE "workspaceId" = {workspaceId:String} AND "timestamp" < {cutoffDate:DateTime64(3)}`, { workspaceId, - cutoffDate: formatDateForClickHouse(cutoffDate), + cutoffDate: formatDateTimeForClickHouse(cutoffDate), }, ); diff --git a/packages/twenty-server/src/engine/core-modules/event-logs/event-logs.service.ts b/packages/twenty-server/src/engine/core-modules/event-logs/event-logs.service.ts index 9e78bcdaac8..4f26975ec32 100644 --- a/packages/twenty-server/src/engine/core-modules/event-logs/event-logs.service.ts +++ b/packages/twenty-server/src/engine/core-modules/event-logs/event-logs.service.ts @@ -8,7 +8,7 @@ import { isDefined } from 'twenty-shared/utils'; import { Repository } from 'typeorm'; import { ClickHouseService } from 'src/database/clickHouse/clickHouse.service'; -import { formatDateForClickHouse } from 'src/database/clickHouse/clickHouse.util'; +import { formatDateTimeForClickHouse } from 'src/database/clickHouse/clickHouse.util'; import { BillingEntitlementKey } from 'src/engine/core-modules/billing/enums/billing-entitlement-key.enum'; import { BillingService } from 'src/engine/core-modules/billing/services/billing.service'; import { UserWorkspaceEntity } from 'src/engine/core-modules/user-workspace/user-workspace.entity'; @@ -235,12 +235,12 @@ export class EventLogsService { if (isDefined(filters.dateRange?.start)) { whereClauses.push('"timestamp" >= {startDate:DateTime64(3)}'); - params.startDate = formatDateForClickHouse(filters.dateRange.start); + params.startDate = formatDateTimeForClickHouse(filters.dateRange.start); } if (isDefined(filters.dateRange?.end)) { whereClauses.push('"timestamp" <= {endDate:DateTime64(3)}'); - params.endDate = formatDateForClickHouse(filters.dateRange.end); + params.endDate = formatDateTimeForClickHouse(filters.dateRange.end); } if (table === EventLogTable.OBJECT_EVENT) { diff --git a/packages/twenty-server/src/engine/core-modules/logic-function/logic-function-executor/logic-function-executor.module.ts b/packages/twenty-server/src/engine/core-modules/logic-function/logic-function-executor/logic-function-executor.module.ts index 6d933ab2f1f..a5b020283a1 100644 --- a/packages/twenty-server/src/engine/core-modules/logic-function/logic-function-executor/logic-function-executor.module.ts +++ b/packages/twenty-server/src/engine/core-modules/logic-function/logic-function-executor/logic-function-executor.module.ts @@ -2,11 +2,12 @@ import { Module } from '@nestjs/common'; import { TypeOrmModule } from '@nestjs/typeorm'; import { ApplicationRegistrationVariableEntity } from 'src/engine/core-modules/application/application-registration-variable/application-registration-variable.entity'; -import { LogicFunctionExecutorService } from 'src/engine/core-modules/logic-function/logic-function-executor/logic-function-executor.service'; -import { ThrottlerModule } from 'src/engine/core-modules/throttler/throttler.module'; import { AuditModule } from 'src/engine/core-modules/audit/audit.module'; import { TokenModule } from 'src/engine/core-modules/auth/token/token.module'; +import { BillingModule } from 'src/engine/core-modules/billing/billing.module'; +import { LogicFunctionExecutorService } from 'src/engine/core-modules/logic-function/logic-function-executor/logic-function-executor.service'; import { SecretEncryptionModule } from 'src/engine/core-modules/secret-encryption/secret-encryption.module'; +import { ThrottlerModule } from 'src/engine/core-modules/throttler/throttler.module'; import { SubscriptionsModule } from 'src/engine/subscriptions/subscriptions.module'; import { WorkspaceCacheModule } from 'src/engine/workspace-cache/workspace-cache.module'; @@ -18,6 +19,7 @@ import { WorkspaceCacheModule } from 'src/engine/workspace-cache/workspace-cache SecretEncryptionModule, SubscriptionsModule, WorkspaceCacheModule, + BillingModule, TypeOrmModule.forFeature([ApplicationRegistrationVariableEntity]), ], providers: [LogicFunctionExecutorService], diff --git a/packages/twenty-server/src/engine/core-modules/logic-function/logic-function-executor/logic-function-executor.service.ts b/packages/twenty-server/src/engine/core-modules/logic-function/logic-function-executor/logic-function-executor.service.ts index 378d3f129c5..952f48aa46f 100644 --- a/packages/twenty-server/src/engine/core-modules/logic-function/logic-function-executor/logic-function-executor.service.ts +++ b/packages/twenty-server/src/engine/core-modules/logic-function/logic-function-executor/logic-function-executor.service.ts @@ -24,6 +24,8 @@ import { FlatApplication } from 'src/engine/core-modules/application/types/flat- import { AuditService } from 'src/engine/core-modules/audit/services/audit.service'; import { LOGIC_FUNCTION_EXECUTED_EVENT } from 'src/engine/core-modules/audit/utils/events/workspace-event/logic-function/logic-function-executed'; import { ApplicationTokenService } from 'src/engine/core-modules/auth/token/services/application-token.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 { LogicFunctionDriverFactory } from 'src/engine/core-modules/logic-function/logic-function-drivers/logic-function-driver.factory'; import { buildEnvVar } from 'src/engine/core-modules/logic-function/logic-function-executor/utils/build-env-var'; import { SecretEncryptionService } from 'src/engine/core-modules/secret-encryption/secret-encryption.service'; @@ -72,6 +74,8 @@ export class LogicFunctionExecutorService { private readonly auditService: AuditService, private readonly applicationLogsService: ApplicationLogsService, private readonly workspaceEventEmitter: WorkspaceEventEmitter, + private readonly billingService: BillingService, + private readonly billingUsageService: BillingUsageService, @InjectRepository(ApplicationRegistrationVariableEntity) private readonly applicationRegistrationVariableRepository: Repository, ) {} @@ -352,5 +356,12 @@ export class LogicFunctionExecutorService { ], workspaceId, ); + + if (this.billingService.isBillingEnabled()) { + await this.billingUsageService.decrementAvailableCredits({ + workspaceId, + usedCredits: 100, + }); + } } } diff --git a/packages/twenty-server/src/engine/core-modules/usage/services/usage-analytics.service.ts b/packages/twenty-server/src/engine/core-modules/usage/services/usage-analytics.service.ts index b4438dc07f0..f0fdf897f0d 100644 --- a/packages/twenty-server/src/engine/core-modules/usage/services/usage-analytics.service.ts +++ b/packages/twenty-server/src/engine/core-modules/usage/services/usage-analytics.service.ts @@ -3,7 +3,7 @@ import { Injectable } from '@nestjs/common'; import { ClickHouseService } from 'src/database/clickHouse/clickHouse.service'; -import { formatDateForClickHouse } from 'src/database/clickHouse/clickHouse.util'; +import { formatDateTimeForClickHouse } from 'src/database/clickHouse/clickHouse.util'; import { fillUsageTimeSeriesGaps } from 'src/engine/core-modules/usage/utils/fill-usage-time-series-gaps.util'; import { toDisplayCredits } from 'src/engine/core-modules/usage/utils/to-display-credits.util'; import { toDollars } from 'src/engine/core-modules/usage/utils/to-dollars.util'; @@ -75,8 +75,8 @@ export class UsageAnalyticsService { `; const rows = await this.clickHouseService.select(query, { - periodStart: formatDateForClickHouse(params.periodStart), - periodEnd: formatDateForClickHouse(params.periodEnd), + periodStart: formatDateTimeForClickHouse(params.periodStart), + periodEnd: formatDateTimeForClickHouse(params.periodEnd), operationTypes: aiOperationTypes, }); @@ -174,8 +174,8 @@ export class UsageAnalyticsService { const rows = await this.clickHouseService.select(query, { workspaceId, - periodStart: formatDateForClickHouse(periodStart), - periodEnd: formatDateForClickHouse(periodEnd), + periodStart: formatDateTimeForClickHouse(periodStart), + periodEnd: formatDateTimeForClickHouse(periodEnd), ...(operationTypes && operationTypes.length > 0 ? { operationTypes } : {}), @@ -222,8 +222,8 @@ export class UsageAnalyticsService { query, { workspaceId, - periodStart: formatDateForClickHouse(periodStart), - periodEnd: formatDateForClickHouse(periodEnd), + periodStart: formatDateTimeForClickHouse(periodStart), + periodEnd: formatDateTimeForClickHouse(periodEnd), ...(operationTypes && operationTypes.length > 0 ? { operationTypes } : {}), diff --git a/packages/twenty-server/src/engine/core-modules/usage/services/usage-event-writer.service.ts b/packages/twenty-server/src/engine/core-modules/usage/services/usage-event-writer.service.ts index 84b5d6b87b2..83584b7f56c 100644 --- a/packages/twenty-server/src/engine/core-modules/usage/services/usage-event-writer.service.ts +++ b/packages/twenty-server/src/engine/core-modules/usage/services/usage-event-writer.service.ts @@ -3,9 +3,9 @@ import { Injectable, Logger } from '@nestjs/common'; import { ClickHouseService } from 'src/database/clickHouse/clickHouse.service'; -import { formatDateForClickHouse } from 'src/database/clickHouse/clickHouse.util'; -import { type UsageEvent } from 'src/engine/core-modules/usage/types/usage-event.type'; +import { formatDateTimeForClickHouse } from 'src/database/clickHouse/clickHouse.util'; import { TwentyConfigService } from 'src/engine/core-modules/twenty-config/twenty-config.service'; +import { type UsageEvent } from 'src/engine/core-modules/usage/types/usage-event.type'; @Injectable() export class UsageEventWriterService { @@ -21,11 +21,14 @@ export class UsageEventWriterService { return; } - const now = formatDateForClickHouse(new Date()); + const now = formatDateTimeForClickHouse(new Date()); const rows = usageEvents.map((usageEvent) => ({ timestamp: now, workspaceId, + periodStart: usageEvent.periodStart + ? formatDateTimeForClickHouse(usageEvent.periodStart) + : undefined, userWorkspaceId: usageEvent.userWorkspaceId ?? '', resourceType: usageEvent.resourceType, operationType: usageEvent.operationType, diff --git a/packages/twenty-server/src/engine/core-modules/usage/types/usage-event.type.ts b/packages/twenty-server/src/engine/core-modules/usage/types/usage-event.type.ts index fd6481a6f88..7c15916f2c5 100644 --- a/packages/twenty-server/src/engine/core-modules/usage/types/usage-event.type.ts +++ b/packages/twenty-server/src/engine/core-modules/usage/types/usage-event.type.ts @@ -1,7 +1,7 @@ /* @license Enterprise */ -import { type UsageResourceType } from 'src/engine/core-modules/usage/enums/usage-resource-type.enum'; import { type UsageOperationType } from 'src/engine/core-modules/usage/enums/usage-operation-type.enum'; +import { type UsageResourceType } from 'src/engine/core-modules/usage/enums/usage-resource-type.enum'; import { type UsageUnit } from 'src/engine/core-modules/usage/enums/usage-unit.enum'; export type UsageEvent = { @@ -10,6 +10,7 @@ export type UsageEvent = { creditsUsedMicro: number; quantity: number; unit: UsageUnit; + periodStart?: Date; resourceId?: string | null; resourceContext?: string | null; userWorkspaceId?: string | null; diff --git a/packages/twenty-server/src/engine/metadata-modules/ai/ai-billing/ai-billing.module.ts b/packages/twenty-server/src/engine/metadata-modules/ai/ai-billing/ai-billing.module.ts index fa0ce4ea2e5..60c8255f7c6 100644 --- a/packages/twenty-server/src/engine/metadata-modules/ai/ai-billing/ai-billing.module.ts +++ b/packages/twenty-server/src/engine/metadata-modules/ai/ai-billing/ai-billing.module.ts @@ -1,11 +1,12 @@ import { Module } from '@nestjs/common'; +import { BillingModule } from 'src/engine/core-modules/billing/billing.module'; import { AiBillingService } from 'src/engine/metadata-modules/ai/ai-billing/services/ai-billing.service'; import { AiModelsModule } from 'src/engine/metadata-modules/ai/ai-models/ai-models.module'; import { WorkspaceEventEmitterModule } from 'src/engine/workspace-event-emitter/workspace-event-emitter.module'; @Module({ - imports: [WorkspaceEventEmitterModule, AiModelsModule], + imports: [WorkspaceEventEmitterModule, AiModelsModule, BillingModule], providers: [AiBillingService], exports: [AiBillingService], }) diff --git a/packages/twenty-server/src/engine/metadata-modules/ai/ai-billing/services/__tests__/ai-billing.service.spec.ts b/packages/twenty-server/src/engine/metadata-modules/ai/ai-billing/services/__tests__/ai-billing.service.spec.ts index 96263c5a5ee..cf5e1e99e94 100644 --- a/packages/twenty-server/src/engine/metadata-modules/ai/ai-billing/services/__tests__/ai-billing.service.spec.ts +++ b/packages/twenty-server/src/engine/metadata-modules/ai/ai-billing/services/__tests__/ai-billing.service.spec.ts @@ -4,6 +4,8 @@ import { USAGE_RECORDED } from 'src/engine/core-modules/usage/constants/usage-re import { UsageOperationType } from 'src/engine/core-modules/usage/enums/usage-operation-type.enum'; import { UsageResourceType } from 'src/engine/core-modules/usage/enums/usage-resource-type.enum'; import { UsageUnit } from 'src/engine/core-modules/usage/enums/usage-unit.enum'; +import { BillingService } from 'src/engine/core-modules/billing/services/billing.service'; +import { BillingUsageService } from 'src/engine/core-modules/billing/services/billing-usage.service'; import { AiBillingService } from 'src/engine/metadata-modules/ai/ai-billing/services/ai-billing.service'; import { ModelFamily } from 'src/engine/metadata-modules/ai/ai-models/types/model-family.enum'; import { AiModelRegistryService } from 'src/engine/metadata-modules/ai/ai-models/services/ai-model-registry.service'; @@ -73,6 +75,18 @@ describe('AiBillingService', () => { provide: AiModelRegistryService, useValue: mockAiModelRegistryMethods, }, + { + provide: BillingService, + useValue: { + isBillingEnabled: jest.fn().mockReturnValue(false), + }, + }, + { + provide: BillingUsageService, + useValue: { + decrementAvailableCredits: jest.fn().mockResolvedValue(undefined), + }, + }, ], }).compile(); diff --git a/packages/twenty-server/src/engine/metadata-modules/ai/ai-billing/services/ai-billing.service.ts b/packages/twenty-server/src/engine/metadata-modules/ai/ai-billing/services/ai-billing.service.ts index 1ddafefee9e..fe5de62ed5a 100644 --- a/packages/twenty-server/src/engine/metadata-modules/ai/ai-billing/services/ai-billing.service.ts +++ b/packages/twenty-server/src/engine/metadata-modules/ai/ai-billing/services/ai-billing.service.ts @@ -1,6 +1,8 @@ import { Injectable, Logger } from '@nestjs/common'; import { type LanguageModelUsage } from 'ai'; +import { BillingUsageService } from 'src/engine/core-modules/billing/services/billing-usage.service'; +import { BillingService } from 'src/engine/core-modules/billing/services/billing.service'; import { USAGE_RECORDED } from 'src/engine/core-modules/usage/constants/usage-recorded.constant'; import { UsageOperationType } from 'src/engine/core-modules/usage/enums/usage-operation-type.enum'; @@ -10,8 +12,8 @@ import { type UsageEvent } from 'src/engine/core-modules/usage/types/usage-event import { NATIVE_WEB_SEARCH_COST_PER_CALL_DOLLARS } from 'src/engine/metadata-modules/ai/ai-billing/constants/native-web-search-cost-per-call-dollars'; import { computeCostBreakdown } from 'src/engine/metadata-modules/ai/ai-billing/utils/compute-cost-breakdown.util'; import { convertDollarsToBillingCredits } from 'src/engine/metadata-modules/ai/ai-billing/utils/convert-dollars-to-billing-credits.util'; -import { type ModelId } from 'src/engine/metadata-modules/ai/ai-models/types/model-id.type'; import { AiModelRegistryService } from 'src/engine/metadata-modules/ai/ai-models/services/ai-model-registry.service'; +import { type ModelId } from 'src/engine/metadata-modules/ai/ai-models/types/model-id.type'; import { WorkspaceEventEmitter } from 'src/engine/workspace-event-emitter/workspace-event-emitter'; export type BillingUsageInput = { @@ -26,6 +28,8 @@ export class AiBillingService { constructor( private readonly workspaceEventEmitter: WorkspaceEventEmitter, private readonly aiModelRegistryService: AiModelRegistryService, + private readonly billingService: BillingService, + private readonly billingUsageService: BillingUsageService, ) {} calculateCost(modelId: ModelId, billingInput: BillingUsageInput): number { @@ -52,14 +56,14 @@ export class AiBillingService { return breakdown.totalCostInDollars; } - calculateAndBillUsage( + async calculateAndBillUsage( modelId: ModelId, billingInput: BillingUsageInput, workspaceId: string, operationType: UsageOperationType, agentId?: string | null, userWorkspaceId?: string | null, - ): void { + ): Promise { const costInDollars = this.calculateCost(modelId, billingInput); const creditsUsedMicro = Math.round( convertDollarsToBillingCredits(costInDollars), @@ -70,7 +74,7 @@ export class AiBillingService { (billingInput.usage.outputTokens ?? 0) + (billingInput.cacheCreationTokens ?? 0); - this.emitAiTokenUsageEvent( + await this.emitAiTokenUsageEvent( workspaceId, creditsUsedMicro, totalTokens, @@ -81,11 +85,11 @@ export class AiBillingService { ); } - billNativeWebSearchUsage( + async billNativeWebSearchUsage( nativeWebSearchCallCount: number, workspaceId: string, userWorkspaceId?: string | null, - ): void { + ): Promise { if (nativeWebSearchCallCount <= 0) { return; } @@ -114,9 +118,16 @@ export class AiBillingService { ], workspaceId, ); + + if (this.billingService.isBillingEnabled()) { + await this.billingUsageService.decrementAvailableCredits({ + workspaceId, + usedCredits: creditsUsedMicro, + }); + } } - private emitAiTokenUsageEvent( + private async emitAiTokenUsageEvent( workspaceId: string, creditsUsedMicro: number, totalTokens: number, @@ -124,7 +135,7 @@ export class AiBillingService { operationType: UsageOperationType, agentId?: string | null, userWorkspaceId?: string | null, - ): void { + ): Promise { this.workspaceEventEmitter.emitCustomBatchEvent( USAGE_RECORDED, [ @@ -141,5 +152,11 @@ export class AiBillingService { ], workspaceId, ); + if (this.billingService.isBillingEnabled()) { + await this.billingUsageService.decrementAvailableCredits({ + workspaceId, + usedCredits: creditsUsedMicro, + }); + } } } diff --git a/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/resolvers/agent-chat.resolver.ts b/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/resolvers/agent-chat.resolver.ts index 5623f3fe954..b7b1237a0e1 100644 --- a/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/resolvers/agent-chat.resolver.ts +++ b/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/resolvers/agent-chat.resolver.ts @@ -9,10 +9,10 @@ import { } from '@nestjs/graphql'; import { InjectRepository } from '@nestjs/typeorm'; +import GraphQLJSON from 'graphql-type-json'; import { PermissionFlagType } from 'twenty-shared/constants'; import { isDefined } from 'twenty-shared/utils'; import { Repository } from 'typeorm'; -import GraphQLJSON from 'graphql-type-json'; import { MetadataResolver } from 'src/engine/api/graphql/graphql-config/decorators/metadata-resolver.decorator'; import { UUIDScalarType } from 'src/engine/api/graphql/workspace-schema-builder/graphql-types/scalars'; @@ -20,8 +20,7 @@ import { BillingException, BillingExceptionCode, } from 'src/engine/core-modules/billing/billing.exception'; -import { BillingProductKey } from 'src/engine/core-modules/billing/enums/billing-product-key.enum'; -import { BillingService } from 'src/engine/core-modules/billing/services/billing.service'; +import { BillingUsageService } from 'src/engine/core-modules/billing/services/billing-usage.service'; import { RedisClientService } from 'src/engine/core-modules/redis-client/redis-client.service'; import { TwentyConfigService } from 'src/engine/core-modules/twenty-config/twenty-config.service'; import { toDisplayCredits } from 'src/engine/core-modules/usage/utils/to-display-credits.util'; @@ -30,13 +29,8 @@ import { AuthUserWorkspaceId } from 'src/engine/decorators/auth/auth-user-worksp import { AuthWorkspace } from 'src/engine/decorators/auth/auth-workspace.decorator'; import { SettingsPermissionGuard } from 'src/engine/guards/settings-permission.guard'; import { WorkspaceAuthGuard } from 'src/engine/guards/workspace-auth.guard'; -import { - AiException, - AiExceptionCode, -} from 'src/engine/metadata-modules/ai/ai.exception'; -import { AiGraphqlApiExceptionInterceptor } from 'src/engine/metadata-modules/ai/interceptors/ai-graphql-api-exception.interceptor'; -import { type BrowsingContextType } from 'src/engine/metadata-modules/ai/ai-agent/types/browsingContext.type'; import { AgentMessageDTO } from 'src/engine/metadata-modules/ai/ai-agent-execution/dtos/agent-message.dto'; +import { type BrowsingContextType } from 'src/engine/metadata-modules/ai/ai-agent/types/browsingContext.type'; import { AgentChatThreadDTO } from 'src/engine/metadata-modules/ai/ai-chat/dtos/agent-chat-thread.dto'; import { AiSystemPromptPreviewDTO } from 'src/engine/metadata-modules/ai/ai-chat/dtos/ai-system-prompt-preview.dto'; import { ChatStreamCatchupChunksDTO } from 'src/engine/metadata-modules/ai/ai-chat/dtos/chat-stream-catchup-chunks.dto'; @@ -45,9 +39,14 @@ import { AgentChatThreadEntity } from 'src/engine/metadata-modules/ai/ai-chat/en import { AgentChatEventPublisherService } from 'src/engine/metadata-modules/ai/ai-chat/services/agent-chat-event-publisher.service'; import { AgentChatStreamingService } from 'src/engine/metadata-modules/ai/ai-chat/services/agent-chat-streaming.service'; import { AgentChatService } from 'src/engine/metadata-modules/ai/ai-chat/services/agent-chat.service'; -import { getCancelChannel } from 'src/engine/metadata-modules/ai/ai-chat/utils/get-cancel-channel.util'; import { SystemPromptBuilderService } from 'src/engine/metadata-modules/ai/ai-chat/services/system-prompt-builder.service'; +import { getCancelChannel } from 'src/engine/metadata-modules/ai/ai-chat/utils/get-cancel-channel.util'; import { AiModelRegistryService } from 'src/engine/metadata-modules/ai/ai-models/services/ai-model-registry.service'; +import { + AiException, + AiExceptionCode, +} from 'src/engine/metadata-modules/ai/ai.exception'; +import { AiGraphqlApiExceptionInterceptor } from 'src/engine/metadata-modules/ai/interceptors/ai-graphql-api-exception.interceptor'; @UseGuards(WorkspaceAuthGuard, SettingsPermissionGuard(PermissionFlagType.AI)) @UseInterceptors(AiGraphqlApiExceptionInterceptor) @@ -58,7 +57,7 @@ export class AgentChatResolver { private readonly agentChatStreamingService: AgentChatStreamingService, private readonly eventPublisherService: AgentChatEventPublisherService, private readonly systemPromptBuilderService: SystemPromptBuilderService, - private readonly billingService: BillingService, + private readonly billingUsageService: BillingUsageService, private readonly twentyConfigService: TwentyConfigService, private readonly aiModelRegistryService: AiModelRegistryService, private readonly redisClientService: RedisClientService, @@ -135,9 +134,8 @@ export class AgentChatResolver { ); if (this.twentyConfigService.get('IS_BILLING_ENABLED')) { - const canBill = await this.billingService.canBillMeteredProduct( + const canBill = await this.billingUsageService.hasAvailableCredits( workspace.id, - BillingProductKey.WORKFLOW_NODE_EXECUTION, ); if (!canBill) { diff --git a/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/services/chat-execution.service.ts b/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/services/chat-execution.service.ts index 5cabaac7f5e..3661f0f1518 100644 --- a/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/services/chat-execution.service.ts +++ b/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/services/chat-execution.service.ts @@ -266,7 +266,7 @@ export class ChatExecutionService { const modelMessages = pruningResult.messages; - const billUsageFromSteps = (steps: StepResult[]) => { + const billUsageFromSteps = async (steps: StepResult[]) => { const usage = steps.reduce( (acc, step) => ({ inputTokens: (acc.inputTokens ?? 0) + (step.usage.inputTokens ?? 0), @@ -308,7 +308,7 @@ export class ChatExecutionService { const cacheCreationTokens = extractCacheCreationTokensFromSteps(steps); - this.aiBillingService.calculateAndBillUsage( + await this.aiBillingService.calculateAndBillUsage( registeredModel.modelId, { usage, cacheCreationTokens }, workspace.id, @@ -333,8 +333,8 @@ export class ChatExecutionService { abortSignal, stopWhen: stepCountIs(AGENT_CONFIG.MAX_STEPS), experimental_telemetry: AI_TELEMETRY_CONFIG, - onAbort: ({ steps }) => { - billUsageFromSteps(steps); + onAbort: async ({ steps }) => { + await billUsageFromSteps(steps); }, experimental_repairToolCall: async ({ toolCall, @@ -360,8 +360,8 @@ export class ChatExecutionService { }); Promise.all([stream.usage, stream.steps]) - .then(([, steps]) => { - billUsageFromSteps(steps); + .then(async ([, steps]) => { + await billUsageFromSteps(steps); }) .catch((error) => { if (error?.name === 'AbortError') { diff --git a/packages/twenty-server/src/engine/workspace-cache/types/workspace-cache-key.type.ts b/packages/twenty-server/src/engine/workspace-cache/types/workspace-cache-key.type.ts index 7f6d696d069..3d94e69bc42 100644 --- a/packages/twenty-server/src/engine/workspace-cache/types/workspace-cache-key.type.ts +++ b/packages/twenty-server/src/engine/workspace-cache/types/workspace-cache-key.type.ts @@ -8,6 +8,7 @@ import { type ResolverNameMapEntry } from 'src/engine/api/graphql/direct-executi import { type FlatApiKey } from 'src/engine/core-modules/api-key/types/flat-api-key.type'; import { type ApplicationVariableCacheMaps } from 'src/engine/core-modules/application/application-variable/types/application-variable-cache-maps.type'; import { type FlatApplicationCacheMaps } from 'src/engine/core-modules/application/types/flat-application-cache-maps.type'; +import { type FlatBillingSubscription } from 'src/engine/core-modules/billing/types/flat-billing-subscription.type'; import { type FlatWorkspaceMemberMaps } from 'src/engine/core-modules/user/types/flat-workspace-member-maps.type'; import { type FlatRoleTargetByAgentIdMaps } from 'src/engine/metadata-modules/flat-agent/types/flat-role-target-by-agent-id-maps.type'; import { type AllFlatEntityMaps } from 'src/engine/metadata-modules/flat-entity/types/all-flat-entity-maps.type'; @@ -56,6 +57,7 @@ export const WORKSPACE_CACHE_KEYS_V2 = { apiKeyMap: 'cache:api-key-map', applicationVariableMaps: 'cache:application-variable', graphQLResolverNameMap: 'direct-execution:graphql-resolver-name-map', + billingSubscription: 'billing:subscription', } as const satisfies Record; export type AdditionalCacheDataMaps = { @@ -72,6 +74,7 @@ export type AdditionalCacheDataMaps = { flatWorkspaceMemberMaps: FlatWorkspaceMemberMaps; applicationVariableMaps: ApplicationVariableCacheMaps; graphQLResolverNameMap: Record; + billingSubscription: FlatBillingSubscription; }; export type WorkspaceCacheDataMap = AllFlatEntityMaps & diff --git a/packages/twenty-server/src/modules/workflow/workflow-executor/workspace-services/__tests__/workflow-executor.workspace-service.spec.ts b/packages/twenty-server/src/modules/workflow/workflow-executor/workspace-services/__tests__/workflow-executor.workspace-service.spec.ts index 750278d97c9..8637d772f53 100644 --- a/packages/twenty-server/src/modules/workflow/workflow-executor/workspace-services/__tests__/workflow-executor.workspace-service.spec.ts +++ b/packages/twenty-server/src/modules/workflow/workflow-executor/workspace-services/__tests__/workflow-executor.workspace-service.spec.ts @@ -2,6 +2,8 @@ import { Test, type TestingModule } from '@nestjs/testing'; import { getWorkflowRunContext, StepStatus } from 'twenty-shared/workflow'; +import { BillingUsageService } from 'src/engine/core-modules/billing/services/billing-usage.service'; +import { BillingService } from 'src/engine/core-modules/billing/services/billing.service'; import { ExceptionHandlerService } from 'src/engine/core-modules/exception-handler/exception-handler.service'; import { MessageQueue } from 'src/engine/core-modules/message-queue/message-queue.constants'; import { MetricsService } from 'src/engine/core-modules/metrics/metrics.service'; @@ -69,6 +71,15 @@ describe('WorkflowExecutorWorkspaceService', () => { getWorkflowRunOrFail: jest.fn(), }; + const mockBillingService = { + isBillingEnabled: jest.fn().mockReturnValue(true), + }; + + const mockBillingUsageService = { + hasAvailableCredits: jest.fn().mockResolvedValue(true), + decrementAvailableCredits: jest.fn().mockResolvedValue(undefined), + }; + const mockExceptionHandlerService = { captureExceptions: jest.fn(), }; @@ -101,6 +112,14 @@ describe('WorkflowExecutorWorkspaceService', () => { provide: WorkflowRunWorkspaceService, useValue: mockWorkflowRunWorkspaceService, }, + { + provide: BillingService, + useValue: mockBillingService, + }, + { + provide: BillingUsageService, + useValue: mockBillingUsageService, + }, { provide: ExceptionHandlerService, useValue: mockExceptionHandlerService, @@ -675,8 +694,8 @@ describe('WorkflowExecutorWorkspaceService', () => { }); describe('sendWorkflowNodeRunEvent', () => { - it('should emit a billing event', () => { - service['sendWorkflowNodeRunEvent']('workspace-id', 'workflow-id'); + it('should emit a billing event', async () => { + await service['sendWorkflowNodeRunEvent']('workspace-id', 'workflow-id'); expect(workspaceEventEmitter.emitCustomBatchEvent).toHaveBeenCalledWith( USAGE_RECORDED, diff --git a/packages/twenty-server/src/modules/workflow/workflow-executor/workspace-services/workflow-executor.workspace-service.ts b/packages/twenty-server/src/modules/workflow/workflow-executor/workspace-services/workflow-executor.workspace-service.ts index 133f793231c..9e03b8a23fd 100644 --- a/packages/twenty-server/src/modules/workflow/workflow-executor/workspace-services/workflow-executor.workspace-service.ts +++ b/packages/twenty-server/src/modules/workflow/workflow-executor/workspace-services/workflow-executor.workspace-service.ts @@ -9,6 +9,8 @@ import { WorkflowRunStepInfos, } from 'twenty-shared/workflow'; +import { BillingUsageService } from 'src/engine/core-modules/billing/services/billing-usage.service'; +import { BillingService } from 'src/engine/core-modules/billing/services/billing.service'; import { ExceptionHandlerService } from 'src/engine/core-modules/exception-handler/exception-handler.service'; 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'; @@ -56,6 +58,8 @@ export class WorkflowExecutorWorkspaceService { private readonly workflowActionFactory: WorkflowActionFactory, private readonly workspaceEventEmitter: WorkspaceEventEmitter, private readonly workflowRunWorkspaceService: WorkflowRunWorkspaceService, + private readonly billingService: BillingService, + private readonly billingUsageService: BillingUsageService, private readonly exceptionHandlerService: ExceptionHandlerService, private readonly metricsService: MetricsService, @InjectMessageQueue(MessageQueue.workflowQueue) @@ -178,7 +182,7 @@ export class WorkflowExecutorWorkspaceService { !actionOutput.shouldFailSafely && !actionOutput.shouldSkipStepExecution ) { - this.sendWorkflowNodeRunEvent(workspaceId, workflowRun.workflowId); + await this.sendWorkflowNodeRunEvent(workspaceId, workflowRun.workflowId); } const { shouldProcessNextSteps } = await this.processStepExecutionResult({ @@ -355,7 +359,10 @@ export class WorkflowExecutorWorkspaceService { }); } - private sendWorkflowNodeRunEvent(workspaceId: string, workflowId: string) { + private async sendWorkflowNodeRunEvent( + workspaceId: string, + workflowId: string, + ) { this.workspaceEventEmitter.emitCustomBatchEvent( USAGE_RECORDED, [ @@ -370,6 +377,13 @@ export class WorkflowExecutorWorkspaceService { ], workspaceId, ); + + if (this.billingService.isBillingEnabled()) { + await this.billingUsageService.decrementAvailableCredits({ + workspaceId, + usedCredits: 100, + }); + } } private async processStepExecutionResult({ diff --git a/packages/twenty-server/test/integration/audit/suites/clickHouse-usage-event-writer.integration-spec.ts b/packages/twenty-server/test/integration/audit/suites/clickHouse-usage-event-writer.integration-spec.ts index f10907c263f..4e5d36d70ce 100644 --- a/packages/twenty-server/test/integration/audit/suites/clickHouse-usage-event-writer.integration-spec.ts +++ b/packages/twenty-server/test/integration/audit/suites/clickHouse-usage-event-writer.integration-spec.ts @@ -6,7 +6,7 @@ import { createClient, } from '@clickhouse/client'; -import { formatDateForClickHouse } from 'src/database/clickHouse/clickHouse.util'; +import { formatDateTimeForClickHouse } from 'src/database/clickHouse/clickHouse.util'; import { UsageOperationType } from 'src/engine/core-modules/usage/enums/usage-operation-type.enum'; import { UsageResourceType } from 'src/engine/core-modules/usage/enums/usage-resource-type.enum'; import { UsageUnit } from 'src/engine/core-modules/usage/enums/usage-unit.enum'; @@ -24,7 +24,7 @@ const buildUsageEventRow = ( resourceContext: string; }> = {}, ) => ({ - timestamp: formatDateForClickHouse(new Date()), + timestamp: formatDateTimeForClickHouse(new Date()), workspaceId, userWorkspaceId: overrides.userWorkspaceId ?? '', resourceType: overrides.resourceType ?? UsageResourceType.AI,