diff --git a/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/constants/agent-chat-checkpoint-interval-ms.constant.ts b/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/constants/agent-chat-checkpoint-interval-ms.constant.ts new file mode 100644 index 0000000000..950b4bc4d7 --- /dev/null +++ b/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/constants/agent-chat-checkpoint-interval-ms.constant.ts @@ -0,0 +1 @@ +export const AGENT_CHAT_CHECKPOINT_INTERVAL_MS = 2_000; diff --git a/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/jobs/__tests__/stream-agent-chat.job.spec.ts b/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/jobs/__tests__/stream-agent-chat.job.spec.ts index f443b8014b..34e48ac5b1 100644 --- a/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/jobs/__tests__/stream-agent-chat.job.spec.ts +++ b/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/jobs/__tests__/stream-agent-chat.job.spec.ts @@ -86,11 +86,15 @@ describe('StreamAgentChatJob', () => { chatStream = createFakeChatStream(), streamChatRejection, addMessageRejection, + assistantPersistRejection, + assistantMessageExistedAtStart = false, }: { workspaceFound?: boolean; chatStream?: ReturnType; streamChatRejection?: Error; addMessageRejection?: Error; + assistantPersistRejection?: Error; + assistantMessageExistedAtStart?: boolean; } = {}) => { const publishedEvents: PublishedEvent[] = []; @@ -107,7 +111,12 @@ describe('StreamAgentChatJob', () => { addMessage: addMessageRejection ? jest.fn().mockRejectedValue(addMessageRejection) : jest.fn().mockResolvedValue({ id: 'assistant-message-id' }), - hasMessageById: jest.fn().mockResolvedValue(false), + upsertAssistantMessage: assistantPersistRejection + ? jest.fn().mockRejectedValue(assistantPersistRejection) + : jest.fn().mockResolvedValue(undefined), + hasMessageById: jest + .fn() + .mockResolvedValue(assistantMessageExistedAtStart), generateTitleIfNeeded: jest.fn().mockResolvedValue(null), notifyThreadUsageUpdated: jest.fn().mockResolvedValue(undefined), }; @@ -190,8 +199,14 @@ describe('StreamAgentChatJob', () => { expect(publishedEvents[publishedEvents.length - 1]).toMatchObject({ type: 'message-persisted', }); - expect(agentChatService.addMessage).toHaveBeenCalledWith( - expect.objectContaining({ turnId: 'turn-id' }), + expect(agentChatService.upsertAssistantMessage).toHaveBeenCalledWith( + expect.objectContaining({ + id: expect.any(String), + turnId: 'turn-id', + parts: expect.arrayContaining([ + expect.objectContaining({ type: 'text' }), + ]), + }), ); expect(threadRepository.update).toHaveBeenCalledWith( 'workspace-id', @@ -206,6 +221,28 @@ describe('StreamAgentChatJob', () => { expect(agentChatStreamingService.flushNextQueuedMessage).toHaveBeenCalled(); }); + it('persists the assistant message but does not re-apply thread totals when a prior execution already persisted it', async () => { + const { job, agentChatService, threadRepository } = buildJob({ + assistantMessageExistedAtStart: true, + }); + + await job.handle(jobData); + + expect(agentChatService.upsertAssistantMessage).toHaveBeenCalledWith( + expect.objectContaining({ turnId: 'turn-id' }), + ); + // The thread-totals accumulation must not run twice for the same stream. + const totalsUpdate = threadRepository.update.mock.calls.find( + ([, criteria]) => + criteria && + typeof criteria === 'object' && + !('activeStreamId' in criteria), + ); + + expect(totalsUpdate).toBeUndefined(); + expect(agentChatService.notifyThreadUsageUpdated).not.toHaveBeenCalled(); + }); + it('never publishes the opaque error chunk to subscribers', async () => { const { job, publishedEvents } = buildJob({ chatStream: createFakeChatStream({ @@ -289,7 +326,7 @@ describe('StreamAgentChatJob', () => { it('terminates the stream with an error when assistant persistence fails after draining chunks', async () => { const { job, publishedEvents } = buildJob({ - addMessageRejection: new Error('insert failed'), + assistantPersistRejection: new Error('insert failed'), }); await expect(job.handle(jobData)).rejects.toThrow('insert failed'); 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 2ab12c30da..84aef08ec8 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 @@ -2,7 +2,7 @@ import { Logger, Scope } from '@nestjs/common'; import { InjectRepository } from '@nestjs/typeorm'; import { isNonEmptyString } from '@sniptt/guards'; -import { createUIMessageStream } from 'ai'; +import { createUIMessageStream, readUIMessageStream } from 'ai'; import type { CodeExecutionData, ExtendedUIMessage, @@ -33,6 +33,7 @@ import { AgentChatStreamingService } from 'src/engine/metadata-modules/ai/ai-cha import { AgentChatService } from 'src/engine/metadata-modules/ai/ai-chat/services/agent-chat.service'; import { ChatExecutionService } from 'src/engine/metadata-modules/ai/ai-chat/services/chat-execution.service'; import { findPendingQuestionPart } from 'src/engine/metadata-modules/ai/ai-chat/utils/find-pending-question-part.util'; +import { AGENT_CHAT_CHECKPOINT_INTERVAL_MS } from 'src/engine/metadata-modules/ai/ai-chat/constants/agent-chat-checkpoint-interval-ms.constant'; import { getCancelChannel } from 'src/engine/metadata-modules/ai/ai-chat/utils/get-cancel-channel.util'; import { mapErrorToStreamError } from 'src/engine/metadata-modules/ai/ai-chat/utils/map-error-to-stream-error.util'; import type { AiModelConfig } from 'src/engine/metadata-modules/ai/ai-models/types/ai-model-config.type'; @@ -221,12 +222,18 @@ export class StreamAgentChatJob { titlePromise: Promise; abortSignal: AbortSignal; }): Promise { - return new Promise((resolve, reject) => { - const assistantMessageId = uuidv5( - data.streamId, - ASSISTANT_MESSAGE_ID_NAMESPACE, - ); + const assistantMessageId = uuidv5( + data.streamId, + ASSISTANT_MESSAGE_ID_NAMESPACE, + ); + const assistantMessageExistedAtStart = + await this.agentChatService.hasMessageById({ + id: assistantMessageId, + workspaceId: data.workspaceId, + }); + + return new Promise((resolve, reject) => { let streamUsage = { inputTokens: 0, outputTokens: 0, @@ -240,6 +247,22 @@ export class StreamAgentChatJob { let streamFinishError: unknown; let checkHasNoMoreAvailableCredits: () => boolean = () => false; + let persistChain: Promise = Promise.resolve(); + let lastCheckpointAt = 0; + let isFinalizingPersist = false; + + const enqueueAssistantPersist = ( + persist: () => Promise, + ): Promise => { + persistChain = persistChain.then(persist).catch((error) => { + this.logger.warn( + `Failed to checkpoint assistant message for stream ${data.streamId}: ${error instanceof Error ? error.message : String(error)}`, + ); + }); + + return persistChain; + }; + // onFinish fires before the uiStream is fully drained. We use this // promise to coordinate: the IIFE waits for DB persist to complete // before publishing message-persisted (after all chunks). @@ -325,8 +348,11 @@ export class StreamAgentChatJob { onFinish: async ({ responseMessage, isAborted }) => { // Rejecting here would race chunks still draining. try { + isFinalizingPersist = true; + await persistChain; await this.handleStreamFinish({ assistantMessageId, + assistantMessageExistedAtStart, responseMessage, isAborted, streamError, @@ -360,11 +386,57 @@ export class StreamAgentChatJob { }, }); + const [publishStream, checkpointStream] = uiStream.tee(); + + void (async () => { + try { + for await (const message of readUIMessageStream({ + stream: checkpointStream, + terminateOnError: false, + })) { + if (isFinalizingPersist || message.parts.length === 0) { + continue; + } + + const now = Date.now(); + + if (now - lastCheckpointAt < AGENT_CHAT_CHECKPOINT_INTERVAL_MS) { + continue; + } + + lastCheckpointAt = now; + const parts = message.parts; + + void enqueueAssistantPersist(async () => { + if (isFinalizingPersist) { + return; + } + + const { turnId } = await userMessagePromise; + + if (!isDefined(turnId)) { + return; + } + + await this.agentChatService.upsertAssistantMessage({ + id: assistantMessageId, + threadId: data.threadId, + turnId, + parts, + workspaceId: data.workspaceId, + }); + }); + } + } catch { + // best-effort; the authoritative persist runs onFinish + } + })(); + // Publish all chunks first, then signal completion. This guarantees // message-persisted arrives after every stream-chunk on the client. void (async () => { try { - for await (const chunk of uiStream) { + for await (const chunk of publishStream) { if ((chunk as { type?: string }).type === 'error') { continue; } @@ -500,6 +572,7 @@ export class StreamAgentChatJob { private async handleStreamFinish({ assistantMessageId, + assistantMessageExistedAtStart, responseMessage, isAborted, streamError, @@ -514,6 +587,7 @@ export class StreamAgentChatJob { userMessagePromise, }: { assistantMessageId: string; + assistantMessageExistedAtStart: boolean; responseMessage: Omit; isAborted: boolean; streamError: unknown; @@ -567,26 +641,26 @@ export class StreamAgentChatJob { const userMessage = await userMessagePromise; - // Idempotent per stream: assistantMessageId is derived from the streamId, - // so a retried job for this stream is skipped here while each distinct - // resume in the turn persists its own message. - const assistantMessageAlreadyPersisted = - await this.agentChatService.hasMessageById({ + if (isDefined(userMessage.turnId)) { + await this.agentChatService.upsertAssistantMessage({ + id: assistantMessageId, + threadId, + turnId: userMessage.turnId, + parts: responseMessage.parts, + workspaceId, + }); + } else if (!assistantMessageExistedAtStart) { + await this.agentChatService.addMessage({ + threadId, + uiMessage: responseMessage, id: assistantMessageId, workspaceId, }); - - if (assistantMessageAlreadyPersisted) { - return; } - await this.agentChatService.addMessage({ - threadId, - uiMessage: responseMessage, - id: assistantMessageId, - turnId: userMessage.turnId ?? undefined, - workspaceId, - }); + if (assistantMessageExistedAtStart) { + return; + } await this.threadRepository.update( workspaceId, 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 5a7ef7b3e4..1b504374ec 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 @@ -266,6 +266,47 @@ export class AgentChatService { } as AgentMessageEntity; } + async upsertAssistantMessage({ + id, + threadId, + turnId, + parts, + workspaceId, + }: { + id: string; + threadId: string; + turnId: string; + parts: ExtendedUIMessage['parts']; + workspaceId: string; + }): Promise { + await this.messageRepository.upsert( + workspaceId, + { + id, + threadId, + turnId, + role: AgentMessageRole.ASSISTANT, + processedAt: new Date(), + }, + ['id'], + ); + + await this.messagePartRepository.delete(workspaceId, { messageId: id }); + + const dbParts = mapUIMessagePartsToDBParts( + finalizeDanglingToolParts(parts), + id, + workspaceId, + ); + + if (dbParts.length > 0) { + await this.messagePartRepository.insert( + workspaceId, + dbParts as QueryDeepPartialEntity[], + ); + } + } + async findLatestSentUserMessage({ threadId, workspaceId,