From 9496a98aa3fd77f23a742a3a91b0eb6e51d088d6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?F=C3=A9lix=20Malfait?= Date: Fri, 3 Jul 2026 18:14:56 +0200 Subject: [PATCH] feat(ai): persist assistant chat messages progressively during streaming (#22524) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Why Today the assistant chat message is written to the DB exactly **once**, at `onFinish` (`handleStreamFinish`). The per-step hook (`onStepFinish` in `chat-execution.service.ts`) only does billing/metrics — no persistence. Content streams to the client live over Redis, but the durable record only lands at the end. The graceful paths already cover partials: normal completion, user cancel, and the shutdown-abort from #22517 all fire `onFinish` with `isAborted` and persist whatever parts exist. The gap is a **true SIGKILL** — OOM, node loss, or a grace-period overrun — where `onFinish` never runs. When that happens mid-turn, the assistant message vanishes from the thread even though its tool calls already executed real CRM mutations. That's the worst failure shape: side effects persisted, the record of them didn't. This closes that gap by materializing the assistant message progressively, so a hard kill leaves the tools-already-run on the thread. It's the app-level piece behind the earlier discussion on #22518 — with this, a retried/interrupted turn also resumes from its own partial (the model reloads history and continues) instead of re-doing completed steps. Scope is **chat only** — the workflow agent path (`AgentAsyncExecutorService`, blocking `generateText`) is a different model with its own step-log persistence and workflow-engine resumption, and is deliberately out of scope here. ## What - `AgentChatService.upsertAssistantMessage`: idempotent message+parts write keyed on the deterministic `uuidv5(streamId)` id (upsert the row, replace its parts), reusing the existing `mapUIMessagePartsToDBParts` / `finalizeDanglingToolParts`. - `stream-agent-chat.job.ts`: tee the assembled UI stream — one branch keeps publishing chunks unchanged; the other drives the SDK's own `readUIMessageStream` and, throttled to `AGENT_CHAT_CHECKPOINT_INTERVAL_MS` (2s), fires a serialized fire-and-forget `upsertAssistantMessage`. No chunk re-assembly — the parts come straight from the SDK assembler, identical to what `onFinish` produces. - `handleStreamFinish` now upserts (authoritative) instead of insert-then-skip. Two ordering/idempotency guards: - Checkpoints are serialized through one promise chain and gated off (`isFinalizingPersist`) before the final write, which drains the chain first — so the authoritative write always lands last and never races a checkpoint on the parts table. - The old `hasMessageById`→skip protected the thread-totals accumulation from double-counting on a re-executed job. Since checkpoints now make the row exist mid-stream, that signal is captured **once at stream start** (`assistantMessageExistedAtStart`) and used to gate the totals update — preserving the exact prior idempotency while allowing progressive writes. `readUIMessageStream` runs with `terminateOnError: false` and the checkpoint consumer swallows errors: checkpoints are best-effort and must never affect the stream or the authoritative persist. ## User impact A worker that dies hard mid-turn no longer erases the assistant message. Combined with #22434's Retry, the user sees the partial turn (including executed tools) and can continue, rather than a turn that silently disappeared while its side effects stuck. Note: this is insurance against true SIGKILL specifically — graceful shutdown (#22514/#22517) already persists partials — so it's most valuable for OOM/node-loss/grace-overrun. Framed that way deliberately; happy to drop it if you'd rather not touch this path for that scope. ## Validation - Unit (`stream-agent-chat.job.spec.ts`, 59 green): the success path persists via `upsertAssistantMessage` with the assembled parts + turnId; a re-executed job whose message existed at start still upserts but does **not** re-apply thread totals; all existing flows (mid-stream error, user cancel, shutdown-abort, missing workspace) unchanged. - Local runtime (isolated instance, real OpenAI stream, checkpoint interval shortened for the test): - **SIGKILL mid-stream** (no graceful onFinish) → the assistant message row (deterministic id) is present afterward with a partial text part (~381 chars) that would otherwise have been lost. - **Normal completion** → the final upsert converges to the full message ("Hello, Tim."), `activeStreamId` cleared, `totalOutputTokens` applied exactly once. Review in cubic --- ...nt-chat-checkpoint-interval-ms.constant.ts | 1 + .../__tests__/stream-agent-chat.job.spec.ts | 45 ++++++- .../ai/ai-chat/jobs/stream-agent-chat.job.ts | 118 ++++++++++++++---- .../ai/ai-chat/services/agent-chat.service.ts | 41 ++++++ 4 files changed, 179 insertions(+), 26 deletions(-) create mode 100644 packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/constants/agent-chat-checkpoint-interval-ms.constant.ts 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,