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 0f1ad0a937..24e6388957 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 @@ -4,10 +4,12 @@ import { InjectRepository } from '@nestjs/typeorm'; import { generateText, jsonSchema, + type LanguageModelUsage, Output, stepCountIs, type ToolSet, } from 'ai'; +import { AUTO_SELECT_SMART_MODEL_ID } from 'twenty-shared/constants'; import { type ActorMetadata } from 'twenty-shared/types'; import { isDefined } from 'twenty-shared/utils'; import { type Repository } from 'typeorm'; @@ -17,9 +19,11 @@ import { type WorkspaceAuthContext } from 'src/engine/core-modules/auth/types/wo 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 { UsageOperationType } from 'src/engine/core-modules/usage/enums/usage-operation-type.enum'; 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'; +import { AiBillingService } from 'src/engine/metadata-modules/ai/ai-billing/services/ai-billing.service'; import { countNativeWebSearchCallsFromSteps } from 'src/engine/metadata-modules/ai/ai-billing/utils/count-native-web-search-calls-from-steps.util'; import { extractCacheCreationTokensFromSteps } from 'src/engine/metadata-modules/ai/ai-billing/utils/extract-cache-creation-tokens.util'; import { mergeLanguageModelUsage } from 'src/engine/metadata-modules/ai/ai-billing/utils/merge-language-model-usage.util'; @@ -37,6 +41,21 @@ import { AiModelRegistryService } from 'src/engine/metadata-modules/ai/ai-models import { RoleTargetEntity } from 'src/engine/metadata-modules/role-target/role-target.entity'; import { type RolePermissionConfig } from 'src/engine/twenty-orm/types/role-permission-config'; +const EMPTY_USAGE: LanguageModelUsage = { + inputTokens: 0, + outputTokens: 0, + totalTokens: 0, + inputTokenDetails: { + noCacheTokens: 0, + cacheReadTokens: 0, + cacheWriteTokens: 0, + }, + outputTokenDetails: { + textTokens: 0, + reasoningTokens: 0, + }, +}; + // Agent execution within workflows uses registry tools plus native model tools. // Workflow registry tools are intentionally excluded to avoid circular // dependencies and recursive workflow execution. @@ -49,6 +68,7 @@ export class AgentAsyncExecutorService { private readonly aiModelConfigService: AiModelConfigService, private readonly toolRegistry: ToolRegistryService, private readonly nativeToolBinder: NativeToolBinderService, + private readonly aiBillingService: AiBillingService, @InjectRepository(RoleTargetEntity) private readonly roleTargetRepository: Repository, @InjectRepository(WorkspaceEntity) @@ -106,13 +126,23 @@ export class AgentAsyncExecutorService { actorContext, rolePermissionConfig, authContext, + workspaceId, + userWorkspaceId, + operationType = UsageOperationType.AI_WORKFLOW_TOKEN, }: { agent: AgentEntity | null; userPrompt: string; actorContext?: ActorMetadata; rolePermissionConfig?: RolePermissionConfig; authContext?: WorkspaceAuthContext; + workspaceId: string; + userWorkspaceId?: string | null; + operationType?: UsageOperationType; }): Promise { + let accumulatedUsage: LanguageModelUsage = EMPTY_USAGE; + let cacheCreationTokens = 0; + let nativeWebSearchCallCount = 0; + try { if (agent) { const workspace = await this.workspaceRepository.findOneBy({ @@ -212,11 +242,11 @@ export class AgentAsyncExecutorService { }, }); - const cacheCreationTokens = extractCacheCreationTokensFromSteps( + accumulatedUsage = textResponse.usage; + cacheCreationTokens = extractCacheCreationTokensFromSteps( textResponse.steps, ); - - const nativeWebSearchCallCount = countNativeWebSearchCallsFromSteps( + nativeWebSearchCallCount = countNativeWebSearchCallsFromSteps( textResponse.steps, ); @@ -246,6 +276,11 @@ export class AgentAsyncExecutorService { experimental_telemetry: AI_TELEMETRY_CONFIG, }); + accumulatedUsage = mergeLanguageModelUsage( + textResponse.usage, + structuredResult.usage, + ); + if (structuredResult.output == null) { throw new AiException( 'Failed to generate structured output from execution results', @@ -255,10 +290,7 @@ export class AgentAsyncExecutorService { return { result: structuredResult.output as object, - usage: mergeLanguageModelUsage( - textResponse.usage, - structuredResult.usage, - ), + usage: accumulatedUsage, cacheCreationTokens, nativeWebSearchCallCount, }; @@ -270,6 +302,21 @@ export class AgentAsyncExecutorService { error instanceof Error ? error.message : 'Agent execution failed', AiExceptionCode.AGENT_EXECUTION_FAILED, ); + } finally { + this.aiBillingService.calculateAndBillUsage( + agent?.modelId ?? AUTO_SELECT_SMART_MODEL_ID, + { usage: accumulatedUsage, cacheCreationTokens }, + workspaceId, + operationType, + agent?.id ?? null, + userWorkspaceId, + ); + + this.aiBillingService.billNativeWebSearchUsage( + nativeWebSearchCallCount, + workspaceId, + userWorkspaceId, + ); } } } diff --git a/packages/twenty-server/src/engine/metadata-modules/ai/ai-agent-monitor/jobs/run-evaluation-input.job.ts b/packages/twenty-server/src/engine/metadata-modules/ai/ai-agent-monitor/jobs/run-evaluation-input.job.ts index 75ba76c24c..424723046a 100644 --- a/packages/twenty-server/src/engine/metadata-modules/ai/ai-agent-monitor/jobs/run-evaluation-input.job.ts +++ b/packages/twenty-server/src/engine/metadata-modules/ai/ai-agent-monitor/jobs/run-evaluation-input.job.ts @@ -58,6 +58,8 @@ export class RunEvaluationInputJob { const executionResult = await this.aiAgentExecutorService.executeAgent({ agent, userPrompt: data.input, + workspaceId: data.workspaceId, + userWorkspaceId: null, }); await this.agentChatService.addMessage({ diff --git a/packages/twenty-server/src/modules/workflow/workflow-executor/workflow-actions/ai-agent/ai-agent-action.module.ts b/packages/twenty-server/src/modules/workflow/workflow-executor/workflow-actions/ai-agent/ai-agent-action.module.ts index 27c0974de5..a1380c921f 100644 --- a/packages/twenty-server/src/modules/workflow/workflow-executor/workflow-actions/ai-agent/ai-agent-action.module.ts +++ b/packages/twenty-server/src/modules/workflow/workflow-executor/workflow-actions/ai-agent/ai-agent-action.module.ts @@ -5,7 +5,6 @@ import { ApplicationModule } from 'src/engine/core-modules/application/applicati import { UserWorkspaceModule } from 'src/engine/core-modules/user-workspace/user-workspace.module'; import { AiAgentExecutionModule } from 'src/engine/metadata-modules/ai/ai-agent-execution/ai-agent-execution.module'; import { AgentEntity } from 'src/engine/metadata-modules/ai/ai-agent/entities/agent.entity'; -import { AiBillingModule } from 'src/engine/metadata-modules/ai/ai-billing/ai-billing.module'; import { RoleModule } from 'src/engine/metadata-modules/role/role.module'; import { UserRoleModule } from 'src/engine/metadata-modules/user-role/user-role.module'; import { WorkflowExecutionContextService } from 'src/modules/workflow/workflow-executor/services/workflow-execution-context.service'; @@ -17,7 +16,6 @@ import { AiAgentWorkflowAction } from './ai-agent.workflow-action'; imports: [ ApplicationModule, AiAgentExecutionModule, - AiBillingModule, TypeOrmModule.forFeature([AgentEntity]), WorkflowRunModule, UserWorkspaceModule, diff --git a/packages/twenty-server/src/modules/workflow/workflow-executor/workflow-actions/ai-agent/ai-agent.workflow-action.ts b/packages/twenty-server/src/modules/workflow/workflow-executor/workflow-actions/ai-agent/ai-agent.workflow-action.ts index b16453fc95..11cec77b20 100644 --- a/packages/twenty-server/src/modules/workflow/workflow-executor/workflow-actions/ai-agent/ai-agent.workflow-action.ts +++ b/packages/twenty-server/src/modules/workflow/workflow-executor/workflow-actions/ai-agent/ai-agent.workflow-action.ts @@ -6,11 +6,9 @@ import { type Repository } from 'typeorm'; import { type WorkflowAction } from 'src/modules/workflow/workflow-executor/interfaces/workflow-action.interface'; +import { UsageOperationType } from 'src/engine/core-modules/usage/enums/usage-operation-type.enum'; import { AgentAsyncExecutorService } from 'src/engine/metadata-modules/ai/ai-agent-execution/services/agent-async-executor.service'; import { AgentEntity } from 'src/engine/metadata-modules/ai/ai-agent/entities/agent.entity'; -import { AiBillingService } from 'src/engine/metadata-modules/ai/ai-billing/services/ai-billing.service'; -import { UsageOperationType } from 'src/engine/core-modules/usage/enums/usage-operation-type.enum'; -import { AUTO_SELECT_SMART_MODEL_ID } from 'twenty-shared/constants'; import { WorkflowStepExecutorException, WorkflowStepExecutorExceptionCode, @@ -26,7 +24,6 @@ import { isWorkflowAiAgentAction } from './guards/is-workflow-ai-agent-action.gu export class AiAgentWorkflowAction implements WorkflowAction { constructor( private readonly aiAgentExecutionService: AgentAsyncExecutorService, - private readonly aiBillingService: AiBillingService, private readonly workflowExecutionContextService: WorkflowExecutionContextService, @InjectRepository(AgentEntity) private readonly agentRepository: Repository, @@ -79,33 +76,18 @@ export class AiAgentWorkflowAction implements WorkflowAction { ? executionContext.authContext.userWorkspaceId : null; - const { result, usage, cacheCreationTokens, nativeWebSearchCallCount } = - await this.aiAgentExecutionService.executeAgent({ - agent, - userPrompt: resolveInput(prompt, context) as string, - actorContext: executionContext.isActingOnBehalfOfUser - ? executionContext.initiator - : undefined, - rolePermissionConfig: executionContext.rolePermissionConfig, - authContext: executionContext.authContext, - }); - - await this.aiBillingService.calculateAndBillUsage( - agent?.modelId ?? AUTO_SELECT_SMART_MODEL_ID, - { usage, cacheCreationTokens }, - workspaceId, - UsageOperationType.AI_WORKFLOW_TOKEN, - agent?.id || null, - userWorkspaceId, - ); - - // billNativeWebSearchUsage short-circuits when count <= 0, so calling - // unconditionally is safe regardless of whether native search fired. - this.aiBillingService.billNativeWebSearchUsage( - nativeWebSearchCallCount, + const { result } = await this.aiAgentExecutionService.executeAgent({ + agent, + userPrompt: resolveInput(prompt, context) as string, + actorContext: executionContext.isActingOnBehalfOfUser + ? executionContext.initiator + : undefined, + rolePermissionConfig: executionContext.rolePermissionConfig, + authContext: executionContext.authContext, workspaceId, userWorkspaceId, - ); + operationType: UsageOperationType.AI_WORKFLOW_TOKEN, + }); return { result,