From 85effd67508492bbc74a424b1429359eebb5bece Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?F=C3=A9lix=20Malfait?= Date: Mon, 27 Apr 2026 22:15:11 +0200 Subject: [PATCH] fix(billing): gate AI credit-cap at entry points instead of workflow executor Restores the credit-cap enforcement that was lifted in #19904 (which had correctly identified that gating per-workflow-step was too coarse), but moves it to the AI entry points where the actual cost lives. - agent-async-executor.service: gate executeAgent. - ai-generate-text.controller: gate the REST handler. - agent-title-generation.service: gate generateThreadTitle. - workflow-executor.workspace-service: replace #19904's TODO with an absolute-behavior comment documenting why the gate is not here. Chat resolver was already gated; no change there. Repair-tool-call util is a sub-call inside an already-gated flow and is left ungated. Net effect: a workspace that exhausts credits via chat or AI agent stops making AI calls, but its non-AI workflows (DB CRUD, branching, action steps) continue running normally. Addresses the root cause of the 2026-04-26 token-usage incident. Co-Authored-By: Claude Opus 4.7 (1M context) --- .../ai-agent-execution.module.ts | 2 ++ .../services/agent-async-executor.service.ts | 23 +++++++++++++++ .../ai/ai-chat/services/agent-chat.service.ts | 6 ++-- .../agent-title-generation.service.ts | 28 ++++++++++++++++++- .../ai-generate-text.module.ts | 8 +++++- .../ai-generate-text.controller.ts | 23 +++++++++++++++ .../workflow-executor.workspace-service.ts | 7 +++-- 7 files changed, 90 insertions(+), 7 deletions(-) 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 9a95d6629f7..05988b1a003 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 @@ -14,9 +14,16 @@ 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 { + 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 { 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'; +import { TwentyConfigService } from 'src/engine/core-modules/twenty-config/twenty-config.service'; import { WorkspaceEntity } from 'src/engine/core-modules/workspace/workspace.entity'; import { WORKFLOW_AGENT_REGISTRY_TOOL_CATEGORIES } from 'src/engine/metadata-modules/ai/ai-agent-execution/constants/workflow-agent-registry-tool-categories.const'; import { type AgentExecutionResult } from 'src/engine/metadata-modules/ai/ai-agent-execution/types/agent-execution-result.type'; @@ -49,6 +56,8 @@ export class AgentAsyncExecutorService { private readonly aiModelConfigService: AiModelConfigService, private readonly toolRegistry: ToolRegistryService, private readonly nativeToolBinder: NativeToolBinderService, + private readonly billingService: BillingService, + private readonly twentyConfigService: TwentyConfigService, @InjectRepository(RoleTargetEntity) private readonly roleTargetRepository: Repository, @InjectRepository(WorkspaceEntity) @@ -113,6 +122,20 @@ export class AgentAsyncExecutorService { rolePermissionConfig?: RolePermissionConfig; authContext?: WorkspaceAuthContext; }): Promise { + if (agent && this.twentyConfigService.get('IS_BILLING_ENABLED')) { + const canBill = await this.billingService.canBillMeteredProduct( + agent.workspaceId, + BillingProductKey.WORKFLOW_NODE_EXECUTION, + ); + + if (!canBill) { + throw new BillingException( + 'Credits exhausted', + BillingExceptionCode.BILLING_CREDITS_EXHAUSTED, + ); + } + } + try { if (agent) { const workspace = await this.workspaceRepository.findOneBy({ diff --git a/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/services/agent-chat.service.ts b/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/services/agent-chat.service.ts index 621970a50a4..50985a2aa43 100644 --- a/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/services/agent-chat.service.ts +++ b/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/services/agent-chat.service.ts @@ -328,8 +328,10 @@ export class AgentChatService { return null; } - const title = - await this.titleGenerationService.generateThreadTitle(messageContent); + const title = await this.titleGenerationService.generateThreadTitle( + messageContent, + workspaceId, + ); await this.threadRepository.update(threadId, { title }); 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 8fed08f1775..8512bc96edf 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 @@ -2,6 +2,13 @@ import { Injectable, Logger } from '@nestjs/common'; import { generateText } from 'ai'; +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 { TwentyConfigService } from 'src/engine/core-modules/twenty-config/twenty-config.service'; import { AI_TELEMETRY_CONFIG } from 'src/engine/metadata-modules/ai/ai-models/constants/ai-telemetry.const'; import { AiModelRegistryService } from 'src/engine/metadata-modules/ai/ai-models/services/ai-model-registry.service'; @@ -11,9 +18,28 @@ export class AgentTitleGenerationService { constructor( private readonly aiModelRegistryService: AiModelRegistryService, + private readonly billingService: BillingService, + private readonly twentyConfigService: TwentyConfigService, ) {} - async generateThreadTitle(messageContent: string): Promise { + async generateThreadTitle( + messageContent: string, + workspaceId: string, + ): Promise { + if (this.twentyConfigService.get('IS_BILLING_ENABLED')) { + const canBill = await this.billingService.canBillMeteredProduct( + workspaceId, + BillingProductKey.WORKFLOW_NODE_EXECUTION, + ); + + if (!canBill) { + throw new BillingException( + 'Credits exhausted', + BillingExceptionCode.BILLING_CREDITS_EXHAUSTED, + ); + } + } + try { const defaultModel = this.aiModelRegistryService.getDefaultSpeedModel(); 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 198625aa4dc..75d602719c9 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,13 +1,19 @@ 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 { PermissionsModule } from 'src/engine/metadata-modules/permissions/permissions.module'; import { WorkspaceCacheStorageModule } from 'src/engine/workspace-cache-storage/workspace-cache-storage.module'; import { AiGenerateTextController } from './controllers/ai-generate-text.controller'; @Module({ - imports: [TokenModule, WorkspaceCacheStorageModule, PermissionsModule], + imports: [ + TokenModule, + WorkspaceCacheStorageModule, + PermissionsModule, + BillingModule, + ], controllers: [AiGenerateTextController], }) export class AiGenerateTextModule {} 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 175de197604..d99cf5db085 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,13 @@ import { generateText } from 'ai'; import { PermissionFlagType } from 'twenty-shared/constants'; import { RestApiExceptionFilter } from 'src/engine/api/rest/rest-api-exception.filter'; +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 { TwentyConfigService } from 'src/engine/core-modules/twenty-config/twenty-config.service'; import type { WorkspaceEntity } from 'src/engine/core-modules/workspace/workspace.entity'; import { AuthWorkspace } from 'src/engine/decorators/auth/auth-workspace.decorator'; import { JwtAuthGuard } from 'src/engine/guards/jwt-auth.guard'; @@ -23,6 +30,8 @@ import { AiModelRegistryService } from 'src/engine/metadata-modules/ai/ai-models export class AiGenerateTextController { constructor( private readonly aiModelRegistryService: AiModelRegistryService, + private readonly billingService: BillingService, + private readonly twentyConfigService: TwentyConfigService, ) {} @Post('generate-text') @@ -38,6 +47,20 @@ export class AiGenerateTextController { ); } + if (this.twentyConfigService.get('IS_BILLING_ENABLED')) { + const canBill = await this.billingService.canBillMeteredProduct( + workspace.id, + BillingProductKey.WORKFLOW_NODE_EXECUTION, + ); + + if (!canBill) { + throw new BillingException( + 'Credits exhausted', + BillingExceptionCode.BILLING_CREDITS_EXHAUSTED, + ); + } + } + 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 133f793231c..175e76e8814 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 @@ -453,9 +453,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);