fix(ai-chat): write the stream heartbeat before the DB claim (#23198)

## Problem Answering an `ask_questions` (select) prompt sometimes killed the turn with "Failed to get response. The response was interrupted before it could finish." The answer was swallowed and Retry rewound the whole turn. It was intermittent, worse on long threads and when coming back from another tab. Reported in [discord quality issue](https://discord.com/channels/1130383047699738754/1526875783170097172). Confirmed in prod: ~28 `ai_chat_turn_failed_total{failure_phase="interrupted"}` over the last 7 days (the only failure phase firing), plus matching `the thread no longer holds this claim` worker logs around the report time. ## Root cause A stream is tracked by two records: the claim (`activeStreamId` in Postgres) and the heartbeat (a Redis key refreshed while the worker runs). `reapDeadStream` treats "claim set but no heartbeat" as a crashed worker and kills the turn. On the answer path the ordering left a window where that was falsely true: 1. `resolvePendingQuestion` writes `activeStreamId` to Postgres (claim set) 2. `enqueueResumeStream` reloads the thread and runs `loadMessagesFromDB` (reads every message and part, signs a URL per file, hundreds of ms on long threads) 3. only then `markClaimed` writes the heartbeat Between 1 and 3 the thread looks dead to the reaper. Worse, `question-answered` was published inside that window, so the client refetched, and the refetch's `chatStreamCatchupChunks` query runs the reaper, racing the server into its own setup window. The keepalive reap tick could land there too. ## Fix Enforce one invariant everywhere: the heartbeat exists before any DB row carries the `activeStreamId`, so "claim without heartbeat" can only ever mean a genuinely dead worker. - New `answerPendingQuestionAndResumeStream` owns the answer flow: `markClaimed` first, then the DB claim, then enqueue, then publish `question-answered` (moved after the enqueue so client refetches can't race the setup, and so we don't tell the client "answered" when the enqueue failed and rolled back). - Both failure paths clean up: clear the heartbeat if resolving fails; restore the pending question and clear the heartbeat if enqueueing fails. - `tryClaimStream` (send / retry / queue-flush) reordered the same way: heartbeat before the claim, cleared if the claim is lost. - `releaseStreamClaim` now also clears the heartbeat so failed claims leave no orphan key. No grace period or schema change needed: the ordering closes the race structurally. The Retry-rewinds-the-turn behavior is unrelated and left as a separate follow-up. ## Testing - New `agent-chat-streaming.service.answer.spec.ts`: heartbeat marked before the claim, publish only after enqueue, both failure paths restore state and clear the key. - Extended `agent-chat-streaming.service.claim.spec.ts`: heartbeat-before-claim ordering and key cleanup on lost claim / failed enqueue. - Full ai-chat suite green (74 tests), lint and typecheck clean. After deploy, `sum(increase(ai_chat_turn_failed_total{failure_phase="interrupted"}[1d]))` trending to zero confirms the fix.  <a href="https://cubic.dev/pr/twentyhq/twenty/pull/23198?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-23 14:17:50 +02:00
parent 9d8ce7c325
commit 5c825e8712
4 changed files with 271 additions and 46 deletions
@@ -8,7 +8,6 @@ import {
  ResolveField,
 } from '@nestjs/graphql';

-import { generateId } from 'ai';
 import GraphQLJSON from 'graphql-type-json';
 import { PermissionFlagType } from 'twenty-shared/constants';
 import { isDefined } from 'twenty-shared/utils';
@@ -352,44 +351,17 @@ export class AgentChatResolver {
      );
    }

-    const streamId = generateId();
-
-    const { turnId, rollback } =
-      await this.agentChatService.resolvePendingQuestion({
-        threadId,
-        messageId,
-        answers,
-        streamId,
-        workspaceId: workspace.id,
-      });
-
-    await this.eventPublisherService
-      .publish({
-        threadId,
-        workspaceId: workspace.id,
-        event: { type: 'question-answered' },
-      })
-      .catch(() => {});
-
-    try {
-      await this.agentChatStreamingService.enqueueResumeStream({
-        threadId,
-        userWorkspaceId,
-        workspace,
-        turnId,
-        streamId,
-        modelId,
-      });
-    } catch (error) {
-      await this.agentChatService.restorePendingQuestion({
-        threadId,
-        messageId,
-        streamId,
-        workspaceId: workspace.id,
-        rollback,
-      });
-      throw error;
-    }
+    const { streamId, turnId } =
+      await this.agentChatStreamingService.answerPendingQuestionAndResumeStream(
+        {
+          threadId,
+          messageId,
+          answers,
+          userWorkspaceId,
+          workspace,
+          modelId,
+        },
+      );

    tagAiChatStreamScope({
      streamId,
@@ -0,0 +1,167 @@
+import { type WorkspaceEntity } from 'src/engine/core-modules/workspace/workspace.entity';
+import { AgentChatStreamingService } from 'src/engine/metadata-modules/ai/ai-chat/services/agent-chat-streaming.service';
+
+describe('AgentChatStreamingService answerPendingQuestionAndResumeStream', () => {
+  const workspace = { id: 'workspace-id' } as WorkspaceEntity;
+
+  const thread = {
+    id: 'thread-id',
+    title: 'Thread',
+    conversationSize: 0,
+    activeStreamId: null,
+    lastStreamError: null,
+    pendingQuestionMessageId: 'question-message-id',
+  };
+
+  const buildService = () => {
+    const publishedEvents: Array<{ type: string }> = [];
+    const threadRepository = {
+      findOne: jest.fn().mockResolvedValue(thread),
+      findOneOrFail: jest.fn().mockResolvedValue(thread),
+      update: jest.fn().mockResolvedValue({ affected: 1 }),
+    };
+    const messageQueueService = { add: jest.fn().mockResolvedValue(undefined) };
+    const agentChatService = {
+      resolvePendingQuestion: jest.fn().mockResolvedValue({
+        turnId: 'turn-id',
+        rollback: { partId: 'part-id', previousOutput: {} },
+      }),
+      restorePendingQuestion: jest.fn().mockResolvedValue(undefined),
+      getMessagesForThread: jest.fn().mockResolvedValue([]),
+    };
+    const eventPublisherService = {
+      publish: jest.fn().mockImplementation(({ event }) => {
+        publishedEvents.push(event);
+
+        return Promise.resolve();
+      }),
+      resetStreamState: jest.fn().mockResolvedValue(undefined),
+    };
+    const streamHeartbeatService = {
+      markClaimed: jest.fn().mockResolvedValue(undefined),
+      isAlive: jest.fn().mockResolvedValue(true),
+      clear: jest.fn().mockResolvedValue(undefined),
+    };
+
+    const service = new AgentChatStreamingService(
+      threadRepository as never,
+      { find: jest.fn().mockResolvedValue([]) } as never,
+      messageQueueService as never,
+      agentChatService as never,
+      eventPublisherService as never,
+      { signFileByIdUrl: jest.fn() } as never,
+      streamHeartbeatService as never,
+      { incrementCounterBy: jest.fn() } as never,
+    );
+
+    return {
+      service,
+      threadRepository,
+      messageQueueService,
+      agentChatService,
+      eventPublisherService,
+      streamHeartbeatService,
+      publishedEvents,
+    };
+  };
+
+  const answerArguments = {
+    threadId: 'thread-id',
+    messageId: 'question-message-id',
+    answers: [{ questionIndex: 0, selectedOptionIndices: [0] }],
+    userWorkspaceId: 'user-workspace-id',
+    workspace,
+  };
+
+  it('marks the heartbeat before the pending question claims the thread', async () => {
+    const { service, agentChatService, streamHeartbeatService } =
+      buildService();
+
+    const result =
+      await service.answerPendingQuestionAndResumeStream(answerArguments);
+
+    expect(result).toEqual({
+      streamId: expect.any(String),
+      turnId: 'turn-id',
+    });
+    expect(streamHeartbeatService.markClaimed).toHaveBeenCalledWith(
+      result.streamId,
+    );
+    expect(
+      streamHeartbeatService.markClaimed.mock.invocationCallOrder[0],
+    ).toBeLessThan(
+      agentChatService.resolvePendingQuestion.mock.invocationCallOrder[0],
+    );
+    expect(streamHeartbeatService.clear).not.toHaveBeenCalled();
+  });
+
+  it('publishes question-answered only after the resume job is enqueued', async () => {
+    const { service, messageQueueService, eventPublisherService } =
+      buildService();
+
+    await service.answerPendingQuestionAndResumeStream(answerArguments);
+
+    expect(eventPublisherService.publish).toHaveBeenCalledWith(
+      expect.objectContaining({ event: { type: 'question-answered' } }),
+    );
+    expect(messageQueueService.add.mock.invocationCallOrder[0]).toBeLessThan(
+      eventPublisherService.publish.mock.invocationCallOrder[0],
+    );
+    expect(messageQueueService.add).toHaveBeenCalledWith(
+      expect.any(String),
+      expect.objectContaining({
+        threadId: 'thread-id',
+        existingTurnId: 'turn-id',
+        isResume: true,
+      }),
+    );
+  });
+
+  it('clears the heartbeat and rethrows when there is no pending question', async () => {
+    const {
+      service,
+      agentChatService,
+      messageQueueService,
+      streamHeartbeatService,
+      publishedEvents,
+    } = buildService();
+
+    agentChatService.resolvePendingQuestion.mockRejectedValue(
+      new Error('No pending question to answer'),
+    );
+
+    await expect(
+      service.answerPendingQuestionAndResumeStream(answerArguments),
+    ).rejects.toThrow('No pending question to answer');
+
+    expect(streamHeartbeatService.clear).toHaveBeenCalled();
+    expect(messageQueueService.add).not.toHaveBeenCalled();
+    expect(publishedEvents).toHaveLength(0);
+  });
+
+  it('restores the question and clears the heartbeat when enqueueing fails', async () => {
+    const {
+      service,
+      agentChatService,
+      messageQueueService,
+      streamHeartbeatService,
+      publishedEvents,
+    } = buildService();
+
+    messageQueueService.add.mockRejectedValue(new Error('redis down'));
+
+    await expect(
+      service.answerPendingQuestionAndResumeStream(answerArguments),
+    ).rejects.toThrow('redis down');
+
+    expect(agentChatService.restorePendingQuestion).toHaveBeenCalledWith(
+      expect.objectContaining({
+        threadId: 'thread-id',
+        messageId: 'question-message-id',
+        rollback: { partId: 'part-id', previousOutput: {} },
+      }),
+    );
+    expect(streamHeartbeatService.clear).toHaveBeenCalled();
+    expect(publishedEvents).toHaveLength(0);
+  });
+});
@@ -105,10 +105,18 @@ describe('AgentChatStreamingService claim & reap', () => {
        expect.objectContaining({ lastStreamError: null }),
      );
      expect(streamHeartbeatService.markClaimed).toHaveBeenCalled();
+      expect(
+        streamHeartbeatService.markClaimed.mock.invocationCallOrder[0],
+      ).toBeLessThan(threadRepository.update.mock.invocationCallOrder[0]);
    });

    it('queues the message when another stream wins the claim race', async () => {
-      const { service, agentChatService, messageQueueService } = buildService({
+      const {
+        service,
+        agentChatService,
+        messageQueueService,
+        streamHeartbeatService,
+      } = buildService({
        claimAffected: 0,
      });

@@ -117,6 +125,7 @@ describe('AgentChatStreamingService claim & reap', () => {
      expect(result.queued).toBe(true);
      expect(agentChatService.queueMessage).toHaveBeenCalled();
      expect(messageQueueService.add).not.toHaveBeenCalled();
+      expect(streamHeartbeatService.clear).toHaveBeenCalled();
    });

    it('queues behind a halted backlog and kicks the drain from the front', async () => {
@@ -139,7 +148,12 @@ describe('AgentChatStreamingService claim & reap', () => {
    });

    it('releases the claim when enqueueing the job fails', async () => {
-      const { service, threadRepository, messageQueueService } = buildService();
+      const {
+        service,
+        threadRepository,
+        messageQueueService,
+        streamHeartbeatService,
+      } = buildService();

      messageQueueService.add.mockRejectedValue(new Error('redis down'));

@@ -152,6 +166,7 @@ describe('AgentChatStreamingService claim & reap', () => {
        { id: 'thread-id', activeStreamId: expect.any(String) },
        { activeStreamId: null },
      );
+      expect(streamHeartbeatService.clear).toHaveBeenCalled();
    });
  });

@@ -2,6 +2,7 @@ import { Injectable, Logger } from '@nestjs/common';

 import { generateId } from 'ai';
 import {
+  type AskQuestionAnswer,
  type ExtendedFileUIPart,
  type ExtendedUIMessagePart,
  isExtendedFileUIPart,
@@ -135,6 +136,8 @@ export class AgentChatStreamingService {
    streamId: string;
    where: FindOptionsWhere<AgentChatThreadEntity>;
  }): Promise<boolean> {
+    await this.streamHeartbeatService.markClaimed(streamId);
+
    const claim = await this.threadRepository.update(
      workspaceId,
      { id: threadId, activeStreamId: IsNull(), ...where },
@@ -142,11 +145,11 @@ export class AgentChatStreamingService {
    );

    if (!claim.affected) {
+      await this.streamHeartbeatService.clear(streamId);
+
      return false;
    }

-    await this.streamHeartbeatService.markClaimed(streamId);
-
    return true;
  }

@@ -422,7 +425,76 @@ export class AgentChatStreamingService {
    }
  }

-  async enqueueResumeStream({
+  async answerPendingQuestionAndResumeStream({
+    threadId,
+    messageId,
+    answers,
+    userWorkspaceId,
+    workspace,
+    modelId,
+  }: {
+    threadId: string;
+    messageId: string;
+    answers: AskQuestionAnswer[];
+    userWorkspaceId: string;
+    workspace: WorkspaceEntity;
+    modelId?: string;
+  }): Promise<{ streamId: string; turnId: string | null }> {
+    const streamId = generateId();
+
+    await this.streamHeartbeatService.markClaimed(streamId);
+
+    let resolved: {
+      turnId: string | null;
+      rollback: { partId: string; previousOutput: Record<string, unknown> };
+    };
+
+    try {
+      resolved = await this.agentChatService.resolvePendingQuestion({
+        threadId,
+        messageId,
+        answers,
+        streamId,
+        workspaceId: workspace.id,
+      });
+    } catch (error) {
+      await this.streamHeartbeatService.clear(streamId);
+      throw error;
+    }
+
+    try {
+      await this.enqueueResumeStream({
+        threadId,
+        userWorkspaceId,
+        workspace,
+        turnId: resolved.turnId,
+        streamId,
+        modelId,
+      });
+    } catch (error) {
+      await this.agentChatService.restorePendingQuestion({
+        threadId,
+        messageId,
+        streamId,
+        workspaceId: workspace.id,
+        rollback: resolved.rollback,
+      });
+      await this.streamHeartbeatService.clear(streamId);
+      throw error;
+    }
+
+    await this.eventPublisherService
+      .publish({
+        threadId,
+        workspaceId: workspace.id,
+        event: { type: 'question-answered' },
+      })
+      .catch(() => {});
+
+    return { streamId, turnId: resolved.turnId };
+  }
+
+  private async enqueueResumeStream({
    threadId,
    userWorkspaceId,
    workspace,
@@ -447,8 +519,6 @@ export class AgentChatStreamingService {
      workspace.id,
    );

-    await this.streamHeartbeatService.markClaimed(streamId);
-
    await this.messageQueueService.add<StreamAgentChatJobData>(
      STREAM_AGENT_CHAT_JOB_NAME,
      {
@@ -624,6 +694,7 @@ export class AgentChatStreamingService {
          `Failed to release stream claim for thread ${threadId}: ${error instanceof Error ? error.message : String(error)}`,
        );
      });
+    await this.streamHeartbeatService.clear(streamId);
  }

  private async loadMessagesFromDB(