feat(ai): persist assistant chat messages progressively during streaming (#22524)

## 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.  <a href="https://cubic.dev/pr/twentyhq/twenty/pull/22524?utm_source=github" target="_blank" rel="noopener noreferrer" data-no-image-dialog="true"><picture><source media="(prefers-color-scheme: dark)" srcset="https://www.cubic.dev/buttons/review-in-cubic-dark.svg"><source media="(prefers-color-scheme: light)" srcset="https://www.cubic.dev/buttons/review-in-cubic-light.svg"><img alt="Review in cubic" src="https://www.cubic.dev/buttons/review-in-cubic-dark.svg"></picture></a>
2026-07-03 18:14:56 +02:00
parent 43730d7748
commit 9496a98aa3
4 changed files with 179 additions and 26 deletions
@@ -0,0 +1 @@
+export const AGENT_CHAT_CHECKPOINT_INTERVAL_MS = 2_000;
@@ -86,11 +86,15 @@ describe('StreamAgentChatJob', () => {
    chatStream = createFakeChatStream(),
    streamChatRejection,
    addMessageRejection,
+    assistantPersistRejection,
+    assistantMessageExistedAtStart = false,
  }: {
    workspaceFound?: boolean;
    chatStream?: ReturnType<typeof createFakeChatStream>;
    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');
@@ -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<string | null>;
    abortSignal: AbortSignal;
  }): Promise<void> {
-    return new Promise<void>((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<void>((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<void> = Promise.resolve();
+      let lastCheckpointAt = 0;
+      let isFinalizingPersist = false;
+
+      const enqueueAssistantPersist = (
+        persist: () => Promise<void>,
+      ): Promise<void> => {
+        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<ExtendedUIMessage>({
+            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<ExtendedUIMessage, 'id'>;
    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,
@@ -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<void> {
+    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<AgentMessagePartEntity>[],
+      );
+    }
+  }
+
  async findLatestSentUserMessage({
    threadId,
    workspaceId,
				`@@ -0,0 +1 @@`
				`export const AGENT_CHAT_CHECKPOINT_INTERVAL_MS = 2_000;`