diff --git a/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/jobs/stream-agent-chat.job.ts b/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/jobs/stream-agent-chat.job.ts index 89ab2d06f9..d7dd95a2b5 100644 --- a/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/jobs/stream-agent-chat.job.ts +++ b/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/jobs/stream-agent-chat.job.ts @@ -281,6 +281,7 @@ export class StreamAgentChatJob { responseMessage, threadId: data.threadId, workspaceId: data.workspaceId, + userWorkspaceId: data.userWorkspaceId, streamUsage, lastStepConversationSize, totalCacheCreationTokens, @@ -434,6 +435,7 @@ export class StreamAgentChatJob { responseMessage, threadId, workspaceId, + userWorkspaceId, streamUsage, lastStepConversationSize, totalCacheCreationTokens, @@ -443,6 +445,7 @@ export class StreamAgentChatJob { responseMessage: Omit; threadId: string; workspaceId: string; + userWorkspaceId: string; streamUsage: { inputTokens: number; outputTokens: number; @@ -492,5 +495,10 @@ export class StreamAgentChatJob { contextWindowTokens: modelConfig.contextWindowTokens, conversationSize: lastStepConversationSize, }); + + await this.agentChatService.notifyThreadUsageUpdated({ + threadId, + userWorkspaceId, + }); } } 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 e40fbc25eb..a05ad658c0 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 @@ -23,8 +23,10 @@ import { } from 'src/engine/metadata-modules/ai/ai.exception'; import { WorkspaceEventBroadcaster } from 'src/engine/subscriptions/workspace-event-broadcaster/workspace-event-broadcaster.service'; +import { toDisplayCredits } from 'src/engine/core-modules/usage/utils/to-display-credits.util'; import { AiChatFileAttachment } from 'src/engine/metadata-modules/ai/ai-chat/types/ai-chat-file-attachment.type'; import { AgentTitleGenerationService } from './agent-title-generation.service'; +import { AgentChatThreadDTO } from '../dtos/agent-chat-thread.dto'; const serializeThreadForBroadcast = ( thread: AgentChatThreadEntity, @@ -38,8 +40,8 @@ const serializeThreadForBroadcast = ( totalCacheCreationTokens: thread.totalCacheCreationTokens, contextWindowTokens: thread.contextWindowTokens, conversationSize: thread.conversationSize, - totalInputCredits: thread.totalInputCredits, - totalOutputCredits: thread.totalOutputCredits, + totalInputCredits: toDisplayCredits(thread.totalInputCredits), + totalOutputCredits: toDisplayCredits(thread.totalOutputCredits), deletedAt: thread.deletedAt, lastMessageAt, createdAt: thread.createdAt, @@ -538,9 +540,32 @@ export class AgentChatService { ); } + async notifyThreadUsageUpdated({ + threadId, + userWorkspaceId, + }: { + threadId: string; + userWorkspaceId: string; + }): Promise { + const thread = await this.getThreadById(threadId, userWorkspaceId); + + await this.broadcastThreadUpdated( + thread, + [ + 'totalInputTokens', + 'totalOutputTokens', + 'totalInputCredits', + 'totalOutputCredits', + 'conversationSize', + 'contextWindowTokens', + ], + userWorkspaceId, + ); + } + private async broadcastThreadUpdated( thread: AgentChatThreadEntity, - updatedFields: string[], + updatedFields: (keyof AgentChatThreadDTO)[], userWorkspaceId: string, ): Promise { const lastMessageAt = await this.getLastMessageAtForThread(thread.id);