diff --git a/packages/twenty-server/src/engine/core-modules/billing/billing.exception.ts b/packages/twenty-server/src/engine/core-modules/billing/billing.exception.ts index d08c217c91d..479f97d4695 100644 --- a/packages/twenty-server/src/engine/core-modules/billing/billing.exception.ts +++ b/packages/twenty-server/src/engine/core-modules/billing/billing.exception.ts @@ -32,6 +32,7 @@ export enum BillingExceptionCode { BILLING_SUBSCRIPTION_PHASE_NOT_FOUND = 'BILLING_SUBSCRIPTION_PHASE_NOT_FOUND', BILLING_TOO_MUCH_SUBSCRIPTIONS_FOUND = 'BILLING_TOO_MUCH_SUBSCRIPTIONS_FOUND', BILLING_CREDITS_EXHAUSTED = 'BILLING_CREDITS_EXHAUSTED', + BILLING_USER_BUDGET_EXCEEDED = 'BILLING_USER_BUDGET_EXCEEDED', } const getBillingExceptionUserFriendlyMessage = (code: BillingExceptionCode) => { @@ -86,6 +87,8 @@ const getBillingExceptionUserFriendlyMessage = (code: BillingExceptionCode) => { return msg`Multiple subscriptions found where one was expected.`; case BillingExceptionCode.BILLING_CREDITS_EXHAUSTED: return msg`You have exhausted your credits. Please upgrade your plan to continue.`; + case BillingExceptionCode.BILLING_USER_BUDGET_EXCEEDED: + return msg`You have exceeded your AI chat budget. Please contact your workspace administrator.`; default: assertUnreachable(code); } diff --git a/packages/twenty-server/src/engine/core-modules/billing/services/billing-budget-guard.service.ts b/packages/twenty-server/src/engine/core-modules/billing/services/billing-budget-guard.service.ts new file mode 100644 index 00000000000..9ceb923403d --- /dev/null +++ b/packages/twenty-server/src/engine/core-modules/billing/services/billing-budget-guard.service.ts @@ -0,0 +1,133 @@ +/* @license Enterprise */ + +import { Injectable } from '@nestjs/common'; +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 { formatDateForClickHouse } from 'src/database/clickHouse/clickHouse.util'; +import { SubscriptionStatus } from 'src/engine/core-modules/billing/enums/billing-subscription-status.enum'; +import { BillingSubscriptionService } from 'src/engine/core-modules/billing/services/billing-subscription.service'; +import { buildUserUsageCacheKey } from 'src/engine/core-modules/billing/utils/build-user-usage-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 { UserWorkspaceEntity } from 'src/engine/core-modules/user-workspace/user-workspace.entity'; +import { WorkspaceEntity } from 'src/engine/core-modules/workspace/workspace.entity'; + +const USER_USAGE_CACHE_TTL_MS = 60_000; + +@Injectable() +export class BillingBudgetGuardService { + constructor( + private readonly clickHouseService: ClickHouseService, + private readonly twentyConfigService: TwentyConfigService, + @InjectCacheStorage(CacheStorageNamespace.EngineWorkspace) + private readonly cacheStorage: CacheStorageService, + @InjectRepository(UserWorkspaceEntity) + private readonly userWorkspaceRepository: Repository, + @InjectRepository(WorkspaceEntity) + private readonly workspaceRepository: Repository, + private readonly billingSubscriptionService: BillingSubscriptionService, + ) {} + + async canUserSpend( + workspaceId: string, + userWorkspaceId: string, + ): Promise { + if (!this.twentyConfigService.get('CLICKHOUSE_URL')) { + return true; + } + + const maxCredits = await this.resolveEffectiveMaxCredits( + workspaceId, + userWorkspaceId, + ); + + if (maxCredits === null) { + return true; + } + + const currentUsage = await this.getUserPeriodUsage( + workspaceId, + userWorkspaceId, + ); + + return currentUsage < maxCredits; + } + + private async getUserPeriodUsage( + workspaceId: string, + userWorkspaceId: string, + ): Promise { + const cacheKey = buildUserUsageCacheKey(workspaceId, userWorkspaceId); + const cached = await this.cacheStorage.get(cacheKey); + + if (cached !== undefined) { + return cached; + } + + const periodStart = await this.getCurrentPeriodStart(workspaceId); + + const result = await this.clickHouseService.select<{ total: string }>( + `SELECT SUM(creditsUsed) as total FROM billingEvent + WHERE workspaceId = {workspaceId:String} + AND userWorkspaceId = {userWorkspaceId:String} + AND timestamp >= {periodStart:DateTime64}`, + { + workspaceId, + userWorkspaceId, + periodStart: formatDateForClickHouse(periodStart), + }, + ); + + const usage = Number(result[0]?.total ?? 0); + + await this.cacheStorage.set(cacheKey, usage, USER_USAGE_CACHE_TTL_MS); + + return usage; + } + + private async resolveEffectiveMaxCredits( + workspaceId: string, + userWorkspaceId: string, + ): Promise { + const userWorkspace = await this.userWorkspaceRepository.findOne({ + where: { id: userWorkspaceId, workspaceId }, + select: ['maxAiChatCreditsPerPeriod'], + }); + + if (isDefined(userWorkspace?.maxAiChatCreditsPerPeriod)) { + return userWorkspace.maxAiChatCreditsPerPeriod; + } + + const workspace = await this.workspaceRepository.findOne({ + where: { id: workspaceId }, + select: ['defaultUserAiChatMaxCreditsPerPeriod'], + }); + + return workspace?.defaultUserAiChatMaxCreditsPerPeriod ?? null; + } + + private async getCurrentPeriodStart(workspaceId: string): Promise { + const subscription = + await this.billingSubscriptionService.getCurrentBillingSubscription({ + workspaceId, + }); + + if (!subscription) { + return new Date(0); + } + + const isTrialing = + subscription.status === SubscriptionStatus.Trialing && + isDefined(subscription.trialStart); + + return isTrialing + ? subscription.trialStart! + : subscription.currentPeriodStart; + } +} diff --git a/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/controllers/agent-chat.controller.ts b/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/controllers/agent-chat.controller.ts index 610cc7b4925..b33c6271e29 100644 --- a/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/controllers/agent-chat.controller.ts +++ b/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/controllers/agent-chat.controller.ts @@ -19,6 +19,7 @@ import { } from 'src/engine/core-modules/billing/billing.exception'; import { BillingProductKey } from 'src/engine/core-modules/billing/enums/billing-product-key.enum'; import { BillingRestApiExceptionFilter } from 'src/engine/core-modules/billing/filters/billing-api-exception.filter'; +import { BillingBudgetGuardService } from 'src/engine/core-modules/billing/services/billing-budget-guard.service'; import { BillingService } from 'src/engine/core-modules/billing/services/billing.service'; import { TwentyConfigService } from 'src/engine/core-modules/twenty-config/twenty-config.service'; import type { WorkspaceEntity } from 'src/engine/core-modules/workspace/workspace.entity'; @@ -47,6 +48,7 @@ export class AgentChatController { constructor( private readonly agentStreamingService: AgentChatStreamingService, private readonly billingService: BillingService, + private readonly billingBudgetGuardService: BillingBudgetGuardService, private readonly twentyConfigService: TwentyConfigService, private readonly aiModelRegistryService: AiModelRegistryService, ) {} @@ -90,6 +92,18 @@ export class AgentChatController { BillingExceptionCode.BILLING_CREDITS_EXHAUSTED, ); } + + const canUserSpend = await this.billingBudgetGuardService.canUserSpend( + workspace.id, + userWorkspaceId, + ); + + if (!canUserSpend) { + throw new BillingException( + 'User AI chat budget exceeded', + BillingExceptionCode.BILLING_USER_BUDGET_EXCEEDED, + ); + } } this.agentStreamingService.streamAgentChat({