From 5c825e8712d907cb9d8a922022a918c5cf019ca2 Mon Sep 17 00:00:00 2001 From: Etienne <45695613+etiennejouan@users.noreply.github.com> Date: Thu, 23 Jul 2026 14:17:50 +0200 Subject: [PATCH] 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. Review in cubic --- .../ai-chat/resolvers/agent-chat.resolver.ts | 50 ++---- ...gent-chat-streaming.service.answer.spec.ts | 167 ++++++++++++++++++ ...agent-chat-streaming.service.claim.spec.ts | 19 +- .../services/agent-chat-streaming.service.ts | 81 ++++++++- 4 files changed, 271 insertions(+), 46 deletions(-) create mode 100644 packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/services/__tests__/agent-chat-streaming.service.answer.spec.ts diff --git a/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/resolvers/agent-chat.resolver.ts b/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/resolvers/agent-chat.resolver.ts index 7088240a6d..31de99a8b1 100644 --- a/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/resolvers/agent-chat.resolver.ts +++ b/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/resolvers/agent-chat.resolver.ts @@ -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, diff --git a/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/services/__tests__/agent-chat-streaming.service.answer.spec.ts b/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/services/__tests__/agent-chat-streaming.service.answer.spec.ts new file mode 100644 index 0000000000..aafa4807eb --- /dev/null +++ b/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/services/__tests__/agent-chat-streaming.service.answer.spec.ts @@ -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); + }); +}); diff --git a/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/services/__tests__/agent-chat-streaming.service.claim.spec.ts b/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/services/__tests__/agent-chat-streaming.service.claim.spec.ts index 65c6d48d6c..15028d9ce0 100644 --- a/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/services/__tests__/agent-chat-streaming.service.claim.spec.ts +++ b/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/services/__tests__/agent-chat-streaming.service.claim.spec.ts @@ -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(); }); }); diff --git a/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/services/agent-chat-streaming.service.ts b/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/services/agent-chat-streaming.service.ts index dfc3a3f2c7..5be71c186f 100644 --- a/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/services/agent-chat-streaming.service.ts +++ b/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/services/agent-chat-streaming.service.ts @@ -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; }): Promise { + 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 }; + }; + + 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( 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(