feat(ai-chat) - Stop ai thinking if credits exhausted (#20526)

Billing is now decremented per-step, not per-turn. The onStepFinish callback in chat-execution.service.ts calls a new decrementAndCheckAvailableCredits method on each model step, so Redis is debited incrementally as the agent runs rather than all at once at the end. Credit exhaustion stops the stream mid-run. When a step depletes the remaining credits, a hasNoMoreAvailableCredits flag is set and passed into the stopWhen predicate of streamText, causing the agent to halt before starting the next step. A new credits-exhausted event is introduced. After the stream drains and the response is persisted, if credits ran out the job publishes a dedicated credits-exhausted event to the frontend instead of the normal message-persisted event. The frontend handles this new event. useAgentChatSubscription has a new credits-exhausted case that sets a BILLING_CREDITS_EXHAUSTED-coded error on the atom, closes the writer, and stops the streaming state — triggering the existing AiChatCreditsExhaustedMessage UI.
2026-05-13 17:09:08 +02:00
parent 068fc9363d
commit aecfe699f4
14 changed files with 212 additions and 52 deletions
@@ -24,21 +24,25 @@ import { UsageOperationType } from 'src/engine/core-modules/usage/enums/usage-op
 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';
-import {
-  AiException,
-  AiExceptionCode,
-} from 'src/engine/metadata-modules/ai/ai.exception';
 import { AGENT_CONFIG } from 'src/engine/metadata-modules/ai/ai-agent/constants/agent-config.const';
 import { WORKFLOW_SYSTEM_PROMPTS } from 'src/engine/metadata-modules/ai/ai-agent/constants/agent-system-prompts.const';
 import { type AgentEntity } from 'src/engine/metadata-modules/ai/ai-agent/entities/agent.entity';
 import { repairToolCall } from 'src/engine/metadata-modules/ai/ai-agent/utils/repair-tool-call.util';
+import { AiBillingService } from 'src/engine/metadata-modules/ai/ai-billing/services/ai-billing.service';
+import { convertDollarsToBillingCredits } from 'src/engine/metadata-modules/ai/ai-billing/utils/convert-dollars-to-billing-credits.util';
+import { countNativeWebSearchCallsFromSteps } from 'src/engine/metadata-modules/ai/ai-billing/utils/count-native-web-search-calls-from-steps.util';
+import {
+  extractCacheCreationTokens,
+  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';
 import { AI_TELEMETRY_CONFIG } from 'src/engine/metadata-modules/ai/ai-models/constants/ai-telemetry.const';
 import { AiModelConfigService } from 'src/engine/metadata-modules/ai/ai-models/services/ai-model-config.service';
 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 { RoleTargetEntity } from 'src/engine/metadata-modules/role-target/role-target.entity';
 import { type RolePermissionConfig } from 'src/engine/twenty-orm/types/role-permission-config';

@@ -222,14 +226,35 @@ export class AgentAsyncExecutorService {

      this.logger.log(`Generated ${Object.keys(tools).length} tools for agent`);

+      let hasNoMoreAvailableCredits = false;
+
      const textResponse = await generateText({
        system: `${WORKFLOW_SYSTEM_PROMPTS.BASE}\n\n${agent ? agent.prompt : ''}`,
        tools,
        model: registeredModel.model,
        prompt: userPrompt,
-        stopWhen: stepCountIs(AGENT_CONFIG.MAX_STEPS),
+        stopWhen: (step) =>
+          stepCountIs(AGENT_CONFIG.MAX_STEPS)(step) ||
+          hasNoMoreAvailableCredits,
        providerOptions,
        experimental_telemetry: AI_TELEMETRY_CONFIG,
+        onStepFinish: async (step) => {
+          const { hasNoMoreAvailableCredits: stepHasNoMoreAvailableCredits } =
+            await this.aiBillingService.decrementAndCheckAvailableCredits(
+              registeredModel.modelId,
+              {
+                usage: step.usage,
+                cacheCreationTokens: extractCacheCreationTokens(
+                  step.providerMetadata,
+                ),
+              },
+              workspaceId,
+            );
+
+          if (stepHasNoMoreAvailableCredits) {
+            hasNoMoreAvailableCredits = true;
+          }
+        },
        experimental_repairToolCall: async ({
          toolCall,
          tools: toolsForRepair,
@@ -265,6 +290,7 @@ export class AgentAsyncExecutorService {
          usage: textResponse.usage,
          cacheCreationTokens,
          nativeWebSearchCallCount,
+          hasNoMoreAvailableCredits,
        };
      }

@@ -278,6 +304,23 @@ export class AgentAsyncExecutorService {
                 Please generate the structured output based on the execution results and context above.`,
        output: Output.object({ schema: jsonSchema(agentSchema) }),
        experimental_telemetry: AI_TELEMETRY_CONFIG,
+        onStepFinish: async (step) => {
+          const { hasNoMoreAvailableCredits: stepHasNoMoreAvailableCredits } =
+            await this.aiBillingService.decrementAndCheckAvailableCredits(
+              registeredModel.modelId,
+              {
+                usage: step.usage,
+                cacheCreationTokens: extractCacheCreationTokens(
+                  step.providerMetadata,
+                ),
+              },
+              workspaceId,
+            );
+
+          if (stepHasNoMoreAvailableCredits) {
+            hasNoMoreAvailableCredits = true;
+          }
+        },
      });

      accumulatedUsage = mergeLanguageModelUsage(
@@ -297,6 +340,7 @@ export class AgentAsyncExecutorService {
        usage: accumulatedUsage,
        cacheCreationTokens,
        nativeWebSearchCallCount,
+        hasNoMoreAvailableCredits,
      };
    } catch (error) {
      if (error instanceof AiException) {
@@ -307,10 +351,24 @@ export class AgentAsyncExecutorService {
        AiExceptionCode.AGENT_EXECUTION_FAILED,
      );
    } finally {
-      void this.aiBillingService.calculateAndBillUsage(
-        agent?.modelId ?? AUTO_SELECT_SMART_MODEL_ID,
-        { usage: accumulatedUsage, cacheCreationTokens },
+      const modelId = agent?.modelId ?? AUTO_SELECT_SMART_MODEL_ID;
+      const costInDollars = this.aiBillingService.calculateCost(modelId, {
+        usage: accumulatedUsage,
+        cacheCreationTokens,
+      });
+      const creditsUsedMicro = Math.round(
+        convertDollarsToBillingCredits(costInDollars),
+      );
+      const totalTokens =
+        (accumulatedUsage.inputTokens ?? 0) +
+        (accumulatedUsage.outputTokens ?? 0) +
+        cacheCreationTokens;
+
+      void this.aiBillingService.emitAiTokenUsageEvent(
        workspaceId,
+        creditsUsedMicro,
+        totalTokens,
+        modelId,
        operationType,
        agent?.id ?? null,
        userWorkspaceId,
@@ -5,4 +5,5 @@ export interface AgentExecutionResult {
  usage: LanguageModelUsage;
  cacheCreationTokens: number;
  nativeWebSearchCallCount: number;
+  hasNoMoreAvailableCredits: boolean;
 }
@@ -85,7 +85,9 @@ describe('AiBillingService', () => {
        {
          provide: BillingUsageService,
          useValue: {
-            decrementAvailableCredits: jest.fn().mockResolvedValue(undefined),
+            decrementAvailableCreditsInCache: jest
+              .fn()
+              .mockResolvedValue(undefined),
          },
        },
        {
@@ -76,6 +76,13 @@ export class AiBillingService {
      (billingInput.usage.outputTokens ?? 0) +
      (billingInput.cacheCreationTokens ?? 0);

+    if (this.billingService.isBillingEnabled()) {
+      await this.billingUsageService.decrementAvailableCreditsInCache({
+        workspaceId,
+        usedCredits: creditsUsedMicro,
+      });
+    }
+
    await this.emitAiTokenUsageEvent(
      workspaceId,
      creditsUsedMicro,
@@ -87,6 +94,29 @@ export class AiBillingService {
    );
  }

+  async decrementAndCheckAvailableCredits(
+    modelId: ModelId,
+    billingInput: BillingUsageInput,
+    workspaceId: string,
+  ): Promise<{ hasNoMoreAvailableCredits: boolean }> {
+    if (!this.billingService.isBillingEnabled()) {
+      return { hasNoMoreAvailableCredits: false };
+    }
+
+    const costInDollars = this.calculateCost(modelId, billingInput);
+    const creditsUsedMicro = Math.round(
+      convertDollarsToBillingCredits(costInDollars),
+    );
+
+    const remainingCredits =
+      await this.billingUsageService.decrementAvailableCreditsInCache({
+        workspaceId,
+        usedCredits: creditsUsedMicro,
+      });
+
+    return { hasNoMoreAvailableCredits: remainingCredits <= 0 };
+  }
+
  async billNativeWebSearchUsage(
    nativeWebSearchCallCount: number,
    workspaceId: string,
@@ -117,7 +147,7 @@ export class AiBillingService {

      periodStart = currentPeriodStart;

-      await this.billingUsageService.decrementAvailableCredits({
+      await this.billingUsageService.decrementAvailableCreditsInCache({
        workspaceId,
        usedCredits: creditsUsedMicro,
      });
@@ -140,7 +170,7 @@ export class AiBillingService {
    );
  }

-  private async emitAiTokenUsageEvent(
+  async emitAiTokenUsageEvent(
    workspaceId: string,
    creditsUsedMicro: number,
    totalTokens: number,
@@ -159,11 +189,6 @@ export class AiBillingService {
      ]);

      periodStart = currentPeriodStart;
-
-      await this.billingUsageService.decrementAvailableCredits({
-        workspaceId,
-        usedCredits: creditsUsedMicro,
-      });
    }

    this.workspaceEventEmitter.emitCustomBatchEvent<UsageEvent>(
@@ -193,6 +193,7 @@ export class StreamAgentChatJob {
      let lastStepConversationSize = 0;
      let totalCacheCreationTokens = 0;
      let streamError: unknown;
+      let checkHasNoMoreAvailableCredits: () => boolean = () => false;

      // onFinish fires before the uiStream is fully drained. We use this
      // promise to coordinate: the IIFE waits for DB persist to complete
@@ -224,7 +225,7 @@ export class StreamAgentChatJob {
            });
          };

-          const { stream, modelConfig } =
+          const { stream, modelConfig, hasNoMoreAvailableCredits } =
            await this.chatExecutionService.streamChat({
              workspace,
              userWorkspaceId: data.userWorkspaceId,
@@ -237,6 +238,8 @@ export class StreamAgentChatJob {
              conversationSizeTokens: data.conversationSizeTokens,
            });

+          checkHasNoMoreAvailableCredits = hasNoMoreAvailableCredits;
+
          const titleWritePromise = titlePromise.then((generatedTitle) => {
            if (generatedTitle) {
              writer.write({
@@ -315,6 +318,13 @@ export class StreamAgentChatJob {

          if (streamError) {
            reject(streamError);
+          } else if (checkHasNoMoreAvailableCredits()) {
+            await this.eventPublisherService.publish({
+              threadId: data.threadId,
+              workspaceId: data.workspaceId,
+              event: { type: 'credits-exhausted' },
+            });
+            resolve();
          } else {
            await this.eventPublisherService.publish({
              threadId: data.threadId,
@@ -38,8 +38,12 @@ import { AGENT_CONFIG } from 'src/engine/metadata-modules/ai/ai-agent/constants/
 import { type BrowsingContextType } from 'src/engine/metadata-modules/ai/ai-agent/types/browsingContext.type';
 import { repairToolCall } from 'src/engine/metadata-modules/ai/ai-agent/utils/repair-tool-call.util';
 import { AiBillingService } from 'src/engine/metadata-modules/ai/ai-billing/services/ai-billing.service';
+import { convertDollarsToBillingCredits } from 'src/engine/metadata-modules/ai/ai-billing/utils/convert-dollars-to-billing-credits.util';
 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 {
+  extractCacheCreationTokens,
+  extractCacheCreationTokensFromSteps,
+} from 'src/engine/metadata-modules/ai/ai-billing/utils/extract-cache-creation-tokens.util';
 import { AI_CHAT_TOOL_NAMES_TO_PRELOAD } from 'src/engine/metadata-modules/ai/ai-chat/constants/ai-chat-tool-names-to-preload.const';
 import { MessagePruningService } from 'src/engine/metadata-modules/ai/ai-chat/services/message-pruning.service';
 import { SystemPromptBuilderService } from 'src/engine/metadata-modules/ai/ai-chat/services/system-prompt-builder.service';
@@ -48,9 +52,9 @@ import {
  type ExtractedFile,
 } from 'src/engine/metadata-modules/ai/ai-chat/utils/extract-code-interpreter-files.util';
 import {
-  injectCacheBreakpoint,
  getCacheProviderOptions,
  getCallLevelCacheProviderOptions,
+  injectCacheBreakpoint,
 } from 'src/engine/metadata-modules/ai/ai-chat/utils/inject-cache-breakpoint.util';
 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';
@@ -72,6 +76,7 @@ export type ChatExecutionOptions = {
 export type ChatExecutionResult = {
  stream: ReturnType<typeof streamText>;
  modelConfig: AiModelConfig;
+  hasNoMoreAvailableCredits: () => boolean;
 };

@Injectable()
@@ -262,7 +267,9 @@ export class ChatExecutionService {

    const modelMessages = pruningResult.messages;

-    const billUsageFromSteps = async (steps: StepResult<ToolSet>[]) => {
+    let hasNoMoreAvailableCredits = false;
+
+    const emitTurnUsageEvent = async (steps: StepResult<ToolSet>[]) => {
      const usage = steps.reduce<LanguageModelUsage>(
        (acc, step) => ({
          inputTokens: (acc.inputTokens ?? 0) + (step.usage.inputTokens ?? 0),
@@ -303,11 +310,24 @@ export class ChatExecutionService {
      );

      const cacheCreationTokens = extractCacheCreationTokensFromSteps(steps);
+      const totalTokens =
+        (usage.inputTokens ?? 0) +
+        (usage.outputTokens ?? 0) +
+        cacheCreationTokens;

-      await this.aiBillingService.calculateAndBillUsage(
+      const costInDollars = this.aiBillingService.calculateCost(
        registeredModel.modelId,
        { usage, cacheCreationTokens },
+      );
+      const creditsUsedMicro = Math.round(
+        convertDollarsToBillingCredits(costInDollars),
+      );
+
+      await this.aiBillingService.emitAiTokenUsageEvent(
        workspace.id,
+        creditsUsedMicro,
+        totalTokens,
+        registeredModel.modelId,
        UsageOperationType.AI_CHAT_TOKEN,
        null,
        userWorkspaceId,
@@ -327,7 +347,8 @@ export class ChatExecutionService {
      messages: [systemMessage, ...modelMessages],
      tools: activeTools,
      abortSignal,
-      stopWhen: stepCountIs(AGENT_CONFIG.MAX_STEPS),
+      stopWhen: (step) =>
+        stepCountIs(AGENT_CONFIG.MAX_STEPS)(step) || hasNoMoreAvailableCredits,
      experimental_telemetry: AI_TELEMETRY_CONFIG,
      providerOptions: getCallLevelCacheProviderOptions(
        registeredModel.sdkPackage,
@@ -335,8 +356,25 @@ export class ChatExecutionService {
      prepareStep: ({ messages }) => ({
        messages: injectCacheBreakpoint(messages, registeredModel.sdkPackage),
      }),
+      onStepFinish: async (step) => {
+        const { hasNoMoreAvailableCredits: stepHasNoMoreAvailableCredits } =
+          await this.aiBillingService.decrementAndCheckAvailableCredits(
+            registeredModel.modelId,
+            {
+              usage: step.usage,
+              cacheCreationTokens: extractCacheCreationTokens(
+                step.providerMetadata,
+              ),
+            },
+            workspace.id,
+          );
+
+        if (stepHasNoMoreAvailableCredits) {
+          hasNoMoreAvailableCredits = true;
+        }
+      },
      onAbort: async ({ steps }) => {
-        await billUsageFromSteps(steps);
+        await emitTurnUsageEvent(steps);
      },
      experimental_repairToolCall: async ({
        toolCall,
@@ -363,7 +401,7 @@ export class ChatExecutionService {

    Promise.all([stream.usage, stream.steps])
      .then(async ([, steps]) => {
-        await billUsageFromSteps(steps);
+        await emitTurnUsageEvent(steps);
      })
      .catch((error) => {
        if (error?.name === 'AbortError') {
@@ -375,6 +413,7 @@ export class ChatExecutionService {
    return {
      stream,
      modelConfig,
+      hasNoMoreAvailableCredits: () => hasNoMoreAvailableCredits,
    };
  }