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.
This commit is contained in:
Etienne
2026-05-13 17:09:08 +02:00
committed by GitHub
parent 068fc9363d
commit aecfe699f4
14 changed files with 212 additions and 52 deletions
@@ -19,10 +19,11 @@ import { agentChatIsStreamingComponentFamilyState } from '@/ai/states/agentChatI
import { agentChatMessagesComponentFamilyState } from '@/ai/states/agentChatMessagesComponentFamilyState';
import { agentChatUsageComponentFamilyState } from '@/ai/states/agentChatUsageComponentFamilyState';
import { currentAiChatThreadTitleComponentFamilyState } from '@/ai/states/currentAiChatThreadTitleComponentFamilyState';
import { AiChatErrorCode } from '@/ai/utils/aiChatErrorCode';
import { dispatchBrowserEvent } from '@/browser-event/utils/dispatchBrowserEvent';
import { sseClientState } from '@/sse-db-event/states/sseClientState';
import { useAtomComponentFamilyStateCallbackState } from '@/ui/utilities/state/jotai/hooks/useAtomComponentFamilyStateCallbackState';
import { useAtomStateValue } from '@/ui/utilities/state/jotai/hooks/useAtomStateValue';
import { sseClientState } from '@/sse-db-event/states/sseClientState';
const THROTTLE_MS = 100;
@@ -313,6 +314,20 @@ export const useAgentChatSubscription = (threadId: string | null) => {
store.set(isStreamingAtom, false);
break;
}
case 'credits-exhausted': {
const noMoreCreditsError = new Error(
'Chat stopped: no more available credits.',
) as Error & { code?: string };
noMoreCreditsError.code = AiChatErrorCode.BILLING_CREDITS_EXHAUSTED;
store.set(errorAtom, noMoreCreditsError);
closeWriter();
dispatchBrowserEvent(AGENT_CHAT_REFETCH_MESSAGES_EVENT_NAME);
store.set(isStreamingAtom, false);
break;
}
}
};
@@ -1,5 +1,5 @@
import { useEndSubscriptionTrialPeriod } from '@/settings/billing/hooks/useEndSubscriptionTrialPeriod';
import { InformationBanner } from '@/information-banner/components/InformationBanner';
import { useEndSubscriptionTrialPeriod } from '@/settings/billing/hooks/useEndSubscriptionTrialPeriod';
import { usePermissionFlagMap } from '@/settings/roles/hooks/usePermissionFlagMap';
import { useLingui } from '@lingui/react/macro';
import { PermissionFlagType } from '~/generated-metadata/graphql';
@@ -285,7 +285,7 @@ export class BillingUsageService {
);
}
private async warmAvailableCredits(
private async warmAvailableCreditsInCache(
workspaceId: string,
periodStart: Date | string,
periodEnd: Date | string,
@@ -421,13 +421,13 @@ export class BillingUsageService {
return Number(resourceCreditPrice.metadata?.credit_amount ?? 0);
}
async decrementAvailableCredits({
async decrementAvailableCreditsInCache({
workspaceId,
usedCredits,
}: {
workspaceId: string;
usedCredits: number;
}): Promise<void> {
}): Promise<number> {
const {
billingSubscription: { currentPeriodStart, currentPeriodEnd },
} = await this.workspaceCacheService.getOrRecompute(workspaceId, [
@@ -447,7 +447,7 @@ export class BillingUsageService {
});
if (!isDefined(cachedAvailableCredits)) {
await this.warmAvailableCredits(
await this.warmAvailableCreditsInCache(
workspaceId,
currentPeriodStart,
currentPeriodEnd,
@@ -470,9 +470,11 @@ export class BillingUsageService {
true,
);
}
return decrementedAvailableCredits;
}
async invalidateAvailableCredits(
async invalidateAvailableCreditsInCache(
workspaceId: string,
periodStart: Date,
): Promise<void> {
@@ -504,7 +506,7 @@ export class BillingUsageService {
currentPeriodStart: subscription.currentPeriodStart,
});
await this.warmAvailableCredits(
await this.warmAvailableCreditsInCache(
subscription.workspaceId,
subscription.currentPeriodStart,
subscription.currentPeriodEnd,
@@ -376,7 +376,7 @@ export class LogicFunctionExecutorService {
periodStart = currentPeriodStart;
await this.billingUsageService.decrementAvailableCredits({
await this.billingUsageService.decrementAvailableCreditsInCache({
workspaceId,
usedCredits: 100,
});
@@ -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,
};
}
@@ -76,18 +76,25 @@ export class AiAgentWorkflowAction implements WorkflowAction {
? executionContext.authContext.userWorkspaceId
: null;
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,
});
const { result, hasNoMoreAvailableCredits } =
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,
});
if (hasNoMoreAvailableCredits) {
return {
error: 'AI agent stopped: no more available credits.',
};
}
return {
result,
@@ -78,7 +78,7 @@ describe('WorkflowExecutorWorkspaceService', () => {
const mockBillingUsageService = {
hasAvailableCredits: jest.fn().mockResolvedValue(true),
decrementAvailableCredits: jest.fn().mockResolvedValue(undefined),
decrementAvailableCreditsInCache: jest.fn().mockResolvedValue(undefined),
};
const mockExceptionHandlerService = {
@@ -375,7 +375,7 @@ export class WorkflowExecutorWorkspaceService {
periodStart = currentPeriodStart;
await this.billingUsageService.decrementAvailableCredits({
await this.billingUsageService.decrementAvailableCreditsInCache({
workspaceId,
usedCredits: 100,
});
@@ -5,4 +5,5 @@ export type AgentChatSubscriptionEvent =
| { type: 'stream-chunk'; chunk: Record<string, unknown>; seq?: number }
| { type: 'message-persisted'; messageId: string }
| { type: 'queue-updated' }
| { type: 'stream-error'; code: string; message: string };
| { type: 'stream-error'; code: string; message: string }
| { type: 'credits-exhausted' };