From 842e679cc67ed92ccc6db85d2f5aa1ef0585aeee Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?F=C3=A9lix=20Malfait?= Date: Wed, 29 Apr 2026 18:48:06 +0200 Subject: [PATCH] fix(billing): gate AI credit-cap at entry points instead of workflow executor (#20096) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Background The 2026-04-26 incident saw 716M Sonnet 4.6 tokens consumed in a single trial workspace. Two causes: failed agent executions weren't billed (addressed by #20065) and the credit-cap gate had been removed from `WorkflowExecutorWorkspaceService.executeStep` in #19904, leaving no enforcement point at all. ## Why not just revert #19904 #19904 was right that gating at the workflow executor is too coarse. When one user exhausted a workspace's credits via chat, *all* workflows hard-failed mid-run — including cheap DB/CRUD/branch automations costing essentially nothing. Reverting would re-introduce that cliff. ## New design: gate at the AI entry points The chat resolver already gates this way (`agent-chat.resolver.ts:137-148`). This PR replicates the same pattern at every other point where the workspace can incur real AI cost: - `executeAgent` in `agent-async-executor.service.ts` - the REST handler in `ai-generate-text.controller.ts` - `generateThreadTitle` in `agent-title-generation.service.ts` In each, after auth/validation: skip if `IS_BILLING_ENABLED` is false; otherwise call `BillingService.canBillMeteredProduct(workspaceId, BillingProductKey.WORKFLOW_NODE_EXECUTION)`; on `false`, throw `BillingException(BILLING_CREDITS_EXHAUSTED)`. No new method, no new exception code, no new product key. This matches industry convention (Lovable/Replit also gate at the expensive-operation boundary, not at every cheap step). ## Deliberately not gated - `WorkflowExecutorWorkspaceService.executeStep` — the design choice is now intentional, so the #19904 TODO is replaced by a one-line absolute-behavior comment explaining why the gate isn't here. Cheap workflow steps (DB CRUD, branching, action steps) are not gated, so a chat-driven cap exhaustion does not block non-AI automations. - `repair-tool-call.util` — repair is a sub-call inside an already-gated AI flow. If the parent is gated, repair will naturally not run. Adding a gate here adds complexity without value. ## Net effect A workspace that exhausts credits via chat or AI agent stops making AI calls. Its non-AI workflows continue running normally. A workflow with both AI and non-AI steps fails at the AI step with `BILLING_CREDITS_EXHAUSTED`, but downstream non-AI steps that don't depend on the AI output still run. ## Conflicts This PR overlaps with three other in-flight PRs in the same files. None of them touch the gate logic; rebasing on top of any of them is trivial: - #20065 (agent-async-executor): adds `workspaceId` to `executeAgent` args and bills in `finally`. The gate at the top of `executeAgent` from this PR sits naturally above that. - #20066 (REST controller): adds usage billing to the controller. - #20067 (title gen): adds usage billing to title generation and tool-call repair. Recommend landing #20065/#20066/#20067 first; this PR rebases trivially on top. ## Tests Out of scope per the PR series convention. The existing chat-resolver gate isn't unit-tested either; this PR follows the same precedent. Follow-up: add integration coverage that exercises a workspace at `hasReachedCurrentPeriodCap=true` against each of the three new gates plus the pre-existing chat-resolver gate. ## Future follow-ups - Per-user soft cap inside a workspace (the Lovable Business-tier pattern), so one user can't exhaust the workspace's cap. - Pre-flight cost estimate so the user sees an "approaching cap" warning before the hard stop. - Rename `BillingProductKey.WORKFLOW_NODE_EXECUTION` — the name predates this design choice and is misleading now that it gates AI entry points rather than workflow nodes. ## Test plan - [ ] Trigger a workspace into `hasReachedCurrentPeriodCap=true`. - [ ] Send a chat message — expect failure with `BILLING_CREDITS_EXHAUSTED`. - [ ] Run a workflow whose only AI step is an `ai-agent` action — expect that step to fail with `BILLING_CREDITS_EXHAUSTED`, downstream non-AI steps still run. - [ ] POST to `/rest/ai/generate-text` — expect `BILLING_CREDITS_EXHAUSTED`. - [ ] Create a new chat thread (which kicks off `generateThreadTitle`) — expect `BILLING_CREDITS_EXHAUSTED`. - [ ] Run a workflow with no AI step (only DB CRUD/branching/actions) — expect it to run unaffected. --------- Co-authored-by: Claude Opus 4.7 (1M context) --- .../billing/services/billing-usage.service.ts | 15 +++++++++++++++ .../ai-agent-execution.module.ts | 2 ++ .../services/agent-async-executor.service.ts | 4 ++++ .../ai-chat/resolvers/agent-chat.resolver.ts | 19 +------------------ .../agent-title-generation.service.ts | 4 ++++ .../ai-generate-text.module.ts | 2 ++ .../ai-generate-text.controller.ts | 4 ++++ .../workflow-executor.workspace-service.ts | 7 ++++--- 8 files changed, 36 insertions(+), 21 deletions(-) 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 1f4ac7bf09e..43143ff1a5b 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 @@ -341,6 +341,10 @@ export class BillingUsageService { } async hasAvailableCredits(workspaceId: string): Promise { + if (!this.twentyConfigService.get('IS_BILLING_ENABLED')) { + return true; + } + const { billingSubscription: subscription } = await this.workspaceCacheService.getOrRecompute(workspaceId, [ 'billingSubscription', @@ -369,6 +373,17 @@ export class BillingUsageService { return availableCredits > 0; } + async hasAvailableCreditsOrThrow(workspaceId: string): Promise { + const hasCredits = await this.hasAvailableCredits(workspaceId); + + if (!hasCredits) { + throw new BillingException( + 'Credits exhausted', + BillingExceptionCode.BILLING_CREDITS_EXHAUSTED, + ); + } + } + async getCurrentPeriodCreditsUsed( workspaceId: string, periodStart: Date, diff --git a/packages/twenty-server/src/engine/metadata-modules/ai/ai-agent-execution/ai-agent-execution.module.ts b/packages/twenty-server/src/engine/metadata-modules/ai/ai-agent-execution/ai-agent-execution.module.ts index 93ec1d7f84f..f87e25152f0 100644 --- a/packages/twenty-server/src/engine/metadata-modules/ai/ai-agent-execution/ai-agent-execution.module.ts +++ b/packages/twenty-server/src/engine/metadata-modules/ai/ai-agent-execution/ai-agent-execution.module.ts @@ -1,6 +1,7 @@ import { forwardRef, Module } from '@nestjs/common'; import { TypeOrmModule } from '@nestjs/typeorm'; +import { BillingModule } from 'src/engine/core-modules/billing/billing.module'; import { WorkspaceDomainsModule } from 'src/engine/core-modules/domain/workspace-domains/workspace-domains.module'; import { FileUrlModule } from 'src/engine/core-modules/file/file-url/file-url.module'; import { ToolProviderModule } from 'src/engine/core-modules/tool-provider/tool-provider.module'; @@ -27,6 +28,7 @@ import { AgentAsyncExecutorService } from './services/agent-async-executor.servi AiBillingModule, AiModelsModule, AiAgentModule, + BillingModule, FileUrlModule, WorkspaceDomainsModule, UserWorkspaceModule, diff --git a/packages/twenty-server/src/engine/metadata-modules/ai/ai-agent-execution/services/agent-async-executor.service.ts b/packages/twenty-server/src/engine/metadata-modules/ai/ai-agent-execution/services/agent-async-executor.service.ts index 24e63889570..35cbb9ef739 100644 --- a/packages/twenty-server/src/engine/metadata-modules/ai/ai-agent-execution/services/agent-async-executor.service.ts +++ b/packages/twenty-server/src/engine/metadata-modules/ai/ai-agent-execution/services/agent-async-executor.service.ts @@ -16,6 +16,7 @@ import { type Repository } from 'typeorm'; import { isUserAuthContext } from 'src/engine/core-modules/auth/guards/is-user-auth-context.guard'; import { type WorkspaceAuthContext } from 'src/engine/core-modules/auth/types/workspace-auth-context.type'; +import { BillingUsageService } from 'src/engine/core-modules/billing/services/billing-usage.service'; import { type ToolProviderContext } from 'src/engine/core-modules/tool-provider/interfaces/tool-provider-context.type'; import { NativeToolBinderService } from 'src/engine/core-modules/tool-provider/native/native-tool-binder.service'; import { ToolRegistryService } from 'src/engine/core-modules/tool-provider/services/tool-registry.service'; @@ -69,6 +70,7 @@ export class AgentAsyncExecutorService { private readonly toolRegistry: ToolRegistryService, private readonly nativeToolBinder: NativeToolBinderService, private readonly aiBillingService: AiBillingService, + private readonly billingUsageService: BillingUsageService, @InjectRepository(RoleTargetEntity) private readonly roleTargetRepository: Repository, @InjectRepository(WorkspaceEntity) @@ -139,6 +141,8 @@ export class AgentAsyncExecutorService { userWorkspaceId?: string | null; operationType?: UsageOperationType; }): Promise { + await this.billingUsageService.hasAvailableCreditsOrThrow(workspaceId); + let accumulatedUsage: LanguageModelUsage = EMPTY_USAGE; let cacheCreationTokens = 0; let nativeWebSearchCallCount = 0; 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 b7b1237a0e1..6547664b817 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 @@ -16,13 +16,8 @@ import { Repository } from 'typeorm'; 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'; -import { - BillingException, - BillingExceptionCode, -} from 'src/engine/core-modules/billing/billing.exception'; 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'; import { type WorkspaceEntity } from 'src/engine/core-modules/workspace/workspace.entity'; import { AuthUserWorkspaceId } from 'src/engine/decorators/auth/auth-user-workspace-id.decorator'; @@ -58,7 +53,6 @@ export class AgentChatResolver { private readonly eventPublisherService: AgentChatEventPublisherService, private readonly systemPromptBuilderService: SystemPromptBuilderService, private readonly billingUsageService: BillingUsageService, - private readonly twentyConfigService: TwentyConfigService, private readonly aiModelRegistryService: AiModelRegistryService, private readonly redisClientService: RedisClientService, @InjectRepository(AgentChatThreadEntity) @@ -133,18 +127,7 @@ export class AgentChatResolver { workspace, ); - if (this.twentyConfigService.get('IS_BILLING_ENABLED')) { - const canBill = await this.billingUsageService.hasAvailableCredits( - workspace.id, - ); - - if (!canBill) { - throw new BillingException( - 'Credits exhausted', - BillingExceptionCode.BILLING_CREDITS_EXHAUSTED, - ); - } - } + await this.billingUsageService.hasAvailableCreditsOrThrow(workspace.id); const thread = await this.threadRepository.findOne({ where: { id: threadId, userWorkspaceId }, diff --git a/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/services/agent-title-generation.service.ts b/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/services/agent-title-generation.service.ts index 736a0c82f81..33fb6faee79 100644 --- a/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/services/agent-title-generation.service.ts +++ b/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/services/agent-title-generation.service.ts @@ -7,6 +7,7 @@ import { generateText, } from 'ai'; +import { BillingUsageService } from 'src/engine/core-modules/billing/services/billing-usage.service'; import { UsageOperationType } from 'src/engine/core-modules/usage/enums/usage-operation-type.enum'; import { AiBillingService } from 'src/engine/metadata-modules/ai/ai-billing/services/ai-billing.service'; import { extractCacheCreationTokensFromSteps } from 'src/engine/metadata-modules/ai/ai-billing/utils/extract-cache-creation-tokens.util'; @@ -20,6 +21,7 @@ export class AgentTitleGenerationService { constructor( private readonly aiModelRegistryService: AiModelRegistryService, private readonly aiBillingService: AiBillingService, + private readonly billingUsageService: BillingUsageService, ) {} async generateThreadTitle( @@ -27,6 +29,8 @@ export class AgentTitleGenerationService { workspaceId: string, userWorkspaceId: string | null, ): Promise { + await this.billingUsageService.hasAvailableCreditsOrThrow(workspaceId); + const defaultModel = this.aiModelRegistryService.getDefaultSpeedModel(); if (!defaultModel) { diff --git a/packages/twenty-server/src/engine/metadata-modules/ai/ai-generate-text/ai-generate-text.module.ts b/packages/twenty-server/src/engine/metadata-modules/ai/ai-generate-text/ai-generate-text.module.ts index e1be39dd361..963117dccfe 100644 --- a/packages/twenty-server/src/engine/metadata-modules/ai/ai-generate-text/ai-generate-text.module.ts +++ b/packages/twenty-server/src/engine/metadata-modules/ai/ai-generate-text/ai-generate-text.module.ts @@ -1,6 +1,7 @@ import { Module } from '@nestjs/common'; import { TokenModule } from 'src/engine/core-modules/auth/token/token.module'; +import { BillingModule } from 'src/engine/core-modules/billing/billing.module'; import { AiBillingModule } from 'src/engine/metadata-modules/ai/ai-billing/ai-billing.module'; import { PermissionsModule } from 'src/engine/metadata-modules/permissions/permissions.module'; import { WorkspaceCacheStorageModule } from 'src/engine/workspace-cache-storage/workspace-cache-storage.module'; @@ -12,6 +13,7 @@ import { AiGenerateTextController } from './controllers/ai-generate-text.control TokenModule, WorkspaceCacheStorageModule, PermissionsModule, + BillingModule, AiBillingModule, ], controllers: [AiGenerateTextController], diff --git a/packages/twenty-server/src/engine/metadata-modules/ai/ai-generate-text/controllers/ai-generate-text.controller.ts b/packages/twenty-server/src/engine/metadata-modules/ai/ai-generate-text/controllers/ai-generate-text.controller.ts index bbc67bb5782..fe24899dc47 100644 --- a/packages/twenty-server/src/engine/metadata-modules/ai/ai-generate-text/controllers/ai-generate-text.controller.ts +++ b/packages/twenty-server/src/engine/metadata-modules/ai/ai-generate-text/controllers/ai-generate-text.controller.ts @@ -4,6 +4,7 @@ import { generateText } from 'ai'; import { PermissionFlagType } from 'twenty-shared/constants'; import { RestApiExceptionFilter } from 'src/engine/api/rest/rest-api-exception.filter'; +import { BillingUsageService } from 'src/engine/core-modules/billing/services/billing-usage.service'; import { UsageOperationType } from 'src/engine/core-modules/usage/enums/usage-operation-type.enum'; import type { WorkspaceEntity } from 'src/engine/core-modules/workspace/workspace.entity'; import { AuthUserWorkspaceId } from 'src/engine/decorators/auth/auth-user-workspace-id.decorator'; @@ -27,6 +28,7 @@ export class AiGenerateTextController { constructor( private readonly aiModelRegistryService: AiModelRegistryService, private readonly aiBillingService: AiBillingService, + private readonly billingUsageService: BillingUsageService, ) {} @Post('generate-text') @@ -43,6 +45,8 @@ export class AiGenerateTextController { ); } + await this.billingUsageService.hasAvailableCreditsOrThrow(workspace.id); + const resolvedModelId = body.modelId ?? workspace.fastModel; this.aiModelRegistryService.validateModelAvailability( 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 9e03b8a23fd..15a3a6db12e 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 @@ -467,9 +467,10 @@ export class WorkflowExecutorWorkspaceService { workflowRunId: string; workspaceId: string; }) { - // TODO: re-enable workflow node execution credit cap once billing limits are revisited. - // Previously gated on BillingService.canBillMeteredProduct(WORKFLOW_NODE_EXECUTION); - // temporarily disabled so workflows keep running when the period cap is reached. + // Credit-cap enforcement lives at the AI entry points (chat resolver, + // executeAgent, generate-text controller, title generation). Cheap + // workflow steps (DB CRUD, branching, actions) are not gated here so a + // chat-driven cap exhaustion does not block non-AI automations. const stepId = step.id; const workflowAction = this.workflowActionFactory.get(step.type);